From 494aa623a50e238d3e0244f2b859a2df4e96c322 Mon Sep 17 00:00:00 2001 From: mac Date: Wed, 26 Aug 2026 23:15:01 +0300 Subject: [PATCH 01/94] server: validate sampling params and rebuild timeout at the API layer resolve_sampling forwarded arbitrary client floats into SamplingParams: negative, NaN and inf temperature, top_p outside (0, 1], top_k 0 -- all reached the engine and died (or misbehaved) deep in the sampler instead of failing the request. Validate the RESOLVED values (checkpoint defaults included) and raise ValueError in the style of the existing max_tokens check; every adapter already maps that to a clean 400. CacheRebuildRequest.timeout gets the matching pydantic bound (0 < t <= 3600) so a NaN/inf/negative timeout cannot reach asyncio.wait_for. Found via the freetoken-mlx downstream audit (docs/AUDIT.md, defects 3-4); the rebuild TOCTOU race from that audit is already structurally fixed upstream, so only the timeout bound is ported. --- python/freetoken/server/api_server.py | 4 +- python/freetoken/server/generation.py | 16 ++++-- tests/server/test_openai_api.py | 63 ++++++++++++++++++++++++ tests/server/test_rebuild_maintenance.py | 15 ++++++ 4 files changed, 93 insertions(+), 5 deletions(-) diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index 3e2acc8542..94b3e10efd 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -33,7 +33,7 @@ init_logger, load_generation_sampling, ) -from pydantic import BaseModel +from pydantic import BaseModel, Field from .args import ServerArgs from .anthropic_api import register_anthropic_routes @@ -494,7 +494,7 @@ class CacheRebuildRequest(BaseModel): # is deferred (needs the drain-gate machinery); constraining the Literal makes an # unsupported value fail fast with a 422 at the API layer instead of a generic 503. mode: Literal["if_idle"] = "if_idle" - timeout: float = 300.0 + timeout: float = Field(default=300.0, gt=0, le=3600) async def dispatch_rebuild( diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index be05d908a0..a60119ecfb 100644 --- a/python/freetoken/server/generation.py +++ b/python/freetoken/server/generation.py @@ -15,6 +15,7 @@ import asyncio import json +import math import time from collections.abc import AsyncIterator from dataclasses import dataclass, field @@ -177,12 +178,21 @@ def pick(value, key, framework): # non-positive value is a client error. if max_tokens is not None and max_tokens < 1: raise ValueError(f"max_tokens must be at least 1, got {max_tokens}") + resolved_temperature = pick(temperature, "temperature", 0.0) + if not math.isfinite(resolved_temperature) or resolved_temperature < 0: + raise ValueError(f"temperature must be a finite number >= 0, got {resolved_temperature}") + resolved_top_p = pick(top_p, "top_p", 1.0) + if not math.isfinite(resolved_top_p) or not 0 < resolved_top_p <= 1: + raise ValueError(f"top_p must be in (0, 1], got {resolved_top_p}") + resolved_top_k = pick(top_k, "top_k", -1) + if resolved_top_k != -1 and resolved_top_k < 1: + raise ValueError(f"top_k must be -1 (disabled) or >= 1, got {resolved_top_k}") return SamplingParams( ignore_eos=ignore_eos, max_tokens=DEFAULT_MAX_OUTPUT_TOKENS if max_tokens is None else max_tokens, - temperature=pick(temperature, "temperature", 0.0), - top_k=pick(top_k, "top_k", -1), - top_p=pick(top_p, "top_p", 1.0), + temperature=resolved_temperature, + top_k=resolved_top_k, + top_p=resolved_top_p, stop_strs=[s for s in stop_list if s], # drop empty strings (would match everything) ) diff --git a/tests/server/test_openai_api.py b/tests/server/test_openai_api.py index facd469b35..f6808a7d33 100644 --- a/tests/server/test_openai_api.py +++ b/tests/server/test_openai_api.py @@ -676,3 +676,66 @@ def test_minimax_http_non_stream_forces_implicit_reasoning_without_request_knob( message = response["choices"][0]["message"] assert message["reasoning_content"] == "private thought" assert message["content"] == "visible answer" + + +import pytest # noqa: E402 + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("temperature", -0.5), + ("top_p", 0.0), + ("top_k", 0), + ], +) +def test_chat_completion_rejects_invalid_sampling(field, value): + app = FastAPI() + state = FakeState([]) + + @app.post("/v1/chat/completions") + async def chat_completion(req: ChatCompletionRequest): + return await handle_chat_completion(req, request=None, state=state, model_sampling={}) + + body = { + "model": "unit-model", + "messages": [{"role": "user", "content": "hi"}], + "stream": True, + field: value, + } + response = TestClient(app).post("/v1/chat/completions", json=body) + + assert response.status_code == 400 + error = response.json()["error"] + assert error["type"] == "invalid_request_error" + assert field in error["message"] + assert state.sent is None + + +@pytest.mark.parametrize( + "sampling", + [ + {"temperature": 0.0, "top_p": 1.0, "top_k": -1}, + {"temperature": 0.0, "top_p": 1.0, "top_k": 1}, + ], +) +def test_chat_completion_accepts_sampling_boundaries(sampling): + app = FastAPI() + state = FakeState([]) + + @app.post("/v1/chat/completions") + async def chat_completion(req: ChatCompletionRequest): + return await handle_chat_completion(req, request=None, state=state, model_sampling={}) + + response = TestClient(app).post( + "/v1/chat/completions", + json={ + "model": "unit-model", + "messages": [{"role": "user", "content": "hi"}], + "stream": True, + **sampling, + }, + ) + + assert response.status_code == 200 + assert state.sent is not None diff --git a/tests/server/test_rebuild_maintenance.py b/tests/server/test_rebuild_maintenance.py index a1d10a7e56..0df3d6bfef 100644 --- a/tests/server/test_rebuild_maintenance.py +++ b/tests/server/test_rebuild_maintenance.py @@ -327,3 +327,18 @@ def test_cache_rebuild_request_rejects_unknown_mode(): assert CacheRebuildRequest(mode="if_idle").mode == "if_idle" with pytest.raises(ValidationError): CacheRebuildRequest(mode="drain") + + +from pydantic import ValidationError # noqa: E402 + +from freetoken.server.api_server import CacheRebuildRequest # noqa: E402 + + +def test_rebuild_timeout_bounds(): + import pytest + + for bad in (0, -5, 3601, float("nan")): + with pytest.raises(ValidationError): + CacheRebuildRequest(timeout=bad) + assert CacheRebuildRequest(timeout=3600).timeout == 3600 + assert CacheRebuildRequest().timeout == 300.0 From 2010271266f1a0253403a20f897d39d85e70cfe2 Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Tue, 13 Oct 2026 20:12:58 -0700 Subject: [PATCH 02/94] refactor(cpu-moe): split the executor into weight formats and gpu sync backends (#670) * refactor(cpu-moe): split the executor into weight formats and gpu sync backends - move each expert weight format into kernels/.cpp behind the MoeFormat interface in format.h; the executor keeps the threads, the work split and the pass 2 route sum - pass each format its banks and integer options by name; MoeParams keeps only the shared shape and activation fields - key the formats by name in one table that weight_formats() exports, replacing the integer format ids and the fp8_block_isa_name hook - replace the per-(layer, batch size) flag slots, pinned IO buffers and task table with one doorbell word MoeTask{layer, num_tokens} and one shared pinned IO block of max_tokens rows - move the GPU/CPU handshake behind GpuSync with a CUDA and a HIP backend picked by setup.py; HIP waits on signal memory and adds explicit batch-memop graph nodes where the runtime does not record the native memops under capture - raise on graph capture when the handshake is off and host functions are not replay-safe - allocate format scratch on 64-byte boundaries - size max_tokens by the explicit cuda_graph_bs as well Co-authored-by: zihaomu * perf(cpu-moe): skip the worker pool for tasks without routed experts * perf(cpu-moe): publish and combine the cpu handshake with kernels * fix(cpu-moe): clean up failed hip capability probes --------- Co-authored-by: zihaomu --- docs/install_amd.md | 8 + python/freetoken/engine/engine.py | 13 +- python/freetoken/kernel/csrc/cpu_moe/common.h | 235 ++ .../kernel/csrc/cpu_moe/cpu_moe_ext.cpp | 2307 ++--------------- python/freetoken/kernel/csrc/cpu_moe/format.h | 131 + .../freetoken/kernel/csrc/cpu_moe/gpu_sync.h | 78 + .../kernel/csrc/cpu_moe/gpu_sync_cuda.cpp | 166 ++ .../kernel/csrc/cpu_moe/gpu_sync_hip.cpp | 275 ++ .../kernel/csrc/cpu_moe/gpu_sync_io.cu | 199 ++ .../kernel/csrc/cpu_moe/gpu_sync_io.h | 39 + .../kernel/csrc/cpu_moe/kernels/bf16.cpp | 159 ++ .../kernel/csrc/cpu_moe/kernels/ds_fp4.cpp | 269 ++ .../kernel/csrc/cpu_moe/kernels/fp8_block.cpp | 203 ++ .../kernel/csrc/cpu_moe/kernels/mxfp4.cpp | 247 ++ .../kernel/csrc/cpu_moe/kernels/nvfp4.cpp | 466 ++++ .../kernel/csrc/cpu_moe/kernels/q4_0.cpp | 216 ++ .../csrc/include/freetoken/hip_compat.h | 12 + python/freetoken/layers/moe.py | 5 +- python/freetoken/moe/benchbw.py | 9 +- python/freetoken/moe/cpu_executor.py | 474 ++-- setup.py | 17 +- tests/moe/hip_probe_faults.cpp | 75 + tests/moe/test_cpu_moe.py | 115 +- tests/moe/test_cpu_moe_fp8_block.py | 25 +- tests/moe/test_cpu_moe_q4_0.py | 2 +- tests/moe/test_dsfp4_prequant.py | 10 +- 26 files changed, 3372 insertions(+), 2383 deletions(-) create mode 100644 python/freetoken/kernel/csrc/cpu_moe/common.h create mode 100644 python/freetoken/kernel/csrc/cpu_moe/format.h create mode 100644 python/freetoken/kernel/csrc/cpu_moe/gpu_sync.h create mode 100644 python/freetoken/kernel/csrc/cpu_moe/gpu_sync_cuda.cpp create mode 100644 python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp create mode 100644 python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.cu create mode 100644 python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.h create mode 100644 python/freetoken/kernel/csrc/cpu_moe/kernels/bf16.cpp create mode 100644 python/freetoken/kernel/csrc/cpu_moe/kernels/ds_fp4.cpp create mode 100644 python/freetoken/kernel/csrc/cpu_moe/kernels/fp8_block.cpp create mode 100644 python/freetoken/kernel/csrc/cpu_moe/kernels/mxfp4.cpp create mode 100644 python/freetoken/kernel/csrc/cpu_moe/kernels/nvfp4.cpp create mode 100644 python/freetoken/kernel/csrc/cpu_moe/kernels/q4_0.cpp create mode 100644 tests/moe/hip_probe_faults.cpp diff --git a/docs/install_amd.md b/docs/install_amd.md index 81d1ad5009..ed934ad657 100644 --- a/docs/install_amd.md +++ b/docs/install_amd.md @@ -45,3 +45,11 @@ On ROCm, these paths use the built-in fallbacks; do not install the `[accel]`, `[fi]`, or `[sgl]` extras. Forcing `--attention-backend fi`, `fa`, or `trtllm`, or NVFP4 Marlin/b12x (e.g. `--quant-backend moe.nvfp4=marlin` or `--quant-backend moe.nvfp4=b12x`), fails with a ROCm-specific error. + +## CPU and hybrid MoE + +`--moe-strategy cpu`, `hybrid` and `--moe-cpu-layers` sync the GPU with the CPU +expert pool through HIP stream memory operations. ROCm 10 records them into CUDA +graphs itself; on ROCm 7.14, which does not, FreeToken adds the graph nodes +explicitly. When the startup probe of that path fails, CUDA graph capture stops +with an error; run those strategies with `--cuda-graph-max-bs 0`. diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 3988caf602..feb33fbb94 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -813,9 +813,8 @@ def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None: """Build the persistent CPU MoE executor (decode-time expert compute). Must run before CUDA graph capture: the worker pool has to be live for the - eager warmup forward, and the pinned IO buffers / host-func task pointers - must be stable for the captured nodes. Buffers/tasks themselves are - allocated lazily on the first (eager) forward at each batch size. + eager warmup forward, and the pinned IO block must be stable for the captured + nodes. """ from freetoken.moe.cpu_executor import CpuMoeExecutor @@ -828,7 +827,7 @@ def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None: ) # Decode batches never exceed max_running_req, but CUDA-graph padding can # round a batch up to the largest captured size; cover both. - max_tokens = max(config.max_running_req, config.cuda_graph_max_bs or 0, 1) + max_tokens = max(config.max_running_req, config.cuda_graph_max_bs or 0, *(config.cuda_graph_bs or []), 1) method = sample.quant_method executor = CpuMoeExecutor( cache, @@ -1365,10 +1364,10 @@ def _cpu_moe_executor_viable(model_config) -> bool: """Whether an automatic CPU-decode decision may target the CPU MoE executor. A default boot must degrade to GPU offload instead of crashing in CpuMoeExecutor after the whole load; explicit cpu/hybrid/--moe-cpu-layers picks still fail loudly.""" - from freetoken.moe.cpu_executor import _WFMT_IDS, compiled_extension_supports + from freetoken.moe.cpu_executor import compiled_extension_supports try: - from freetoken.kernel import _cpu_moe # noqa: F401 + from freetoken.kernel import _cpu_moe except ImportError: return False act = getattr(model_config, "hidden_act", "silu") @@ -1379,7 +1378,7 @@ def _cpu_moe_executor_viable(model_config) -> bool: return False expert_quant = getattr(model_config, "expert_quant", "none") fmt = expert_quant if expert_quant != "none" else (moe_wfmt or "bf16") - return fmt == "mxfp4" or fmt in _WFMT_IDS + return fmt == "mxfp4" or fmt in _cpu_moe.weight_formats() def _pin_budget_bytes(reserved: int = 0) -> int | None: diff --git a/python/freetoken/kernel/csrc/cpu_moe/common.h b/python/freetoken/kernel/csrc/cpu_moe/common.h new file mode 100644 index 0000000000..adf7e14584 --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/common.h @@ -0,0 +1,235 @@ +// Numeric helpers and ISA detection shared by the CPU MoE formats. +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#if defined(__x86_64__) || defined(__i386__) +#include +#define CPU_MOE_X86 1 +#else +#define CPU_MOE_X86 0 +#endif + +// Each SIMD kernel is compiled for its own ISA and picked at runtime (pick_isa), so the build sets no global -march. +#define CPU_MOE_TARGET(...) __attribute__((target(__VA_ARGS__))) + +#if CPU_MOE_X86 && ((defined(__GNUC__) && __GNUC__ >= 10) || defined(__clang__)) +#define CPU_MOE_HAS_AVX512BF16 1 +#define CPU_MOE_HAS_AVX512VNNI 1 +#endif + +namespace freetoken::cpu_moe { + +// Scratch rows are read with 64-byte SIMD loads, so their base is cache-line aligned whatever the heap state; std::vector alone guarantees 16 bytes. +template +struct CacheLineAllocator { + using value_type = T; + CacheLineAllocator() = default; + template + CacheLineAllocator(const CacheLineAllocator&) {} + T* allocate(size_t n) { return static_cast(::operator new(n * sizeof(T), std::align_val_t(64))); } + void deallocate(T* p, size_t) { ::operator delete(p, std::align_val_t(64)); } + template + bool operator==(const CacheLineAllocator&) const { return true; } +}; + +template +using AlignedVector = std::vector>; + +using bf16_t = uint16_t; + +inline float bf16_to_f32(bf16_t v) { + uint32_t u = static_cast(v) << 16; + float f; + std::memcpy(&f, &u, sizeof(f)); + return f; +} + +inline bf16_t f32_to_bf16(float f) { + uint32_t u; + std::memcpy(&u, &f, sizeof(u)); + // round-to-nearest-even + const uint32_t lsb = (u >> 16) & 1u; + u += 0x7fffu + lsb; + return static_cast(u >> 16); +} + +// ACT_SWIGLUOAI is the clamped (up + 1) swiglu (gpt-oss "swigluoai" / +// MiniMax-M3): gate/up are combined jointly with the runtime alpha/limit +// scalars, so it is handled in the gated_rows epilogue (act_apply never sees it; +// the mxfp4 kernel additionally fuses its own copy of the same math). +// ACT_SWIGLU_CLAMP (GLM-5.3 "swiglu_limit") is the same clamped form WITHOUT +// the (up + 1) bias: clamp(gate, max=lim) * sigmoid(alpha*gate) * clamp(up, +-lim). +enum ActKind { + ACT_SILU = 0, + ACT_GELU = 1, + ACT_GELU_TANH = 2, + ACT_SWIGLUOAI = 3, + ACT_SWIGLU_CLAMP = 4, +}; + +inline float act_apply(int act, float x) { + if (act == ACT_SILU) return x / (1.0f + std::exp(-x)); + if (act == ACT_GELU) + return 0.5f * x * (1.0f + std::erf(x * 0.70710678118654752440f)); + // gelu_tanh + const float k0 = 0.79788456080286535588f; // sqrt(2/pi) + const float inner = k0 * (x + 0.044715f * x * x * x); + return 0.5f * x * (1.0f + std::tanh(inner)); +} + +// Software prefetch distance (bytes) ahead of the current weight row stream. The +// weight stream is the bandwidth bottleneck (read once, never reused); nudging the +// HW prefetcher with a few cache lines of lookahead raises sustained throughput. +constexpr int PF_AHEAD = 512; + +const float kE2M1[16] = {0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, + -0.0f, -0.5f, -1.0f, -1.5f, -2.0f, -3.0f, -4.0f, -6.0f}; + +inline float fp16_to_f32(uint16_t h) { + const uint32_t sign = (uint32_t)(h & 0x8000u) << 16; + uint32_t exp = (h >> 10) & 0x1Fu; + uint32_t man = h & 0x3FFu; + uint32_t f; + if (exp == 0) { + if (man == 0) { + f = sign; + } else { + exp = 127 - 15 + 1; + while ((man & 0x400u) == 0) { + man <<= 1; + --exp; + } + man &= 0x3FFu; + f = sign | (exp << 23) | (man << 13); + } + } else if (exp == 0x1Fu) { + f = sign | 0x7F800000u | (man << 13); + } else { + f = sign | ((exp + (127 - 15)) << 23) | (man << 13); + } + float out; + std::memcpy(&out, &f, sizeof(out)); + return out; +} + +// e4m3 (OCP "fn": finite, max-normal 448, exp bias 7). Decoded into a 256-entry LUT. +inline float e4m3_decode(uint8_t v) { + const float sign = (v & 0x80u) ? -1.0f : 1.0f; + const uint32_t exp = (v >> 3) & 0xFu; + const uint32_t man = v & 0x7u; + if (exp == 0) return sign * (man / 8.0f) * 0.015625f; // 2^(1-7) = 2^-6 + return sign * (1.0f + man / 8.0f) * std::ldexp(1.0f, (int)exp - 7); +} + +#if CPU_MOE_X86 +CPU_MOE_TARGET("avx2,fma") +inline float hsum256(__m256 v) { + __m128 lo = _mm256_castps256_ps128(v); + lo = _mm_add_ps(lo, _mm256_extractf128_ps(v, 1)); + lo = _mm_add_ps(lo, _mm_movehl_ps(lo, lo)); + lo = _mm_add_ss(lo, _mm_shuffle_ps(lo, lo, 0x55)); + return _mm_cvtss_f32(lo); +} + +// AVX2 e2m1 nibble decode: codes (0..15) in 8 int32 lanes -> fp32. AVX2 vpermps is +// only 8-wide, so instead of a 16-entry LUT we use the e2m1 sign/magnitude symmetry: +// value = (code&8 ? - : +) * mag8[code&7], mag8 = e2m1[0..7]. The sign is bit 3 of +// the code shifted into the fp32 sign bit (bit 31). Bit-identical to the e2m1 LUT. +CPU_MOE_TARGET("avx2,fma") +inline __m256 e2m1_decode8(__m256i codes, __m256 mag8) { + __m256 mag = _mm256_permutevar8x32_ps(mag8, _mm256_and_si256(codes, _mm256_set1_epi32(7))); + __m256i sgn = _mm256_slli_epi32(_mm256_and_si256(codes, _mm256_set1_epi32(8)), 28); + return _mm256_xor_ps(mag, _mm256_castsi256_ps(sgn)); +} +#endif // CPU_MOE_X86 + +// SIMD tiers, ascending. Each format picks the highest tier <= the one chosen by +// pick_isa() that it implements (fp4 and fp8_block formats have no bf16-specific tier, so the +// avx512bf16 tier maps to their avx512 kernel). +enum IsaTier { ISA_SCALAR = 0, ISA_AVX2 = 1, ISA_AVX512 = 2, ISA_AVX512BF16 = 3 }; + +// Best tier the CPU+build supports, optionally capped DOWN by +// FREETOKEN_CPU_MOE_ISA={scalar,avx2,avx512,avx512bf16} (A/B testing on a machine +// that supports more). FREETOKEN_CPU_MOE_SCALAR=1 forces scalar (legacy alias). +inline IsaTier pick_isa() { +#if CPU_MOE_X86 + if (getenv("FREETOKEN_CPU_MOE_SCALAR")) return ISA_SCALAR; + IsaTier best = ISA_SCALAR; + if (__builtin_cpu_supports("avx2") && __builtin_cpu_supports("fma")) best = ISA_AVX2; + if (best >= ISA_AVX2 && __builtin_cpu_supports("avx512f")) best = ISA_AVX512; +#ifdef CPU_MOE_HAS_AVX512BF16 + if (best >= ISA_AVX512 && __builtin_cpu_supports("avx512bf16")) best = ISA_AVX512BF16; +#endif + if (const char* f = getenv("FREETOKEN_CPU_MOE_ISA")) { + IsaTier want = best; + if (!std::strcmp(f, "scalar")) want = ISA_SCALAR; + else if (!std::strcmp(f, "avx2")) want = ISA_AVX2; + else if (!std::strcmp(f, "avx512")) want = ISA_AVX512; + else if (!std::strcmp(f, "avx512bf16")) want = ISA_AVX512BF16; + if (want < best) best = want; // cap downward; never force above hw/build support + } + return best; +#else + return ISA_SCALAR; +#endif +} + +inline const char* tier_name(IsaTier t) { +#if CPU_MOE_X86 +#ifdef CPU_MOE_HAS_AVX512BF16 + if (t >= ISA_AVX512BF16) return "avx512bf16"; +#endif + if (t >= ISA_AVX512) return "avx512f"; + if (t >= ISA_AVX2) return "avx2"; +#endif + (void)t; + return "scalar"; +} + +// AVX-VNNI (VEX-256 VPDPBUSD) availability: Alder/Raptor Lake, Sapphire Rapids+, Zen5. +// Distinct from AVX-512 VNNI. Opt out with FREETOKEN_CPU_MOE_NO_VNNI=1 (A/B the W4A8 path). +inline bool cpu_has_avxvnni() { +#if CPU_MOE_X86 + const char* no = getenv("FREETOKEN_CPU_MOE_NO_VNNI"); + if (no && no[0] && no[0] != '0') return false; // ignore unset/empty/"0" + return __builtin_cpu_supports("avx2") && __builtin_cpu_supports("avxvnni"); +#else + return false; +#endif +} + +// AVX-512 VNNI (512-bit VPDPBUSD): Cascade Lake+, Ice Lake, Sapphire/Emerald Rapids, +// Zen4+. 2x the 256-bit AVX-VNNI width. FREETOKEN_CPU_MOE_NO_AVX512VNNI=1 forces the +// 256-bit path (A/B the two W4A8 kernels on the same box); FREETOKEN_CPU_MOE_NO_VNNI=1 +// still disables the whole W4A8 family (back to the faithful fp32 nvdot), so it is +// honored here too. +inline bool cpu_has_avx512vnni() { +#if CPU_MOE_X86 && defined(CPU_MOE_HAS_AVX512VNNI) + const char* no = getenv("FREETOKEN_CPU_MOE_NO_AVX512VNNI"); + if (no && no[0] && no[0] != '0') return false; + const char* no_vnni = getenv("FREETOKEN_CPU_MOE_NO_VNNI"); + if (no_vnni && no_vnni[0] && no_vnni[0] != '0') return false; + return __builtin_cpu_supports("avx512vnni"); +#else + return false; +#endif +} + +// Split an interleaved bf16 row into fp32 even/odd halves (even[m]=src[2m]). +// bf16->fp32 is exact, so this only reorders -- done once per token/route and +// reused across every output row of the GEMV. +inline void deinterleave_bf16_f32(const bf16_t* src, float* even, float* odd, int K) { + for (int m = 0; m < K / 2; ++m) { + even[m] = bf16_to_f32(src[2 * m]); + odd[m] = bf16_to_f32(src[2 * m + 1]); + } +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp index 6f85971e9c..c8b4d1b465 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp @@ -1,32 +1,27 @@ -// CPU-compute MoE executor for the "cpu" offload backend. +// CPU-compute MoE executor for the cpu and hybrid MoE strategies, and its Python binding. // // Decode ships activations to the CPU, computes the routed experts here (reading -// the pinned host expert banks at full RAM bandwidth), and ships the results -// back. To keep the whole decode path inside a single CUDA graph we expose -// submit/sync as host nodes via cudaLaunchHostFunc -- the callbacks only touch a -// CPU worker pool + pinned host buffers and never call any CUDA API. +// the pinned host expert banks at full RAM bandwidth), and ships the results back. +// One task (one MoE layer) is in flight at a time: submit() wakes the persistent worker +// pool, the pool runs pass 1 (gate_up) and pass 2 (down), and sync() waits for it. +// The GPU rings a doorbell to submit and waits on it to sync (GpuSync, gpu_sync.h), so the +// whole decode step stays in one CUDA graph; cudaLaunchHostFunc nodes are the fallback. // -// One task is in flight at a time (per MoE layer): submit() wakes the pool, -// sync() blocks the host-func thread until the pool drains. The heavy GEMV runs -// on the persistent worker threads, not the host-func thread. -// -// Weight formats: bf16, NVFP4, MXFP4, ds_fp4, Q4_0 and block-FP8 expert banks (see WFmt and -// the per-format bank schemas). Compute is FP32-accumulate; the intermediate is -// stored bf16 to match the GPU decode path. ISA is chosen once at construction -// (AVX-512-BF16 dpbf16 -> AVX-512F widening -> AVX2+FMA -> scalar). +// Weight formats live in kernels/, one file each behind the MoeFormat interface (format.h). +// Compute is FP32-accumulate; the intermediate is stored bf16 to match the GPU decode path. #include #include #include -#include #include #include -#include #include +#include #include #include #include #include +#include #include #include @@ -40,1379 +35,99 @@ #define CPU_MOE_HAS_AFFINITY 0 #endif -#if defined(__x86_64__) || defined(__i386__) -#include -#define CPU_MOE_X86 1 -#else -#define CPU_MOE_X86 0 -#endif - -namespace { - -inline void cpu_relax() { -#if CPU_MOE_X86 - _mm_pause(); -#elif defined(__aarch64__) - asm volatile("yield" ::: "memory"); -#endif -} - -using bf16_t = uint16_t; - -inline float bf16_to_f32(bf16_t v) { - uint32_t u = static_cast(v) << 16; - float f; - std::memcpy(&f, &u, sizeof(f)); - return f; -} - -inline bf16_t f32_to_bf16(float f) { - uint32_t u; - std::memcpy(&u, &f, sizeof(u)); - // round-to-nearest-even - const uint32_t lsb = (u >> 16) & 1u; - u += 0x7fffu + lsb; - return static_cast(u >> 16); -} - -// ACT_SWIGLUOAI is the clamped (up + 1) swiglu (gpt-oss "swigluoai" / -// MiniMax-M3): gate/up are combined jointly with the runtime alpha/limit -// scalars, so it is handled in the do_pass1 epilogue (act_apply never sees it; -// the mxfp4 kernel additionally fuses its own copy of the same math). -// ACT_SWIGLU_CLAMP (GLM-5.3 "swiglu_limit") is the same clamped form WITHOUT -// the (up + 1) bias: clamp(gate, max=lim) * sigmoid(alpha*gate) * clamp(up, +-lim). -enum ActKind { - ACT_SILU = 0, - ACT_GELU = 1, - ACT_GELU_TANH = 2, - ACT_SWIGLUOAI = 3, - ACT_SWIGLU_CLAMP = 4, -}; - -inline float act_apply(int act, float x) { - if (act == ACT_SILU) return x / (1.0f + std::exp(-x)); - if (act == ACT_GELU) - return 0.5f * x * (1.0f + std::erf(x * 0.70710678118654752440f)); - // gelu_tanh - const float k0 = 0.79788456080286535588f; // sqrt(2/pi) - const float inner = k0 * (x + 0.044715f * x * x * x); - return 0.5f * x * (1.0f + std::tanh(inner)); -} - -// ------------------------------- dot products ------------------------------- -// dot(weight[bf16], act[bf16], n) -> fp32. The selected impl is a function -// pointer chosen at runtime; the per-row call overhead is negligible vs n. - -using dot_fn = float (*)(const bf16_t*, const bf16_t*, int); - -float dot_scalar(const bf16_t* w, const bf16_t* x, int n) { - float acc = 0.0f; - for (int i = 0; i < n; ++i) acc += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); - return acc; -} - -// Software prefetch distance (bytes) ahead of the current weight row stream. The -// weight stream is the bandwidth bottleneck (read once, never reused); nudging the -// HW prefetcher with a few cache lines of lookahead raises sustained throughput. -constexpr int PF_AHEAD = 512; - -#if CPU_MOE_X86 -__attribute__((target("avx512f"))) -float dot_avx512f(const bf16_t* w, const bf16_t* x, int n) { - // 4 independent accumulators -> more in-flight loads (memory-level parallelism), - // which is what lifts a bandwidth-bound GEMV toward peak. - __m512 a0 = _mm512_setzero_ps(), a1 = _mm512_setzero_ps(); - __m512 a2 = _mm512_setzero_ps(), a3 = _mm512_setzero_ps(); - int i = 0; - for (; i + 64 <= n; i += 64) { - _mm_prefetch(reinterpret_cast(w + i) + PF_AHEAD, _MM_HINT_T0); - for (int j = 0; j < 64; j += 16) { - __m256i wi = _mm256_loadu_si256(reinterpret_cast(w + i + j)); - __m256i xi = _mm256_loadu_si256(reinterpret_cast(x + i + j)); - __m512 wf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(wi), 16)); - __m512 xf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(xi), 16)); - __m512& acc = (j == 0) ? a0 : (j == 16) ? a1 : (j == 32) ? a2 : a3; - acc = _mm512_fmadd_ps(wf, xf, acc); - } - } - for (; i + 16 <= n; i += 16) { - __m256i wi = _mm256_loadu_si256(reinterpret_cast(w + i)); - __m256i xi = _mm256_loadu_si256(reinterpret_cast(x + i)); - __m512 wf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(wi), 16)); - __m512 xf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(xi), 16)); - a0 = _mm512_fmadd_ps(wf, xf, a0); - } - float s = _mm512_reduce_add_ps(_mm512_add_ps(_mm512_add_ps(a0, a1), _mm512_add_ps(a2, a3))); - for (; i < n; ++i) s += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); - return s; -} - -#if (defined(__GNUC__) && __GNUC__ >= 10) || defined(__clang__) -#define CPU_MOE_HAS_AVX512BF16 1 -__attribute__((target("avx512bf16,avx512f"))) -static inline __m512bh load_bh(const bf16_t* p) { - __m512i raw = _mm512_loadu_si512(reinterpret_cast(p)); - __m512bh out; - std::memcpy(&out, &raw, sizeof(out)); - return out; -} - -__attribute__((target("avx512bf16,avx512f"))) -float dot_avx512bf16(const bf16_t* w, const bf16_t* x, int n) { - // 4 accumulators (128 bf16/iter) for memory-level parallelism + a prefetch nudge. - __m512 a0 = _mm512_setzero_ps(), a1 = _mm512_setzero_ps(); - __m512 a2 = _mm512_setzero_ps(), a3 = _mm512_setzero_ps(); - int i = 0; - for (; i + 128 <= n; i += 128) { - _mm_prefetch(reinterpret_cast(w + i) + PF_AHEAD, _MM_HINT_T0); - a0 = _mm512_dpbf16_ps(a0, load_bh(w + i), load_bh(x + i)); - a1 = _mm512_dpbf16_ps(a1, load_bh(w + i + 32), load_bh(x + i + 32)); - a2 = _mm512_dpbf16_ps(a2, load_bh(w + i + 64), load_bh(x + i + 64)); - a3 = _mm512_dpbf16_ps(a3, load_bh(w + i + 96), load_bh(x + i + 96)); - } - for (; i + 32 <= n; i += 32) { - a0 = _mm512_dpbf16_ps(a0, load_bh(w + i), load_bh(x + i)); - } - float s = _mm512_reduce_add_ps(_mm512_add_ps(_mm512_add_ps(a0, a1), _mm512_add_ps(a2, a3))); - for (; i < n; ++i) s += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); - return s; -} -#endif // avx512bf16 available - -// --------- AVX2 + FMA (256-bit fallback for CPUs without AVX-512) ---------- -// Covers Intel 12-14th gen / Arrow Lake (AVX-512 fused off) and AMD Zen<4. The -// bf16->fp32 widen is a zero-extend + <<16; with 4 independent accumulators the -// GEMV is memory-bandwidth bound, same as the AVX-512 path (just half the width). -__attribute__((target("avx2,fma"))) -inline float hsum256(__m256 v) { - __m128 lo = _mm256_castps256_ps128(v); - lo = _mm_add_ps(lo, _mm256_extractf128_ps(v, 1)); - lo = _mm_add_ps(lo, _mm_movehl_ps(lo, lo)); - lo = _mm_add_ss(lo, _mm_shuffle_ps(lo, lo, 0x55)); - return _mm_cvtss_f32(lo); -} - -__attribute__((target("avx2,fma"))) -float dot_avx2(const bf16_t* w, const bf16_t* x, int n) { - __m256 a0 = _mm256_setzero_ps(), a1 = _mm256_setzero_ps(); - __m256 a2 = _mm256_setzero_ps(), a3 = _mm256_setzero_ps(); - int i = 0; - for (; i + 32 <= n; i += 32) { - _mm_prefetch(reinterpret_cast(w + i) + PF_AHEAD, _MM_HINT_T0); - for (int j = 0; j < 32; j += 8) { - __m128i wi = _mm_loadu_si128(reinterpret_cast(w + i + j)); - __m128i xi = _mm_loadu_si128(reinterpret_cast(x + i + j)); - __m256 wf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(wi), 16)); - __m256 xf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(xi), 16)); - __m256& acc = (j == 0) ? a0 : (j == 8) ? a1 : (j == 16) ? a2 : a3; - acc = _mm256_fmadd_ps(wf, xf, acc); - } - } - for (; i + 8 <= n; i += 8) { - __m128i wi = _mm_loadu_si128(reinterpret_cast(w + i)); - __m128i xi = _mm_loadu_si128(reinterpret_cast(x + i)); - __m256 wf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(wi), 16)); - __m256 xf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(xi), 16)); - a0 = _mm256_fmadd_ps(wf, xf, a0); - } - float s = hsum256(_mm256_add_ps(_mm256_add_ps(a0, a1), _mm256_add_ps(a2, a3))); - for (; i < n; ++i) s += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); - return s; -} -#endif // CPU_MOE_X86 +#include "format.h" +#include "gpu_sync.h" +#include "gpu_sync_io.h" -// --------------------------- NVFP4 (W4A16) dequant --------------------------- -// Weights: e2m1 4-bit codes (2/byte, low nibble first), per-16 block scale in -// fp8-e4m3, per-output-row global scale in fp16. Dequant matches the GPU kernels -// (freetoken/kernel/triton/nvfp4_dequant.py): w = E2M1[code] * e4m3(scale) * global. -// Activations stay bf16 (W4A16); the GEMV dequantizes weights inside the K-loop. +namespace freetoken::cpu_moe { -const float kE2M1[16] = {0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, - -0.0f, -0.5f, -1.0f, -1.5f, -2.0f, -3.0f, -4.0f, -6.0f}; +std::unique_ptr make_bf16_format(const MoeParams& p, const FormatArgs& a); +std::unique_ptr make_nvfp4_format(const MoeParams& p, const FormatArgs& a); +std::unique_ptr make_mxfp4_format(const MoeParams& p, const FormatArgs& a); +std::unique_ptr make_ds_fp4_format(const MoeParams& p, const FormatArgs& a); +std::unique_ptr make_q4_0_format(const MoeParams& p, const FormatArgs& a); +std::unique_ptr make_fp8_block_format(const MoeParams& p, const FormatArgs& a); -// e2m1 * 2 as exact int8 (all e2m1 values are multiples of 0.5). Used by the AVX-VNNI -// W4A8 path: nibble -> int8 weight via PSHUFB LUT, then VPDPBUSD against int8 activations; -// the *2 is undone by a 0.5 folded into the final scale. Mirrors ggml's kvalues_mxfp4. -alignas(16) const int8_t kE2M1x2[16] = {0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12}; - -inline float fp16_to_f32(uint16_t h) { - const uint32_t sign = (uint32_t)(h & 0x8000u) << 16; - uint32_t exp = (h >> 10) & 0x1Fu; - uint32_t man = h & 0x3FFu; - uint32_t f; - if (exp == 0) { - if (man == 0) { - f = sign; - } else { - exp = 127 - 15 + 1; - while ((man & 0x400u) == 0) { - man <<= 1; - --exp; - } - man &= 0x3FFu; - f = sign | (exp << 23) | (man << 13); - } - } else if (exp == 0x1Fu) { - f = sign | 0x7F800000u | (man << 13); - } else { - f = sign | ((exp + (127 - 15)) << 23) | (man << 13); - } - float out; - std::memcpy(&out, &f, sizeof(out)); - return out; -} - -// e4m3 (OCP "fn": finite, max-normal 448, exp bias 7). Decoded into a 256-entry LUT. -inline float e4m3_decode(uint8_t v) { - const float sign = (v & 0x80u) ? -1.0f : 1.0f; - const uint32_t exp = (v >> 3) & 0xFu; - const uint32_t man = v & 0x7u; - if (exp == 0) return sign * (man / 8.0f) * 0.015625f; // 2^(1-7) = 2^-6 - return sign * (1.0f + man / 8.0f) * std::ldexp(1.0f, (int)exp - 7); -} - -// ------------------------ block-FP8 (W8A16) dequant ------------------------ -// fp8-e4m3 weights row-major [rows, K] plus one bf16 scale per 128x128 block. The scale tensor is named weight_scale_inv but multiplies, as in kernel/triton/fp8_blockscale_moe.py. -// Whole 128-blocks read the activations as fp32 even/odd halves (xe[m] = x[2m], xo[m] = x[2m+1]); a ragged tail reads the bf16 row x. -// 0x7F/0xFF (NaN in torch float8_e4m3fn) decode to +-480 in every kernel, as e4m3_decode does. -using fp8_block_dot_fn = float (*)(const uint8_t*, const float*, const float*, const bf16_t*, - const bf16_t*, int, const float*); - -constexpr int FP8_BLK = 128; // K elements covered by one weight scale - -inline float fp8_block_tail(const uint8_t* w, const bf16_t* x, int k0, int K, const float* e4m3) { - float blk = 0.0f; - for (int k = k0; k < K; ++k) blk += e4m3[w[k]] * bf16_to_f32(x[k]); - return blk; -} - -float dot_fp8_block_scalar(const uint8_t* w, const float* xe, const float* xo, const bf16_t* x, - const bf16_t* s, int K, const float* e4m3) { - float acc = 0.0f; - int b = 0, k0 = 0; - for (; k0 + FP8_BLK <= K; k0 += FP8_BLK, ++b) { - // 8 independent sums: a single running sum is bound by the add latency - float a[8] = {}; - const uint8_t* wb = w + k0; - const float* pe = xe + k0 / 2; - const float* po = xo + k0 / 2; - for (int j = 0; j < FP8_BLK / 2; j += 4) - for (int l = 0; l < 4; ++l) { - a[2 * l] += e4m3[wb[2 * (j + l)]] * pe[j + l]; - a[2 * l + 1] += e4m3[wb[2 * (j + l) + 1]] * po[j + l]; - } - acc += (((a[0] + a[1]) + (a[2] + a[3])) + ((a[4] + a[5]) + (a[6] + a[7]))) * - bf16_to_f32(s[b]); - } - if (k0 < K) acc += fp8_block_tail(w, x, k0, K, e4m3) * bf16_to_f32(s[b]); - return acc; -} - -#if CPU_MOE_X86 -// A byte b read as fp16 bits (b & 0x80) << 8 | (b & 0x7F) << 7 is its e4m3 value times 2^-8, subnormals included, so the block scale carries the 256. -// In a 16-bit lane holding bytes (2m, 2m+1), srai 1 then & 0xBF80 builds that for byte 2m+1; a left shift by 8 first does byte 2m. Every AVX2 CPU has F16C. -__attribute__((target("avx2,fma,f16c"))) -float dot_fp8_block_avx2(const uint8_t* w, const float* xe, const float* xo, const bf16_t* x, - const bf16_t* s, int K, const float* e4m3) { - const __m256i keep = _mm256_set1_epi16((short)0xBF80); - __m256 total = _mm256_setzero_ps(); - int b = 0, k0 = 0; - for (; k0 + FP8_BLK <= K; k0 += FP8_BLK, ++b) { - _mm_prefetch(reinterpret_cast(w + k0) + PF_AHEAD, _MM_HINT_T0); - __m256 a0 = _mm256_setzero_ps(), a1 = _mm256_setzero_ps(); - __m256 a2 = _mm256_setzero_ps(), a3 = _mm256_setzero_ps(); - for (int j = 0; j < FP8_BLK; j += 32) { - const __m256i v = _mm256_loadu_si256(reinterpret_cast(w + k0 + j)); - const __m256i odd = _mm256_and_si256(_mm256_srai_epi16(v, 1), keep); - const __m256i even = _mm256_and_si256(_mm256_srai_epi16(_mm256_slli_epi16(v, 8), 1), keep); - const float* pe = xe + (k0 + j) / 2; - const float* po = xo + (k0 + j) / 2; - a0 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_castsi256_si128(even)), _mm256_loadu_ps(pe), a0); - a1 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_extracti128_si256(even, 1)), _mm256_loadu_ps(pe + 8), a1); - a2 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_castsi256_si128(odd)), _mm256_loadu_ps(po), a2); - a3 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_extracti128_si256(odd, 1)), _mm256_loadu_ps(po + 8), a3); - } - const __m256 blk = _mm256_add_ps(_mm256_add_ps(a0, a1), _mm256_add_ps(a2, a3)); - total = _mm256_fmadd_ps(blk, _mm256_set1_ps(bf16_to_f32(s[b]) * 256.0f), total); - } - float acc = hsum256(total); - if (k0 < K) acc += fp8_block_tail(w, x, k0, K, e4m3) * bf16_to_f32(s[b]); - return acc; -} - -__attribute__((target("avx512f,avx2,fma,f16c"))) -float dot_fp8_block_avx512f(const uint8_t* w, const float* xe, const float* xo, const bf16_t* x, - const bf16_t* s, int K, const float* e4m3) { - const __m256i keep = _mm256_set1_epi16((short)0xBF80); - __m512 total = _mm512_setzero_ps(); - int b = 0, k0 = 0; - for (; k0 + FP8_BLK <= K; k0 += FP8_BLK, ++b) { - _mm_prefetch(reinterpret_cast(w + k0) + PF_AHEAD, _MM_HINT_T0); - __m512 a0 = _mm512_setzero_ps(), a1 = _mm512_setzero_ps(); - __m512 a2 = _mm512_setzero_ps(), a3 = _mm512_setzero_ps(); - for (int j = 0; j < FP8_BLK; j += 64) { - const __m256i v0 = _mm256_loadu_si256(reinterpret_cast(w + k0 + j)); - const __m256i v1 = _mm256_loadu_si256(reinterpret_cast(w + k0 + j + 32)); - const __m256i o0 = _mm256_and_si256(_mm256_srai_epi16(v0, 1), keep); - const __m256i e0 = _mm256_and_si256(_mm256_srai_epi16(_mm256_slli_epi16(v0, 8), 1), keep); - const __m256i o1 = _mm256_and_si256(_mm256_srai_epi16(v1, 1), keep); - const __m256i e1 = _mm256_and_si256(_mm256_srai_epi16(_mm256_slli_epi16(v1, 8), 1), keep); - const float* pe = xe + (k0 + j) / 2; - const float* po = xo + (k0 + j) / 2; - a0 = _mm512_fmadd_ps(_mm512_cvtph_ps(e0), _mm512_loadu_ps(pe), a0); - a1 = _mm512_fmadd_ps(_mm512_cvtph_ps(o0), _mm512_loadu_ps(po), a1); - a2 = _mm512_fmadd_ps(_mm512_cvtph_ps(e1), _mm512_loadu_ps(pe + 16), a2); - a3 = _mm512_fmadd_ps(_mm512_cvtph_ps(o1), _mm512_loadu_ps(po + 16), a3); - } - const __m512 blk = _mm512_add_ps(_mm512_add_ps(a0, a1), _mm512_add_ps(a2, a3)); - total = _mm512_fmadd_ps(blk, _mm512_set1_ps(bf16_to_f32(s[b]) * 256.0f), total); - } - float acc = _mm512_reduce_add_ps(total); - if (k0 < K) acc += fp8_block_tail(w, x, k0, K, e4m3) * bf16_to_f32(s[b]); - return acc; -} -#endif // CPU_MOE_X86 - -// Activations pre-deinterleaved to fp32 (xe[m]=x[2m], xo[m]=x[2m+1]); see the -// ds_fp4 dot below for why (drops the hot loop to ~1.5 shuffle ops / 16 weights). -using nvdot_fn = float (*)(const uint8_t*, const uint8_t*, float, const float*, const float*, - int, const float*, const float*); - -float dot_nvfp4_scalar(const uint8_t* packed, const uint8_t* scale, float global, - const float* xe, const float* xo, int K, const float* e2m1, - const float* e4m3) { - float acc = 0.0f; - const int nb = K / 16; - for (int b = 0; b < nb; ++b) { - const float bs = e4m3[scale[b]]; - const uint8_t* pk = packed + (size_t)b * 8; - const float* xeb = xe + (size_t)b * 8; // 16 K -> 8 even + 8 odd - const float* xob = xo + (size_t)b * 8; - float bsum = 0.0f; - for (int j = 0; j < 8; ++j) { - const uint8_t byte = pk[j]; - bsum += e2m1[byte & 0xF] * xeb[j]; - bsum += e2m1[byte >> 4] * xob[j]; - } - acc += bs * bsum; - } - return acc * global; -} - -// ---- NVFP4 W4A8 (int8 activations) dot: nibble->int8 via LUT, per-16 act scale ---- -// asi8: int8 activations laid out per-16 block as [even(8), odd(8)]; asb[b] = per-block -// activation scale (absmax/127). Result folds the e2m1*2 -> *0.5 into the scale. -using nvi8dot_fn = float (*)(const uint8_t*, const uint8_t*, float, const int8_t*, int, - const float*, const float*); - -[[maybe_unused]] float dot_nvfp4_i8_scalar(const uint8_t* packed, const uint8_t* scale, - float global, const int8_t* asi8, int K, const float* e4m3, - const float* asb) { - float acc = 0.0f; - const int nb = K / 16; - for (int b = 0; b < nb; ++b) { - const uint8_t* pk = packed + (size_t)b * 8; - const int8_t* ae = asi8 + (size_t)b * 16; // even(8) - const int8_t* ao = ae + 8; // odd(8) - int isum = 0; - for (int j = 0; j < 8; ++j) { - isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j]; - isum += (int)kE2M1x2[pk[j] >> 4] * (int)ao[j]; - } - acc += (e4m3[scale[b]] * asb[b]) * (float)isum; - } - return acc * (0.5f * global); -} - -#if CPU_MOE_X86 -// AVX2 e2m1 nibble decode: codes (0..15) in 8 int32 lanes -> fp32. AVX2 vpermps is -// only 8-wide, so instead of a 16-entry LUT we use the e2m1 sign/magnitude symmetry: -// value = (code&8 ? - : +) * mag8[code&7], mag8 = e2m1[0..7]. The sign is bit 3 of -// the code shifted into the fp32 sign bit (bit 31). Bit-identical to the e2m1 LUT. -__attribute__((target("avx2,fma"))) -inline __m256 e2m1_decode8(__m256i codes, __m256 mag8) { - __m256 mag = _mm256_permutevar8x32_ps(mag8, _mm256_and_si256(codes, _mm256_set1_epi32(7))); - __m256i sgn = _mm256_slli_epi32(_mm256_and_si256(codes, _mm256_set1_epi32(8)), 28); - return _mm256_xor_ps(mag, _mm256_castsi256_ps(sgn)); -} - -// Two 16-K blocks (16 packed bytes) per iter: lo nibbles -> even-K, hi -> odd-K, -// gathered via two vpermps. The per-16 e4m3 scale differs across the two blocks, so -// it is applied per lane (low 8 lanes = block b, high 8 = block b+1). -__attribute__((target("avx512f"))) -inline __m512 nvfp4_blk2(const uint8_t* pk, const float* xeb, const float* xob, __m512 lut, - __m512i loma, float s0, float s1) { - __m512i wi = _mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast(pk))); - __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut); - __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut); - __m512 prod = _mm512_fmadd_ps(vlo, _mm512_loadu_ps(xeb), _mm512_mul_ps(vhi, _mm512_loadu_ps(xob))); - // lanes 0-7 (block b) -> s0, lanes 8-15 (block b+1) -> s1 (pure AVX512F mask move) - __m512 scv = _mm512_mask_mov_ps(_mm512_set1_ps(s0), 0xFF00, _mm512_set1_ps(s1)); - return _mm512_mul_ps(prod, scv); -} - -__attribute__((target("avx512f"))) -float dot_nvfp4_avx512(const uint8_t* packed, const uint8_t* scale, float global, - const float* xe, const float* xo, int K, const float* e2m1, - const float* e4m3) { - const __m512 lut = _mm512_loadu_ps(e2m1); // 16 e2m1 values - const __m512i loma = _mm512_set1_epi32(0xF); - __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps(); - const int nb = K / 16; // 8 packed bytes + one e4m3 scale per block - int b = 0; - for (; b + 4 <= nb; b += 4) { // two blk2 calls -> 4 blocks, 2 accumulators - acc0 = _mm512_add_ps(acc0, nvfp4_blk2(packed + (size_t)b * 8, xe + (size_t)b * 8, - xo + (size_t)b * 8, lut, loma, e4m3[scale[b]], e4m3[scale[b + 1]])); - acc1 = _mm512_add_ps(acc1, nvfp4_blk2(packed + (size_t)(b + 2) * 8, xe + (size_t)(b + 2) * 8, - xo + (size_t)(b + 2) * 8, lut, loma, e4m3[scale[b + 2]], e4m3[scale[b + 3]])); - } - for (; b + 2 <= nb; b += 2) - acc0 = _mm512_add_ps(acc0, nvfp4_blk2(packed + (size_t)b * 8, xe + (size_t)b * 8, - xo + (size_t)b * 8, lut, loma, e4m3[scale[b]], e4m3[scale[b + 1]])); - float s = _mm512_reduce_add_ps(_mm512_add_ps(acc0, acc1)); - for (; b < nb; ++b) { // odd final 16-K block - const uint8_t* pk = packed + (size_t)b * 8; - const float* xeb = xe + (size_t)b * 8; - const float* xob = xo + (size_t)b * 8; - float bsum = 0.0f; - for (int j = 0; j < 8; ++j) { - bsum += e2m1[pk[j] & 0xF] * xeb[j] + e2m1[pk[j] >> 4] * xob[j]; - } - s += e4m3[scale[b]] * bsum; - } - return s * global; -} - -// AVX2: one 16-K block (8 packed bytes) per call, 8 even + 8 odd lanes. -__attribute__((target("avx2,fma"))) -inline __m256 nvfp4_blk_avx2(const uint8_t* pk, const float* xeb, const float* xob, - __m256 mag8, float sc) { - __m256i wi = _mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cast(pk))); - __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); - __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); - __m256 prod = _mm256_fmadd_ps(vlo, _mm256_loadu_ps(xeb), _mm256_mul_ps(vhi, _mm256_loadu_ps(xob))); - return _mm256_mul_ps(prod, _mm256_set1_ps(sc)); -} - -__attribute__((target("avx2,fma"))) -float dot_nvfp4_avx2(const uint8_t* packed, const uint8_t* scale, float global, - const float* xe, const float* xo, int K, const float* e2m1, - const float* e4m3) { - const __m256 mag8 = _mm256_loadu_ps(e2m1); // e2m1[0..7] magnitudes - __m256 acc0 = _mm256_setzero_ps(), acc1 = _mm256_setzero_ps(); - const int nb = K / 16; - int b = 0; - for (; b + 2 <= nb; b += 2) { - acc0 = _mm256_add_ps(acc0, nvfp4_blk_avx2(packed + (size_t)b * 8, xe + (size_t)b * 8, - xo + (size_t)b * 8, mag8, e4m3[scale[b]])); - acc1 = _mm256_add_ps(acc1, nvfp4_blk_avx2(packed + (size_t)(b + 1) * 8, xe + (size_t)(b + 1) * 8, - xo + (size_t)(b + 1) * 8, mag8, e4m3[scale[b + 1]])); - } - for (; b < nb; ++b) - acc0 = _mm256_add_ps(acc0, nvfp4_blk_avx2(packed + (size_t)b * 8, xe + (size_t)b * 8, - xo + (size_t)b * 8, mag8, e4m3[scale[b]])); - return hsum256(_mm256_add_ps(acc0, acc1)) * global; -} - -// AVX-VNNI W4A8: decode 8 packed bytes (16 nibbles) of one 16-block to int8 [lo(8),hi(8)] -// via PSHUFB against the e2m1*2 LUT (replaces the 2 vpermps fp32 expands -- ~4x less -// port-5 traffic). lo=even-K weights, hi=odd-K, matching the [even(8),odd(8)] act layout. -__attribute__((target("avx2,avxvnni,fma"))) -inline __m128i nvfp4_decode_block_i8(const uint8_t* pk, __m128i lut) { - __m128i b = _mm_loadl_epi64(reinterpret_cast(pk)); // 8 bytes - __m128i lo = _mm_and_si128(b, _mm_set1_epi8(0x0F)); - __m128i hi = _mm_and_si128(_mm_srli_epi16(b, 4), _mm_set1_epi8(0x0F)); - return _mm_shuffle_epi8(lut, _mm_unpacklo_epi64(lo, hi)); // [lo(8),hi(8)] -> int8 -} - -// Two 16-blocks per VPDPBUSD (32 int8). Sign trick (ggml mul_add_epi8): |w|*(sign(w)*a)=w*a, -// so u8*s8 needs no offset/correction term. Per-block scale (e4m3 * act-scale) folded in fp32 -// (lanes 0-3 -> block b, 4-7 -> block b+1). Bit-faithful weight; only the int8 activation quant -// (W4A8) differs from the bf16 reference. -__attribute__((target("avx2,avxvnni,fma"))) -float dot_nvfp4_i8_vnni(const uint8_t* packed, const uint8_t* scale, float global, - const int8_t* asi8, int K, const float* e4m3, const float* asb) { - const __m128i lut = _mm_loadu_si128(reinterpret_cast(kE2M1x2)); - __m256 accF = _mm256_setzero_ps(); - const int nb = K / 16; - int b = 0; - for (; b + 2 <= nb; b += 2) { - __m128i wb = nvfp4_decode_block_i8(packed + (size_t)b * 8, lut); - __m128i wb1 = nvfp4_decode_block_i8(packed + (size_t)(b + 1) * 8, lut); - __m256i w = _mm256_set_m128i(wb1, wb); // [blk b | blk b+1] - __m256i a = _mm256_loadu_si256(reinterpret_cast(asi8 + (size_t)b * 16)); - __m256i aw = _mm256_sign_epi8(w, w); // |w| (u8 operand) - __m256i sa = _mm256_sign_epi8(a, w); // sign(w)*a (s8 operand) - __m256i di = _mm256_dpbusd_avx_epi32(_mm256_setzero_si256(), aw, sa); - __m256 scv = _mm256_blend_ps(_mm256_set1_ps(e4m3[scale[b]] * asb[b]), - _mm256_set1_ps(e4m3[scale[b + 1]] * asb[b + 1]), 0xF0); - accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(di), scv, accF); - } - float s = hsum256(accF); - for (; b < nb; ++b) { // tail (odd block count) - const uint8_t* pk = packed + (size_t)b * 8; - const int8_t* ae = asi8 + (size_t)b * 16; const int8_t* ao = ae + 8; - int isum = 0; - for (int j = 0; j < 8; ++j) - isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j] + (int)kE2M1x2[pk[j] >> 4] * (int)ao[j]; - s += (e4m3[scale[b]] * asb[b]) * (float)isum; - } - return s * (0.5f * global); -} - -#if (defined(__GNUC__) && __GNUC__ >= 10) || defined(__clang__) -#define CPU_MOE_HAS_AVX512VNNI 1 - -// Software-prefetch distance for the W4A8 weight stream, in 16-K blocks (8 packed -// bytes each). Returns -1 when FREETOKEN_CPU_MOE_PF_BLOCKS is unset: the kernel then -// uses the built-in default min(512 blocks = 4 KB, 2 rows) -- 4 KB is the empirical -// optimum on large-row machines (Emerald Rapids sweep), while the 2-row cap keeps a small-row -// model's overshoot bounded (the executor works in 32-row tiles, so a fixed byte -// distance otherwise prefetches another worker's tile: duplicated DRAM traffic that -// regresses at the bandwidth ceiling). An EXPLICIT env value is honored verbatim -// (no clamp; 0 disables): the per-machine optimum can sit past the safe default (+20% -// at 4 KB on a 24-thread Ice Lake with 256B rows), so the escape hatch must reach it. -// Prefetch never faults, so overshooting a row/bank tail is safe. -static int nvfp4_pf_blocks() { - static const int v = [] { - const char* s = getenv("FREETOKEN_CPU_MOE_PF_BLOCKS"); - return (s && s[0]) ? atoi(s) : -1; - }(); - return v; -} - -// AVX-512 VNNI W4A8: FOUR 16-K blocks per VPDPBUSD (64 int8) -- 2x the AVX-VNNI -// (256-bit) path. Decode 32 packed bytes -> 64 int8 with a single _mm512_shuffle_epi8 -// (e2m1*2 LUT replicated to all 4 128-bit lanes). AVX-512 has no _mm512_sign_epi8, so -// the u8*s8 sign trick (|w| as u8, sign(w)*a as s8) uses abs_epi8 + a masked negate. -// Bit-faithful weight; only the int8 activation quant (W4A8) differs from bf16. -// One 4-block group (64 int8) -> scaled fp32 partial (16 lanes). Isolated as a helper so -// the caller can run several independent chains into separate accumulators (the decode -> -// dpbusd -> scale chain is long, so a single accumulator leaves the core latency-bound). -__attribute__((target("avx512f,avx512bw,avx512vnni,avx2"))) -static inline __m512 nvfp4_i8_grp4(const uint8_t* packed, const uint8_t* scale, - const int8_t* asi8, const float* e4m3, const float* asb, - int b, __m512i lut, __m512i idx, __m512i mask0F, - __m512i idxsc) { - const __mmask64 hi_half = 0xFF00FF00FF00FF00ULL; // bytes 8-15 of each 128b lane -> hi nibbles - __m256i raw = _mm256_loadu_si256(reinterpret_cast(packed + (size_t)b * 8)); - __m512i src = _mm512_permutexvar_epi64(idx, _mm512_castsi256_si512(raw)); - __m512i lo = _mm512_and_si512(src, mask0F); - __m512i hi = _mm512_and_si512(_mm512_srli_epi16(src, 4), mask0F); - __m512i comb = _mm512_mask_blend_epi8(hi_half, lo, hi); - __m512i w = _mm512_shuffle_epi8(lut, comb); // 64 int8 weights (e2m1*2) - // u8*s8 sign trick without _mm512_sign_epi8: aw=|w|, sa = (w<0 ? -a : a) (a moot at w==0) - __m512i a = _mm512_loadu_si512(reinterpret_cast(asi8 + (size_t)b * 16)); - __m512i aw = _mm512_abs_epi8(w); - __mmask64 neg = _mm512_movepi8_mask(w); - __m512i sa = _mm512_mask_sub_epi8(a, neg, _mm512_setzero_si512(), a); - __m512i di = _mm512_dpbusd_epi32(_mm512_setzero_si512(), aw, sa); // 16 int32 (groups of 4) - // int32[0..3]->blk b, [4..7]->b+1, [8..11]->b+2, [12..15]->b+3. The 4 block scales - // (e4m3 LUT x per-block act scale) are computed vectorized: a 4-byte load + epu8->epi32 - // widen + one 4-lane LUT gather + one mul replaces 8 scalar loads + 4 scalar muls + - // a set_ps assembly, which otherwise dominates the per-group op count. - int sc_raw; - memcpy(&sc_raw, scale + b, 4); - __m128i sc4 = _mm_cvtepu8_epi32(_mm_cvtsi32_si128(sc_raw)); - __m128 s4 = _mm_mul_ps(_mm_i32gather_ps(e4m3, sc4, 4), _mm_loadu_ps(asb + b)); - __m512 scv = _mm512_permutexvar_ps(idxsc, _mm512_castps128_ps512(s4)); - return _mm512_mul_ps(_mm512_cvtepi32_ps(di), scv); -} - -__attribute__((target("avx512f,avx512bw,avx512vnni,avx2"))) -float dot_nvfp4_i8_avx512vnni(const uint8_t* packed, const uint8_t* scale, float global, - const int8_t* asi8, int K, const float* e4m3, const float* asb) { - const __m512i lut = _mm512_broadcast_i32x4( - _mm_loadu_si128(reinterpret_cast(kE2M1x2))); - const __m512i idx = _mm512_set_epi64(3, 3, 2, 2, 1, 1, 0, 0); // block i -> 128b lane i - const __m512i mask0F = _mm512_set1_epi8(0x0F); - const __m512i idxsc = _mm512_set_epi32(3, 3, 3, 3, 2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0); - __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps(); - __m512 acc2 = _mm512_setzero_ps(), acc3 = _mm512_setzero_ps(); - const int nb = K / 16; - // Built-in default: min(4 KB, 2 rows) -- see nvfp4_pf_blocks(). An explicitly set - // env value is used verbatim so a per-machine sweep can reach operating points past - // the conservative default. - const int pfb = nvfp4_pf_blocks(); - const int pf = (pfb < 0) ? std::min(512, 2 * nb) : pfb; - int b = 0; - // Four independent 4-block groups per iter -> four accumulator chains hide the ~10-op - // decode->dpbusd->scale latency (1-2 chains leave the loop latency-bound: the per-core - // rate sat at ~60% of the core's achievable DRAM stream rate). - for (; b + 16 <= nb; b += 16) { - // Prefetch the weight stream ahead: the interleaved decode lowers the L1-miss - // concurrency the HW prefetcher sustains on its own (134 -> 167 GB/s on Emerald Rapids). - if (pf > 0) { - _mm_prefetch(reinterpret_cast(packed + ((size_t)b + (size_t)pf) * 8), - _MM_HINT_T0); - _mm_prefetch(reinterpret_cast(packed + ((size_t)b + (size_t)pf) * 8 + 64), - _MM_HINT_T0); - } - acc0 = _mm512_add_ps(acc0, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b, lut, idx, - mask0F, idxsc)); - acc1 = _mm512_add_ps(acc1, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 4, lut, idx, - mask0F, idxsc)); - acc2 = _mm512_add_ps(acc2, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 8, lut, idx, - mask0F, idxsc)); - acc3 = _mm512_add_ps(acc3, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 12, lut, idx, - mask0F, idxsc)); - } - for (; b + 4 <= nb; b += 4) - acc0 = _mm512_add_ps(acc0, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b, lut, idx, - mask0F, idxsc)); - float s = _mm512_reduce_add_ps( - _mm512_add_ps(_mm512_add_ps(acc0, acc1), _mm512_add_ps(acc2, acc3))); - for (; b < nb; ++b) { // tail (<4 remaining 16-K blocks) - const uint8_t* pk = packed + (size_t)b * 8; - const int8_t* ae = asi8 + (size_t)b * 16; const int8_t* ao = ae + 8; - int isum = 0; - for (int j = 0; j < 8; ++j) - isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j] + (int)kE2M1x2[pk[j] >> 4] * (int)ao[j]; - s += (e4m3[scale[b]] * asb[b]) * (float)isum; - } - return s * (0.5f * global); -} -#endif // avx512vnni available -#endif - -// ===================================================================================== -// CUDA stream memory operations (driver API, resolved via dlopen -- no link-time or -// toolchain dependence). The GPU side of the flag handshake: submit = WRITE_VALUE -// (done[slot]=0 then ready[slot]=1), sync = WAIT_VALUE(done[slot] >= 1). The wait is -// executed by the GPU front-end (no SM-resident kernel), so GPU "utilization" stays -// truthful during CPU compute windows -- a resident spin kernel pinned it at 99%, -// which laptop CPU/GPU dynamic power schedulers answered by clamping the CPU's max -// frequency (GEMV workers -1.5x: the reported edge regression). Availability is -// probed functionally at startup (memops_probe); anything unsupported (Windows WDDM, -// vGPU, old drivers) falls back to the cudaLaunchHostFunc path. -#if defined(_WIN32) -#include -static void* cumemop_dlopen() { return (void*)::LoadLibraryA("nvcuda.dll"); } -static void* cumemop_dlsym(void* h, const char* n) { - return (void*)::GetProcAddress((HMODULE)h, n); -} -#else -#include -static void* cumemop_dlopen() { - void* h = dlopen("libcuda.so.1", RTLD_LAZY | RTLD_LOCAL); - if (h == nullptr) h = dlopen("libcuda.so", RTLD_LAZY | RTLD_LOCAL); - return h; -} -static void* cumemop_dlsym(void* h, const char* n) { return dlsym(h, n); } -#endif - -using cuMemOp64_fn = int (*)(void* stream, unsigned long long addr, unsigned long long value, - unsigned int flags); -static cuMemOp64_fn g_cu_write64 = nullptr; -static cuMemOp64_fn g_cu_wait64 = nullptr; -static constexpr unsigned int kCuWaitValueGeq = 0x0; // CU_STREAM_WAIT_VALUE_GEQ -static constexpr unsigned int kCuWriteDefault = 0x0; // CU_STREAM_WRITE_VALUE_DEFAULT - -static bool cumemop_resolve() { - static bool resolved = [] { - void* h = cumemop_dlopen(); - if (h == nullptr) return false; - // 11.7+ made the v2 entry points the default; older drivers export only the v1 - // names with the same signature. - g_cu_write64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWriteValue64_v2")); - if (g_cu_write64 == nullptr) - g_cu_write64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWriteValue64")); - g_cu_wait64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWaitValue64_v2")); - if (g_cu_wait64 == nullptr) - g_cu_wait64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWaitValue64")); - return g_cu_write64 != nullptr && g_cu_wait64 != nullptr; - }(); - return resolved; -} - -// Functional probe on a scratch pinned int64: enqueue WRITE(7) + WAIT(>=7) + sync. -// Returns true only if the whole memop path works on THIS stream/device/driver. -static bool cumemops_probe(uintptr_t stream, uintptr_t scratch_addr) { - if (!cumemop_resolve()) return false; - auto* s = reinterpret_cast(stream); - if (g_cu_write64(s, (unsigned long long)scratch_addr, 7ULL, kCuWriteDefault) != 0) return false; - if (g_cu_wait64(s, (unsigned long long)scratch_addr, 7ULL, kCuWaitValueGeq) != 0) return false; - return cudaStreamSynchronize(reinterpret_cast(stream)) == cudaSuccess; -} - -// GPU side of the flag handshake (see the block comment above): enqueued on the -// caller's (possibly capturing) stream; the WAIT immediate is the constant 1, -// replay-safe under CUDA graphs. -// The startup probe validates EAGER memops; a driver could still reject them at graph -// capture time. Those enqueue errors would otherwise be swallowed here and surface only -// as a later EndCapture failure -- log the first CUresult so triage is one step. -static void cumemop_check(int rc, const char* what) { - static std::atomic warned{false}; - if (rc != 0 && !warned.exchange(true)) { - std::fprintf(stderr, - "[freetoken/cpu_moe] %s failed with CUresult=%d (first occurrence; " - "subsequent errors are not repeated). If this happened during CUDA " - "graph capture, the driver lacks capture support for stream memops -- " - "set FREETOKEN_CPU_MOE_FLAG_SYNC=0.\n", - what, rc); - } -} - -static void cumemop_submit(uintptr_t stream, uintptr_t done_addr, uintptr_t ready_addr, - int64_t slot) { - auto* s = reinterpret_cast(stream); - // Order matters and is preserved by the front end: reset done BEFORE raising ready, - // so the coordinator's completion write for THIS step can never be wiped. - cumemop_check(g_cu_write64(s, (unsigned long long)(done_addr + (size_t)slot * 8), 0ULL, - kCuWriteDefault), - "cuStreamWriteValue64(done)"); - cumemop_check(g_cu_write64(s, (unsigned long long)(ready_addr + (size_t)slot * 8), 1ULL, - kCuWriteDefault), - "cuStreamWriteValue64(ready)"); -} - -static void cumemop_sync(uintptr_t stream, uintptr_t done_addr, int64_t slot) { - cumemop_check(g_cu_wait64(reinterpret_cast(stream), - (unsigned long long)(done_addr + (size_t)slot * 8), 1ULL, - kCuWaitValueGeq), - "cuStreamWaitValue64(done)"); -} - -struct DotChoice { - dot_fn fn; - const char* name; -}; - -// SIMD tiers, ascending. Each format picks the highest tier <= the one chosen by -// pick_isa() that it implements (fp4 and fp8_block formats have no bf16-specific tier, so the -// avx512bf16 tier maps to their avx512 kernel). -enum IsaTier { ISA_SCALAR = 0, ISA_AVX2 = 1, ISA_AVX512 = 2, ISA_AVX512BF16 = 3 }; - -// Best tier the CPU+build supports, optionally capped DOWN by -// FREETOKEN_CPU_MOE_ISA={scalar,avx2,avx512,avx512bf16} (A/B testing on a machine -// that supports more). FREETOKEN_CPU_MOE_SCALAR=1 forces scalar (legacy alias). -inline IsaTier pick_isa() { -#if CPU_MOE_X86 - if (getenv("FREETOKEN_CPU_MOE_SCALAR")) return ISA_SCALAR; - IsaTier best = ISA_SCALAR; - if (__builtin_cpu_supports("avx2") && __builtin_cpu_supports("fma")) best = ISA_AVX2; - if (best >= ISA_AVX2 && __builtin_cpu_supports("avx512f")) best = ISA_AVX512; -#ifdef CPU_MOE_HAS_AVX512BF16 - if (best >= ISA_AVX512 && __builtin_cpu_supports("avx512bf16")) best = ISA_AVX512BF16; -#endif - if (const char* f = getenv("FREETOKEN_CPU_MOE_ISA")) { - IsaTier want = best; - if (!std::strcmp(f, "scalar")) want = ISA_SCALAR; - else if (!std::strcmp(f, "avx2")) want = ISA_AVX2; - else if (!std::strcmp(f, "avx512")) want = ISA_AVX512; - else if (!std::strcmp(f, "avx512bf16")) want = ISA_AVX512BF16; - if (want < best) best = want; // cap downward; never force above hw/build support - } - return best; -#else - return ISA_SCALAR; -#endif -} - -DotChoice select_dot() { - const IsaTier t = pick_isa(); -#if CPU_MOE_X86 -#ifdef CPU_MOE_HAS_AVX512BF16 - if (t >= ISA_AVX512BF16) return {dot_avx512bf16, "avx512bf16"}; -#endif - if (t >= ISA_AVX512) return {dot_avx512f, "avx512f"}; - if (t >= ISA_AVX2) return {dot_avx2, "avx2"}; -#endif - (void)t; - return {dot_scalar, "scalar"}; -} - -nvdot_fn select_nvdot() { - const IsaTier t = pick_isa(); -#if CPU_MOE_X86 - if (t >= ISA_AVX512) return dot_nvfp4_avx512; - if (t >= ISA_AVX2) return dot_nvfp4_avx2; -#endif - (void)t; - return dot_nvfp4_scalar; -} - -// AVX-VNNI (VEX-256 VPDPBUSD) availability: Alder/Raptor Lake, Sapphire Rapids+, Zen5. -// Distinct from AVX-512 VNNI. Opt out with FREETOKEN_CPU_MOE_NO_VNNI=1 (A/B the W4A8 path). -inline bool cpu_has_avxvnni() { -#if CPU_MOE_X86 - const char* no = getenv("FREETOKEN_CPU_MOE_NO_VNNI"); - if (no && no[0] && no[0] != '0') return false; // ignore unset/empty/"0" - return __builtin_cpu_supports("avx2") && __builtin_cpu_supports("avxvnni"); -#else - return false; -#endif -} +namespace { -// AVX-512 VNNI (512-bit VPDPBUSD): Cascade Lake+, Ice Lake, Sapphire/Emerald Rapids, -// Zen4+. 2x the 256-bit AVX-VNNI width. FREETOKEN_CPU_MOE_NO_AVX512VNNI=1 forces the -// 256-bit path (A/B the two W4A8 kernels on the same box); FREETOKEN_CPU_MOE_NO_VNNI=1 -// still disables the whole W4A8 family (back to the faithful fp32 nvdot), so it is -// honored here too. -inline bool cpu_has_avx512vnni() { -#if CPU_MOE_X86 && defined(CPU_MOE_HAS_AVX512VNNI) - const char* no = getenv("FREETOKEN_CPU_MOE_NO_AVX512VNNI"); - if (no && no[0] && no[0] != '0') return false; - const char* no_vnni = getenv("FREETOKEN_CPU_MOE_NO_VNNI"); - if (no_vnni && no_vnni[0] && no_vnni[0] != '0') return false; - return __builtin_cpu_supports("avx512vnni"); +inline void pin_this_thread(int cpu) { +#if CPU_MOE_HAS_AFFINITY + cpu_set_t set; + CPU_ZERO(&set); + CPU_SET(cpu, &set); + pthread_setaffinity_np(pthread_self(), sizeof(set), &set); #else - return false; -#endif -} - -// Best W4A8 (int8-activation) nvfp4 dot, or nullptr if no SIMD VNNI (caller keeps the -// faithful fp32 nvdot path). The scalar i8 dot exists only as a correctness reference. -nvi8dot_fn select_nvi8dot() { -#if CPU_MOE_X86 -#if defined(CPU_MOE_HAS_AVX512VNNI) - if (cpu_has_avx512vnni()) return dot_nvfp4_i8_avx512vnni; -#endif - if (cpu_has_avxvnni()) return dot_nvfp4_i8_vnni; -#endif - return nullptr; -} - -// ----------------------- DeepSeek-V4 ds_fp4 (W4A8) --------------------------- -// Row-major e2m1 (2/byte, low nibble first) + e8m0 per-32 block scale, no global -// (w = E2M1[code] * 2^(e8m0-127)); activations are FP8-e4m3 round-tripped (per-128 -// block, ue8m0 scale) before each GEMM. Matches kernel/triton/dsv4 (fused_moe + -// fp8_linear): silu(clamp(gate,max=lim)) * clamp(up,-lim,lim), router weight on the -// down output. - -// Activations are pre-deinterleaved to fp32 (xe[m]=x[2m], xo[m]=x[2m+1]) once per -// token/route and reused across all output rows. This drops the hot dot to one -// vpmovzxbd + two vpermps per 32 weights (1.5 shuffle ops / 16 vs 4 for a bf16, -// dup-permute, per-element gather), so the row-major fp4 GEMV stops being port-5 -// bound and approaches the bf16 memory-bandwidth ceiling. -using dsdot_fn = float (*)(const uint8_t*, const uint8_t*, const float*, const float*, int, - const float*, const float*); - -float dot_dsfp4_scalar(const uint8_t* packed, const uint8_t* scale, const float* xe, - const float* xo, int K, const float* e2m1, const float* e8m0) { - float acc = 0.0f; - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const float sc = e8m0[scale[b]]; - const uint8_t* pk = packed + (size_t)b * 16; // 16 bytes = 32 codes - const float* xeb = xe + (size_t)b * 16; - const float* xob = xo + (size_t)b * 16; - float bsum = 0.0f; - for (int j = 0; j < 16; ++j) { - const uint8_t byte = pk[j]; - bsum += e2m1[byte & 0xF] * xeb[j]; // low nibble -> even-K activation - bsum += e2m1[byte >> 4] * xob[j]; // high nibble -> odd-K activation - } - acc += sc * bsum; - } - return acc; -} - -#if CPU_MOE_X86 -// One 32-block: 16 bytes -> 16 low + 16 high nibble values via two vpermps, times -// the pre-split even/odd fp32 activations, folded by the per-32 e8m0 scale. -__attribute__((target("avx512f"))) -inline __m512 dsfp4_blk(const uint8_t* pk, const float* xeb, const float* xob, __m512 lut, - __m512i loma, float sc) { - __m512i wi = _mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast(pk))); - __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut); - __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut); - __m512 prod = _mm512_fmadd_ps(vlo, _mm512_loadu_ps(xeb), _mm512_mul_ps(vhi, _mm512_loadu_ps(xob))); - return _mm512_mul_ps(prod, _mm512_set1_ps(sc)); -} - -__attribute__((target("avx512f"))) -float dot_dsfp4_avx512(const uint8_t* packed, const uint8_t* scale, const float* xe, - const float* xo, int K, const float* e2m1, const float* e8m0) { - const __m512 lut = _mm512_loadu_ps(e2m1); - const __m512i loma = _mm512_set1_epi32(0xF); - __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps(); - const int nb = K / 32; // 16 packed bytes + one e8m0 scale per block - int b = 0; - for (; b + 2 <= nb; b += 2) { // two independent accumulators hide FMA latency - acc0 = _mm512_add_ps(acc0, dsfp4_blk(packed + (size_t)b * 16, xe + (size_t)b * 16, - xo + (size_t)b * 16, lut, loma, e8m0[scale[b]])); - acc1 = _mm512_add_ps(acc1, dsfp4_blk(packed + (size_t)(b + 1) * 16, xe + (size_t)(b + 1) * 16, - xo + (size_t)(b + 1) * 16, lut, loma, e8m0[scale[b + 1]])); - } - for (; b < nb; ++b) - acc0 = _mm512_add_ps(acc0, dsfp4_blk(packed + (size_t)b * 16, xe + (size_t)b * 16, - xo + (size_t)b * 16, lut, loma, e8m0[scale[b]])); - return _mm512_reduce_add_ps(_mm512_add_ps(acc0, acc1)); -} - -// AVX2: a 32-K block is 16 bytes -> two 8-lane halves (8 even + 8 odd each). -__attribute__((target("avx2,fma"))) -inline __m256 dsfp4_half_avx2(const uint8_t* pk, const float* xeb, const float* xob, __m256 mag8) { - __m256i wi = _mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cast(pk))); - __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); - __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); - return _mm256_fmadd_ps(vlo, _mm256_loadu_ps(xeb), _mm256_mul_ps(vhi, _mm256_loadu_ps(xob))); -} - -__attribute__((target("avx2,fma"))) -float dot_dsfp4_avx2(const uint8_t* packed, const uint8_t* scale, const float* xe, - const float* xo, int K, const float* e2m1, const float* e8m0) { - const __m256 mag8 = _mm256_loadu_ps(e2m1); - __m256 acc0 = _mm256_setzero_ps(), acc1 = _mm256_setzero_ps(); - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const uint8_t* pk = packed + (size_t)b * 16; - const float* xeb = xe + (size_t)b * 16; - const float* xob = xo + (size_t)b * 16; - const __m256 sc = _mm256_set1_ps(e8m0[scale[b]]); - acc0 = _mm256_fmadd_ps(dsfp4_half_avx2(pk, xeb, xob, mag8), sc, acc0); - acc1 = _mm256_fmadd_ps(dsfp4_half_avx2(pk + 8, xeb + 8, xob + 8, mag8), sc, acc1); - } - return hsum256(_mm256_add_ps(acc0, acc1)); -} -#endif - -dsdot_fn select_dsdot() { - const IsaTier t = pick_isa(); -#if CPU_MOE_X86 - if (t >= ISA_AVX512) return dot_dsfp4_avx512; - if (t >= ISA_AVX2) return dot_dsfp4_avx2; -#endif - (void)t; - return dot_dsfp4_scalar; -} - -// ------------------------- mxfp4 (gpt-oss) GEMV ----------------------------- -// Transposed split-K layout: blk[Kpairs, N2] (N innermost), scl[Kpairs/16, N2] -// e8m0 per 32-K. Computes out[c] = sum_kb (E2M1[lo]*x[2kb] + E2M1[hi]*x[2kb+1]) -// * 2^(e8m0-127) for a contiguous column tile (blk/scl already offset to col 0 of -// the tile). Vectorized over N (16 columns / __m512), K stays the outer (cache- -// sequential) loop. Used by both gate_up (K=H) and down (K=I). -using mxgemv_fn = void (*)(float*, const uint8_t*, const uint8_t*, const bf16_t*, int, int, - int, const float*, const float*); - -void mxfp4_gemv_scalar(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, - int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { - for (int c = 0; c < ncol; ++c) out[c] = 0.0f; - for (int kb = 0; kb < Kpairs; ++kb) { - const uint8_t* w = blk + (size_t)kb * N2; - const uint8_t* s = scl + (size_t)(kb >> 4) * N2; - const float xl = bf16_to_f32(x[2 * kb]); - const float xh = bf16_to_f32(x[2 * kb + 1]); - for (int c = 0; c < ncol; ++c) { - const uint8_t byte = w[c]; - out[c] += (e2m1[byte & 0xF] * xl + e2m1[byte >> 4] * xh) * e8m0[s[c]]; - } - } -} - -#if CPU_MOE_X86 -__attribute__((target("avx512f"))) -void mxfp4_gemv_avx512(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, - int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { - (void)e8m0; // e8m0[c]=2^(c-127) computed via bit construction (no gather) - const __m512 lut = _mm512_loadu_ps(e2m1); - const __m512i loma = _mm512_set1_epi32(0xF); - // K-outer / N-inner: each kb cache line is read once and all live column chunks - // (up to 4 -> 64 cols) accumulate from registers, so DRAM/L2 stream the tile once. - int c0 = 0; - for (; c0 + 16 <= ncol; c0 += 64) { - const int nchunk = std::min(4, (ncol - c0) / 16); - __m512 acc[4]; - for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm512_setzero_ps(); - for (int kblk = 0; kblk < Kpairs; kblk += 16) { // 16 K-pairs = 32 K = one scale row - __m512 sc[4]; - for (int ci = 0; ci < nchunk; ++ci) { - __m128i sraw = _mm_loadu_si128(reinterpret_cast( - scl + (size_t)(kblk >> 4) * N2 + c0 + ci * 16)); - sc[ci] = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu8_epi32(sraw), 23)); - } - __m512 blk_acc[4]; - for (int ci = 0; ci < nchunk; ++ci) blk_acc[ci] = _mm512_setzero_ps(); - for (int kk = 0; kk < 16; ++kk) { - const int kb = kblk + kk; - const uint8_t* wbase = blk + (size_t)kb * N2 + c0; - // The transposed layout strides K by N2 bytes; prefetch ahead so the strided - // reads are not exposed to DRAM latency (the HW streamer misses big strides). - constexpr int PFD = 8; - if (kb + PFD < Kpairs) - _mm_prefetch(reinterpret_cast(blk + (size_t)(kb + PFD) * N2 + c0), - _MM_HINT_T0); - const __m512 xl = _mm512_set1_ps(bf16_to_f32(x[2 * kb])); - const __m512 xh = _mm512_set1_ps(bf16_to_f32(x[2 * kb + 1])); - for (int ci = 0; ci < nchunk; ++ci) { - __m512i wi = _mm512_cvtepu8_epi32( - _mm_loadu_si128(reinterpret_cast(wbase + ci * 16))); - __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut); - __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut); - blk_acc[ci] = _mm512_fmadd_ps(vlo, xl, blk_acc[ci]); - blk_acc[ci] = _mm512_fmadd_ps(vhi, xh, blk_acc[ci]); - } - } - for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm512_fmadd_ps(blk_acc[ci], sc[ci], acc[ci]); - } - for (int ci = 0; ci < nchunk; ++ci) _mm512_storeu_ps(out + c0 + ci * 16, acc[ci]); - } - for (int c = c0; c < ncol; ++c) { // tail columns (< 16) - float o = 0.0f; - for (int kb = 0; kb < Kpairs; ++kb) { - const uint8_t byte = blk[(size_t)kb * N2 + c]; - uint32_t bits = (uint32_t)scl[(size_t)(kb >> 4) * N2 + c] << 23; - float sc; - std::memcpy(&sc, &bits, 4); - o += (e2m1[byte & 0xF] * bf16_to_f32(x[2 * kb]) + - e2m1[byte >> 4] * bf16_to_f32(x[2 * kb + 1])) * sc; - } - out[c] = o; - } -} - -__attribute__((target("avx2,fma"))) -void mxfp4_gemv_avx2(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, - int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { - (void)e8m0; // e8m0[s]=2^(s-127) built via s<<23 (no gather) - const __m256 mag8 = _mm256_loadu_ps(e2m1); - int c0 = 0; - for (; c0 + 8 <= ncol; c0 += 32) { // up to 4 chunks of 8 = 32 cols - const int nchunk = std::min(4, (ncol - c0) / 8); - __m256 acc[4]; - for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm256_setzero_ps(); - for (int kblk = 0; kblk < Kpairs; kblk += 16) { // 16 K-pairs = one scale row - __m256 sc[4]; - for (int ci = 0; ci < nchunk; ++ci) { - __m128i sraw = _mm_loadl_epi64(reinterpret_cast( - scl + (size_t)(kblk >> 4) * N2 + c0 + ci * 8)); - sc[ci] = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu8_epi32(sraw), 23)); - } - __m256 blk_acc[4]; - for (int ci = 0; ci < nchunk; ++ci) blk_acc[ci] = _mm256_setzero_ps(); - for (int kk = 0; kk < 16; ++kk) { - const int kb = kblk + kk; - const uint8_t* wbase = blk + (size_t)kb * N2 + c0; - constexpr int PFD = 8; - if (kb + PFD < Kpairs) - _mm_prefetch(reinterpret_cast(blk + (size_t)(kb + PFD) * N2 + c0), - _MM_HINT_T0); - const __m256 xl = _mm256_set1_ps(bf16_to_f32(x[2 * kb])); - const __m256 xh = _mm256_set1_ps(bf16_to_f32(x[2 * kb + 1])); - for (int ci = 0; ci < nchunk; ++ci) { - __m256i wi = _mm256_cvtepu8_epi32( - _mm_loadl_epi64(reinterpret_cast(wbase + ci * 8))); - __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); - __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); - blk_acc[ci] = _mm256_fmadd_ps(vlo, xl, blk_acc[ci]); - blk_acc[ci] = _mm256_fmadd_ps(vhi, xh, blk_acc[ci]); - } - } - for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm256_fmadd_ps(blk_acc[ci], sc[ci], acc[ci]); - } - for (int ci = 0; ci < nchunk; ++ci) _mm256_storeu_ps(out + c0 + ci * 8, acc[ci]); - } - for (int c = c0; c < ncol; ++c) { // tail columns (< 8); none when ncol%8==0 - float o = 0.0f; - for (int kb = 0; kb < Kpairs; ++kb) { - const uint8_t byte = blk[(size_t)kb * N2 + c]; - uint32_t bits = (uint32_t)scl[(size_t)(kb >> 4) * N2 + c] << 23; - float sc; - std::memcpy(&sc, &bits, 4); - o += (e2m1[byte & 0xF] * bf16_to_f32(x[2 * kb]) + - e2m1[byte >> 4] * bf16_to_f32(x[2 * kb + 1])) * sc; - } - out[c] = o; - } -} -#endif - -fp8_block_dot_fn select_fp8_block_dot() { - const IsaTier t = pick_isa(); -#if CPU_MOE_X86 - // the avx512bf16 tier keeps the fp32 kernel: on Sapphire Rapids it is faster than widening to bf16 for dpbf16 - if (t >= ISA_AVX512) return dot_fp8_block_avx512f; - if (t >= ISA_AVX2) return dot_fp8_block_avx2; + (void)cpu; #endif - (void)t; - return dot_fp8_block_scalar; } -mxgemv_fn select_mxgemv() { - const IsaTier t = pick_isa(); +inline void cpu_relax() { #if CPU_MOE_X86 - if (t >= ISA_AVX512) return mxfp4_gemv_avx512; - if (t >= ISA_AVX2) return mxfp4_gemv_avx2; + _mm_pause(); +#elif defined(__aarch64__) + asm volatile("yield" ::: "memory"); #endif - (void)t; - return mxfp4_gemv_scalar; -} - -// Round a clamped |x|<=448 to nearest float8-e4m3 (RNE), back to fp32. Matches -// torch.float8_e4m3fn / triton .to(float8e4nv). -inline float e4m3_round(float x) { - const float sign = x < 0.0f ? -1.0f : 1.0f; - const float a = std::fabs(x); - if (a == 0.0f) return 0.0f; - if (a >= 448.0f) return sign * 448.0f; - int e; - std::frexp(a, &e); // a in [2^(e-1), 2^e) - float step = std::ldexp(1.0f, e - 4); - const float min_step = std::ldexp(1.0f, -9); // e4m3 subnormal step (2^-9) - if (step < min_step) step = min_step; - float r = std::nearbyint(a / step) * step; - if (r > 448.0f) r = 448.0f; - return sign * r; -} - -// IEEE ceil(log2(v)) for v>0 (matches dsv4 _log2_ceil / fast_round_scale). -inline int ceil_log2_pos(float v) { - uint32_t bits; - std::memcpy(&bits, &v, sizeof(bits)); - const int exp = (int)((bits >> 23) & 0xFF); - const int man = (int)(bits & 0x7FFFFF); - return exp - 127 + (man != 0 ? 1 : 0); -} - -// Split an interleaved bf16 row into fp32 even/odd halves (even[m]=src[2m]). -// bf16->fp32 is exact, so this only reorders -- done once per token/route and -// reused across every output row of the GEMV. -inline void deinterleave_bf16_f32(const bf16_t* src, float* even, float* odd, int K) { - for (int m = 0; m < K / 2; ++m) { - even[m] = bf16_to_f32(src[2 * m]); - odd[m] = bf16_to_f32(src[2 * m + 1]); - } -} - -// DeepSeek-V4 / V4.1 activation FP8 round-trip (bf16 in/out): per `block` (the checkpoint's -// fp8 block: 128 on V4, 32 on V4.1), s = 2^ceil(log2(max(|x|,1e-4)/448)); -// y = round_e4m3(clamp(x/s,+-448)) * s. -void fp8_roundtrip_bf16(const bf16_t* src, bf16_t* dst, int K, int block) { - for (int b0 = 0; b0 < K; b0 += block) { - const int b1 = std::min(K, b0 + block); - float amax = 1e-4f; - for (int i = b0; i < b1; ++i) amax = std::max(amax, std::fabs(bf16_to_f32(src[i]))); - const float s = std::ldexp(1.0f, ceil_log2_pos(amax * (1.0f / 448.0f))); - const float inv_s = 1.0f / s; - for (int i = b0; i < b1; ++i) { - float q = bf16_to_f32(src[i]) * inv_s; - q = std::min(448.0f, std::max(-448.0f, q)); - dst[i] = f32_to_bf16(e4m3_round(q) * s); - } - } } -// --------------------------------- executor --------------------------------- - -struct CpuMoeExecutor; +constexpr int kMaxTopK = 64; +// One MoE layer on the first num_tokens rows of the pinned IO. +// It is also the doorbell word the GPU writes, a constant baked into each captured graph; num_tokens >= 1 keeps it nonzero, since 0 means no doorbell. struct MoeTask { - CpuMoeExecutor* exec; - int layer_id; - int num_tokens; - const bf16_t* x; // [num_tokens, H] - const int32_t* ids; // [num_tokens, top_k] (raw expert ids; <0 = skip) - const float* w; // [num_tokens, top_k] - bf16_t* y; // [num_tokens, H] + uint64_t layer : 32; + uint64_t num_tokens : 32; }; +static_assert(sizeof(MoeTask) == sizeof(uint64_t), "MoeTask must pack into one 64-bit word"); -// Output-row tiling. Small enough to give every worker independent work even at -// batch size 1; large enough to amortize the atomic work-grab. -// -// Bandwidth notes (Sapphire Rapids 8480+, 13 cores): the two passes already read -// every expert weight byte exactly once per token (each output row block is owned -// by one worker), and x stays hot in L1 across a (token,expert)'s rows -- so the -// kernel is single-read bandwidth-optimal at bs=1 (~205 GB/s vs ~55 GB/s PCIe). -// One worker per *physical* core, pinned, is the sweet spot; SMT oversubscription -// thrashes the spin-barrier. Deferred (not worth it here / for this workload): -// - AMX-bf16: a GEMM tile engine; decode is M=1 GEMV so tiles sit idle. It would -// only pay off in a grouped/batched (dedup) path. -// - expert dedup for bs>1: read each distinct expert once and GEMM its tokens. -// Helps locality+bytes when bs is large; decode batches here are tiny (<=4). -// - NUMA: a single node is assumed. Multi-socket machines would split each -// expert's K dimension per node (banks are already per-row contiguous). -constexpr int IBLK = 32; -constexpr int HBLK = 32; - -// -------------------------------- Q4_0 (W4A8) -------------------------------- -// Native GGUF Q4_0 experts (gemma4 GGUF): per-32 block = fp16 scale d + 16 packed -// bytes; byte j holds element j in its low nibble and j+16 in its high nibble, so a -// block's storage order is [lo0..lo15, hi0..hi15] and w = (nibble - 8) * d. Matches -// the reference dequant (models/gguf/dequant.py) and the packed banks the GPU offload -// path streams. -// -// llama.cpp ggml_vec_dot_q4_0_q8_0: W4A8. The activation is pre-quantized to Q8_0 -// (per-32-block int8 ``aq`` + fp32 scale ``asb``); each block unpacks its 16 bytes to -// 32 int8 weights in [-8,7] (bytes_from_nibbles_32: low nibbles -> elems 0..15, high -// -> 16..31) and runs an integer block dot -- VPDPBUSD (AVX-VNNI) or VPMADDUBSW+VPMADDWD -// (AVX2) with the ggml sign trick |w|*(sign(w)*a)=w*a, or a scalar int loop -- then -// scales the block sum by wd*xd in fp32. No fp weight dequant / shuffle chain. The GPU -// offload path (ggml_moe_a8_vec / MMVQ) is also W4A8, so cpu and hybrid stay close. -using q4dot_fn = float (*)(const uint8_t*, const int8_t*, const float*, int); - -float q4_0_dot_i8_scalar(const uint8_t* w, const int8_t* aq, const float* asb, int K) { - float acc = 0.0f; - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const uint8_t* blk = w + (size_t)b * 18; - uint16_t dh; - std::memcpy(&dh, blk, sizeof(dh)); - const uint8_t* q = blk + 2; // 16 nibble bytes - const int8_t* a = aq + (size_t)b * 32; - int isum = 0; - for (int j = 0; j < 16; ++j) { - isum += ((int)(q[j] & 0x0F) - 8) * (int)a[j]; // elem j - isum += ((int)(q[j] >> 4) - 8) * (int)a[16 + j]; // elem 16+j - } - acc += fp16_to_f32(dh) * asb[b] * (float)isum; - } - return acc; -} - -#if CPU_MOE_X86 -// fp16 block scale -> fp32 via HW F16C (single value in lane 0). -__attribute__((target("f16c"))) -static inline float q4_scale(uint16_t h) { - return _mm_cvtss_f32(_mm_cvtph_ps(_mm_cvtsi32_si128((int)h))); +inline uint64_t to_word(MoeTask t) { + uint64_t w; + std::memcpy(&w, &t, sizeof(w)); + return w; } -// Unpack one Q4_0 block's 16 bytes -> 32 int8 weights in [-8,7] (elems 0..15 = low -// nibbles, 16..31 = high nibbles). ``eight`` = _mm256_set1_epi8(8). -__attribute__((target("avx2"))) -static inline __m256i q4_unpack32(const uint8_t* blk, __m128i mask, __m256i eight) { - const __m128i qb = _mm_loadu_si128(reinterpret_cast(blk + 2)); - const __m128i lo = _mm_and_si128(qb, mask); - const __m128i hi = _mm_and_si128(_mm_srli_epi16(qb, 4), mask); - return _mm256_sub_epi8(_mm256_set_m128i(hi, lo), eight); +inline MoeTask from_word(uint64_t w) { + MoeTask t; + std::memcpy(&t, &w, sizeof(t)); + return t; } -// AVX2 W4A8 (llama.cpp non-VNNI mul_sum_i8_pairs): integer block dot via VPMADDUBSW + -// VPMADDWD (sign trick), scaled by wd*xd. |aw*sa| pair sums <= 8*127*2 < 32767 -> no -// int16 saturation. This is the fast path on AVX2 CPUs without AVX-VNNI (and the -// avx512-tier fallback, since the block dot is 256-bit either way). -__attribute__((target("avx2,fma,f16c"))) -float q4_0_dot_i8_avx2(const uint8_t* w, const int8_t* aq, const float* asb, int K) { - const __m128i mask = _mm_set1_epi8(0x0F); - const __m256i eight = _mm256_set1_epi8(8); - const __m256i ones16 = _mm256_set1_epi16(1); - __m256 accF = _mm256_setzero_ps(); - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const uint8_t* blk = w + (size_t)b * 18; - _mm_prefetch(reinterpret_cast(blk) + 512, _MM_HINT_T0); - uint16_t dh; - std::memcpy(&dh, blk, sizeof(dh)); - __m256i wq = q4_unpack32(blk, mask, eight); - __m256i a = _mm256_loadu_si256(reinterpret_cast(aq + (size_t)b * 32)); - __m256i aw = _mm256_sign_epi8(wq, wq); // |wq| (unsigned operand) - __m256i sa = _mm256_sign_epi8(a, wq); // sign(wq) * a (signed operand) - __m256i d32 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, sa), ones16); // 8 int32 - accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d32), _mm256_set1_ps(q4_scale(dh) * asb[b]), accF); - } - return hsum256(accF); -} - -// AVX-VNNI W4A8: one VPDPBUSD per block (the fast path on modern CPUs). -__attribute__((target("avx2,avxvnni,fma,f16c"))) -float q4_0_dot_i8_vnni(const uint8_t* w, const int8_t* aq, const float* asb, int K) { - const __m128i mask = _mm_set1_epi8(0x0F); - const __m256i eight = _mm256_set1_epi8(8); - __m256 accF = _mm256_setzero_ps(); - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const uint8_t* blk = w + (size_t)b * 18; - _mm_prefetch(reinterpret_cast(blk) + 512, _MM_HINT_T0); - uint16_t dh; - std::memcpy(&dh, blk, sizeof(dh)); - __m256i wq = q4_unpack32(blk, mask, eight); - __m256i a = _mm256_loadu_si256(reinterpret_cast(aq + (size_t)b * 32)); - __m256i aw = _mm256_sign_epi8(wq, wq); // |wq| (unsigned operand) - __m256i sa = _mm256_sign_epi8(a, wq); // sign(wq) * a (signed operand) - __m256i di = _mm256_dpbusd_avx_epi32(_mm256_setzero_si256(), aw, sa); - // All 32 elems of the block share wd*xd; distribute over di's 8 partial sums and - // reduce at the end (equivalent to scale * block_total). - accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(di), _mm256_set1_ps(q4_scale(dh) * asb[b]), accF); - } - return hsum256(accF); -} -#endif // CPU_MOE_X86 - -// All tiers are W4A8 (int8 activations pre-quantized to Q8_0). AVX-VNNI is orthogonal to -// the ISA tier (gated by cpu_has_avxvnni() / FREETOKEN_CPU_MOE_NO_VNNI), so it wins when -// present; otherwise the 256-bit VPMADDUBSW kernel covers both the avx2 and avx512 tiers. -q4dot_fn select_q4dot() { - const IsaTier t = pick_isa(); -#if CPU_MOE_X86 - if (cpu_has_avxvnni()) return q4_0_dot_i8_vnni; - if (t >= ISA_AVX2) return q4_0_dot_i8_avx2; -#endif - (void)t; - return q4_0_dot_i8_scalar; -} - -enum WFmt { WF_BF16 = 0, WF_NVFP4 = 1, WF_MXFP4 = 2, WF_DSFP4 = 3, WF_Q4_0 = 4, - WF_FP8_BLOCK = 5 }; +// Every weight format this module serves, keyed by the name cpu_executor.py passes; weight_formats() exports the names. +const std::pair (*)(const MoeParams&, const FormatArgs&)> kFormats[] = { + {"bf16", make_bf16_format}, {"nvfp4", make_nvfp4_format}, {"mxfp4_triton", make_mxfp4_format}, + {"ds_fp4", make_ds_fp4_format}, {"q4_0", make_q4_0_format}, {"fp8_block", make_fp8_block_format}, +}; -// Each ctor pointer arg is the address of a CPU int64 array of length -// num_layers (one base address per layer, built by cpu_executor.py's -// _make_table), not a single flat bank. tbl_at resolves -// tbl[layer_id] once per task/pass; a null table (bank unused by this fmt, ptr -// arg 0) resolves to nullptr without dereferencing. -inline const void* tbl_at(const uint64_t* tbl, int layer_id) { - return tbl ? reinterpret_cast(tbl[layer_id]) : nullptr; +std::unique_ptr make_format(const std::string& name, const MoeParams& p, const FormatArgs& a) { + for (const auto& [n, make] : kFormats) + if (name == n) return make(p, a); + throw std::invalid_argument("unknown CPU MoE weight format " + name); } struct CpuMoeExecutor { + // A host-function callback gets one pointer, so it carries the executor with the task. + struct HostTask { + CpuMoeExecutor* exec; + MoeTask task; + }; + + const int num_layers, num_experts; + const bool apply_on_input; + MoeParams params; // the format keeps a reference; declared before fmt + std::unique_ptr fmt; + bool preps_inter = false; + bool rounds_routes = false; + int n_iblk = 0, n_hblk = 0; // kRowBlock blocks per intermediate row and per output row + + // Pinned IO shared by every layer and batch size (one task in flight): [max_tokens, H] x and y, [max_tokens, top_k] ids and w. + const bf16_t* io_x; + const int32_t* io_ids; + const float* io_w; + bf16_t* io_y; + AlignedVector g_scratch; // [max_tokens * top_k * I] intermediate + + // Worker pool. int num_threads; - int num_layers, num_experts, top_k; - int H, I; - int act, apply_on_input; - int fmt; // WFmt - bool needs_di = false; // pre-deinterleave activations to fp32 (nvfp4/ds_fp4/fp8_block) - // Per-layer pointer tables (one base address per layer, see tbl_at). gate_up_tbl - // doubles as the bf16 gate_up table and the nvfp4/mxfp4/q4_0/ds_fp4 packed-gate_up - // table (down_tbl likewise for down); which reinterpretation applies is picked by - // fmt at each resolve site (see gemm1_dot/gemm2_dot/do_pass1_mxfp4/do_pass1_dsfp4). - const uint64_t* gate_up_tbl; // bf16: [E,2I,H] rows; else: packed e2m1/mxfp4-blocks - const uint64_t* down_tbl; // bf16: [E,H,I] rows; else: packed e2m1/mxfp4-blocks - const uint64_t* gu_scale_tbl; // nvfp4/mxfp4/ds_fp4: [E,2I,*] block scales - const uint64_t* gu_global_tbl; // nvfp4: [E,2I] fp16 row globals - const uint64_t* dn_scale_tbl; // nvfp4/mxfp4/ds_fp4: [E,H,*] block scales - const uint64_t* dn_global_tbl; // nvfp4: [E,H] fp16 row globals - const uint64_t* gu_bias_tbl; // mxfp4: [E,2I] bf16 biases - const uint64_t* dn_bias_tbl; // mxfp4: [E,H] bf16 biases - float swiglu_alpha; - float swiglu_limit; // +inf == no clamp - dot_fn dot; - nvdot_fn nvdot; - nvi8dot_fn nvi8dot = nullptr; // AVX-VNNI W4A8 nvfp4 dot (nullptr -> use fp32 nvdot) - bool use_vnni = false; // nvfp4 + AVX-VNNI: decode via int8 VPDPBUSD (W4A8) - bool use_q4a8 = false; // q4_0: always W4A8 (llama.cpp Q4_0 x Q8_0); int8 pre-quant - dsdot_fn dsdot; - mxgemv_fn mxgemv; - q4dot_fn q4dot; - fp8_block_dot_fn fp8_block_dot; - // fp8_block: bf16 scale per 128x128 weight block. These are the scale tensor's - // trailing strides, including any padding used by the host bank layout. - int fp8_block_gu_scale_row = 0, fp8_block_dn_scale_row = 0; - // ds_fp4: the caller already FP8-round-tripped the input activations on the GPU - // (same reference grid), so submit() must not repeat it on the host-callback - // thread. That scalar per-element pass is single-threaded ON THE DECODE CRITICAL - // PATH (~0.3ms/layer at H=4096, every worker and the GPU waiting on it); moving - // it to a captured GPU elementwise kernel removes it while keeping the official - // W4A8 numerics bit-exact. Set via set_input_prequant (see cpu_executor.py). - bool input_prequant = false; - // ds_fp4 activation round-trip block (the checkpoint's fp8 block); set via set_act_block. - int act_block = 128; - // Q4_0 packed-row byte strides (H/32*18 for gate_up over K=H, I/32*18 for down over K=I). - int q4_gu_row_bytes = 0, q4_dn_row_bytes = 0; - float e2m1_lut[16]; - float e4m3_lut[256]; - float e8m0_lut[256]; // mxfp4 block scale: 2^(s-127), s clamped to [0,254] - const char* isa; - - std::vector g_scratch; // [max_tokens * top_k * I] intermediate - std::vector xq_scratch; // [max_tokens * H] ds_fp4 fp8-roundtripped input - // nvfp4/ds_fp4/fp8_block activations pre-deinterleaved to fp32 (even/odd K) for the row-major dot. - std::vector xe_scratch, xo_scratch; // [max_tokens * H/2] (input) - std::vector ge_scratch, go_scratch; // [max_tokens*top_k*I/2] (intermediate) - // AVX-VNNI W4A8: per-16-block int8 activations [even(8),odd(8)] + per-block scale. - std::vector xi8_scratch, gi8_scratch; // [max_tokens*H], [max_tokens*top_k*I] - std::vector xas_scratch, gas_scratch; // [max_tokens*H/16], [..*top_k*I/16] - std::string isa_str; - + std::vector core_ids; // worker tid -> logical CPU to pin to (may be empty) std::vector workers; std::mutex task_mtx; std::condition_variable task_cv; @@ -1421,7 +136,7 @@ struct CpuMoeExecutor { std::atomic stop{false}; // spinning workers read it without the lock uint64_t cur_gen = 0; - MoeTask* cur_task = nullptr; + MoeTask cur_task{}; // Separate cache lines: idle workers poll `submitted` while busy workers hit p1_next/p2_next. alignas(64) std::atomic submitted{0}; alignas(64) std::atomic completed{0}; @@ -1431,264 +146,56 @@ struct CpuMoeExecutor { std::atomic spin_wait{false}; bool coord_pinned = false; + // Work split of the current task. alignas(64) std::atomic p1_next{0}; std::atomic p2_next{0}; - std::atomic prt_next{0}; // intermediate-row prep phase (prep_g_row) + std::atomic prt_next{0}; // intermediate-row prep phase (prep_inter) int64_t p1_total = 0, p2_total = 0, prt_total = 0; - int n_iblk = 0, n_hblk = 0; std::atomic done_count{0}; std::atomic bar_count{0}; std::atomic bar_sense{0}; - std::vector owned_tasks; // persistent task descriptors (graph-stable) - std::vector core_ids; // worker tid -> logical CPU to pin to (may be empty) - // ---- Flag-based GPU<->CPU handshake (replaces the per-layer cudaLaunchHostFunc pair) ---- - // A tiny GPU kernel bumps ready_flags[slot] at submit; this coordinator thread busy-polls - // it, runs the slot's task on the worker pool, and sets done_flags[slot], which a GPU - // spin-wait kernel polls at sync. This removes the ~2x30-50us host-func dispatch round - // trips per MoE layer per decode step that otherwise idle the GPU (~6 ms/step on a - // 75-layer model). One slot per (layer, decode batch size) pair -- the Python side - // allocates slots as tasks are created. Flags live in mapped-pinned host memory (UVA: - // the same pointers are used by the GPU kernels and by this thread). + // The GPU rings the doorbell with a task word at submit; this coordinator thread busy-polls it, runs the task on the worker pool, and completes it, which the GPU waits on at sync. + // This removes the ~2x30-50us host-func dispatch round trips per MoE layer per decode step that otherwise idle the GPU (~6 ms/step on a 75-layer model). + std::unique_ptr gpu_sync; + std::unique_ptr kernel_io; // replaces the copies around the doorbell once enable_kernel_io() runs std::thread coord_thread; std::atomic coord_stop{false}; - volatile int64_t* ready_flags = nullptr; // GPU increments, this thread polls - volatile int64_t* done_flags = nullptr; // this thread sets, GPU spin-waits - int coord_num_slots = 0; - std::vector flag_task; // slot -> task (registered lazily) - std::vector flag_served; // slot -> completed dispatch count (tests/debug) - std::mutex flag_task_mtx; - - // Portable ordering for the flag handshake: "ready observed => the DMA'd inputs that - // preceded the bump are visible" and "y stores are visible before done". Plain - // volatile loads lean on x86 TSO; acquire/release makes it hold on aarch64 too - // (GH200/Jetson) at zero x86 cost. (MSVC branch is x86-only today: compiler barrier - // + TSO.) - static int64_t flag_load_acquire(const volatile int64_t* p) { -#if defined(_MSC_VER) - const int64_t v = *p; - _ReadWriteBarrier(); - return v; -#else - return __atomic_load_n(const_cast(p), __ATOMIC_ACQUIRE); -#endif - } - - static void flag_store_release(volatile int64_t* p, int64_t v) { -#if defined(_MSC_VER) - _ReadWriteBarrier(); - *p = v; -#else - __atomic_store_n(const_cast(p), v, __ATOMIC_RELEASE); -#endif - } + std::atomic flag_served{0}; // completed doorbells (tests, watchdog) + std::map host_tasks; // host-function userData, keyed by doorbell word (graph-stable) CpuMoeExecutor(int num_threads_, int num_layers_, int num_experts_, int top_k_, - int hidden_size, int inter_size, int max_tokens, int activation_id, - int apply_router_weight_on_input, int weight_format, - uintptr_t gate_up_ptr, uintptr_t down_ptr, uintptr_t gate_up_scale_ptr, - uintptr_t gate_up_global_ptr, uintptr_t down_scale_ptr, - uintptr_t down_global_ptr, uintptr_t gate_up_bias_ptr, - uintptr_t down_bias_ptr, double swiglu_alpha_, double swiglu_limit_, - std::vector core_ids_, int fp8_block_gu_scale_stride_, - int fp8_block_dn_scale_stride_) - : num_threads(num_threads_ > 0 ? num_threads_ : 1), - num_layers(num_layers_), + int hidden_size, int inter_size, int max_tokens_, int activation_id, + int apply_router_weight_on_input, const std::string& weight_format, + std::map banks, std::map options, + double swiglu_alpha_, double swiglu_limit_, + std::vector core_ids_, uintptr_t x_ptr, uintptr_t ids_ptr, uintptr_t w_ptr, + uintptr_t y_ptr) + : num_layers(num_layers_), num_experts(num_experts_), - top_k(top_k_), - H(hidden_size), - I(inter_size), - act(activation_id), - apply_on_input(apply_router_weight_on_input), - fmt(weight_format), - gate_up_tbl(reinterpret_cast(gate_up_ptr)), - down_tbl(reinterpret_cast(down_ptr)), - gu_scale_tbl(reinterpret_cast(gate_up_scale_ptr)), - gu_global_tbl(reinterpret_cast(gate_up_global_ptr)), - dn_scale_tbl(reinterpret_cast(down_scale_ptr)), - dn_global_tbl(reinterpret_cast(down_global_ptr)), - gu_bias_tbl(reinterpret_cast(gate_up_bias_ptr)), - dn_bias_tbl(reinterpret_cast(down_bias_ptr)), - swiglu_alpha(static_cast(swiglu_alpha_)), - swiglu_limit(static_cast(swiglu_limit_)), - fp8_block_gu_scale_row(fp8_block_gu_scale_stride_), - fp8_block_dn_scale_row(fp8_block_dn_scale_stride_), + apply_on_input(apply_router_weight_on_input != 0), + params{top_k_, max_tokens_, hidden_size, inter_size, activation_id, static_cast(swiglu_alpha_), + static_cast(swiglu_limit_)}, + io_x(reinterpret_cast(x_ptr)), + io_ids(reinterpret_cast(ids_ptr)), + io_w(reinterpret_cast(w_ptr)), + io_y(reinterpret_cast(y_ptr)), + num_threads(num_threads_ > 0 ? num_threads_ : 1), core_ids(std::move(core_ids_)) { - DotChoice c = select_dot(); - dot = c.fn; - nvdot = select_nvdot(); - dsdot = select_dsdot(); - mxgemv = select_mxgemv(); - q4dot = select_q4dot(); - fp8_block_dot = select_fp8_block_dot(); - if (weight_format == WF_Q4_0) { - if (H % 32 != 0 || I % 32 != 0) - throw std::runtime_error("Q4_0 CPU MoE requires H and I to be multiples of 32"); - q4_gu_row_bytes = (H / 32) * 18; // K = H (gate_up rows) - q4_dn_row_bytes = (I / 32) * 18; // K = I (down rows) - } - isa = c.name; - // nvfp4 (AVX-VNNI only): W4A8 int8 decode when the CPU supports it. q4_0 is always - // W4A8 (activations pre-quantized to Q8_0); select_q4dot picks VPDPBUSD / VPMADDUBSW - // / scalar for the tier, so the tag reflects which of those q4dot resolved to. - nvi8dot = select_nvi8dot(); - use_vnni = (weight_format == WF_NVFP4) && (nvi8dot != nullptr); - use_q4a8 = (weight_format == WF_Q4_0); - const char* q4tag = use_q4a8 ? (cpu_has_avxvnni() ? "+vnni(q4_0-w4a8)" : "+q4_0-w4a8") : ""; - const char* vnni_tag = - cpu_has_avx512vnni() ? "+avx512vnni(nvfp4-w4a8)" : "+vnni(nvfp4-w4a8)"; - isa_str = std::string(c.name) + (use_vnni ? vnni_tag : "") + q4tag; - isa = isa_str.c_str(); - for (int i = 0; i < 16; ++i) e2m1_lut[i] = kE2M1[i]; - for (int i = 0; i < 256; ++i) e4m3_lut[i] = e4m3_decode((uint8_t)i); - // e8m0 (mxfp4 block scale) = 2^(s-127); the GPU GEMV clamps s to [0,254]. - for (int i = 0; i < 256; ++i) e8m0_lut[i] = std::ldexp(1.0f, std::min(i, 254) - 127); - g_scratch.assign(static_cast(max_tokens) * top_k * I, 0); - // Row-major fp4 (nvfp4/ds_fp4) and fp8_block pre-deinterleave activations to fp32 even/odd. - needs_di = (fmt == WF_NVFP4 || fmt == WF_DSFP4 || fmt == WF_FP8_BLOCK); - if (needs_di) { - if (fmt == WF_DSFP4) xq_scratch.assign(static_cast(max_tokens) * H, 0); - xe_scratch.assign(static_cast(max_tokens) * (H / 2), 0); - xo_scratch.assign(static_cast(max_tokens) * (H / 2), 0); - ge_scratch.assign(static_cast(max_tokens) * top_k * (I / 2), 0); - go_scratch.assign(static_cast(max_tokens) * top_k * (I / 2), 0); - if (use_vnni) { - xi8_scratch.assign(static_cast(max_tokens) * H, 0); - xas_scratch.assign(static_cast(max_tokens) * (H / 16), 0); - gi8_scratch.assign(static_cast(max_tokens) * top_k * I, 0); - gas_scratch.assign(static_cast(max_tokens) * top_k * (I / 16), 0); - } - } - // q4_0 W4A8: per-32-block Q8_0 activations (int8 + fp32 scale) for input + intermediate. - if (use_q4a8) { - xi8_scratch.assign(static_cast(max_tokens) * H, 0); - xas_scratch.assign(static_cast(max_tokens) * (H / 32), 0); - gi8_scratch.assign(static_cast(max_tokens) * top_k * I, 0); - gas_scratch.assign(static_cast(max_tokens) * top_k * (I / 32), 0); - } + if (params.top_k > kMaxTopK) throw std::invalid_argument("CPU MoE supports top_k <= " + std::to_string(kMaxTopK)); + fmt = make_format(weight_format, params, FormatArgs{weight_format, std::move(banks), std::move(options)}); + preps_inter = fmt->preps_inter(); + rounds_routes = fmt->rounds_routes(); + n_iblk = (params.I + kRowBlock - 1) / kRowBlock; + n_hblk = (params.H + kRowBlock - 1) / kRowBlock; + g_scratch.assign(static_cast(params.max_tokens) * params.top_k * params.I, 0); for (int t = 0; t < num_threads; ++t) workers.emplace_back([this, t] { worker_loop(t); }); } - // Quantize a bf16 activation row to Q8_0 (llama.cpp): per-32-block symmetric int8 in - // natural order + one fp32 scale (amax/127) per block. Done once per token/route, - // amortized over every output row that the q4_0 W4A8 GEMV reads. K % 32 == 0. - void quant_q8_0(const bf16_t* x, int K, int8_t* aq, float* asb) { - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const bf16_t* xb = x + (size_t)b * 32; - float xf[32], amax = 0.0f; - for (int j = 0; j < 32; ++j) { - xf[j] = bf16_to_f32(xb[j]); - amax = std::max(amax, std::fabs(xf[j])); - } - const float d = amax > 0.0f ? amax / 127.0f : 1.0f; - asb[b] = d; - const float inv = amax > 0.0f ? 1.0f / d : 0.0f; - int8_t* o = aq + (size_t)b * 32; - for (int j = 0; j < 32; ++j) - o[j] = (int8_t)std::max(-127, std::min(127, (int)std::lround(xf[j] * inv))); - } - } - - // Quantize the pre-deinterleaved fp32 even/odd activations to per-16-block int8 in the - // [even(8),odd(8)] layout the VNNI dot expects. Done once per token/route (amortized over - // every output row), so a scalar pass is fine relative to the GEMV. - void quant_i8_pg16(const float* xe, const float* xo, int K, int8_t* asi8, float* asb) { - const int nb = K / 16; - for (int b = 0; b < nb; ++b) { - const float* xeb = xe + (size_t)b * 8; - const float* xob = xo + (size_t)b * 8; - float amax = 0.0f; - for (int j = 0; j < 8; ++j) - amax = std::max(amax, std::max(std::fabs(xeb[j]), std::fabs(xob[j]))); - const float s = amax > 0.0f ? amax / 127.0f : 1.0f; - asb[b] = s; - const float inv = 1.0f / s; - int8_t* ae = asi8 + (size_t)b * 16; - for (int j = 0; j < 8; ++j) { - int qe = (int)std::lround(xeb[j] * inv), qo = (int)std::lround(xob[j] * inv); - ae[j] = (int8_t)std::max(-127, std::min(127, qe)); - ae[8 + j] = (int8_t)std::max(-127, std::min(127, qo)); - } - } - } - - // gate_up output row `row` (in [0, 2I)) dotted with activation over K = H. ``e`` is - // the layer-local expert row (0..num_experts); the layer bases (already resolved - // once per task/pass by the caller via tbl_at) pick the layer's own tensors. - // bf16 uses the interleaved bf16 row; nvfp4 uses the pre-split fp32 even/odd halves - // (or, with AVX-VNNI, the per-16-block int8 activations). - // fp8_block uses the halves for whole 128-blocks and the bf16 row for a ragged tail. - inline float gemm1_dot(const bf16_t* gate_up_l, const uint8_t* gu_packed_l, - const uint8_t* gu_scale_l, const uint16_t* gu_global_l, int e, int row, - const bf16_t* x, const float* xe, const float* xo, const int8_t* xi8, - const float* xas) { - if (fmt == WF_BF16) { - const bf16_t* w = gate_up_l + ((size_t)e * (2 * I) + row) * H; - return dot(w, x, H); - } - if (fmt == WF_Q4_0) { - const uint8_t* w = - gu_packed_l + ((size_t)e * (2 * I) + row) * (size_t)q4_gu_row_bytes; - return q4dot(w, xi8, xas, H); // W4A8: int8 activations (Q8_0), scale in xas - } - if (fmt == WF_FP8_BLOCK) { - const uint8_t* w = gu_packed_l + ((size_t)e * (2 * I) + row) * (size_t)H; - const bf16_t* sc = reinterpret_cast(gu_scale_l) + - ((size_t)e * ((2 * I + FP8_BLK - 1) / FP8_BLK) + row / FP8_BLK) * - (size_t)fp8_block_gu_scale_row; - return fp8_block_dot(w, xe, xo, x, sc, H, e4m3_lut); - } - const size_t r = (size_t)e * (2 * I) + row; - if (use_vnni) - return nvi8dot(gu_packed_l + r * (size_t)(H / 2), gu_scale_l + r * (size_t)(H / 16), - fp16_to_f32(gu_global_l[r]), xi8, H, e4m3_lut, xas); - return nvdot(gu_packed_l + r * (size_t)(H / 2), gu_scale_l + r * (size_t)(H / 16), - fp16_to_f32(gu_global_l[r]), xe, xo, H, e2m1_lut, e4m3_lut); - } - - // down output row `row` (in [0, H)) dotted with the intermediate over K = I. Same - // layer-local-base convention as gemm1_dot. - inline float gemm2_dot(const bf16_t* down_l, const uint8_t* dn_packed_l, - const uint8_t* dn_scale_l, const uint16_t* dn_global_l, int e, int row, - const bf16_t* g, const float* ge, const float* go, const int8_t* gi8, - const float* gas) { - if (fmt == WF_BF16) { - const bf16_t* w = down_l + ((size_t)e * H + row) * I; - return dot(w, g, I); - } - if (fmt == WF_Q4_0) { - const uint8_t* w = dn_packed_l + ((size_t)e * H + row) * (size_t)q4_dn_row_bytes; - return q4dot(w, gi8, gas, I); // W4A8: int8 activations (Q8_0), scale in gas - } - if (fmt == WF_FP8_BLOCK) { - const uint8_t* w = dn_packed_l + ((size_t)e * H + row) * (size_t)I; - const bf16_t* sc = reinterpret_cast(dn_scale_l) + - ((size_t)e * ((H + FP8_BLK - 1) / FP8_BLK) + row / FP8_BLK) * - (size_t)fp8_block_dn_scale_row; - return fp8_block_dot(w, ge, go, g, sc, I, e4m3_lut); - } - const size_t r = (size_t)e * H + row; - if (use_vnni) - return nvi8dot(dn_packed_l + r * (size_t)(I / 2), dn_scale_l + r * (size_t)(I / 16), - fp16_to_f32(dn_global_l[r]), gi8, I, e4m3_lut, gas); - return nvdot(dn_packed_l + r * (size_t)(I / 2), dn_scale_l + r * (size_t)(I / 16), - fp16_to_f32(dn_global_l[r]), ge, go, I, e2m1_lut, e4m3_lut); - } - void pin_self(int tid) { -#if CPU_MOE_HAS_AFFINITY - if (core_ids.empty()) return; - const int cpu = core_ids[tid % static_cast(core_ids.size())]; - cpu_set_t set; - CPU_ZERO(&set); - CPU_SET(cpu, &set); - pthread_setaffinity_np(pthread_self(), sizeof(set), &set); -#else - (void)tid; -#endif + if (!core_ids.empty()) pin_this_thread(core_ids[tid % static_cast(core_ids.size())]); } ~CpuMoeExecutor() { @@ -1701,30 +208,23 @@ struct CpuMoeExecutor { task_cv.notify_all(); for (auto& th : workers) if (th.joinable()) th.join(); - for (MoeTask* t : owned_tasks) delete t; } - uintptr_t create_task(int layer_id, int num_tokens, uintptr_t x_ptr, - uintptr_t ids_ptr, uintptr_t w_ptr, uintptr_t y_ptr) { - MoeTask* t = new MoeTask{this, - layer_id, - num_tokens, - reinterpret_cast(x_ptr), - reinterpret_cast(ids_ptr), - reinterpret_cast(w_ptr), - reinterpret_cast(y_ptr)}; - owned_tasks.push_back(t); - return reinterpret_cast(t); + MoeTask task(int layer_id, int num_tokens) const { + if (layer_id < 0 || layer_id >= num_layers || num_tokens < 1 || num_tokens > params.max_tokens) + throw std::out_of_range("CPU MoE task (layer " + std::to_string(layer_id) + ", " + + std::to_string(num_tokens) + " tokens) is outside the executor's layers or IO"); + return MoeTask{static_cast(layer_id), static_cast(num_tokens)}; } - const char* isa_name() const { return isa; } + std::string isa_name() const { return fmt->isa_name(); } - const char* fp8_block_isa_name() const { -#if CPU_MOE_X86 - if (fp8_block_dot == dot_fp8_block_avx512f) return "avx512f"; - if (fp8_block_dot == dot_fp8_block_avx2) return "avx2"; -#endif - return "scalar"; + // Whether any route of the task names an expert; do_pass1 and do_pass2 skip the others. + bool has_routes(const MoeTask& t) const { + const int64_t n = static_cast(t.num_tokens) * params.top_k; + for (int64_t i = 0; i < n; ++i) + if (io_ids[i] >= 0 && io_ids[i] < num_experts) return true; + return false; } void barrier(int& local_sense) { @@ -1733,275 +233,58 @@ struct CpuMoeExecutor { bar_count.store(0); bar_sense.store(local_sense); } else { - while (bar_sense.load() != local_sense) { -#if CPU_MOE_X86 - _mm_pause(); -#endif - } + while (bar_sense.load() != local_sense) cpu_relax(); } } - void do_pass1(const MoeTask* t, int64_t p) { - if (fmt == WF_MXFP4) { - do_pass1_mxfp4(t, p); - return; - } - if (fmt == WF_DSFP4) { - do_pass1_dsfp4(t, p); - return; - } + // Pass 1 item p = one kRowBlock block of one (token, route)'s intermediate row. + void do_pass1(const MoeTask& t, int64_t p) { + const int top_k = params.top_k, H = params.H, I = params.I; const int64_t ib = p % n_iblk; const int64_t tk = p / n_iblk; const int k = static_cast(tk % top_k); const int tok = static_cast(tk / top_k); - const int e = t->ids[static_cast(tok) * top_k + k]; + const int e = io_ids[static_cast(tok) * top_k + k]; if (e < 0 || e >= num_experts) return; - const float w_in = apply_on_input ? t->w[static_cast(tok) * top_k + k] : 1.0f; - // Resolve this task's layer base once; row indexing below is layer-local (e). - const bf16_t* gate_up_l = reinterpret_cast(tbl_at(gate_up_tbl, t->layer_id)); - const uint8_t* gu_packed_l = reinterpret_cast(tbl_at(gate_up_tbl, t->layer_id)); - const uint8_t* gu_scale_l = reinterpret_cast(tbl_at(gu_scale_tbl, t->layer_id)); - const uint16_t* gu_global_l = - reinterpret_cast(tbl_at(gu_global_tbl, t->layer_id)); - const bf16_t* x_row = t->x + (size_t)tok * H; - const float* xe = needs_di ? xe_scratch.data() + (size_t)tok * (H / 2) : nullptr; - const float* xo = needs_di ? xo_scratch.data() + (size_t)tok * (H / 2) : nullptr; - const int8_t* xi8 = - (use_vnni || use_q4a8) ? xi8_scratch.data() + (size_t)tok * H : nullptr; - const float* xas = use_vnni ? xas_scratch.data() + (size_t)tok * (H / 16) - : use_q4a8 ? xas_scratch.data() + (size_t)tok * (H / 32) - : nullptr; - bf16_t* g_row = g_scratch.data() + ((size_t)tok * top_k + k) * I; - const int i0 = static_cast(ib) * IBLK; - const int i1 = std::min(I, i0 + IBLK); - const bool clamped = act == ACT_SWIGLUOAI || act == ACT_SWIGLU_CLAMP; - const float up_bias = act == ACT_SWIGLUOAI ? 1.0f : 0.0f; - const float lim = swiglu_limit, alpha = swiglu_alpha; - for (int i = i0; i < i1; ++i) { - // gate = row i, up = row I+i - float gate = - gemm1_dot(gate_up_l, gu_packed_l, gu_scale_l, gu_global_l, e, i, x_row, xe, xo, xi8, xas) * w_in; - float up = gemm1_dot(gate_up_l, gu_packed_l, gu_scale_l, gu_global_l, e, I + i, x_row, - xe, xo, xi8, xas) * w_in; - if (clamped) { - // clamp(gate, max=lim) * sigmoid(alpha * gate) * (clamp(up, +-lim) + up_bias) - // -- swigluoai carries the +1 up bias (gpt-oss/MiniMax); swiglu_clamp - // (GLM-5.3) does not. lim == +inf: no clamp. - if (gate > lim) gate = lim; - if (up > lim) up = lim; - else if (up < -lim) up = -lim; - const float glu = gate / (1.0f + std::exp(-gate * alpha)); - g_row[i] = f32_to_bf16(glu * (up + up_bias)); - } else { - g_row[i] = f32_to_bf16(act_apply(act, gate) * up); - } - } + const float w_in = apply_on_input ? io_w[static_cast(tok) * top_k + k] : 1.0f; + const int i0 = static_cast(ib) * kRowBlock; + const int i1 = std::min(I, i0 + kRowBlock); + fmt->gate_up(static_cast(t.layer), e, tok, i0, i1, io_x + (size_t)tok * H, w_in, + g_scratch.data() + ((size_t)tok * top_k + k) * I); } - void do_pass2(const MoeTask* t, int64_t p) { - if (fmt == WF_MXFP4) { - do_pass2_mxfp4(t, p); - return; - } - if (fmt == WF_DSFP4) { - do_pass2_dsfp4(t, p); - return; - } - const int64_t hb = p % n_hblk; - const int tok = static_cast(p / n_hblk); - const int h0 = static_cast(hb) * HBLK; - const int h1 = std::min(H, h0 + HBLK); - // Resolve this task's layer base once; row indexing below is layer-local (e). - const bf16_t* down_l = reinterpret_cast(tbl_at(down_tbl, t->layer_id)); - const uint8_t* dn_packed_l = reinterpret_cast(tbl_at(down_tbl, t->layer_id)); - const uint8_t* dn_scale_l = reinterpret_cast(tbl_at(dn_scale_tbl, t->layer_id)); - const uint16_t* dn_global_l = - reinterpret_cast(tbl_at(dn_global_tbl, t->layer_id)); - bf16_t* y_row = t->y + (size_t)tok * H; - for (int h = h0; h < h1; ++h) { - float acc = 0.0f; - for (int k = 0; k < top_k; ++k) { - const int e = t->ids[static_cast(tok) * top_k + k]; - if (e < 0 || e >= num_experts) continue; - const float w_out = apply_on_input ? 1.0f : t->w[static_cast(tok) * top_k + k]; - const size_t gr = (size_t)tok * top_k + k; - const bf16_t* g_row = g_scratch.data() + gr * I; - const float* ge = needs_di ? ge_scratch.data() + gr * (I / 2) : nullptr; - const float* go = needs_di ? go_scratch.data() + gr * (I / 2) : nullptr; - const int8_t* gi8 = (use_vnni || use_q4a8) ? gi8_scratch.data() + gr * I : nullptr; - const float* gas = use_vnni ? gas_scratch.data() + gr * (I / 16) - : use_q4a8 ? gas_scratch.data() + gr * (I / 32) - : nullptr; - acc += gemm2_dot(down_l, dn_packed_l, dn_scale_l, dn_global_l, e, h, g_row, ge, go, gi8, - gas) * w_out; - } - y_row[h] = f32_to_bf16(acc); - } - } - - // ----------------------------- mxfp4 (gpt-oss) ----------------------------- - // Transposed split-K layout (N innermost), so the GEMV streams K and accumulates - // a contiguous block of N output columns -> cache-line-efficient, no repack, no - // extra host memory. Pass 1 fuses gate_up + clamped-swiglu(+bias); pass 2 fuses - // down(+bias) * router-weight, summed over the token's routes. - // - // Dequant: w = E2M1[code] * 2^(e8m0_scale - 127); two codes per byte (low nibble - // first), one e8m0 scale per 32 contiguous K. Matches kernel/triton/mxfp4_moe.py. - - void do_pass1_mxfp4(const MoeTask* t, int64_t p) { - const int64_t ib = p % n_iblk; - const int64_t tk = p / n_iblk; - const int k = static_cast(tk % top_k); - const int tok = static_cast(tk / top_k); - const int e = t->ids[static_cast(tok) * top_k + k]; - if (e < 0 || e >= num_experts) return; - // Resolve this task's layer base once; row indexing below is layer-local (e). - const uint8_t* gu_packed_l = reinterpret_cast(tbl_at(gate_up_tbl, t->layer_id)); - const uint8_t* gu_scale_l = reinterpret_cast(tbl_at(gu_scale_tbl, t->layer_id)); - const bf16_t* gu_bias_l = reinterpret_cast(tbl_at(gu_bias_tbl, t->layer_id)); - const int N2 = 2 * I; // gate_up output width (gate/up interleaved) - const int Hh = H / 2; // packed-K rows - const int i0 = static_cast(ib) * IBLK; - const int i1 = std::min(I, i0 + IBLK); - const int nunit = i1 - i0; // intermediate units owned by this tile - const int col0 = 2 * i0; // first gate_up column - const int ncol = 2 * nunit; // gate_up columns owned by this tile - const bf16_t* x_row = t->x + (size_t)tok * H; - const uint8_t* blk_e = gu_packed_l + (size_t)e * Hh * N2; - const uint8_t* scl_e = gu_scale_l + (size_t)e * (size_t)(H / 32) * N2; - float gu[2 * IBLK]; - mxgemv(gu, blk_e + col0, scl_e + col0, x_row, Hh, N2, ncol, e2m1_lut, e8m0_lut); - const bf16_t* bias_e = gu_bias_l + (size_t)e * N2 + col0; - bf16_t* g_row = g_scratch.data() + ((size_t)tok * top_k + k) * I; - const float lim = swiglu_limit, alpha = swiglu_alpha; - for (int j = 0; j < nunit; ++j) { - float gate = gu[2 * j] + bf16_to_f32(bias_e[2 * j]); - float up = gu[2 * j + 1] + bf16_to_f32(bias_e[2 * j + 1]); - if (gate > lim) gate = lim; - if (up > lim) up = lim; - else if (up < -lim) up = -lim; - const float glu = gate / (1.0f + std::exp(-gate * alpha)); // gate * sigmoid(alpha*gate) - g_row[i0 + j] = f32_to_bf16(glu * (up + 1.0f)); - } - } - - void do_pass2_mxfp4(const MoeTask* t, int64_t p) { + // Pass 2 item p = one kRowBlock block of one token's output row, summed over its routes. + void do_pass2(const MoeTask& t, int64_t p) { + const int top_k = params.top_k, H = params.H, I = params.I; const int64_t hb = p % n_hblk; const int tok = static_cast(p / n_hblk); - const int Ih = I / 2; // packed-K rows - const int h0 = static_cast(hb) * HBLK; - const int h1 = std::min(H, h0 + HBLK); - const int nh = h1 - h0; - // Resolve this task's layer base once; row indexing below is layer-local (e). - const uint8_t* dn_packed_l = reinterpret_cast(tbl_at(down_tbl, t->layer_id)); - const uint8_t* dn_scale_l = reinterpret_cast(tbl_at(dn_scale_tbl, t->layer_id)); - const bf16_t* dn_bias_l = reinterpret_cast(tbl_at(dn_bias_tbl, t->layer_id)); - float acc[HBLK]; - for (int c = 0; c < nh; ++c) acc[c] = 0.0f; + const int h0 = static_cast(hb) * kRowBlock; + const int h1 = std::min(H, h0 + kRowBlock); + DownRoute routes[kMaxTopK]; + float w_out[kMaxTopK]; + int n = 0; for (int k = 0; k < top_k; ++k) { - const int e = t->ids[static_cast(tok) * top_k + k]; + const int e = io_ids[static_cast(tok) * top_k + k]; if (e < 0 || e >= num_experts) continue; - const float wt = t->w[static_cast(tok) * top_k + k]; - const bf16_t* g_row = g_scratch.data() + ((size_t)tok * top_k + k) * I; - const uint8_t* blk_e = dn_packed_l + (size_t)e * Ih * H; - const uint8_t* scl_e = dn_scale_l + (size_t)e * (size_t)(I / 32) * H; - float part[HBLK]; - mxgemv(part, blk_e + h0, scl_e + h0, g_row, Ih, H, nh, e2m1_lut, e8m0_lut); - const bf16_t* bias_e = dn_bias_l + (size_t)e * H + h0; - for (int c = 0; c < nh; ++c) acc[c] += (part[c] + bf16_to_f32(bias_e[c])) * wt; - } - bf16_t* y_row = t->y + (size_t)tok * H; - for (int c = 0; c < nh; ++c) y_row[h0 + c] = f32_to_bf16(acc[c]); - } - - // ----------------------------- ds_fp4 (DSV4) ------------------------------- - // Row-major e2m1 + e8m0/32 (no global); silu-swiglu with clamp; FP8-roundtripped - // activations (x once in submit -> xq_scratch; the intermediate g in a dedicated - // round-trip phase between the two passes). Router weight applies on the down output. - - void do_pass1_dsfp4(const MoeTask* t, int64_t p) { - const int64_t ib = p % n_iblk; - const int64_t tk = p / n_iblk; - const int k = static_cast(tk % top_k); - const int tok = static_cast(tk / top_k); - const int e = t->ids[static_cast(tok) * top_k + k]; - if (e < 0 || e >= num_experts) return; - // Resolve this task's layer base once; row indexing below is layer-local (e). - const uint8_t* gu_packed_l = reinterpret_cast(tbl_at(gate_up_tbl, t->layer_id)); - const uint8_t* gu_scale_l = reinterpret_cast(tbl_at(gu_scale_tbl, t->layer_id)); - const int N2 = 2 * I, Hh = H / 2, Hs = H / 32; - // fp8-roundtripped input, pre-deinterleaved to fp32 even/odd halves. - const float* xe = xe_scratch.data() + (size_t)tok * (H / 2); - const float* xo = xo_scratch.data() + (size_t)tok * (H / 2); - const uint8_t* gp = gu_packed_l + (size_t)e * N2 * Hh; - const uint8_t* gs = gu_scale_l + (size_t)e * N2 * Hs; - bf16_t* g_row = g_scratch.data() + ((size_t)tok * top_k + k) * I; - const int i0 = static_cast(ib) * IBLK; - const int i1 = std::min(I, i0 + IBLK); - const float lim = swiglu_limit; - for (int i = i0; i < i1; ++i) { - // gate_up is stored bf16 by the reference GEMV before swiglu; round to match. - float gate = bf16_to_f32(f32_to_bf16( - dsdot(gp + (size_t)i * Hh, gs + (size_t)i * Hs, xe, xo, H, e2m1_lut, e8m0_lut))); - float up = bf16_to_f32(f32_to_bf16(dsdot( - gp + (size_t)(I + i) * Hh, gs + (size_t)(I + i) * Hs, xe, xo, H, e2m1_lut, e8m0_lut))); - if (lim > 0.0f) { - if (gate > lim) gate = lim; - if (up > lim) up = lim; - else if (up < -lim) up = -lim; - } - const float glu = gate / (1.0f + std::exp(-gate)); // silu(gate) - g_row[i] = f32_to_bf16(glu * up); - } - } - - // Prepare one intermediate row (token,route) for the down GEMV, reused across every down output row. - // ds_fp4 first FP8 round-trips it (DSV4 act_quant); nvfp4/ds_fp4/fp8_block deinterleave it to fp32 even/odd; q4_0 Q8_0-quantizes it. - void prep_g_row(int64_t r) { - bf16_t* g = g_scratch.data() + (size_t)r * I; - if (use_q4a8) { // q4_0 W4A8: Q8_0-quantize the intermediate row for the down GEMV. - quant_q8_0(g, I, gi8_scratch.data() + (size_t)r * I, - gas_scratch.data() + (size_t)r * (I / 32)); - return; - } - if (fmt == WF_DSFP4) fp8_roundtrip_bf16(g, g, I, act_block); - float* ge = ge_scratch.data() + (size_t)r * (I / 2); - float* go = go_scratch.data() + (size_t)r * (I / 2); - deinterleave_bf16_f32(g, ge, go, I); - if (use_vnni) - quant_i8_pg16(ge, go, I, gi8_scratch.data() + (size_t)r * I, - gas_scratch.data() + (size_t)r * (I / 16)); - } - - void do_pass2_dsfp4(const MoeTask* t, int64_t p) { - const int64_t hb = p % n_hblk; - const int tok = static_cast(p / n_hblk); - const int Ih = I / 2, Is = I / 32; - const int h0 = static_cast(hb) * HBLK; - const int h1 = std::min(H, h0 + HBLK); - // Resolve this task's layer base once; row indexing below is layer-local (e). - const uint8_t* dn_packed_l = reinterpret_cast(tbl_at(down_tbl, t->layer_id)); - const uint8_t* dn_scale_l = reinterpret_cast(tbl_at(dn_scale_tbl, t->layer_id)); - bf16_t* y_row = t->y + (size_t)tok * H; + const size_t r = (size_t)tok * top_k + k; + routes[n] = {e, static_cast(r), g_scratch.data() + r * I}; + w_out[n] = apply_on_input ? 1.0f : io_w[r]; + ++n; + } + float part[kMaxTopK * kRowBlock]; + if (n > 0) fmt->down(static_cast(t.layer), h0, h1, routes, n, part); + bf16_t* y_row = io_y + (size_t)tok * H; for (int h = h0; h < h1; ++h) { float acc = 0.0f; - for (int k = 0; k < top_k; ++k) { - const int e = t->ids[static_cast(tok) * top_k + k]; - if (e < 0 || e >= num_experts) continue; - const float wt = t->w[static_cast(tok) * top_k + k]; - const float* ge = ge_scratch.data() + ((size_t)tok * top_k + k) * (I / 2); - const float* go = go_scratch.data() + ((size_t)tok * top_k + k) * (I / 2); - const uint8_t* dp = dn_packed_l + (size_t)e * (size_t)H * Ih + (size_t)h * Ih; - const uint8_t* ds = dn_scale_l + (size_t)e * (size_t)H * Is + (size_t)h * Is; - // The reference rounds each route's weighted down output to bf16 before the - // fp32 sum over routes (down [T, top_k, H] bf16 -> .sum(dim=1)). - acc += bf16_to_f32(f32_to_bf16(dsdot(dp, ds, ge, go, I, e2m1_lut, e8m0_lut) * wt)); + for (int j = 0; j < n; ++j) { + const float v = part[j * kRowBlock + h - h0] * w_out[j]; + acc += rounds_routes ? bf16_to_f32(f32_to_bf16(v)) : v; } y_row[h] = f32_to_bf16(acc); } } - void run_task_body(const MoeTask* t) { + void run_task_body(const MoeTask& t) { int local_sense = 0; for (;;) { int64_t p = p1_next.fetch_add(1, std::memory_order_relaxed); @@ -2009,13 +292,13 @@ struct CpuMoeExecutor { do_pass1(t, p); } barrier(local_sense); - // Prepare the intermediate rows (per token,route) for the down GEMV, see prep_g_row. + // Prepare the intermediate rows (per token,route) for the down GEMV, see prep_inter. // Needs all of pass1 done (a full row spans every iblk). - if (needs_di || use_q4a8) { + if (preps_inter) { for (;;) { int64_t r = prt_next.fetch_add(1, std::memory_order_relaxed); if (r >= prt_total) break; - prep_g_row(r); + fmt->prep_inter(r, g_scratch.data() + (size_t)r * params.I); } barrier(local_sense); } @@ -2051,8 +334,7 @@ struct CpuMoeExecutor { } // submit() writes cur_task before the release store to `submitted`. my_gen = gen; - MoeTask* t = cur_task; - run_task_body(t); + run_task_body(cur_task); last_active = wclk::now(); if (done_count.fetch_add(1) + 1 == num_threads) { completed.store(my_gen, std::memory_order_release); @@ -2064,69 +346,20 @@ struct CpuMoeExecutor { } } - void submit(MoeTask* t) { - n_iblk = (I + IBLK - 1) / IBLK; - n_hblk = (H + HBLK - 1) / HBLK; - // Grow the per-token intermediate scratch if a larger batch shows up than the - // construction-time hint (CUDA-graph capture warms the largest bs first, so - // this happens at most once, before any capture, while the pool is idle). - const size_t need = static_cast(t->num_tokens) * top_k * I; - if (need > g_scratch.size()) g_scratch.resize(need); - p1_total = static_cast(t->num_tokens) * top_k * n_iblk; - p2_total = static_cast(t->num_tokens) * n_hblk; - prt_total = (needs_di || use_q4a8) ? static_cast(t->num_tokens) * top_k : 0; + void submit(MoeTask t) { + const int num_tokens = static_cast(t.num_tokens); + p1_total = static_cast(num_tokens) * params.top_k * n_iblk; + p2_total = static_cast(num_tokens) * n_hblk; + prt_total = preps_inter ? static_cast(num_tokens) * params.top_k : 0; p1_next.store(0, std::memory_order_relaxed); p2_next.store(0, std::memory_order_relaxed); prt_next.store(0, std::memory_order_relaxed); done_count.store(0, std::memory_order_relaxed); bar_count.store(0, std::memory_order_relaxed); bar_sense.store(0, std::memory_order_relaxed); - // Deinterleave the per-token input once, up front (ds_fp4 FP8-round-trips it first). + // Prepare the per-token input once, up front. // Single-threaded: tiny for decode, and done before the workers are woken below. - if (needs_di) { - const size_t xn = static_cast(t->num_tokens) * H; - if (xn / 2 > xe_scratch.size()) { - xe_scratch.resize(xn / 2); - xo_scratch.resize(xn / 2); - ge_scratch.resize(static_cast(t->num_tokens) * top_k * (I / 2)); - go_scratch.resize(static_cast(t->num_tokens) * top_k * (I / 2)); - if (use_vnni) { - xi8_scratch.resize(xn); - xas_scratch.resize(xn / 16); - gi8_scratch.resize(static_cast(t->num_tokens) * top_k * I); - gas_scratch.resize(static_cast(t->num_tokens) * top_k * (I / 16)); - } - } - const bool ds = (fmt == WF_DSFP4) && !input_prequant; - if (ds && xn > xq_scratch.size()) xq_scratch.resize(xn); - for (int tok = 0; tok < t->num_tokens; ++tok) { - const bf16_t* src = t->x + (size_t)tok * H; - if (ds) { // DSV4 FP8-round-trips the input before the gate_up GEMV - bf16_t* xq = xq_scratch.data() + (size_t)tok * H; - fp8_roundtrip_bf16(src, xq, H, act_block); - src = xq; - } - float* xe = xe_scratch.data() + (size_t)tok * (H / 2); - float* xo = xo_scratch.data() + (size_t)tok * (H / 2); - deinterleave_bf16_f32(src, xe, xo, H); - if (use_vnni) - quant_i8_pg16(xe, xo, H, xi8_scratch.data() + (size_t)tok * H, - xas_scratch.data() + (size_t)tok * (H / 16)); - } - } - // q4_0 W4A8: Q8_0-quantize the per-token input once (single-threaded, tiny for decode). - if (use_q4a8) { - const size_t xn = static_cast(t->num_tokens) * H; - if (xn > xi8_scratch.size()) { - xi8_scratch.resize(xn); - xas_scratch.resize(xn / 32); - gi8_scratch.resize(static_cast(t->num_tokens) * top_k * I); - gas_scratch.resize(static_cast(t->num_tokens) * top_k * (I / 32)); - } - for (int tok = 0; tok < t->num_tokens; ++tok) - quant_q8_0(t->x + (size_t)tok * H, H, xi8_scratch.data() + (size_t)tok * H, - xas_scratch.data() + (size_t)tok * (H / 32)); - } + for (int tok = 0; tok < num_tokens; ++tok) fmt->prep_input(tok, io_x + (size_t)tok * params.H); { std::lock_guard lk(task_mtx); cur_task = t; @@ -2153,62 +386,56 @@ struct CpuMoeExecutor { sync_cv.wait(lk, [&] { return completed.load(std::memory_order_acquire) >= target; }); } - void submit_with_cuda_stream(uintptr_t stream, uintptr_t task) { - cudaLaunchHostFunc(reinterpret_cast(stream), &CpuMoeExecutor::submit_cb, - reinterpret_cast(task)); + void submit_with_cuda_stream(uintptr_t stream, int layer_id, int num_tokens) { + const MoeTask t = task(layer_id, num_tokens); + HostTask& h = host_tasks.try_emplace(to_word(t), HostTask{this, t}).first->second; + cudaLaunchHostFunc(reinterpret_cast(stream), &CpuMoeExecutor::submit_cb, &h); } - void sync_with_cuda_stream(uintptr_t stream, uintptr_t task) { - cudaLaunchHostFunc(reinterpret_cast(stream), &CpuMoeExecutor::sync_cb, - reinterpret_cast(task)); + void sync_with_cuda_stream(uintptr_t stream) { + cudaLaunchHostFunc(reinterpret_cast(stream), &CpuMoeExecutor::sync_cb, this); } - // Register a (layer, batch-size) slot's task so the coordinator can dispatch it on a - // flag bump. - void register_flag_task(int slot, uintptr_t task) { - std::lock_guard lk(flag_task_mtx); - if (static_cast(flag_task.size()) <= slot) flag_task.resize(slot + 1, nullptr); - flag_task[slot] = reinterpret_cast(task); - } - - int64_t flag_served_count(int slot) const { - return (slot >= 0 && slot < static_cast(flag_served.size())) ? flag_served[slot] : 0; - } - - // Start the busy-poll coordinator over the mapped-pinned flag arrays. ``pin_core`` >= 0 + // Start the busy-poll coordinator over the doorbell. ``pin_core`` >= 0 // pins the coordinator to that logical CPU (the worker auto-sizing reserves it), so its // polling never migrates onto / contends with a GEMV worker's core. - void start_flag_coordinator(uintptr_t ready_ptr, uintptr_t done_ptr, int num_slots, - int pin_core) { - ready_flags = reinterpret_cast(ready_ptr); - done_flags = reinterpret_cast(done_ptr); - coord_num_slots = num_slots; - { - std::lock_guard lk(flag_task_mtx); - if (static_cast(flag_task.size()) < num_slots) flag_task.resize(num_slots, nullptr); - } - flag_served.assign(num_slots, 0); + void start_flag_coordinator(int pin_core) { + gpu_sync = std::make_unique(); coord_stop.store(false); coord_pinned = CPU_MOE_HAS_AFFINITY && pin_core >= 0; coord_thread = std::thread([this, pin_core] { -#if CPU_MOE_HAS_AFFINITY - if (pin_core >= 0) { - cpu_set_t set; - CPU_ZERO(&set); - CPU_SET(pin_core, &set); - pthread_setaffinity_np(pthread_self(), sizeof(set), &set); - } -#endif + if (pin_core >= 0) pin_this_thread(pin_core); coordinator_loop(); }); } + void flag_signal(uintptr_t stream, int layer_id, int num_tokens) { + gpu_sync->gpu_submit(stream, to_word(task(layer_id, num_tokens))); + } + void flag_wait(uintptr_t stream) { gpu_sync->gpu_wait(stream); } + void enable_kernel_io() { + kernel_io = std::make_unique(gpu_sync->ready_word(), gpu_sync->done_word(), io_x, io_ids, io_w, io_y, params.H, + params.top_k); + } + void flag_publish(uintptr_t stream, uintptr_t x, uintptr_t ids, uintptr_t w, int layer_id, int num_tokens) { + kernel_io->publish(stream, reinterpret_cast(x), reinterpret_cast(ids), + reinterpret_cast(w), num_tokens, to_word(task(layer_id, num_tokens))); + } + void flag_combine(uintptr_t stream, uintptr_t out, int num_tokens, bool accumulate) { + kernel_io->combine(stream, reinterpret_cast(out), num_tokens, accumulate); + } + bool flag_pending() const { return gpu_sync && gpu_sync->pending(); } + // Complete a doorbell the coordinator never answered so the stream can make progress. + void poison_flag() { + if (gpu_sync) gpu_sync->cpu_done(); + } + void coordinator_loop() { // Idle backoff (edge-device friendliness): spin hot only while decode traffic is - // flowing, with a TIME-based hot window (an iteration count varies 2-6x with slot - // count and pause cost, so some configs dozed every token). 50 ms since the last - // flag comfortably covers intra- and inter-token gaps -- including host-func-only - // stretches (prefill bursts, --moe-cpu-layers subsets) -- so steady decode never + // flowing, with a TIME-based hot window (an iteration count varies 2-6x with pause + // cost, so some configs dozed every token). 50 ms since the last + // flag comfortably covers intra- and inter-token gaps -- including prefill bursts and + // --moe-cpu-layers gaps -- so steady decode never // sleeps and keeps the ~us-level wakeup. Past it, i.e. the engine is actually idle, // escalate to timed sleeps capped at 2 ms: a dozing coordinator costs <0.1% of a // core instead of 100%, and the only price is a <=2 ms discovery delay on the FIRST @@ -2222,30 +449,19 @@ struct CpuMoeExecutor { int64_t sleep_us = 100; bool dozing = false; while (!coord_stop.load(std::memory_order_relaxed)) { - bool any = false; - for (int L = 0; L < coord_num_slots; ++L) { - // Binary handshake (memop-compatible: the GPU-side WAIT compares against an - // immediate baked at graph capture, so the protocol resets per step instead of - // counting). Acquire: everything the GPU made visible before setting ready -- - // the D2H input copies -- is visible to the worker pool after this read. - if (flag_load_acquire(&ready_flags[L]) != 0) { - flag_store_release(&ready_flags[L], 0); // consume this step's doorbell - MoeTask* t; - { - std::lock_guard lk(flag_task_mtx); - t = (L < static_cast(flag_task.size())) ? flag_task[L] : nullptr; - } - if (t != nullptr) { - submit(t); - sync(coord_pinned && spin_wait.load(std::memory_order_relaxed)); - } - // Release: the workers' y stores are visible before the GPU sees done. - flag_store_release(&done_flags[L], 1); - if (L < static_cast(flag_served.size())) ++flag_served[L]; - any = true; + const uint64_t word = gpu_sync->cpu_poll(); + if (word != 0) { + const MoeTask t = from_word(word); + if (has_routes(t)) { + submit(t); + sync(coord_pinned && spin_wait.load(std::memory_order_relaxed)); + } else { + // What the pool would write for a task without routes, minus its fork and join. + std::memset(io_y, 0, sizeof(bf16_t) * t.num_tokens * params.H); } - } - if (any) { + // Release: the workers' y stores are visible before the GPU sees done. + gpu_sync->cpu_done(); + flag_served.fetch_add(1, std::memory_order_relaxed); last_active = coord_clock::now(); empty_polls = 0; sleep_us = 100; @@ -2254,15 +470,13 @@ struct CpuMoeExecutor { } // Hot phase: pause-spin, consulting the clock only every 1024 empty polls (the // amortization must gate the CLOCK, not the sleep -- gating the sleep decision - // let 1023/1024 idle iterations hot-scan all slots, ~20% of a core at 992 slots). + // let 1023/1024 idle iterations hot-poll). // Doze phase: sleep on EVERY iteration until traffic returns; one wake-up scan // (~0.5 us) per 2 ms sleep is ~0.03% duty. if (!dozing) { if ((++empty_polls & 1023u) != 0 || coord_clock::now() - last_active < kHotWindow) { -#if CPU_MOE_X86 - _mm_pause(); -#endif + cpu_relax(); continue; } dozing = true; @@ -2270,92 +484,79 @@ struct CpuMoeExecutor { std::this_thread::sleep_for(std::chrono::microseconds(sleep_us)); sleep_us = std::min(sleep_us * 2, kSleepCapUs); } - // Teardown: release any in-flight (or future) spin-wait immediately so a replay - // caught mid-shutdown exits its sync kernel now instead of owning the watchdog - // stall. Runs before the destructor's join() returns, while the flag arrays are - // still alive on the Python side. - for (int L = 0; L < coord_num_slots; ++L) { - flag_store_release(&done_flags[L], INT64_MAX); - } + // Teardown: complete any in-flight wait so a replay caught mid-shutdown does not stall + // until the watchdog fires. Runs before the destructor's join() returns, while the doorbell is still alive. + gpu_sync->cpu_done(); } // Eager (non-graph) path: run one task to completion on the pool. - void run_task(uintptr_t task) { - MoeTask* t = reinterpret_cast(task); - submit(t); + void run(int layer_id, int num_tokens) { + submit(task(layer_id, num_tokens)); sync(); } static void CUDART_CB submit_cb(void* ud) { - MoeTask* t = reinterpret_cast(ud); - t->exec->submit(t); - } - static void CUDART_CB sync_cb(void* ud) { - MoeTask* t = reinterpret_cast(ud); - t->exec->sync(); + HostTask* h = reinterpret_cast(ud); + h->exec->submit(h->task); } + static void CUDART_CB sync_cb(void* ud) { reinterpret_cast(ud)->sync(); } }; } // namespace +} // namespace freetoken::cpu_moe PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { namespace py = pybind11; + using freetoken::cpu_moe::CpuMoeExecutor; + using freetoken::cpu_moe::GpuSync; py::class_(m, "CpuMoeExecutor") - .def(py::init, int, int>(), + .def(py::init, std::map, + double, double, std::vector, uintptr_t, uintptr_t, uintptr_t, uintptr_t>(), py::arg("num_threads"), py::arg("num_layers"), py::arg("num_experts"), py::arg("top_k"), py::arg("hidden_size"), py::arg("inter_size"), py::arg("max_tokens"), py::arg("activation_id"), py::arg("apply_router_weight_on_input"), py::arg("weight_format"), - py::arg("gate_up_ptr"), py::arg("down_ptr"), py::arg("gate_up_scale_ptr"), - py::arg("gate_up_global_ptr"), py::arg("down_scale_ptr"), - py::arg("down_global_ptr"), py::arg("gate_up_bias_ptr"), - py::arg("down_bias_ptr"), py::arg("swiglu_alpha"), py::arg("swiglu_limit"), - py::arg("core_ids"), py::arg("fp8_block_gu_scale_stride") = 0, - py::arg("fp8_block_dn_scale_stride") = 0) - .def("create_task", &CpuMoeExecutor::create_task, py::arg("layer_id"), - py::arg("num_tokens"), py::arg("x_ptr"), py::arg("ids_ptr"), py::arg("w_ptr"), + py::arg("banks"), py::arg("options"), py::arg("swiglu_alpha"), py::arg("swiglu_limit"), + py::arg("core_ids"), py::arg("x_ptr"), py::arg("ids_ptr"), py::arg("w_ptr"), py::arg("y_ptr")) .def("submit_with_cuda_stream", &CpuMoeExecutor::submit_with_cuda_stream, - py::arg("stream"), py::arg("task"), py::call_guard()) + py::arg("stream"), py::arg("layer_id"), py::arg("num_tokens"), + py::call_guard()) .def("sync_with_cuda_stream", &CpuMoeExecutor::sync_with_cuda_stream, - py::arg("stream"), py::arg("task"), py::call_guard()) - .def("run_task", &CpuMoeExecutor::run_task, py::arg("task"), + py::arg("stream"), py::call_guard()) + .def("run", &CpuMoeExecutor::run, py::arg("layer_id"), py::arg("num_tokens"), py::call_guard()) - .def("register_flag_task", &CpuMoeExecutor::register_flag_task, - py::arg("slot"), py::arg("task")) - .def("flag_served_count", &CpuMoeExecutor::flag_served_count, py::arg("slot")) - .def("start_flag_coordinator", &CpuMoeExecutor::start_flag_coordinator, - py::arg("ready_ptr"), py::arg("done_ptr"), py::arg("num_slots"), - py::arg("pin_core")) - .def("set_input_prequant", - [](CpuMoeExecutor& e, bool v) { e.input_prequant = v; }, - py::arg("value")) - .def("set_act_block", - [](CpuMoeExecutor& e, int v) { e.act_block = v; }, - py::arg("value")) + .def("flag_served", [](const CpuMoeExecutor& e) { return e.flag_served.load(); }) + .def("start_flag_coordinator", &CpuMoeExecutor::start_flag_coordinator, py::arg("pin_core")) + .def("flag_signal", &CpuMoeExecutor::flag_signal, py::arg("stream"), py::arg("layer_id"), + py::arg("num_tokens")) + .def("flag_wait", &CpuMoeExecutor::flag_wait, py::arg("stream")) + .def("flag_pending", &CpuMoeExecutor::flag_pending) + .def("enable_kernel_io", &CpuMoeExecutor::enable_kernel_io) + .def("flag_publish", &CpuMoeExecutor::flag_publish, py::arg("stream"), py::arg("x"), py::arg("ids"), py::arg("w"), + py::arg("layer_id"), py::arg("num_tokens")) + .def("flag_combine", &CpuMoeExecutor::flag_combine, py::arg("stream"), py::arg("out"), py::arg("num_tokens"), + py::arg("accumulate")) + .def("poison_flag", &CpuMoeExecutor::poison_flag) .def("isa_name", &CpuMoeExecutor::isa_name) - .def("fp8_block_isa_name", &CpuMoeExecutor::fp8_block_isa_name) .def("set_spin_wait", [](CpuMoeExecutor& e, bool v) { e.spin_wait.store(v, std::memory_order_relaxed); }, py::arg("enabled")) .def("get_spin_wait", [](const CpuMoeExecutor& e) { return e.spin_wait.load(); }); - m.def("memops_probe", &cumemops_probe, py::arg("stream"), py::arg("scratch_addr")); - m.def("memop_submit", &cumemop_submit, py::arg("stream"), py::arg("done_addr"), - py::arg("ready_addr"), py::arg("slot")); - m.def("memop_sync", &cumemop_sync, py::arg("stream"), py::arg("done_addr"), - py::arg("slot")); + m.def("memops_probe", &GpuSync::probe, py::arg("stream")); + m.def("kernel_io_probe", &freetoken::cpu_moe::KernelIo::probe, py::arg("stream")); + m.def("host_func_capture_safe", &GpuSync::host_func_capture_safe); // ABI capability marker: the highest ActKind this build implements in the // GENERIC epilogue. CpuMoeExecutor.__init__ probes it before requesting an act // id the epilogue must handle -- a prebuilt .so from before ACT_SWIGLUOAI // accepts id 3 without error and silently computes the wrong activation // (act_apply falls through to gelu_tanh); the probe turns a stale extension // into a loud rebuild instruction instead of wrong model outputs. - m.def("max_generic_act_id", []() { return static_cast(ACT_SWIGLU_CLAMP); }); - // Weight-format ABI probe. A stale prebuilt .so accepts a newer WFmt id without - // complaint and then falls through to the wrong dequant branch -- silently wrong - // numbers, not a crash. The Python executor refuses to build against an extension - // that predates the format it was asked for. - m.def("max_weight_format_id", []() { return static_cast(WF_FP8_BLOCK); }); + m.def("max_generic_act_id", []() { return static_cast(freetoken::cpu_moe::ACT_SWIGLU_CLAMP); }); + m.def("weight_formats", []() { + py::list names; + for (const auto& f : freetoken::cpu_moe::kFormats) names.append(f.first); + return names; + }); } diff --git a/python/freetoken/kernel/csrc/cpu_moe/format.h b/python/freetoken/kernel/csrc/cpu_moe/format.h new file mode 100644 index 0000000000..286ae110f7 --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/format.h @@ -0,0 +1,131 @@ +// The interface between the CPU MoE executor and one expert weight format. +#pragma once + +#include +#include +#include +#include + +#include "common.h" + +namespace freetoken::cpu_moe { + +// Output-row tiling. Small enough to give every worker independent work even at +// batch size 1; large enough to amortize the atomic work-grab. +// +// Bandwidth notes (Sapphire Rapids 8480+, 13 cores): the two passes already read +// every expert weight byte exactly once per token (each output row block is owned +// by one worker), and x stays hot in L1 across a (token,expert)'s rows -- so the +// kernel is single-read bandwidth-optimal at bs=1 (~205 GB/s vs ~55 GB/s PCIe). +// One worker per *physical* core, pinned, is the sweet spot; SMT oversubscription +// thrashes the spin-barrier. Deferred (not worth it here / for this workload): +// - AMX-bf16: a GEMM tile engine; decode is M=1 GEMV so tiles sit idle. It would +// only pay off in a grouped/batched (dedup) path. +// - expert dedup for bs>1: read each distinct expert once and GEMM its tokens. +// Helps locality+bytes when bs is large; decode batches here are tiny (<=4). +// - NUMA: a single node is assumed. Multi-socket machines would split each +// expert's K dimension per node (banks are already per-row contiguous). +constexpr int kRowBlock = 32; + +struct MoeParams { + int top_k = 0, max_tokens = 0; + int H = 0, I = 0; + int act = ACT_SILU; + float swiglu_alpha = 1.0f; + float swiglu_limit = INFINITY; // +inf == no clamp +}; + +// What cpu_executor.py passes for one format, by name: banks map canonical roles (gate_up, gate_up_scale, ...) to per-layer tables, options hold small integers. +// A format reads what it needs once, in its constructor. +struct FormatArgs { + std::string format; + std::map banks; + std::map options; + + const uint64_t* bank(const std::string& role) const { + const auto it = banks.find(role); + if (it == banks.end() || it->second == 0) + throw std::invalid_argument("CPU MoE " + format + " experts need the " + role + " bank"); + return reinterpret_cast(it->second); + } + int64_t option(const std::string& name) const { + const auto it = options.find(name); + if (it == options.end()) throw std::invalid_argument("CPU MoE " + format + " experts need the " + name + " option"); + return it->second; + } + int64_t option(const std::string& name, int64_t fallback) const { + const auto it = options.find(name); + return it == options.end() ? fallback : it->second; + } +}; + +// A bank table is the address of a CPU int64 array of length +// num_layers (one base address per layer, built by cpu_executor.py's +// _make_table), not a single flat bank. tbl_at resolves +// tbl[layer_id] once per task/pass. +inline const void* tbl_at(const uint64_t* tbl, int layer_id) { + return reinterpret_cast(tbl[layer_id]); +} + +// One route of a pass 2 item: expert e on intermediate row r = tok * top_k + k, whose bf16 values are g. +struct DownRoute { + int e; + int64_t r; + const bf16_t* g; +}; + +// g[i] for i in [i0, i1) from the gate row i and the up row I + i; dot(row) is the row's GEMV. +// ACT_SWIGLUOAI and ACT_SWIGLU_CLAMP clamp both inputs; the other kinds go through act_apply. +template +inline void gated_rows(const MoeParams& p, int i0, int i1, float w_in, bf16_t* g, RowDot dot) { + const bool clamped = p.act == ACT_SWIGLUOAI || p.act == ACT_SWIGLU_CLAMP; + const float up_bias = p.act == ACT_SWIGLUOAI ? 1.0f : 0.0f; + const float lim = p.swiglu_limit, alpha = p.swiglu_alpha; + for (int i = i0; i < i1; ++i) { + // gate = row i, up = row I+i + float gate = dot(i) * w_in; + float up = dot(p.I + i) * w_in; + if (clamped) { + // clamp(gate, max=lim) * sigmoid(alpha * gate) * (clamp(up, +-lim) + up_bias) + // -- swigluoai carries the +1 up bias (gpt-oss/MiniMax); swiglu_clamp + // (GLM-5.3) does not. lim == +inf: no clamp. + if (gate > lim) gate = lim; + if (up > lim) up = lim; + else if (up < -lim) up = -lim; + const float glu = gate / (1.0f + std::exp(-gate * alpha)); + g[i] = f32_to_bf16(glu * (up + up_bias)); + } else { + g[i] = f32_to_bf16(act_apply(p.act, gate) * up); + } + } +} + +// One expert weight format: its kernels, its activation layout and the scratch for it. +// The executor owns the threads and the work split. +// Scratch is sized once from max_tokens; gate_up, down and prep_inter run on the worker threads, prep_input before they wake. +class MoeFormat { + public: + explicit MoeFormat(const MoeParams& p) : p(p) {} + virtual ~MoeFormat() = default; + + virtual std::string isa_name() const = 0; + // Once per input row x[tok] before pass 1. + virtual void prep_input(int /*tok*/, const bf16_t* /*x*/) {} + // Whether pass 2 needs prep_inter on every intermediate row. + virtual bool preps_inter() const { return false; } + virtual void prep_inter(int64_t /*r*/, bf16_t* /*g*/) {} + // Pass 1: g[i] for i in [i0, i1) (at most kRowBlock rows), expert e of `layer` on token tok. + // w_in scales the gate and up dots: the router weight when it applies on the input, else 1. + virtual void gate_up(int layer, int e, int tok, int i0, int i1, const bf16_t* x, float w_in, + bf16_t* g) const = 0; + // Pass 2: out[j * kRowBlock + h - h0] for h in [h0, h1) (at most kRowBlock rows) and each of the n routes, before the router weight. + // Visiting the routes' rows h together keeps n weight streams in flight, which the down GEMV needs to reach DRAM bandwidth. + virtual void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const = 0; + // Round each router-weighted down output to bf16 before the sum over routes. + virtual bool rounds_routes() const { return false; } + + protected: + const MoeParams& p; +}; + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync.h b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync.h new file mode 100644 index 0000000000..829a8f34c3 --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync.h @@ -0,0 +1,78 @@ +// GPU <-> CPU handshake for the CPU MoE coordinator, one implementation per GPU runtime (gpu_sync_cuda.cpp, gpu_sync_hip.cpp). +// +// A backend keeps two contracts, in eager mode and in graph replay: +// 1. submit, GPU -> CPU: after gpu_submit(stream, task), cpu_poll() returns task once, and every earlier write of the stream is visible to the CPU. +// 2. wait, CPU -> GPU: gpu_wait(stream) holds later work of the stream until cpu_done(), and that work sees every CPU write made before cpu_done(). +// task is a nonzero word, opaque here; the executor encodes the layer and token count in it, as a constant baked into each captured graph. +// One task is in flight at a time: decode runs on one stream and every layer waits before the next one signals. +// KernelIo (gpu_sync_io.h) may replace the copies around submit and wait: its kernel copies the inputs into the IO rows and does gpu_submit's writes itself. +// It skips a step without a CPU route: done = 1 and no ready, so cpu_poll() never sees the step and gpu_wait() passes at once. +#pragma once + +#include +#include + +#if defined(_MSC_VER) +#include +#endif + +namespace freetoken::cpu_moe { + +// Portable ordering for the flag handshake: "ready observed => the DMA'd inputs that +// preceded the bump are visible" and "y stores are visible before done". Plain +// volatile loads lean on x86 TSO; acquire/release makes it hold on aarch64 too +// (GH200/Jetson) at zero x86 cost. (MSVC branch is x86-only today: compiler barrier +// + TSO.) +inline int64_t flag_load_acquire(const volatile int64_t* p) { +#if defined(_MSC_VER) + const int64_t v = *p; + _ReadWriteBarrier(); + return v; +#else + return __atomic_load_n(const_cast(p), __ATOMIC_ACQUIRE); +#endif +} + +inline void flag_store_release(volatile int64_t* p, int64_t v) { +#if defined(_MSC_VER) + _ReadWriteBarrier(); + *p = v; +#else + __atomic_store_n(const_cast(p), v, __ATOMIC_RELEASE); +#endif +} + +class GpuSync { + public: + // Allocates the handshake memory; throws if this runtime cannot. + GpuSync(); + ~GpuSync(); + GpuSync(const GpuSync&) = delete; + GpuSync& operator=(const GpuSync&) = delete; + + // Whether both contracts hold on this device and stream, including graph replay where the runtime needs a real check. + static bool probe(uintptr_t stream); + // Whether the cudaLaunchHostFunc fallback is replayed correctly from a captured graph. + static bool host_func_capture_safe(); + + // 1. submit. gpu_submit() is enqueued on `stream`, which may be capturing; cpu_poll() returns 0 when nothing is submitted. + void gpu_submit(uintptr_t stream, uint64_t task); + uint64_t cpu_poll(); + + // 2. wait. cpu_done() also unblocks the stream on teardown and from the watchdog. + void gpu_wait(uintptr_t stream); + void cpu_done(); + + // Submitted but not polled yet (watchdog only). + bool pending() const; + + // Device addresses of the two words, for KernelIo. + uintptr_t ready_word() const; + uintptr_t done_word() const; + + private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_cuda.cpp b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_cuda.cpp new file mode 100644 index 0000000000..a2a089bf67 --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_cuda.cpp @@ -0,0 +1,166 @@ +#include +#include +#include +#include +#include +#include + +#if defined(_WIN32) +#include +#else +#include +#endif + +#include + +#include "gpu_sync.h" + +namespace freetoken::cpu_moe { +namespace { + +// ===================================================================================== +// CUDA stream memory operations (driver API, resolved via dlopen -- no link-time or +// toolchain dependence). The GPU side of the flag handshake: submit = WRITE_VALUE +// (done=0 then ready=task word), sync = WAIT_VALUE(done >= 1). The wait is +// executed by the GPU front-end (no SM-resident kernel), so GPU "utilization" stays +// truthful during CPU compute windows -- a resident spin kernel pinned it at 99%, +// which laptop CPU/GPU dynamic power schedulers answered by clamping the CPU's max +// frequency (GEMV workers -1.5x: the reported edge regression). Availability is +// probed functionally at startup (memops_probe); anything unsupported (some driver models, +// vGPU, old drivers) falls back to the cudaLaunchHostFunc path. +#if defined(_WIN32) +static void* cumemop_dlopen() { return (void*)::LoadLibraryA("nvcuda.dll"); } +static void* cumemop_dlsym(void* h, const char* n) { + return (void*)::GetProcAddress((HMODULE)h, n); +} +#else +static void* cumemop_dlopen() { + void* h = dlopen("libcuda.so.1", RTLD_LAZY | RTLD_LOCAL); + if (h == nullptr) h = dlopen("libcuda.so", RTLD_LAZY | RTLD_LOCAL); + return h; +} +static void* cumemop_dlsym(void* h, const char* n) { return dlsym(h, n); } +#endif + +using cuMemOp64_fn = int (*)(void* stream, unsigned long long addr, unsigned long long value, + unsigned int flags); +static cuMemOp64_fn g_cu_write64 = nullptr; +static cuMemOp64_fn g_cu_wait64 = nullptr; +static constexpr unsigned int kCuWaitValueGeq = 0x0; // CU_STREAM_WAIT_VALUE_GEQ +static constexpr unsigned int kCuWriteDefault = 0x0; // CU_STREAM_WRITE_VALUE_DEFAULT + +static bool cumemop_resolve() { + static bool resolved = [] { + void* h = cumemop_dlopen(); + if (h == nullptr) return false; + // 11.7+ made the v2 entry points the default; older drivers export only the v1 + // names with the same signature. + g_cu_write64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWriteValue64_v2")); + if (g_cu_write64 == nullptr) + g_cu_write64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWriteValue64")); + g_cu_wait64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWaitValue64_v2")); + if (g_cu_wait64 == nullptr) + g_cu_wait64 = reinterpret_cast(cumemop_dlsym(h, "cuStreamWaitValue64")); + return g_cu_write64 != nullptr && g_cu_wait64 != nullptr; + }(); + return resolved; +} + +// GPU side of the flag handshake (see the block comment above): enqueued on the +// caller's (possibly capturing) stream; the WAIT immediate is the constant 1, +// replay-safe under CUDA graphs. +// The startup probe validates EAGER memops; a driver could still reject them at graph +// capture time. Those enqueue errors would otherwise be swallowed here and surface only +// as a later EndCapture failure -- log the first CUresult so triage is one step. +static void cumemop_check(int rc, const char* what) { + static std::atomic warned{false}; + if (rc != 0 && !warned.exchange(true)) { + std::fprintf(stderr, + "[freetoken/cpu_moe] %s failed with CUresult=%d (first occurrence; " + "subsequent errors are not repeated). If this happened during CUDA " + "graph capture, the driver lacks capture support for stream memops -- " + "set FREETOKEN_CPU_MOE_FLAG_SYNC=0.\n", + what, rc); + } +} + +} // namespace + +// ready and done live on separate cache lines of one mapped-pinned block; with UVA the GPU memops and the coordinator use the same addresses. +struct GpuSync::Impl { + void* block = nullptr; + void* device_block = nullptr; + volatile int64_t* ready = nullptr; + volatile int64_t* done = nullptr; +}; + +GpuSync::GpuSync() : impl_(std::make_unique()) { + if (!cumemop_resolve()) throw std::runtime_error("CPU MoE flag sync: CUDA stream memops are unavailable"); + const cudaError_t rc = cudaHostAlloc(&impl_->block, 128, cudaHostAllocPortable | cudaHostAllocMapped); + if (rc != cudaSuccess) + throw std::runtime_error(std::string("CPU MoE flag allocation failed: ") + cudaGetErrorString(rc)); + std::memset(impl_->block, 0, 128); + impl_->ready = static_cast(impl_->block); + impl_->done = impl_->ready + 8; + const cudaError_t dev_rc = cudaHostGetDevicePointer(&impl_->device_block, impl_->block, 0); + if (dev_rc != cudaSuccess) { + cudaFreeHost(impl_->block); + throw std::runtime_error(std::string("CPU MoE flag device address: ") + cudaGetErrorString(dev_rc)); + } +} + +GpuSync::~GpuSync() { cudaFreeHost(impl_->block); } + +// Functional probe on a scratch pinned int64: enqueue WRITE(7) + WAIT(>=7) + sync. +// Returns true only if the whole memop path works on THIS stream/device/driver. +bool GpuSync::probe(uintptr_t stream) { + if (!cumemop_resolve()) return false; + void* scratch = nullptr; + if (cudaHostAlloc(&scratch, sizeof(int64_t), cudaHostAllocPortable | cudaHostAllocMapped) != cudaSuccess) + return false; + *static_cast(scratch) = 0; + auto* s = reinterpret_cast(stream); + const auto addr = reinterpret_cast(scratch); + const bool ok = g_cu_write64(s, addr, 7ULL, kCuWriteDefault) == 0 && g_cu_wait64(s, addr, 7ULL, kCuWaitValueGeq) == 0 && + cudaStreamSynchronize(reinterpret_cast(stream)) == cudaSuccess; + // a failed enqueue can leave the WRITE pending on the scratch word + cudaStreamSynchronize(reinterpret_cast(stream)); + cudaFreeHost(scratch); + return ok; +} + +bool GpuSync::host_func_capture_safe() { return true; } + +void GpuSync::gpu_submit(uintptr_t stream, uint64_t task) { + auto* s = reinterpret_cast(stream); + // Order matters and is preserved by the front end: reset done BEFORE raising ready, + // so the coordinator's completion write for THIS step can never be wiped. + cumemop_check(g_cu_write64(s, reinterpret_cast(impl_->done), 0ULL, kCuWriteDefault), + "cuStreamWriteValue64(done)"); + cumemop_check(g_cu_write64(s, reinterpret_cast(impl_->ready), task, kCuWriteDefault), + "cuStreamWriteValue64(ready)"); +} + +void GpuSync::gpu_wait(uintptr_t stream) { + cumemop_check(g_cu_wait64(reinterpret_cast(stream), reinterpret_cast(impl_->done), 1ULL, + kCuWaitValueGeq), + "cuStreamWaitValue64(done)"); +} + +uint64_t GpuSync::cpu_poll() { + // Acquire: everything the GPU made visible before writing ready -- the D2H input copies -- is visible to the worker pool after this read. + const int64_t task = flag_load_acquire(impl_->ready); + if (task == 0) return 0; + flag_store_release(impl_->ready, 0); // consume this step's doorbell + return static_cast(task); +} + +void GpuSync::cpu_done() { flag_store_release(impl_->done, 1); } + +bool GpuSync::pending() const { return flag_load_acquire(impl_->ready) != 0 && flag_load_acquire(impl_->done) == 0; } + +uintptr_t GpuSync::ready_word() const { return reinterpret_cast(impl_->device_block); } + +uintptr_t GpuSync::done_word() const { return reinterpret_cast(impl_->device_block) + 8 * sizeof(int64_t); } + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp new file mode 100644 index 0000000000..31ff9be539 --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp @@ -0,0 +1,275 @@ +// HIP stream memory operations for the CPU MoE flag handshake. +// +// HIP WAIT addresses must come from hipMallocSignalMemory, one 64-bit word per allocation. +// +// ROCm 10 records hipStreamWrite/WaitValue64 into the graph under stream capture, one node per call. +// ROCm 7.14 runs them while the stream is being captured but records nothing, so there we add explicit batch-memop nodes and splice them into the stream's dependency set. +// Which of the two applies is measured once at runtime (records_native_memops), not read from the version number. +// The probe then does a real capture + instantiate + replay through the chosen path. +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "gpu_sync.h" + +namespace freetoken::cpu_moe { + +namespace { + +void hip_check(hipError_t rc, const char* what) { + if (rc != hipSuccess) + throw std::runtime_error(std::string("CPU MoE HIP flag sync: ") + what + " failed: " + hipGetErrorString(rc)); +} + +// Capability failures must not poison eager fallback or hide a caller's earlier error. +struct ProbeErrorScope { + ProbeErrorScope() { hip_check(hipPeekAtLastError(), "pending error before probe"); } + ~ProbeErrorScope() { (void)hipGetLastError(); } +}; + +hipError_t alloc_signal(uint64_t** out) { + void* storage = nullptr; + const hipError_t rc = hipExtMallocWithFlags(&storage, sizeof(uint64_t), hipMallocSignalMemory); + if (rc == hipSuccess) { + *out = static_cast(storage); + __atomic_store_n(*out, 0ULL, __ATOMIC_RELEASE); + } + return rc; +} + +void free_signal(uint64_t* signal) { + if (signal != nullptr) (void)hipFree(signal); +} + +volatile int64_t* as_flag(uint64_t* signal) { return reinterpret_cast(signal); } + +bool capturing(hipStream_t stream) { + hipStreamCaptureStatus status = hipStreamCaptureStatusNone; + hip_check(hipStreamIsCapturing(stream, &status), "capture status query"); + return status == hipStreamCaptureStatusActive; +} + +// Whether this runtime records a native hipStreamWriteValue64 issued under capture as a graph node. +bool records_native_memops() { + static const bool native = [] { + ProbeErrorScope errors; + uint64_t* word = nullptr; + hipStream_t stream = nullptr; + hipGraph_t graph = nullptr; + hipError_t rc = alloc_signal(&word); + if (rc == hipSuccess) rc = hipStreamCreateWithFlags(&stream, hipStreamNonBlocking); + if (rc == hipSuccess) rc = hipStreamBeginCapture(stream, hipStreamCaptureModeRelaxed); + const bool capture_started = rc == hipSuccess; + if (rc == hipSuccess) rc = hipStreamWriteValue64(stream, word, 1, hipStreamWriteValueDefault); + if (capture_started) { + const hipError_t end_rc = hipStreamEndCapture(stream, &graph); + if (rc == hipSuccess) rc = end_rc; + } + size_t nodes = 0; + if (rc == hipSuccess) rc = hipGraphGetNodes(graph, nullptr, &nodes); + if (graph != nullptr) (void)hipGraphDestroy(graph); + if (stream != nullptr) { + (void)hipStreamSynchronize(stream); + (void)hipStreamDestroy(stream); + } + free_signal(word); + return rc == hipSuccess && nodes == 1; + }(); + return native; +} + +// Appends one batch-memop node to the capture of `stream`; the node keeps a pointer to `ops`. +hipError_t add_capture_memop_node(hipStream_t stream, hipStreamBatchMemOpParams* ops, unsigned int count) { + hipStreamCaptureStatus status = hipStreamCaptureStatusNone; + hipGraph_t graph = nullptr; + const hipGraphNode_t* dependencies = nullptr; + size_t dependency_count = 0; + hipError_t rc = hipStreamGetCaptureInfo_v2(stream, &status, nullptr, &graph, &dependencies, &dependency_count); + if (rc != hipSuccess) return rc; + if (status != hipStreamCaptureStatusActive || graph == nullptr) return hipErrorIllegalState; + + hipCtx_t context = nullptr; +#if defined(__GNUC__) +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wdeprecated-declarations" +#endif + rc = hipCtxGetCurrent(&context); +#if defined(__GNUC__) +#pragma GCC diagnostic pop +#endif + if (rc != hipSuccess) return rc; + + hipBatchMemOpNodeParams params{}; + params.ctx = context; + params.count = count; + params.paramArray = ops; + hipGraphNode_t node = nullptr; + rc = hipGraphAddBatchMemOpNode(&node, graph, dependencies, dependency_count, ¶ms); + if (rc != hipSuccess) return rc; + rc = hipStreamUpdateCaptureDependencies(stream, &node, 1, hipStreamSetCaptureDependencies); + // EndCapture rejects a node that was added but never joined to the stream. + if (rc != hipSuccess && node != nullptr) (void)hipGraphDestroyNode(node); + return rc; +} + +hipStreamBatchMemOpParams write_op(uint64_t* address, uint64_t value) { + hipStreamBatchMemOpParams op{}; + op.writeValue.operation = hipStreamMemOpWriteValue64; + op.writeValue.address = reinterpret_cast(address); + op.writeValue.value64 = value; + return op; +} + +hipStreamBatchMemOpParams wait_op(uint64_t* address) { + hipStreamBatchMemOpParams op{}; + op.waitValue.operation = hipStreamMemOpWaitValue64; + op.waitValue.address = reinterpret_cast(address); + op.waitValue.value64 = 1; + op.waitValue.flags = hipStreamWaitValueGte; + return op; +} + +bool needs_explicit_nodes(hipStream_t stream) { return capturing(stream) && !records_native_memops(); } + +// node_op is null for a native call, else the parameter slot of an explicit node, which the node keeps a pointer to. +// The ops of one batch node run in parallel, so each write gets its own node and consecutive writes stay in order. +hipError_t write_value(hipStream_t stream, uint64_t* address, uint64_t value, hipStreamBatchMemOpParams* node_op) { + if (node_op == nullptr) return hipStreamWriteValue64(stream, address, value, hipStreamWriteValueDefault); + *node_op = write_op(address, value); + return add_capture_memop_node(stream, node_op, 1); +} + +hipError_t wait_done(hipStream_t stream, uint64_t* done, hipStreamBatchMemOpParams* node_op) { + if (node_op == nullptr) return hipStreamWaitValue64(stream, done, 1, hipStreamWaitValueGte, UINT64_MAX); + *node_op = wait_op(done); + return add_capture_memop_node(stream, node_op, 1); +} + +} // namespace + +struct GpuSync::Impl { + uint64_t* ready = nullptr; // signal memory + uint64_t* done = nullptr; + // Explicit nodes keep these addresses, so one parameter pair per distinct task, never moved (std::map nodes are stable). + std::map> submit_ops; // done = 0, ready = task + hipStreamBatchMemOpParams wait_op{}; +}; + +GpuSync::GpuSync() : impl_(std::make_unique()) { + hipError_t rc = alloc_signal(&impl_->ready); + if (rc == hipSuccess) rc = alloc_signal(&impl_->done); + if (rc != hipSuccess) { + free_signal(impl_->ready); + free_signal(impl_->done); + hip_check(rc, "signal allocation"); + } + // Measure now: the measurement captures and allocates, which a graph capture in progress would reject. + (void)records_native_memops(); +} + +GpuSync::~GpuSync() { + free_signal(impl_->ready); + free_signal(impl_->done); +} + +bool GpuSync::probe(uintptr_t /*stream*/) { + ProbeErrorScope errors; + int device = 0; + int can_wait = 0; + if (hipGetDevice(&device) != hipSuccess || + hipDeviceGetAttribute(&can_wait, hipDeviceAttributeCanUseStreamWaitValue, device) != hipSuccess || can_wait == 0) + return false; + const bool explicit_nodes = !records_native_memops(); + + uint64_t* ready = nullptr; + uint64_t* done = nullptr; + hipStream_t stream = nullptr; + hipGraph_t graph = nullptr; + hipGraphExec_t graph_exec = nullptr; + std::array node_ops{}; + hipError_t rc = alloc_signal(&ready); + if (rc == hipSuccess) rc = alloc_signal(&done); + if (rc == hipSuccess) rc = hipStreamCreateWithFlags(&stream, hipStreamNonBlocking); + // A runtime that runs the WAIT during capture instead of recording it must not block here. + if (rc == hipSuccess) __atomic_store_n(done, 1ULL, __ATOMIC_RELEASE); + if (rc == hipSuccess) rc = hipStreamBeginCapture(stream, hipStreamCaptureModeRelaxed); + const bool capture_started = rc == hipSuccess; + if (rc == hipSuccess) rc = write_value(stream, ready, 1, explicit_nodes ? &node_ops[0] : nullptr); + if (rc == hipSuccess) rc = wait_done(stream, done, explicit_nodes ? &node_ops[1] : nullptr); + // End the capture even on failure so its stream and graph are not left active during cleanup. + if (capture_started) { + const hipError_t end_rc = hipStreamEndCapture(stream, &graph); + if (rc == hipSuccess) rc = end_rc; + } + if (rc == hipSuccess) rc = hipStreamSynchronize(stream); + size_t nodes = 0; + if (rc == hipSuccess) rc = hipGraphGetNodes(graph, nullptr, &nodes); + if (rc == hipSuccess && nodes != 2) rc = hipErrorNotSupported; + if (rc == hipSuccess) rc = hipGraphInstantiate(&graph_exec, graph, nullptr, nullptr, 0); + + bool served = false; + if (rc == hipSuccess) { + __atomic_store_n(ready, 0ULL, __ATOMIC_RELEASE); + __atomic_store_n(done, 0ULL, __ATOMIC_RELEASE); + std::thread responder([&] { + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(2); + while (__atomic_load_n(ready, __ATOMIC_ACQUIRE) == 0 && std::chrono::steady_clock::now() < deadline) + std::this_thread::yield(); + served = __atomic_load_n(ready, __ATOMIC_ACQUIRE) == 1; + // The probe graph never resets done, so this also releases a launch that arrives after the deadline. + __atomic_store_n(done, 1ULL, __ATOMIC_RELEASE); + }); + rc = hipGraphLaunch(graph_exec, stream); + if (rc == hipSuccess) rc = hipStreamSynchronize(stream); + responder.join(); + } + const bool ok = rc == hipSuccess && served; + + if (graph_exec != nullptr) (void)hipGraphExecDestroy(graph_exec); + if (graph != nullptr) (void)hipGraphDestroy(graph); + if (stream != nullptr) (void)hipStreamDestroy(stream); + free_signal(ready); + free_signal(done); + return ok; +} + +void GpuSync::gpu_submit(uintptr_t stream, uint64_t task) { + auto s = reinterpret_cast(stream); + hipStreamBatchMemOpParams* ops = needs_explicit_nodes(s) ? impl_->submit_ops[task].data() : nullptr; + // Reset done before raising ready, so the coordinator's completion for this step can never be wiped. + hip_check(write_value(s, impl_->done, 0, ops), "write done"); + hip_check(write_value(s, impl_->ready, task, ops ? ops + 1 : nullptr), "write ready"); +} + +void GpuSync::gpu_wait(uintptr_t stream) { + auto s = reinterpret_cast(stream); + hip_check(wait_done(s, impl_->done, needs_explicit_nodes(s) ? &impl_->wait_op : nullptr), "wait done"); +} + +bool GpuSync::host_func_capture_safe() { return false; } + +uint64_t GpuSync::cpu_poll() { + const int64_t task = flag_load_acquire(as_flag(impl_->ready)); + if (task == 0) return 0; + flag_store_release(as_flag(impl_->ready), 0); + return static_cast(task); +} + +void GpuSync::cpu_done() { flag_store_release(as_flag(impl_->done), 1); } + +bool GpuSync::pending() const { + return flag_load_acquire(as_flag(impl_->ready)) != 0 && flag_load_acquire(as_flag(impl_->done)) == 0; +} + +uintptr_t GpuSync::ready_word() const { return reinterpret_cast(impl_->ready); } + +uintptr_t GpuSync::done_word() const { return reinterpret_cast(impl_->done); } + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.cu b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.cu new file mode 100644 index 0000000000..b41961f0bf --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.cu @@ -0,0 +1,199 @@ +#include +#include +#include + +#include + +#include "gpu_sync_io.h" + +namespace freetoken::cpu_moe { +namespace { + +__device__ inline float bf16_bits_to_float(uint16_t b) { return __uint_as_float(static_cast(b) << 16); } + +// Round to nearest even, the same rounding as a torch bf16 add. +__device__ inline uint16_t float_to_bf16_bits(float f) { + const uint32_t u = __float_as_uint(f); + if ((u & 0x7fffffffu) > 0x7f800000u) return 0x7fc0; + return static_cast((u + 0x7fffu + ((u >> 16) & 1u)) >> 16); +} + +// One block. Copies x/ids/w into the IO rows, then done = 0 and ready = task. +// A step without a CPU route writes done = 1 and rings nothing: the coordinator stays asleep and the GPU wait passes at once. +__global__ void publish_kernel(const uint16_t* __restrict__ x, const int32_t* __restrict__ ids, const float* __restrict__ w, + int64_t n_x, int64_t n_ids, uint16_t* io_x, int32_t* io_ids, float* io_w, + volatile int64_t* ready, volatile int64_t* done, int64_t task, int32_t* has_cpu) { + __shared__ int any; + if (threadIdx.x == 0) any = 0; + __syncthreads(); + for (int64_t i = threadIdx.x; i < n_ids; i += blockDim.x) + if (ids[i] >= 0) any = 1; + __syncthreads(); + if (!any) { + if (threadIdx.x == 0) { + *has_cpu = 0; + *done = 1; + __threadfence_system(); + } + return; + } + const bool vec = (reinterpret_cast(x) | reinterpret_cast(io_x)) % 16 == 0 && n_x % 8 == 0; + if (vec) { + for (int64_t i = threadIdx.x; i < n_x / 8; i += blockDim.x) + reinterpret_cast(io_x)[i] = reinterpret_cast(x)[i]; + } else { + for (int64_t i = threadIdx.x; i < n_x; i += blockDim.x) io_x[i] = x[i]; + } + for (int64_t i = threadIdx.x; i < n_ids; i += blockDim.x) { + io_ids[i] = ids[i]; + io_w[i] = w[i]; + } + __threadfence_system(); + __syncthreads(); + if (threadIdx.x == 0) { + *has_cpu = 1; + *done = 0; + __threadfence_system(); + *ready = task; + __threadfence_system(); + } +} + +// out = CPU rows, or out += CPU rows with accumulate; a step publish skipped contributes zeros. +// io_y is volatile: the same host rows hold a new step on every call. +__global__ void combine_kernel(uint16_t* out, const volatile uint16_t* io_y, const int32_t* has_cpu, int64_t n, bool accumulate) { + const bool on = *has_cpu != 0; + if (accumulate && !on) return; + for (int64_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; i < n; i += static_cast(gridDim.x) * blockDim.x) { + float v = accumulate ? bf16_bits_to_float(out[i]) : 0.0f; + if (on) v += bf16_bits_to_float(io_y[i]); + out[i] = float_to_bf16_bits(v); + } +} + +cudaError_t launch_publish(cudaStream_t stream, const uint16_t* x, const int32_t* ids, const float* w, int64_t n_x, int64_t n_ids, + uint16_t* io_x, int32_t* io_ids, float* io_w, volatile int64_t* ready, volatile int64_t* done, int64_t task, + int32_t* has_cpu) { + publish_kernel<<<1, 256, 0, stream>>>(x, ids, w, n_x, n_ids, io_x, io_ids, io_w, ready, done, task, has_cpu); + return cudaGetLastError(); +} + +cudaError_t launch_combine(cudaStream_t stream, uint16_t* out, const uint16_t* io_y, const int32_t* has_cpu, int64_t n, bool accumulate) { + if (n == 0) return cudaSuccess; + const auto blocks = static_cast(std::min((n + 255) / 256, 1024)); + combine_kernel<<>>(out, io_y, has_cpu, n, accumulate); + return cudaGetLastError(); +} + +void check(cudaError_t rc, const char* what) { + if (rc != cudaSuccess) throw std::runtime_error(std::string("CPU MoE kernel IO: ") + what + ": " + cudaGetErrorString(rc)); +} + +template +T* device_address(const void* host) { + void* dev = nullptr; + check(cudaHostGetDevicePointer(&dev, const_cast(host), 0), "device address of the IO rows"); + return static_cast(dev); +} + +} // namespace + +KernelIo::KernelIo(uintptr_t ready, uintptr_t done, const void* x, const int32_t* ids, const float* w, const void* y, int hidden, + int top_k) + : ready_(reinterpret_cast(ready)), + done_(reinterpret_cast(done)), + x_(device_address(x)), + ids_(device_address(ids)), + w_(device_address(w)), + y_(device_address(y)), + hidden_(hidden), + top_k_(top_k) { + check(cudaMalloc(reinterpret_cast(&has_cpu_), sizeof(int32_t)), "has_cpu allocation"); + const cudaError_t rc = cudaMemset(has_cpu_, 0, sizeof(int32_t)); + if (rc != cudaSuccess) { + (void)cudaFree(has_cpu_); + check(rc, "has_cpu reset"); + } +} + +KernelIo::~KernelIo() { (void)cudaFree(has_cpu_); } + +void KernelIo::publish(uintptr_t stream, const void* x, const int32_t* ids, const float* w, int num_tokens, uint64_t task) { + check(launch_publish(reinterpret_cast(stream), static_cast(x), ids, w, + static_cast(num_tokens) * hidden_, static_cast(num_tokens) * top_k_, x_, ids_, w_, ready_, done_, + static_cast(task), has_cpu_), + "publish launch"); +} + +void KernelIo::combine(uintptr_t stream, void* out, int num_tokens, bool accumulate) { + check(launch_combine(reinterpret_cast(stream), static_cast(out), y_, has_cpu_, + static_cast(num_tokens) * hidden_, accumulate), + "combine launch"); +} + +// A ringing step, its combine and a skipped step on scratch mapped memory, checking every word they write. +std::string KernelIo::probe(uintptr_t stream_handle) { + struct alignas(16) Area { + int64_t ready, done; + int32_t has_cpu, pad[3]; + uint16_t x[8], io_x[8], io_y[8], out[8]; + int32_t ids[2], io_ids[2]; + float w[2], io_w[2]; + }; + const auto stream = reinterpret_cast(stream_handle); + void* host = nullptr; + void* dev = nullptr; + if (cudaHostAlloc(&host, sizeof(Area), cudaHostAllocPortable | cudaHostAllocMapped) != cudaSuccess) + return "pinned scratch allocation failed"; + std::string reason; + auto ok = [&](cudaError_t rc, const char* what) { + if (rc != cudaSuccess && reason.empty()) reason = std::string(what) + ": " + cudaGetErrorString(rc); + return rc == cudaSuccess; + }; + auto fail = [&](const char* why) { + if (reason.empty()) reason = why; + }; + if (ok(cudaHostGetDevicePointer(&dev, host, 0), "device address of the scratch")) { + auto* h = static_cast(host); + auto* d = static_cast(dev); + std::memset(h, 0, sizeof(Area)); + for (int i = 0; i < 8; ++i) { + h->x[i] = static_cast(0x3f80 + i); + h->io_y[i] = 0x3f80; // 1.0 + h->out[i] = 0x4000; // 2.0 + } + h->ids[0] = 3; + h->ids[1] = -1; + h->w[0] = 0.25f; + h->w[1] = 0.75f; + auto publish = [&] { + return launch_publish(stream, d->x, d->ids, d->w, 8, 2, d->io_x, d->io_ids, d->io_w, &d->ready, &d->done, 5, &d->has_cpu); + }; + + if (ok(publish(), "publish launch") && ok(cudaStreamSynchronize(stream), "publish")) { + if (h->ready != 5 || h->done != 0 || h->has_cpu != 1) fail("publish wrote wrong handshake words"); + if (std::memcmp(h->io_x, h->x, sizeof(h->x)) || std::memcmp(h->io_ids, h->ids, sizeof(h->ids)) || + std::memcmp(h->io_w, h->w, sizeof(h->w))) + fail("publish wrote wrong IO rows"); + } + if (reason.empty() && ok(launch_combine(stream, d->out, d->io_y, &d->has_cpu, 8, true), "combine launch") && + ok(cudaStreamSynchronize(stream), "combine")) { + for (int i = 0; i < 8; ++i) + if (h->out[i] != 0x4040) fail("combine returned a wrong sum"); // 3.0 + } + if (reason.empty()) { + h->ready = 0; + h->ids[0] = -1; + if (ok(publish(), "skip launch") && ok(launch_combine(stream, d->out, d->io_y, &d->has_cpu, 8, false), "skip combine launch") && + ok(cudaStreamSynchronize(stream), "skipped step")) { + bool zeros = true; + for (int i = 0; i < 8; ++i) zeros = zeros && h->out[i] == 0; + if (h->ready != 0 || h->done != 1 || h->has_cpu != 0 || !zeros) fail("a skipped step wrote wrong words"); + } + } + } + (void)cudaFreeHost(host); + return reason; +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.h b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.h new file mode 100644 index 0000000000..6eacb3936b --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.h @@ -0,0 +1,39 @@ +// Kernel IO for the GPU <-> CPU handshake (gpu_sync.h): kernels on the executor's mapped pinned IO rows in place of the copies around submit and wait. +// One source for both runtimes: nvcc builds gpu_sync_io.cu on CUDA, hipcc on ROCm through freetoken/hip_compat.h. +#pragma once + +#include +#include + +namespace freetoken::cpu_moe { + +class KernelIo { + public: + // ready, done: device addresses of GpuSync's words. x, ids, w, y: host addresses of the pinned IO rows. Throws if setup fails. + KernelIo(uintptr_t ready, uintptr_t done, const void* x, const int32_t* ids, const float* w, const void* y, int hidden, int top_k); + ~KernelIo(); + KernelIo(const KernelIo&) = delete; + KernelIo& operator=(const KernelIo&) = delete; + + // Copies the step's device inputs into the IO rows and does gpu_submit's writes, or skips the step when no route goes to the CPU. + // x [n, hidden] bf16, ids [n, top_k] int32 (< 0 = not a CPU route), w [n, top_k] fp32. + void publish(uintptr_t stream, const void* x, const int32_t* ids, const float* w, int num_tokens, uint64_t task); + // Follows gpu_wait(): writes the CPU rows to out [n, hidden] bf16, or adds them to it; a skipped step gives zeros. + void combine(uintptr_t stream, void* out, int num_tokens, bool accumulate); + + // Whether both kernels run and write the right words on this device: "" or the reason. + static std::string probe(uintptr_t stream); + + private: + volatile int64_t* ready_; + volatile int64_t* done_; + uint16_t* x_; + int32_t* ids_; + float* w_; + const uint16_t* y_; + int hidden_; + int top_k_; + int32_t* has_cpu_ = nullptr; // device word: whether the last published step has CPU routes +}; + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/bf16.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/bf16.cpp new file mode 100644 index 0000000000..273c394609 --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/bf16.cpp @@ -0,0 +1,159 @@ +#include "../format.h" + +namespace freetoken::cpu_moe { +namespace { + +// ------------------------------- dot products ------------------------------- +// dot(weight[bf16], act[bf16], n) -> fp32. The selected impl is a function +// pointer chosen at runtime; the per-row call overhead is negligible vs n. + +using dot_fn = float (*)(const bf16_t*, const bf16_t*, int); + +float dot_scalar(const bf16_t* w, const bf16_t* x, int n) { + float acc = 0.0f; + for (int i = 0; i < n; ++i) acc += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); + return acc; +} + +#if CPU_MOE_X86 +CPU_MOE_TARGET("avx512f") +float dot_avx512f(const bf16_t* w, const bf16_t* x, int n) { + // 4 independent accumulators -> more in-flight loads (memory-level parallelism), + // which is what lifts a bandwidth-bound GEMV toward peak. + __m512 a0 = _mm512_setzero_ps(), a1 = _mm512_setzero_ps(); + __m512 a2 = _mm512_setzero_ps(), a3 = _mm512_setzero_ps(); + int i = 0; + for (; i + 64 <= n; i += 64) { + _mm_prefetch(reinterpret_cast(w + i) + PF_AHEAD, _MM_HINT_T0); + for (int j = 0; j < 64; j += 16) { + __m256i wi = _mm256_loadu_si256(reinterpret_cast(w + i + j)); + __m256i xi = _mm256_loadu_si256(reinterpret_cast(x + i + j)); + __m512 wf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(wi), 16)); + __m512 xf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(xi), 16)); + __m512& acc = (j == 0) ? a0 : (j == 16) ? a1 : (j == 32) ? a2 : a3; + acc = _mm512_fmadd_ps(wf, xf, acc); + } + } + for (; i + 16 <= n; i += 16) { + __m256i wi = _mm256_loadu_si256(reinterpret_cast(w + i)); + __m256i xi = _mm256_loadu_si256(reinterpret_cast(x + i)); + __m512 wf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(wi), 16)); + __m512 xf = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(xi), 16)); + a0 = _mm512_fmadd_ps(wf, xf, a0); + } + float s = _mm512_reduce_add_ps(_mm512_add_ps(_mm512_add_ps(a0, a1), _mm512_add_ps(a2, a3))); + for (; i < n; ++i) s += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); + return s; +} + +#ifdef CPU_MOE_HAS_AVX512BF16 +CPU_MOE_TARGET("avx512bf16,avx512f") +static inline __m512bh load_bh(const bf16_t* p) { + __m512i raw = _mm512_loadu_si512(reinterpret_cast(p)); + __m512bh out; + std::memcpy(&out, &raw, sizeof(out)); + return out; +} + +CPU_MOE_TARGET("avx512bf16,avx512f") +float dot_avx512bf16(const bf16_t* w, const bf16_t* x, int n) { + // 4 accumulators (128 bf16/iter) for memory-level parallelism + a prefetch nudge. + __m512 a0 = _mm512_setzero_ps(), a1 = _mm512_setzero_ps(); + __m512 a2 = _mm512_setzero_ps(), a3 = _mm512_setzero_ps(); + int i = 0; + for (; i + 128 <= n; i += 128) { + _mm_prefetch(reinterpret_cast(w + i) + PF_AHEAD, _MM_HINT_T0); + a0 = _mm512_dpbf16_ps(a0, load_bh(w + i), load_bh(x + i)); + a1 = _mm512_dpbf16_ps(a1, load_bh(w + i + 32), load_bh(x + i + 32)); + a2 = _mm512_dpbf16_ps(a2, load_bh(w + i + 64), load_bh(x + i + 64)); + a3 = _mm512_dpbf16_ps(a3, load_bh(w + i + 96), load_bh(x + i + 96)); + } + for (; i + 32 <= n; i += 32) { + a0 = _mm512_dpbf16_ps(a0, load_bh(w + i), load_bh(x + i)); + } + float s = _mm512_reduce_add_ps(_mm512_add_ps(_mm512_add_ps(a0, a1), _mm512_add_ps(a2, a3))); + for (; i < n; ++i) s += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); + return s; +} +#endif // avx512bf16 available + +// --------- AVX2 + FMA (256-bit fallback for CPUs without AVX-512) ---------- +// Covers Intel 12-14th gen / Arrow Lake (AVX-512 fused off) and AMD Zen<4. The +// bf16->fp32 widen is a zero-extend + <<16; with 4 independent accumulators the +// GEMV is memory-bandwidth bound, same as the AVX-512 path (just half the width). +CPU_MOE_TARGET("avx2,fma") +float dot_avx2(const bf16_t* w, const bf16_t* x, int n) { + __m256 a0 = _mm256_setzero_ps(), a1 = _mm256_setzero_ps(); + __m256 a2 = _mm256_setzero_ps(), a3 = _mm256_setzero_ps(); + int i = 0; + for (; i + 32 <= n; i += 32) { + _mm_prefetch(reinterpret_cast(w + i) + PF_AHEAD, _MM_HINT_T0); + for (int j = 0; j < 32; j += 8) { + __m128i wi = _mm_loadu_si128(reinterpret_cast(w + i + j)); + __m128i xi = _mm_loadu_si128(reinterpret_cast(x + i + j)); + __m256 wf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(wi), 16)); + __m256 xf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(xi), 16)); + __m256& acc = (j == 0) ? a0 : (j == 8) ? a1 : (j == 16) ? a2 : a3; + acc = _mm256_fmadd_ps(wf, xf, acc); + } + } + for (; i + 8 <= n; i += 8) { + __m128i wi = _mm_loadu_si128(reinterpret_cast(w + i)); + __m128i xi = _mm_loadu_si128(reinterpret_cast(x + i)); + __m256 wf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(wi), 16)); + __m256 xf = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu16_epi32(xi), 16)); + a0 = _mm256_fmadd_ps(wf, xf, a0); + } + float s = hsum256(_mm256_add_ps(_mm256_add_ps(a0, a1), _mm256_add_ps(a2, a3))); + for (; i < n; ++i) s += bf16_to_f32(w[i]) * bf16_to_f32(x[i]); + return s; +} +#endif // CPU_MOE_X86 + +dot_fn select_dot() { + const IsaTier t = pick_isa(); +#if CPU_MOE_X86 +#ifdef CPU_MOE_HAS_AVX512BF16 + if (t >= ISA_AVX512BF16) return dot_avx512bf16; +#endif + if (t >= ISA_AVX512) return dot_avx512f; + if (t >= ISA_AVX2) return dot_avx2; +#endif + (void)t; + return dot_scalar; +} + +class Bf16Format final : public MoeFormat { + public: + Bf16Format(const MoeParams& p, const FormatArgs& a) + : MoeFormat(p), dot_(select_dot()), gate_up_(a.bank("gate_up")), down_(a.bank("down")) {} + + std::string isa_name() const override { return tier_name(pick_isa()); } + + void gate_up(int layer, int e, int /*tok*/, int i0, int i1, const bf16_t* x, float w_in, + bf16_t* g) const override { + const bf16_t* w = static_cast(tbl_at(gate_up_, layer)) + (size_t)e * (2 * p.I) * p.H; + gated_rows(p, i0, i1, w_in, g, [&](int row) { return dot_(w + (size_t)row * p.H, x, p.H); }); + } + + void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override { + const bf16_t* w = static_cast(tbl_at(down_, layer)); + for (int h = h0; h < h1; ++h) + for (int j = 0; j < n; ++j) + out[j * kRowBlock + h - h0] = dot_(w + ((size_t)routes[j].e * p.H + h) * p.I, routes[j].g, p.I); + } + + private: + dot_fn dot_; + // bf16 rows: gate_up [E, 2I, H], down [E, H, I]. + const uint64_t* gate_up_; + const uint64_t* down_; +}; + +} // namespace + +std::unique_ptr make_bf16_format(const MoeParams& p, const FormatArgs& a) { + return std::make_unique(p, a); +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/ds_fp4.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/ds_fp4.cpp new file mode 100644 index 0000000000..a7afe9804d --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/ds_fp4.cpp @@ -0,0 +1,269 @@ +#include + +#include "../format.h" + +namespace freetoken::cpu_moe { +namespace { + +// ----------------------- DeepSeek-V4 ds_fp4 (W4A8) --------------------------- +// Row-major e2m1 (2/byte, low nibble first) + e8m0 per-32 block scale, no global +// (w = E2M1[code] * 2^(e8m0-127)); activations are FP8-e4m3 round-tripped (per-128 +// block, ue8m0 scale) before each GEMM. Matches kernel/triton/dsv4 (fused_moe + +// fp8_linear): silu(clamp(gate,max=lim)) * clamp(up,-lim,lim), router weight on the +// down output. + +// Activations are pre-deinterleaved to fp32 (xe[m]=x[2m], xo[m]=x[2m+1]) once per +// token/route and reused across all output rows. This drops the hot dot to one +// vpmovzxbd + two vpermps per 32 weights (1.5 shuffle ops / 16 vs 4 for a bf16, +// dup-permute, per-element gather), so the row-major fp4 GEMV stops being port-5 +// bound and approaches the bf16 memory-bandwidth ceiling. +using dsdot_fn = float (*)(const uint8_t*, const uint8_t*, const float*, const float*, int, + const float*, const float*); + +float dot_dsfp4_scalar(const uint8_t* packed, const uint8_t* scale, const float* xe, + const float* xo, int K, const float* e2m1, const float* e8m0) { + float acc = 0.0f; + const int nb = K / 32; + for (int b = 0; b < nb; ++b) { + const float sc = e8m0[scale[b]]; + const uint8_t* pk = packed + (size_t)b * 16; // 16 bytes = 32 codes + const float* xeb = xe + (size_t)b * 16; + const float* xob = xo + (size_t)b * 16; + float bsum = 0.0f; + for (int j = 0; j < 16; ++j) { + const uint8_t byte = pk[j]; + bsum += e2m1[byte & 0xF] * xeb[j]; // low nibble -> even-K activation + bsum += e2m1[byte >> 4] * xob[j]; // high nibble -> odd-K activation + } + acc += sc * bsum; + } + return acc; +} + +#if CPU_MOE_X86 +// One 32-block: 16 bytes -> 16 low + 16 high nibble values via two vpermps, times +// the pre-split even/odd fp32 activations, folded by the per-32 e8m0 scale. +CPU_MOE_TARGET("avx512f") +inline __m512 dsfp4_blk(const uint8_t* pk, const float* xeb, const float* xob, __m512 lut, + __m512i loma, float sc) { + __m512i wi = _mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast(pk))); + __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut); + __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut); + __m512 prod = _mm512_fmadd_ps(vlo, _mm512_loadu_ps(xeb), _mm512_mul_ps(vhi, _mm512_loadu_ps(xob))); + return _mm512_mul_ps(prod, _mm512_set1_ps(sc)); +} + +CPU_MOE_TARGET("avx512f") +float dot_dsfp4_avx512(const uint8_t* packed, const uint8_t* scale, const float* xe, + const float* xo, int K, const float* e2m1, const float* e8m0) { + const __m512 lut = _mm512_loadu_ps(e2m1); + const __m512i loma = _mm512_set1_epi32(0xF); + __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps(); + const int nb = K / 32; // 16 packed bytes + one e8m0 scale per block + int b = 0; + for (; b + 2 <= nb; b += 2) { // two independent accumulators hide FMA latency + acc0 = _mm512_add_ps(acc0, dsfp4_blk(packed + (size_t)b * 16, xe + (size_t)b * 16, + xo + (size_t)b * 16, lut, loma, e8m0[scale[b]])); + acc1 = _mm512_add_ps(acc1, dsfp4_blk(packed + (size_t)(b + 1) * 16, xe + (size_t)(b + 1) * 16, + xo + (size_t)(b + 1) * 16, lut, loma, e8m0[scale[b + 1]])); + } + for (; b < nb; ++b) + acc0 = _mm512_add_ps(acc0, dsfp4_blk(packed + (size_t)b * 16, xe + (size_t)b * 16, + xo + (size_t)b * 16, lut, loma, e8m0[scale[b]])); + return _mm512_reduce_add_ps(_mm512_add_ps(acc0, acc1)); +} + +// AVX2: a 32-K block is 16 bytes -> two 8-lane halves (8 even + 8 odd each). +CPU_MOE_TARGET("avx2,fma") +inline __m256 dsfp4_half_avx2(const uint8_t* pk, const float* xeb, const float* xob, __m256 mag8) { + __m256i wi = _mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cast(pk))); + __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); + __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); + return _mm256_fmadd_ps(vlo, _mm256_loadu_ps(xeb), _mm256_mul_ps(vhi, _mm256_loadu_ps(xob))); +} + +CPU_MOE_TARGET("avx2,fma") +float dot_dsfp4_avx2(const uint8_t* packed, const uint8_t* scale, const float* xe, + const float* xo, int K, const float* e2m1, const float* e8m0) { + const __m256 mag8 = _mm256_loadu_ps(e2m1); + __m256 acc0 = _mm256_setzero_ps(), acc1 = _mm256_setzero_ps(); + const int nb = K / 32; + for (int b = 0; b < nb; ++b) { + const uint8_t* pk = packed + (size_t)b * 16; + const float* xeb = xe + (size_t)b * 16; + const float* xob = xo + (size_t)b * 16; + const __m256 sc = _mm256_set1_ps(e8m0[scale[b]]); + acc0 = _mm256_fmadd_ps(dsfp4_half_avx2(pk, xeb, xob, mag8), sc, acc0); + acc1 = _mm256_fmadd_ps(dsfp4_half_avx2(pk + 8, xeb + 8, xob + 8, mag8), sc, acc1); + } + return hsum256(_mm256_add_ps(acc0, acc1)); +} +#endif + +dsdot_fn select_dsdot() { + const IsaTier t = pick_isa(); +#if CPU_MOE_X86 + if (t >= ISA_AVX512) return dot_dsfp4_avx512; + if (t >= ISA_AVX2) return dot_dsfp4_avx2; +#endif + (void)t; + return dot_dsfp4_scalar; +} + +// Round a clamped |x|<=448 to nearest float8-e4m3 (RNE), back to fp32. Matches +// torch.float8_e4m3fn / triton .to(float8e4nv). +inline float e4m3_round(float x) { + const float sign = x < 0.0f ? -1.0f : 1.0f; + const float a = std::fabs(x); + if (a == 0.0f) return 0.0f; + if (a >= 448.0f) return sign * 448.0f; + int e; + std::frexp(a, &e); // a in [2^(e-1), 2^e) + float step = std::ldexp(1.0f, e - 4); + const float min_step = std::ldexp(1.0f, -9); // e4m3 subnormal step (2^-9) + if (step < min_step) step = min_step; + float r = std::nearbyint(a / step) * step; + if (r > 448.0f) r = 448.0f; + return sign * r; +} + +// IEEE ceil(log2(v)) for v>0 (matches dsv4 _log2_ceil / fast_round_scale). +inline int ceil_log2_pos(float v) { + uint32_t bits; + std::memcpy(&bits, &v, sizeof(bits)); + const int exp = (int)((bits >> 23) & 0xFF); + const int man = (int)(bits & 0x7FFFFF); + return exp - 127 + (man != 0 ? 1 : 0); +} + +// DeepSeek-V4 / V4.1 activation FP8 round-trip (bf16 in/out): per `block` (the checkpoint's +// fp8 block: 128 on V4, 32 on V4.1), s = 2^ceil(log2(max(|x|,1e-4)/448)); +// y = round_e4m3(clamp(x/s,+-448)) * s. +void fp8_roundtrip_bf16(const bf16_t* src, bf16_t* dst, int K, int block) { + for (int b0 = 0; b0 < K; b0 += block) { + const int b1 = std::min(K, b0 + block); + float amax = 1e-4f; + for (int i = b0; i < b1; ++i) amax = std::max(amax, std::fabs(bf16_to_f32(src[i]))); + const float s = std::ldexp(1.0f, ceil_log2_pos(amax * (1.0f / 448.0f))); + const float inv_s = 1.0f / s; + for (int i = b0; i < b1; ++i) { + float q = bf16_to_f32(src[i]) * inv_s; + q = std::min(448.0f, std::max(-448.0f, q)); + dst[i] = f32_to_bf16(e4m3_round(q) * s); + } + } +} + +class DsFp4Format final : public MoeFormat { + public: + DsFp4Format(const MoeParams& p, const FormatArgs& a) : MoeFormat(p), dot_(select_dsdot()) { + gate_up_ = a.bank("gate_up"); + gu_scale_ = a.bank("gate_up_scale"); + down_ = a.bank("down"); + dn_scale_ = a.bank("down_scale"); + act_block_ = static_cast(a.option("act_block")); + input_prequant_ = a.option("input_prequant", 0) != 0; + const size_t tokens = p.max_tokens, routes = (size_t)p.max_tokens * p.top_k; + xq_.resize(tokens * p.H); + xe_.resize(tokens * (p.H / 2)); + xo_.resize(tokens * (p.H / 2)); + ge_.resize(routes * (p.I / 2)); + go_.resize(routes * (p.I / 2)); + for (int i = 0; i < 16; ++i) e2m1_lut_[i] = kE2M1[i]; + for (int i = 0; i < 256; ++i) e8m0_lut_[i] = std::ldexp(1.0f, std::min(i, 254) - 127); + } + + std::string isa_name() const override { return tier_name(pick_isa()); } + + + // DSV4 FP8-round-trips the input before the gate_up GEMV, unless the GPU already did. + void prep_input(int tok, const bf16_t* x) override { + const bf16_t* src = x; + if (!input_prequant_) { + bf16_t* xq = xq_.data() + (size_t)tok * p.H; + fp8_roundtrip_bf16(x, xq, p.H, act_block_); + src = xq; + } + deinterleave_bf16_f32(src, xe_.data() + (size_t)tok * (p.H / 2), xo_.data() + (size_t)tok * (p.H / 2), p.H); + } + + bool preps_inter() const override { return true; } + + void prep_inter(int64_t r, bf16_t* g) override { + fp8_roundtrip_bf16(g, g, p.I, act_block_); + deinterleave_bf16_f32(g, ge_.data() + (size_t)r * (p.I / 2), go_.data() + (size_t)r * (p.I / 2), p.I); + } + + void gate_up(int layer, int e, int tok, int i0, int i1, const bf16_t* /*x*/, float /*w_in*/, + bf16_t* g) const override { + const int H = p.H, I = p.I; + const uint8_t* gu_packed_l = static_cast(tbl_at(gate_up_, layer)); + const uint8_t* gu_scale_l = static_cast(tbl_at(gu_scale_, layer)); + const int N2 = 2 * I, Hh = H / 2, Hs = H / 32; + const float* xe = xe_.data() + (size_t)tok * (H / 2); + const float* xo = xo_.data() + (size_t)tok * (H / 2); + const uint8_t* gp = gu_packed_l + (size_t)e * N2 * Hh; + const uint8_t* gs = gu_scale_l + (size_t)e * N2 * Hs; + const float lim = p.swiglu_limit; + for (int i = i0; i < i1; ++i) { + // gate_up is stored bf16 by the reference GEMV before swiglu; round to match. + float gate = bf16_to_f32(f32_to_bf16( + dot_(gp + (size_t)i * Hh, gs + (size_t)i * Hs, xe, xo, H, e2m1_lut_, e8m0_lut_))); + float up = bf16_to_f32(f32_to_bf16(dot_( + gp + (size_t)(I + i) * Hh, gs + (size_t)(I + i) * Hs, xe, xo, H, e2m1_lut_, e8m0_lut_))); + if (lim > 0.0f) { + if (gate > lim) gate = lim; + if (up > lim) up = lim; + else if (up < -lim) up = -lim; + } + const float glu = gate / (1.0f + std::exp(-gate)); // silu(gate) + g[i] = f32_to_bf16(glu * up); + } + } + + void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override { + const int H = p.H, I = p.I, Ih = I / 2, Is = I / 32; + const uint8_t* dn_packed_l = static_cast(tbl_at(down_, layer)); + const uint8_t* dn_scale_l = static_cast(tbl_at(dn_scale_, layer)); + for (int h = h0; h < h1; ++h) { + for (int j = 0; j < n; ++j) { + const size_t e = (size_t)routes[j].e, r = (size_t)routes[j].r; + const uint8_t* dp = dn_packed_l + e * (size_t)H * Ih + (size_t)h * Ih; + const uint8_t* ds = dn_scale_l + e * (size_t)H * Is + (size_t)h * Is; + out[j * kRowBlock + h - h0] = dot_(dp, ds, ge_.data() + r * Ih, go_.data() + r * Ih, I, e2m1_lut_, e8m0_lut_); + } + } + } + + // The reference rounds each route's weighted down output to bf16 before the fp32 sum over routes. + bool rounds_routes() const override { return true; } + + private: + dsdot_fn dot_; + // e2m1 rows [E,2I,H/2] with e8m0 scales per 32 along K [E,2I,H/32]; down likewise with K = I. + const uint64_t* gate_up_ = nullptr; + const uint64_t* gu_scale_ = nullptr; + const uint64_t* down_ = nullptr; + const uint64_t* dn_scale_ = nullptr; + // Activation FP8 round-trip block, the checkpoint's fp8 block. + int act_block_ = 128; + // The caller already FP8-round-tripped the input activations on the GPU + // (same reference grid), so prep_input must not repeat it on the host-callback + // thread. That scalar per-element pass is single-threaded ON THE DECODE CRITICAL + // PATH (~0.3ms/layer at H=4096, every worker and the GPU waiting on it); moving + // it to a captured GPU elementwise kernel removes it while keeping the official + // W4A8 numerics bit-exact (see cpu_executor.py). + bool input_prequant_ = false; + float e2m1_lut_[16]; + float e8m0_lut_[256]; + AlignedVector xq_; + AlignedVector xe_, xo_, ge_, go_; +}; + +} // namespace + +std::unique_ptr make_ds_fp4_format(const MoeParams& p, const FormatArgs& a) { + return std::make_unique(p, a); +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/fp8_block.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/fp8_block.cpp new file mode 100644 index 0000000000..d27fb1868e --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/fp8_block.cpp @@ -0,0 +1,203 @@ +#include +#include + +#include "../format.h" + +namespace freetoken::cpu_moe { +namespace { + +// ------------------------ block-FP8 (W8A16) dequant ------------------------ +// fp8-e4m3 weights row-major [rows, K] plus one bf16 scale per 128x128 block. The scale tensor is named weight_scale_inv but multiplies, as in kernel/triton/fp8_blockscale_moe.py. +// Whole 128-blocks read the activations as fp32 even/odd halves (xe[m] = x[2m], xo[m] = x[2m+1]); a ragged tail reads the bf16 row x. +// 0x7F/0xFF (NaN in torch float8_e4m3fn) decode to +-480 in every kernel, as e4m3_decode does. +using fp8_block_dot_fn = float (*)(const uint8_t*, const float*, const float*, const bf16_t*, + const bf16_t*, int, const float*); + +constexpr int FP8_BLK = 128; // K elements covered by one weight scale + +inline float fp8_block_tail(const uint8_t* w, const bf16_t* x, int k0, int K, const float* e4m3) { + float blk = 0.0f; + for (int k = k0; k < K; ++k) blk += e4m3[w[k]] * bf16_to_f32(x[k]); + return blk; +} + +float dot_fp8_block_scalar(const uint8_t* w, const float* xe, const float* xo, const bf16_t* x, + const bf16_t* s, int K, const float* e4m3) { + float acc = 0.0f; + int b = 0, k0 = 0; + for (; k0 + FP8_BLK <= K; k0 += FP8_BLK, ++b) { + // 8 independent sums: a single running sum is bound by the add latency + float a[8] = {}; + const uint8_t* wb = w + k0; + const float* pe = xe + k0 / 2; + const float* po = xo + k0 / 2; + for (int j = 0; j < FP8_BLK / 2; j += 4) + for (int l = 0; l < 4; ++l) { + a[2 * l] += e4m3[wb[2 * (j + l)]] * pe[j + l]; + a[2 * l + 1] += e4m3[wb[2 * (j + l) + 1]] * po[j + l]; + } + acc += (((a[0] + a[1]) + (a[2] + a[3])) + ((a[4] + a[5]) + (a[6] + a[7]))) * + bf16_to_f32(s[b]); + } + if (k0 < K) acc += fp8_block_tail(w, x, k0, K, e4m3) * bf16_to_f32(s[b]); + return acc; +} + +#if CPU_MOE_X86 +// A byte b read as fp16 bits (b & 0x80) << 8 | (b & 0x7F) << 7 is its e4m3 value times 2^-8, subnormals included, so the block scale carries the 256. +// In a 16-bit lane holding bytes (2m, 2m+1), srai 1 then & 0xBF80 builds that for byte 2m+1; a left shift by 8 first does byte 2m. Every AVX2 CPU has F16C. +CPU_MOE_TARGET("avx2,fma,f16c") +float dot_fp8_block_avx2(const uint8_t* w, const float* xe, const float* xo, const bf16_t* x, + const bf16_t* s, int K, const float* e4m3) { + const __m256i keep = _mm256_set1_epi16((short)0xBF80); + __m256 total = _mm256_setzero_ps(); + int b = 0, k0 = 0; + for (; k0 + FP8_BLK <= K; k0 += FP8_BLK, ++b) { + _mm_prefetch(reinterpret_cast(w + k0) + PF_AHEAD, _MM_HINT_T0); + __m256 a0 = _mm256_setzero_ps(), a1 = _mm256_setzero_ps(); + __m256 a2 = _mm256_setzero_ps(), a3 = _mm256_setzero_ps(); + for (int j = 0; j < FP8_BLK; j += 32) { + const __m256i v = _mm256_loadu_si256(reinterpret_cast(w + k0 + j)); + const __m256i odd = _mm256_and_si256(_mm256_srai_epi16(v, 1), keep); + const __m256i even = _mm256_and_si256(_mm256_srai_epi16(_mm256_slli_epi16(v, 8), 1), keep); + const float* pe = xe + (k0 + j) / 2; + const float* po = xo + (k0 + j) / 2; + a0 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_castsi256_si128(even)), _mm256_loadu_ps(pe), a0); + a1 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_extracti128_si256(even, 1)), _mm256_loadu_ps(pe + 8), a1); + a2 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_castsi256_si128(odd)), _mm256_loadu_ps(po), a2); + a3 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm256_extracti128_si256(odd, 1)), _mm256_loadu_ps(po + 8), a3); + } + const __m256 blk = _mm256_add_ps(_mm256_add_ps(a0, a1), _mm256_add_ps(a2, a3)); + total = _mm256_fmadd_ps(blk, _mm256_set1_ps(bf16_to_f32(s[b]) * 256.0f), total); + } + float acc = hsum256(total); + if (k0 < K) acc += fp8_block_tail(w, x, k0, K, e4m3) * bf16_to_f32(s[b]); + return acc; +} + +CPU_MOE_TARGET("avx512f,avx2,fma,f16c") +float dot_fp8_block_avx512f(const uint8_t* w, const float* xe, const float* xo, const bf16_t* x, + const bf16_t* s, int K, const float* e4m3) { + const __m256i keep = _mm256_set1_epi16((short)0xBF80); + __m512 total = _mm512_setzero_ps(); + int b = 0, k0 = 0; + for (; k0 + FP8_BLK <= K; k0 += FP8_BLK, ++b) { + _mm_prefetch(reinterpret_cast(w + k0) + PF_AHEAD, _MM_HINT_T0); + __m512 a0 = _mm512_setzero_ps(), a1 = _mm512_setzero_ps(); + __m512 a2 = _mm512_setzero_ps(), a3 = _mm512_setzero_ps(); + for (int j = 0; j < FP8_BLK; j += 64) { + const __m256i v0 = _mm256_loadu_si256(reinterpret_cast(w + k0 + j)); + const __m256i v1 = _mm256_loadu_si256(reinterpret_cast(w + k0 + j + 32)); + const __m256i o0 = _mm256_and_si256(_mm256_srai_epi16(v0, 1), keep); + const __m256i e0 = _mm256_and_si256(_mm256_srai_epi16(_mm256_slli_epi16(v0, 8), 1), keep); + const __m256i o1 = _mm256_and_si256(_mm256_srai_epi16(v1, 1), keep); + const __m256i e1 = _mm256_and_si256(_mm256_srai_epi16(_mm256_slli_epi16(v1, 8), 1), keep); + const float* pe = xe + (k0 + j) / 2; + const float* po = xo + (k0 + j) / 2; + a0 = _mm512_fmadd_ps(_mm512_cvtph_ps(e0), _mm512_loadu_ps(pe), a0); + a1 = _mm512_fmadd_ps(_mm512_cvtph_ps(o0), _mm512_loadu_ps(po), a1); + a2 = _mm512_fmadd_ps(_mm512_cvtph_ps(e1), _mm512_loadu_ps(pe + 16), a2); + a3 = _mm512_fmadd_ps(_mm512_cvtph_ps(o1), _mm512_loadu_ps(po + 16), a3); + } + const __m512 blk = _mm512_add_ps(_mm512_add_ps(a0, a1), _mm512_add_ps(a2, a3)); + total = _mm512_fmadd_ps(blk, _mm512_set1_ps(bf16_to_f32(s[b]) * 256.0f), total); + } + float acc = _mm512_reduce_add_ps(total); + if (k0 < K) acc += fp8_block_tail(w, x, k0, K, e4m3) * bf16_to_f32(s[b]); + return acc; +} +#endif // CPU_MOE_X86 + +fp8_block_dot_fn select_fp8_block_dot() { + const IsaTier t = pick_isa(); +#if CPU_MOE_X86 + // the avx512bf16 tier keeps the fp32 kernel: on Sapphire Rapids it is faster than widening to bf16 for dpbf16 + if (t >= ISA_AVX512) return dot_fp8_block_avx512f; + if (t >= ISA_AVX2) return dot_fp8_block_avx2; +#endif + (void)t; + return dot_fp8_block_scalar; +} + +class Fp8BlockFormat final : public MoeFormat { + public: + Fp8BlockFormat(const MoeParams& p, const FormatArgs& a) : MoeFormat(p), dot_(select_fp8_block_dot()) { + gate_up_ = a.bank("gate_up"); + gu_scale_ = a.bank("gate_up_scale"); + down_ = a.bank("down"); + dn_scale_ = a.bank("down_scale"); + gu_scale_stride_ = a.option("gu_scale_stride"); + dn_scale_stride_ = a.option("dn_scale_stride"); + const size_t tokens = p.max_tokens, routes = (size_t)p.max_tokens * p.top_k; + xe_.resize(tokens * (p.H / 2)); + xo_.resize(tokens * (p.H / 2)); + ge_.resize(routes * (p.I / 2)); + go_.resize(routes * (p.I / 2)); + for (int i = 0; i < 256; ++i) e4m3_lut_[i] = e4m3_decode((uint8_t)i); + } + + std::string isa_name() const override { return tier_name(pick_isa()); } + + + void prep_input(int tok, const bf16_t* x) override { + deinterleave_bf16_f32(x, xe_.data() + (size_t)tok * (p.H / 2), xo_.data() + (size_t)tok * (p.H / 2), p.H); + } + + bool preps_inter() const override { return true; } + + void prep_inter(int64_t r, bf16_t* g) override { + deinterleave_bf16_f32(g, ge_.data() + (size_t)r * (p.I / 2), go_.data() + (size_t)r * (p.I / 2), p.I); + } + + // The halves cover whole 128-blocks; the bf16 row covers a ragged tail. + void gate_up(int layer, int e, int tok, int i0, int i1, const bf16_t* x, float w_in, + bf16_t* g) const override { + const int H = p.H, I = p.I; + const uint8_t* packed = static_cast(tbl_at(gate_up_, layer)); + const bf16_t* scale = static_cast(tbl_at(gu_scale_, layer)); + const float* xe = xe_.data() + (size_t)tok * (H / 2); + const float* xo = xo_.data() + (size_t)tok * (H / 2); + gated_rows(p, i0, i1, w_in, g, [&](int row) { + const uint8_t* w = packed + ((size_t)e * (2 * I) + row) * (size_t)H; + const bf16_t* sc = scale + ((size_t)e * ((2 * I + FP8_BLK - 1) / FP8_BLK) + row / FP8_BLK) * + (size_t)gu_scale_stride_; + return dot_(w, xe, xo, x, sc, H, e4m3_lut_); + }); + } + + void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override { + const int H = p.H, I = p.I; + const uint8_t* packed = static_cast(tbl_at(down_, layer)); + const bf16_t* scale = static_cast(tbl_at(dn_scale_, layer)); + for (int h = h0; h < h1; ++h) { + for (int j = 0; j < n; ++j) { + const size_t e = (size_t)routes[j].e, r = (size_t)routes[j].r; + const uint8_t* w = packed + (e * H + h) * (size_t)I; + const bf16_t* sc = scale + (e * ((H + FP8_BLK - 1) / FP8_BLK) + h / FP8_BLK) * + (size_t)dn_scale_stride_; + out[j * kRowBlock + h - h0] = + dot_(w, ge_.data() + r * (I / 2), go_.data() + r * (I / 2), routes[j].g, sc, I, e4m3_lut_); + } + } + } + + private: + fp8_block_dot_fn dot_; + // e4m3 rows [E,2I,H] and [E,H,I] with one bf16 scale per 128x128 block. + const uint64_t* gate_up_ = nullptr; + const uint64_t* gu_scale_ = nullptr; + const uint64_t* down_ = nullptr; + const uint64_t* dn_scale_ = nullptr; + // Row strides of the scale tensors, including any padding of the host bank layout. + int64_t gu_scale_stride_ = 0, dn_scale_stride_ = 0; + float e4m3_lut_[256]; + AlignedVector xe_, xo_, ge_, go_; +}; + +} // namespace + +std::unique_ptr make_fp8_block_format(const MoeParams& p, const FormatArgs& a) { + return std::make_unique(p, a); +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/mxfp4.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/mxfp4.cpp new file mode 100644 index 0000000000..484d6dc001 --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/mxfp4.cpp @@ -0,0 +1,247 @@ +#include "../format.h" + +namespace freetoken::cpu_moe { +namespace { + +// ------------------------- mxfp4 (gpt-oss) GEMV ----------------------------- +// Transposed split-K layout: blk[Kpairs, N2] (N innermost), scl[Kpairs/16, N2] +// e8m0 per 32-K. Computes out[c] = sum_kb (E2M1[lo]*x[2kb] + E2M1[hi]*x[2kb+1]) +// * 2^(e8m0-127) for a contiguous column tile (blk/scl already offset to col 0 of +// the tile). Vectorized over N (16 columns / __m512), K stays the outer (cache- +// sequential) loop. Used by both gate_up (K=H) and down (K=I). +using mxgemv_fn = void (*)(float*, const uint8_t*, const uint8_t*, const bf16_t*, int, int, + int, const float*, const float*); + +void mxfp4_gemv_scalar(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, + int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { + for (int c = 0; c < ncol; ++c) out[c] = 0.0f; + for (int kb = 0; kb < Kpairs; ++kb) { + const uint8_t* w = blk + (size_t)kb * N2; + const uint8_t* s = scl + (size_t)(kb >> 4) * N2; + const float xl = bf16_to_f32(x[2 * kb]); + const float xh = bf16_to_f32(x[2 * kb + 1]); + for (int c = 0; c < ncol; ++c) { + const uint8_t byte = w[c]; + out[c] += (e2m1[byte & 0xF] * xl + e2m1[byte >> 4] * xh) * e8m0[s[c]]; + } + } +} + +#if CPU_MOE_X86 +CPU_MOE_TARGET("avx512f") +void mxfp4_gemv_avx512(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, + int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { + (void)e8m0; // e8m0[c]=2^(c-127) computed via bit construction (no gather) + const __m512 lut = _mm512_loadu_ps(e2m1); + const __m512i loma = _mm512_set1_epi32(0xF); + // K-outer / N-inner: each kb cache line is read once and all live column chunks + // (up to 4 -> 64 cols) accumulate from registers, so DRAM/L2 stream the tile once. + int c0 = 0; + for (; c0 + 16 <= ncol; c0 += 64) { + const int nchunk = std::min(4, (ncol - c0) / 16); + __m512 acc[4]; + for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm512_setzero_ps(); + for (int kblk = 0; kblk < Kpairs; kblk += 16) { // 16 K-pairs = 32 K = one scale row + __m512 sc[4]; + for (int ci = 0; ci < nchunk; ++ci) { + __m128i sraw = _mm_loadu_si128(reinterpret_cast( + scl + (size_t)(kblk >> 4) * N2 + c0 + ci * 16)); + sc[ci] = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu8_epi32(sraw), 23)); + } + __m512 blk_acc[4]; + for (int ci = 0; ci < nchunk; ++ci) blk_acc[ci] = _mm512_setzero_ps(); + for (int kk = 0; kk < 16; ++kk) { + const int kb = kblk + kk; + const uint8_t* wbase = blk + (size_t)kb * N2 + c0; + // The transposed layout strides K by N2 bytes; prefetch ahead so the strided + // reads are not exposed to DRAM latency (the HW streamer misses big strides). + constexpr int PFD = 8; + if (kb + PFD < Kpairs) + _mm_prefetch(reinterpret_cast(blk + (size_t)(kb + PFD) * N2 + c0), + _MM_HINT_T0); + const __m512 xl = _mm512_set1_ps(bf16_to_f32(x[2 * kb])); + const __m512 xh = _mm512_set1_ps(bf16_to_f32(x[2 * kb + 1])); + for (int ci = 0; ci < nchunk; ++ci) { + __m512i wi = _mm512_cvtepu8_epi32( + _mm_loadu_si128(reinterpret_cast(wbase + ci * 16))); + __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut); + __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut); + blk_acc[ci] = _mm512_fmadd_ps(vlo, xl, blk_acc[ci]); + blk_acc[ci] = _mm512_fmadd_ps(vhi, xh, blk_acc[ci]); + } + } + for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm512_fmadd_ps(blk_acc[ci], sc[ci], acc[ci]); + } + for (int ci = 0; ci < nchunk; ++ci) _mm512_storeu_ps(out + c0 + ci * 16, acc[ci]); + } + for (int c = c0; c < ncol; ++c) { // tail columns (< 16) + float o = 0.0f; + for (int kb = 0; kb < Kpairs; ++kb) { + const uint8_t byte = blk[(size_t)kb * N2 + c]; + uint32_t bits = (uint32_t)scl[(size_t)(kb >> 4) * N2 + c] << 23; + float sc; + std::memcpy(&sc, &bits, 4); + o += (e2m1[byte & 0xF] * bf16_to_f32(x[2 * kb]) + + e2m1[byte >> 4] * bf16_to_f32(x[2 * kb + 1])) * sc; + } + out[c] = o; + } +} + +CPU_MOE_TARGET("avx2,fma") +void mxfp4_gemv_avx2(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, + int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { + (void)e8m0; // e8m0[s]=2^(s-127) built via s<<23 (no gather) + const __m256 mag8 = _mm256_loadu_ps(e2m1); + int c0 = 0; + for (; c0 + 8 <= ncol; c0 += 32) { // up to 4 chunks of 8 = 32 cols + const int nchunk = std::min(4, (ncol - c0) / 8); + __m256 acc[4]; + for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm256_setzero_ps(); + for (int kblk = 0; kblk < Kpairs; kblk += 16) { // 16 K-pairs = one scale row + __m256 sc[4]; + for (int ci = 0; ci < nchunk; ++ci) { + __m128i sraw = _mm_loadl_epi64(reinterpret_cast( + scl + (size_t)(kblk >> 4) * N2 + c0 + ci * 8)); + sc[ci] = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu8_epi32(sraw), 23)); + } + __m256 blk_acc[4]; + for (int ci = 0; ci < nchunk; ++ci) blk_acc[ci] = _mm256_setzero_ps(); + for (int kk = 0; kk < 16; ++kk) { + const int kb = kblk + kk; + const uint8_t* wbase = blk + (size_t)kb * N2 + c0; + constexpr int PFD = 8; + if (kb + PFD < Kpairs) + _mm_prefetch(reinterpret_cast(blk + (size_t)(kb + PFD) * N2 + c0), + _MM_HINT_T0); + const __m256 xl = _mm256_set1_ps(bf16_to_f32(x[2 * kb])); + const __m256 xh = _mm256_set1_ps(bf16_to_f32(x[2 * kb + 1])); + for (int ci = 0; ci < nchunk; ++ci) { + __m256i wi = _mm256_cvtepu8_epi32( + _mm_loadl_epi64(reinterpret_cast(wbase + ci * 8))); + __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); + __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); + blk_acc[ci] = _mm256_fmadd_ps(vlo, xl, blk_acc[ci]); + blk_acc[ci] = _mm256_fmadd_ps(vhi, xh, blk_acc[ci]); + } + } + for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm256_fmadd_ps(blk_acc[ci], sc[ci], acc[ci]); + } + for (int ci = 0; ci < nchunk; ++ci) _mm256_storeu_ps(out + c0 + ci * 8, acc[ci]); + } + for (int c = c0; c < ncol; ++c) { // tail columns (< 8); none when ncol%8==0 + float o = 0.0f; + for (int kb = 0; kb < Kpairs; ++kb) { + const uint8_t byte = blk[(size_t)kb * N2 + c]; + uint32_t bits = (uint32_t)scl[(size_t)(kb >> 4) * N2 + c] << 23; + float sc; + std::memcpy(&sc, &bits, 4); + o += (e2m1[byte & 0xF] * bf16_to_f32(x[2 * kb]) + + e2m1[byte >> 4] * bf16_to_f32(x[2 * kb + 1])) * sc; + } + out[c] = o; + } +} +#endif + +mxgemv_fn select_mxgemv() { + const IsaTier t = pick_isa(); +#if CPU_MOE_X86 + if (t >= ISA_AVX512) return mxfp4_gemv_avx512; + if (t >= ISA_AVX2) return mxfp4_gemv_avx2; +#endif + (void)t; + return mxfp4_gemv_scalar; +} + +// ----------------------------- mxfp4 (gpt-oss) ----------------------------- +// Transposed split-K layout (N innermost), so the GEMV streams K and accumulates +// a contiguous block of N output columns -> cache-line-efficient, no repack, no +// extra host memory. Pass 1 fuses gate_up + clamped-swiglu(+bias); pass 2 adds the +// down bias, and the executor applies the router weight and sums the routes. +// +// Dequant: w = E2M1[code] * 2^(e8m0_scale - 127); two codes per byte (low nibble +// first), one e8m0 scale per 32 contiguous K. Matches kernel/triton/mxfp4_moe.py. +class Mxfp4Format final : public MoeFormat { + public: + Mxfp4Format(const MoeParams& p, const FormatArgs& a) : MoeFormat(p), gemv_(select_mxgemv()) { + gate_up_ = a.bank("gate_up"); + gu_scale_ = a.bank("gate_up_scale"); + gu_bias_ = a.bank("gate_up_bias"); + down_ = a.bank("down"); + dn_scale_ = a.bank("down_scale"); + dn_bias_ = a.bank("down_bias"); + for (int i = 0; i < 16; ++i) e2m1_lut_[i] = kE2M1[i]; + // e8m0 (mxfp4 block scale) = 2^(s-127); the GPU GEMV clamps s to [0,254]. + for (int i = 0; i < 256; ++i) e8m0_lut_[i] = std::ldexp(1.0f, std::min(i, 254) - 127); + } + + std::string isa_name() const override { return tier_name(pick_isa()); } + + void gate_up(int layer, int e, int /*tok*/, int i0, int i1, const bf16_t* x, float /*w_in*/, + bf16_t* g) const override { + const int H = p.H, I = p.I; + const uint8_t* gu_packed_l = static_cast(tbl_at(gate_up_, layer)); + const uint8_t* gu_scale_l = static_cast(tbl_at(gu_scale_, layer)); + const bf16_t* gu_bias_l = static_cast(tbl_at(gu_bias_, layer)); + const int N2 = 2 * I; // gate_up output width (gate/up interleaved) + const int Hh = H / 2; // packed-K rows + const int nunit = i1 - i0; // intermediate units owned by this tile + const int col0 = 2 * i0; // first gate_up column + const int ncol = 2 * nunit; // gate_up columns owned by this tile + const uint8_t* blk_e = gu_packed_l + (size_t)e * Hh * N2; + const uint8_t* scl_e = gu_scale_l + (size_t)e * (size_t)(H / 32) * N2; + float gu[2 * kRowBlock]; + gemv_(gu, blk_e + col0, scl_e + col0, x, Hh, N2, ncol, e2m1_lut_, e8m0_lut_); + const bf16_t* bias_e = gu_bias_l + (size_t)e * N2 + col0; + const float lim = p.swiglu_limit, alpha = p.swiglu_alpha; + for (int j = 0; j < nunit; ++j) { + float gate = gu[2 * j] + bf16_to_f32(bias_e[2 * j]); + float up = gu[2 * j + 1] + bf16_to_f32(bias_e[2 * j + 1]); + if (gate > lim) gate = lim; + if (up > lim) up = lim; + else if (up < -lim) up = -lim; + const float glu = gate / (1.0f + std::exp(-gate * alpha)); // gate * sigmoid(alpha*gate) + g[i0 + j] = f32_to_bf16(glu * (up + 1.0f)); + } + } + + // The transposed layout already streams a whole column block per call, so the routes go one after another. + void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override { + const int H = p.H, I = p.I; + const int Ih = I / 2; // packed-K rows + const int nh = h1 - h0; + const uint8_t* dn_packed_l = static_cast(tbl_at(down_, layer)); + const uint8_t* dn_scale_l = static_cast(tbl_at(dn_scale_, layer)); + const bf16_t* dn_bias_l = static_cast(tbl_at(dn_bias_, layer)); + for (int j = 0; j < n; ++j) { + const size_t e = (size_t)routes[j].e; + const uint8_t* blk_e = dn_packed_l + e * Ih * H; + const uint8_t* scl_e = dn_scale_l + e * (size_t)(I / 32) * H; + float* part = out + j * kRowBlock; + gemv_(part, blk_e + h0, scl_e + h0, routes[j].g, Ih, H, nh, e2m1_lut_, e8m0_lut_); + const bf16_t* bias_e = dn_bias_l + e * H + h0; + for (int c = 0; c < nh; ++c) part[c] += bf16_to_f32(bias_e[c]); + } + } + + private: + mxgemv_fn gemv_; + // gpt-oss split-K layout, N innermost: blocks [E,H/2,2I], e8m0 scales [E,H/32,2I], bf16 bias [E,2I]; down likewise with K = I, N = H. + const uint64_t* gate_up_ = nullptr; + const uint64_t* gu_scale_ = nullptr; + const uint64_t* gu_bias_ = nullptr; + const uint64_t* down_ = nullptr; + const uint64_t* dn_scale_ = nullptr; + const uint64_t* dn_bias_ = nullptr; + float e2m1_lut_[16]; + float e8m0_lut_[256]; +}; + +} // namespace + +std::unique_ptr make_mxfp4_format(const MoeParams& p, const FormatArgs& a) { + return std::make_unique(p, a); +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/nvfp4.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/nvfp4.cpp new file mode 100644 index 0000000000..429955faf3 --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/nvfp4.cpp @@ -0,0 +1,466 @@ +#include + +#include "../format.h" + +namespace freetoken::cpu_moe { +namespace { + +// --------------------------- NVFP4 (W4A16) dequant --------------------------- +// Weights: e2m1 4-bit codes (2/byte, low nibble first), per-16 block scale in +// fp8-e4m3, per-output-row global scale in fp16. Dequant matches the GPU kernels +// (freetoken/kernel/triton/nvfp4_dequant.py): w = E2M1[code] * e4m3(scale) * global. +// Activations stay bf16 (W4A16); the GEMV dequantizes weights inside the K-loop. + + +// e2m1 * 2 as exact int8 (all e2m1 values are multiples of 0.5). Used by the AVX-VNNI +// W4A8 path: nibble -> int8 weight via PSHUFB LUT, then VPDPBUSD against int8 activations; +// the *2 is undone by a 0.5 folded into the final scale. Mirrors ggml's kvalues_mxfp4. +alignas(16) const int8_t kE2M1x2[16] = {0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12}; + +// Activations pre-deinterleaved to fp32 (xe[m]=x[2m], xo[m]=x[2m+1]); see the +// ds_fp4 dot in ds_fp4.cpp for why (drops the hot loop to ~1.5 shuffle ops / 16 weights). +using nvdot_fn = float (*)(const uint8_t*, const uint8_t*, float, const float*, const float*, + int, const float*, const float*); + +float dot_nvfp4_scalar(const uint8_t* packed, const uint8_t* scale, float global, + const float* xe, const float* xo, int K, const float* e2m1, + const float* e4m3) { + float acc = 0.0f; + const int nb = K / 16; + for (int b = 0; b < nb; ++b) { + const float bs = e4m3[scale[b]]; + const uint8_t* pk = packed + (size_t)b * 8; + const float* xeb = xe + (size_t)b * 8; // 16 K -> 8 even + 8 odd + const float* xob = xo + (size_t)b * 8; + float bsum = 0.0f; + for (int j = 0; j < 8; ++j) { + const uint8_t byte = pk[j]; + bsum += e2m1[byte & 0xF] * xeb[j]; + bsum += e2m1[byte >> 4] * xob[j]; + } + acc += bs * bsum; + } + return acc * global; +} + +// ---- NVFP4 W4A8 (int8 activations) dot: nibble->int8 via LUT, per-16 act scale ---- +// asi8: int8 activations laid out per-16 block as [even(8), odd(8)]; asb[b] = per-block +// activation scale (absmax/127). Result folds the e2m1*2 -> *0.5 into the scale. +using nvi8dot_fn = float (*)(const uint8_t*, const uint8_t*, float, const int8_t*, int, + const float*, const float*); + +[[maybe_unused]] float dot_nvfp4_i8_scalar(const uint8_t* packed, const uint8_t* scale, + float global, const int8_t* asi8, int K, const float* e4m3, + const float* asb) { + float acc = 0.0f; + const int nb = K / 16; + for (int b = 0; b < nb; ++b) { + const uint8_t* pk = packed + (size_t)b * 8; + const int8_t* ae = asi8 + (size_t)b * 16; // even(8) + const int8_t* ao = ae + 8; // odd(8) + int isum = 0; + for (int j = 0; j < 8; ++j) { + isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j]; + isum += (int)kE2M1x2[pk[j] >> 4] * (int)ao[j]; + } + acc += (e4m3[scale[b]] * asb[b]) * (float)isum; + } + return acc * (0.5f * global); +} + +#if CPU_MOE_X86 +// Two 16-K blocks (16 packed bytes) per iter: lo nibbles -> even-K, hi -> odd-K, +// gathered via two vpermps. The per-16 e4m3 scale differs across the two blocks, so +// it is applied per lane (low 8 lanes = block b, high 8 = block b+1). +CPU_MOE_TARGET("avx512f") +inline __m512 nvfp4_blk2(const uint8_t* pk, const float* xeb, const float* xob, __m512 lut, + __m512i loma, float s0, float s1) { + __m512i wi = _mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast(pk))); + __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut); + __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut); + __m512 prod = _mm512_fmadd_ps(vlo, _mm512_loadu_ps(xeb), _mm512_mul_ps(vhi, _mm512_loadu_ps(xob))); + // lanes 0-7 (block b) -> s0, lanes 8-15 (block b+1) -> s1 (pure AVX512F mask move) + __m512 scv = _mm512_mask_mov_ps(_mm512_set1_ps(s0), 0xFF00, _mm512_set1_ps(s1)); + return _mm512_mul_ps(prod, scv); +} + +CPU_MOE_TARGET("avx512f") +float dot_nvfp4_avx512(const uint8_t* packed, const uint8_t* scale, float global, + const float* xe, const float* xo, int K, const float* e2m1, + const float* e4m3) { + const __m512 lut = _mm512_loadu_ps(e2m1); // 16 e2m1 values + const __m512i loma = _mm512_set1_epi32(0xF); + __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps(); + const int nb = K / 16; // 8 packed bytes + one e4m3 scale per block + int b = 0; + for (; b + 4 <= nb; b += 4) { // two blk2 calls -> 4 blocks, 2 accumulators + acc0 = _mm512_add_ps(acc0, nvfp4_blk2(packed + (size_t)b * 8, xe + (size_t)b * 8, + xo + (size_t)b * 8, lut, loma, e4m3[scale[b]], e4m3[scale[b + 1]])); + acc1 = _mm512_add_ps(acc1, nvfp4_blk2(packed + (size_t)(b + 2) * 8, xe + (size_t)(b + 2) * 8, + xo + (size_t)(b + 2) * 8, lut, loma, e4m3[scale[b + 2]], e4m3[scale[b + 3]])); + } + for (; b + 2 <= nb; b += 2) + acc0 = _mm512_add_ps(acc0, nvfp4_blk2(packed + (size_t)b * 8, xe + (size_t)b * 8, + xo + (size_t)b * 8, lut, loma, e4m3[scale[b]], e4m3[scale[b + 1]])); + float s = _mm512_reduce_add_ps(_mm512_add_ps(acc0, acc1)); + for (; b < nb; ++b) { // odd final 16-K block + const uint8_t* pk = packed + (size_t)b * 8; + const float* xeb = xe + (size_t)b * 8; + const float* xob = xo + (size_t)b * 8; + float bsum = 0.0f; + for (int j = 0; j < 8; ++j) { + bsum += e2m1[pk[j] & 0xF] * xeb[j] + e2m1[pk[j] >> 4] * xob[j]; + } + s += e4m3[scale[b]] * bsum; + } + return s * global; +} + +// AVX2: one 16-K block (8 packed bytes) per call, 8 even + 8 odd lanes. +CPU_MOE_TARGET("avx2,fma") +inline __m256 nvfp4_blk_avx2(const uint8_t* pk, const float* xeb, const float* xob, + __m256 mag8, float sc) { + __m256i wi = _mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cast(pk))); + __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); + __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); + __m256 prod = _mm256_fmadd_ps(vlo, _mm256_loadu_ps(xeb), _mm256_mul_ps(vhi, _mm256_loadu_ps(xob))); + return _mm256_mul_ps(prod, _mm256_set1_ps(sc)); +} + +CPU_MOE_TARGET("avx2,fma") +float dot_nvfp4_avx2(const uint8_t* packed, const uint8_t* scale, float global, + const float* xe, const float* xo, int K, const float* e2m1, + const float* e4m3) { + const __m256 mag8 = _mm256_loadu_ps(e2m1); // e2m1[0..7] magnitudes + __m256 acc0 = _mm256_setzero_ps(), acc1 = _mm256_setzero_ps(); + const int nb = K / 16; + int b = 0; + for (; b + 2 <= nb; b += 2) { + acc0 = _mm256_add_ps(acc0, nvfp4_blk_avx2(packed + (size_t)b * 8, xe + (size_t)b * 8, + xo + (size_t)b * 8, mag8, e4m3[scale[b]])); + acc1 = _mm256_add_ps(acc1, nvfp4_blk_avx2(packed + (size_t)(b + 1) * 8, xe + (size_t)(b + 1) * 8, + xo + (size_t)(b + 1) * 8, mag8, e4m3[scale[b + 1]])); + } + for (; b < nb; ++b) + acc0 = _mm256_add_ps(acc0, nvfp4_blk_avx2(packed + (size_t)b * 8, xe + (size_t)b * 8, + xo + (size_t)b * 8, mag8, e4m3[scale[b]])); + return hsum256(_mm256_add_ps(acc0, acc1)) * global; +} + +// AVX-VNNI W4A8: decode 8 packed bytes (16 nibbles) of one 16-block to int8 [lo(8),hi(8)] +// via PSHUFB against the e2m1*2 LUT (replaces the 2 vpermps fp32 expands -- ~4x less +// port-5 traffic). lo=even-K weights, hi=odd-K, matching the [even(8),odd(8)] act layout. +CPU_MOE_TARGET("avx2,avxvnni,fma") +inline __m128i nvfp4_decode_block_i8(const uint8_t* pk, __m128i lut) { + __m128i b = _mm_loadl_epi64(reinterpret_cast(pk)); // 8 bytes + __m128i lo = _mm_and_si128(b, _mm_set1_epi8(0x0F)); + __m128i hi = _mm_and_si128(_mm_srli_epi16(b, 4), _mm_set1_epi8(0x0F)); + return _mm_shuffle_epi8(lut, _mm_unpacklo_epi64(lo, hi)); // [lo(8),hi(8)] -> int8 +} + +// Two 16-blocks per VPDPBUSD (32 int8). Sign trick (ggml mul_add_epi8): |w|*(sign(w)*a)=w*a, +// so u8*s8 needs no offset/correction term. Per-block scale (e4m3 * act-scale) folded in fp32 +// (lanes 0-3 -> block b, 4-7 -> block b+1). Bit-faithful weight; only the int8 activation quant +// (W4A8) differs from the bf16 reference. +CPU_MOE_TARGET("avx2,avxvnni,fma") +float dot_nvfp4_i8_vnni(const uint8_t* packed, const uint8_t* scale, float global, + const int8_t* asi8, int K, const float* e4m3, const float* asb) { + const __m128i lut = _mm_loadu_si128(reinterpret_cast(kE2M1x2)); + __m256 accF = _mm256_setzero_ps(); + const int nb = K / 16; + int b = 0; + for (; b + 2 <= nb; b += 2) { + __m128i wb = nvfp4_decode_block_i8(packed + (size_t)b * 8, lut); + __m128i wb1 = nvfp4_decode_block_i8(packed + (size_t)(b + 1) * 8, lut); + __m256i w = _mm256_set_m128i(wb1, wb); // [blk b | blk b+1] + __m256i a = _mm256_loadu_si256(reinterpret_cast(asi8 + (size_t)b * 16)); + __m256i aw = _mm256_sign_epi8(w, w); // |w| (u8 operand) + __m256i sa = _mm256_sign_epi8(a, w); // sign(w)*a (s8 operand) + __m256i di = _mm256_dpbusd_avx_epi32(_mm256_setzero_si256(), aw, sa); + __m256 scv = _mm256_blend_ps(_mm256_set1_ps(e4m3[scale[b]] * asb[b]), + _mm256_set1_ps(e4m3[scale[b + 1]] * asb[b + 1]), 0xF0); + accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(di), scv, accF); + } + float s = hsum256(accF); + for (; b < nb; ++b) { // tail (odd block count) + const uint8_t* pk = packed + (size_t)b * 8; + const int8_t* ae = asi8 + (size_t)b * 16; const int8_t* ao = ae + 8; + int isum = 0; + for (int j = 0; j < 8; ++j) + isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j] + (int)kE2M1x2[pk[j] >> 4] * (int)ao[j]; + s += (e4m3[scale[b]] * asb[b]) * (float)isum; + } + return s * (0.5f * global); +} + +#ifdef CPU_MOE_HAS_AVX512VNNI + +// Software-prefetch distance for the W4A8 weight stream, in 16-K blocks (8 packed +// bytes each). Returns -1 when FREETOKEN_CPU_MOE_PF_BLOCKS is unset: the kernel then +// uses the built-in default min(512 blocks = 4 KB, 2 rows) -- 4 KB is the empirical +// optimum on large-row machines (Emerald Rapids sweep), while the 2-row cap keeps a small-row +// model's overshoot bounded (the executor works in 32-row tiles, so a fixed byte +// distance otherwise prefetches another worker's tile: duplicated DRAM traffic that +// regresses at the bandwidth ceiling). An EXPLICIT env value is honored verbatim +// (no clamp; 0 disables): the per-machine optimum can sit past the safe default (+20% +// at 4 KB on a 24-thread Ice Lake with 256B rows), so the escape hatch must reach it. +// Prefetch never faults, so overshooting a row/bank tail is safe. +static int nvfp4_pf_blocks() { + static const int v = [] { + const char* s = getenv("FREETOKEN_CPU_MOE_PF_BLOCKS"); + return (s && s[0]) ? atoi(s) : -1; + }(); + return v; +} + +// AVX-512 VNNI W4A8: FOUR 16-K blocks per VPDPBUSD (64 int8) -- 2x the AVX-VNNI +// (256-bit) path. Decode 32 packed bytes -> 64 int8 with a single _mm512_shuffle_epi8 +// (e2m1*2 LUT replicated to all 4 128-bit lanes). AVX-512 has no _mm512_sign_epi8, so +// the u8*s8 sign trick (|w| as u8, sign(w)*a as s8) uses abs_epi8 + a masked negate. +// Bit-faithful weight; only the int8 activation quant (W4A8) differs from bf16. +// One 4-block group (64 int8) -> scaled fp32 partial (16 lanes). Isolated as a helper so +// the caller can run several independent chains into separate accumulators (the decode -> +// dpbusd -> scale chain is long, so a single accumulator leaves the core latency-bound). +CPU_MOE_TARGET("avx512f,avx512bw,avx512vnni,avx2") +static inline __m512 nvfp4_i8_grp4(const uint8_t* packed, const uint8_t* scale, + const int8_t* asi8, const float* e4m3, const float* asb, + int b, __m512i lut, __m512i idx, __m512i mask0F, + __m512i idxsc) { + const __mmask64 hi_half = 0xFF00FF00FF00FF00ULL; // bytes 8-15 of each 128b lane -> hi nibbles + __m256i raw = _mm256_loadu_si256(reinterpret_cast(packed + (size_t)b * 8)); + __m512i src = _mm512_permutexvar_epi64(idx, _mm512_castsi256_si512(raw)); + __m512i lo = _mm512_and_si512(src, mask0F); + __m512i hi = _mm512_and_si512(_mm512_srli_epi16(src, 4), mask0F); + __m512i comb = _mm512_mask_blend_epi8(hi_half, lo, hi); + __m512i w = _mm512_shuffle_epi8(lut, comb); // 64 int8 weights (e2m1*2) + // u8*s8 sign trick without _mm512_sign_epi8: aw=|w|, sa = (w<0 ? -a : a) (a moot at w==0) + __m512i a = _mm512_loadu_si512(reinterpret_cast(asi8 + (size_t)b * 16)); + __m512i aw = _mm512_abs_epi8(w); + __mmask64 neg = _mm512_movepi8_mask(w); + __m512i sa = _mm512_mask_sub_epi8(a, neg, _mm512_setzero_si512(), a); + __m512i di = _mm512_dpbusd_epi32(_mm512_setzero_si512(), aw, sa); // 16 int32 (groups of 4) + // int32[0..3]->blk b, [4..7]->b+1, [8..11]->b+2, [12..15]->b+3. The 4 block scales + // (e4m3 LUT x per-block act scale) are computed vectorized: a 4-byte load + epu8->epi32 + // widen + one 4-lane LUT gather + one mul replaces 8 scalar loads + 4 scalar muls + + // a set_ps assembly, which otherwise dominates the per-group op count. + int sc_raw; + memcpy(&sc_raw, scale + b, 4); + __m128i sc4 = _mm_cvtepu8_epi32(_mm_cvtsi32_si128(sc_raw)); + __m128 s4 = _mm_mul_ps(_mm_i32gather_ps(e4m3, sc4, 4), _mm_loadu_ps(asb + b)); + __m512 scv = _mm512_permutexvar_ps(idxsc, _mm512_castps128_ps512(s4)); + return _mm512_mul_ps(_mm512_cvtepi32_ps(di), scv); +} + +CPU_MOE_TARGET("avx512f,avx512bw,avx512vnni,avx2") +float dot_nvfp4_i8_avx512vnni(const uint8_t* packed, const uint8_t* scale, float global, + const int8_t* asi8, int K, const float* e4m3, const float* asb) { + const __m512i lut = _mm512_broadcast_i32x4( + _mm_loadu_si128(reinterpret_cast(kE2M1x2))); + const __m512i idx = _mm512_set_epi64(3, 3, 2, 2, 1, 1, 0, 0); // block i -> 128b lane i + const __m512i mask0F = _mm512_set1_epi8(0x0F); + const __m512i idxsc = _mm512_set_epi32(3, 3, 3, 3, 2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0); + __m512 acc0 = _mm512_setzero_ps(), acc1 = _mm512_setzero_ps(); + __m512 acc2 = _mm512_setzero_ps(), acc3 = _mm512_setzero_ps(); + const int nb = K / 16; + // Built-in default: min(4 KB, 2 rows) -- see nvfp4_pf_blocks(). An explicitly set + // env value is used verbatim so a per-machine sweep can reach operating points past + // the conservative default. + const int pfb = nvfp4_pf_blocks(); + const int pf = (pfb < 0) ? std::min(512, 2 * nb) : pfb; + int b = 0; + // Four independent 4-block groups per iter -> four accumulator chains hide the ~10-op + // decode->dpbusd->scale latency (1-2 chains leave the loop latency-bound: the per-core + // rate sat at ~60% of the core's achievable DRAM stream rate). + for (; b + 16 <= nb; b += 16) { + // Prefetch the weight stream ahead: the interleaved decode lowers the L1-miss + // concurrency the HW prefetcher sustains on its own (134 -> 167 GB/s on Emerald Rapids). + if (pf > 0) { + _mm_prefetch(reinterpret_cast(packed + ((size_t)b + (size_t)pf) * 8), + _MM_HINT_T0); + _mm_prefetch(reinterpret_cast(packed + ((size_t)b + (size_t)pf) * 8 + 64), + _MM_HINT_T0); + } + acc0 = _mm512_add_ps(acc0, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b, lut, idx, + mask0F, idxsc)); + acc1 = _mm512_add_ps(acc1, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 4, lut, idx, + mask0F, idxsc)); + acc2 = _mm512_add_ps(acc2, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 8, lut, idx, + mask0F, idxsc)); + acc3 = _mm512_add_ps(acc3, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b + 12, lut, idx, + mask0F, idxsc)); + } + for (; b + 4 <= nb; b += 4) + acc0 = _mm512_add_ps(acc0, nvfp4_i8_grp4(packed, scale, asi8, e4m3, asb, b, lut, idx, + mask0F, idxsc)); + float s = _mm512_reduce_add_ps( + _mm512_add_ps(_mm512_add_ps(acc0, acc1), _mm512_add_ps(acc2, acc3))); + for (; b < nb; ++b) { // tail (<4 remaining 16-K blocks) + const uint8_t* pk = packed + (size_t)b * 8; + const int8_t* ae = asi8 + (size_t)b * 16; const int8_t* ao = ae + 8; + int isum = 0; + for (int j = 0; j < 8; ++j) + isum += (int)kE2M1x2[pk[j] & 0xF] * (int)ae[j] + (int)kE2M1x2[pk[j] >> 4] * (int)ao[j]; + s += (e4m3[scale[b]] * asb[b]) * (float)isum; + } + return s * (0.5f * global); +} +#endif // avx512vnni available +#endif + +nvdot_fn select_nvdot() { + const IsaTier t = pick_isa(); +#if CPU_MOE_X86 + if (t >= ISA_AVX512) return dot_nvfp4_avx512; + if (t >= ISA_AVX2) return dot_nvfp4_avx2; +#endif + (void)t; + return dot_nvfp4_scalar; +} + +// Best W4A8 (int8-activation) nvfp4 dot, or nullptr if no SIMD VNNI (caller keeps the +// faithful fp32 nvdot path). The scalar i8 dot exists only as a correctness reference. +nvi8dot_fn select_nvi8dot() { +#if CPU_MOE_X86 +#if defined(CPU_MOE_HAS_AVX512VNNI) + if (cpu_has_avx512vnni()) return dot_nvfp4_i8_avx512vnni; +#endif + if (cpu_has_avxvnni()) return dot_nvfp4_i8_vnni; +#endif + return nullptr; +} + +// Quantize the pre-deinterleaved fp32 even/odd activations to per-16-block int8 in the +// [even(8),odd(8)] layout the VNNI dot expects. Done once per token/route (amortized over +// every output row), so a scalar pass is fine relative to the GEMV. +void quant_i8_pg16(const float* xe, const float* xo, int K, int8_t* asi8, float* asb) { + const int nb = K / 16; + for (int b = 0; b < nb; ++b) { + const float* xeb = xe + (size_t)b * 8; + const float* xob = xo + (size_t)b * 8; + float amax = 0.0f; + for (int j = 0; j < 8; ++j) + amax = std::max(amax, std::max(std::fabs(xeb[j]), std::fabs(xob[j]))); + const float s = amax > 0.0f ? amax / 127.0f : 1.0f; + asb[b] = s; + const float inv = 1.0f / s; + int8_t* ae = asi8 + (size_t)b * 16; + for (int j = 0; j < 8; ++j) { + int qe = (int)std::lround(xeb[j] * inv), qo = (int)std::lround(xob[j] * inv); + ae[j] = (int8_t)std::max(-127, std::min(127, qe)); + ae[8 + j] = (int8_t)std::max(-127, std::min(127, qo)); + } + } +} + +class Nvfp4Format final : public MoeFormat { + public: + Nvfp4Format(const MoeParams& p, const FormatArgs& a) + : MoeFormat(p), nvdot_(select_nvdot()), nvi8dot_(select_nvi8dot()) { + gate_up_ = a.bank("gate_up"); + gu_scale_ = a.bank("gate_up_scale"); + gu_global_ = a.bank("gate_up_global"); + down_ = a.bank("down"); + dn_scale_ = a.bank("down_scale"); + dn_global_ = a.bank("down_global"); + const size_t tokens = p.max_tokens, routes = (size_t)p.max_tokens * p.top_k; + xe_.resize(tokens * (p.H / 2)); + xo_.resize(tokens * (p.H / 2)); + ge_.resize(routes * (p.I / 2)); + go_.resize(routes * (p.I / 2)); + if (nvi8dot_ != nullptr) { + xi8_.resize(tokens * p.H); + xas_.resize(tokens * (p.H / 16)); + gi8_.resize(routes * p.I); + gas_.resize(routes * (p.I / 16)); + } + isa_ = tier_name(pick_isa()); + if (nvi8dot_ != nullptr) isa_ += cpu_has_avx512vnni() ? "+avx512vnni(nvfp4-w4a8)" : "+vnni(nvfp4-w4a8)"; + for (int i = 0; i < 16; ++i) e2m1_lut_[i] = kE2M1[i]; + for (int i = 0; i < 256; ++i) e4m3_lut_[i] = e4m3_decode((uint8_t)i); + } + + std::string isa_name() const override { return isa_; } + + + void prep_input(int tok, const bf16_t* x) override { + float* xe = xe_.data() + (size_t)tok * (p.H / 2); + float* xo = xo_.data() + (size_t)tok * (p.H / 2); + deinterleave_bf16_f32(x, xe, xo, p.H); + if (nvi8dot_ != nullptr) + quant_i8_pg16(xe, xo, p.H, xi8_.data() + (size_t)tok * p.H, xas_.data() + (size_t)tok * (p.H / 16)); + } + + bool preps_inter() const override { return true; } + + void prep_inter(int64_t r, bf16_t* g) override { + float* ge = ge_.data() + (size_t)r * (p.I / 2); + float* go = go_.data() + (size_t)r * (p.I / 2); + deinterleave_bf16_f32(g, ge, go, p.I); + if (nvi8dot_ != nullptr) + quant_i8_pg16(ge, go, p.I, gi8_.data() + (size_t)r * p.I, gas_.data() + (size_t)r * (p.I / 16)); + } + + void gate_up(int layer, int e, int tok, int i0, int i1, const bf16_t* /*x*/, float w_in, + bf16_t* g) const override { + const int H = p.H; + const uint8_t* packed = static_cast(tbl_at(gate_up_, layer)); + const uint8_t* scale = static_cast(tbl_at(gu_scale_, layer)); + const uint16_t* global = static_cast(tbl_at(gu_global_, layer)); + const float* xe = xe_.data() + (size_t)tok * (H / 2); + const float* xo = xo_.data() + (size_t)tok * (H / 2); + const int8_t* xi8 = nvi8dot_ != nullptr ? xi8_.data() + (size_t)tok * H : nullptr; + const float* xas = nvi8dot_ != nullptr ? xas_.data() + (size_t)tok * (H / 16) : nullptr; + gated_rows(p, i0, i1, w_in, g, [&](int row) { + const size_t r = (size_t)e * (2 * p.I) + row; + if (nvi8dot_ != nullptr) + return nvi8dot_(packed + r * (size_t)(H / 2), scale + r * (size_t)(H / 16), + fp16_to_f32(global[r]), xi8, H, e4m3_lut_, xas); + return nvdot_(packed + r * (size_t)(H / 2), scale + r * (size_t)(H / 16), + fp16_to_f32(global[r]), xe, xo, H, e2m1_lut_, e4m3_lut_); + }); + } + + void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override { + const int I = p.I; + const uint8_t* packed = static_cast(tbl_at(down_, layer)); + const uint8_t* scale = static_cast(tbl_at(dn_scale_, layer)); + const uint16_t* global = static_cast(tbl_at(dn_global_, layer)); + for (int h = h0; h < h1; ++h) { + for (int j = 0; j < n; ++j) { + const size_t r = (size_t)routes[j].r; + const size_t row = (size_t)routes[j].e * p.H + h; + out[j * kRowBlock + h - h0] = nvi8dot_ != nullptr + ? nvi8dot_(packed + row * (size_t)(I / 2), scale + row * (size_t)(I / 16), fp16_to_f32(global[row]), + gi8_.data() + r * I, I, e4m3_lut_, gas_.data() + r * (I / 16)) + : nvdot_(packed + row * (size_t)(I / 2), scale + row * (size_t)(I / 16), fp16_to_f32(global[row]), + ge_.data() + r * (I / 2), go_.data() + r * (I / 2), I, e2m1_lut_, e4m3_lut_); + } + } + } + + private: + nvdot_fn nvdot_; + nvi8dot_fn nvi8dot_; // AVX-VNNI W4A8 nvfp4 dot (nullptr -> use fp32 nvdot) + // ModelOpt rows: packed e2m1 [E,2I,H/2], e4m3 scales per 16 along K [E,2I,H/16], fp16 row globals [E,2I]; down likewise with K = I. + const uint64_t* gate_up_ = nullptr; + const uint64_t* gu_scale_ = nullptr; + const uint64_t* gu_global_ = nullptr; + const uint64_t* down_ = nullptr; + const uint64_t* dn_scale_ = nullptr; + const uint64_t* dn_global_ = nullptr; + std::string isa_; + float e2m1_lut_[16]; + float e4m3_lut_[256]; + // Activations pre-deinterleaved to fp32 even/odd for the row-major dot; with VNNI also per-16-block int8 + scale. + AlignedVector xe_, xo_, ge_, go_; + AlignedVector xi8_, gi8_; + AlignedVector xas_, gas_; +}; + +} // namespace + +std::unique_ptr make_nvfp4_format(const MoeParams& p, const FormatArgs& a) { + return std::make_unique(p, a); +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/cpu_moe/kernels/q4_0.cpp b/python/freetoken/kernel/csrc/cpu_moe/kernels/q4_0.cpp new file mode 100644 index 0000000000..600a194ebc --- /dev/null +++ b/python/freetoken/kernel/csrc/cpu_moe/kernels/q4_0.cpp @@ -0,0 +1,216 @@ +#include +#include + +#include "../format.h" + +namespace freetoken::cpu_moe { +namespace { + +// -------------------------------- Q4_0 (W4A8) -------------------------------- +// Native GGUF Q4_0 experts (gemma4 GGUF): per-32 block = fp16 scale d + 16 packed +// bytes; byte j holds element j in its low nibble and j+16 in its high nibble, so a +// block's storage order is [lo0..lo15, hi0..hi15] and w = (nibble - 8) * d. Matches +// the reference dequant (models/gguf/dequant.py) and the packed banks the GPU offload +// path streams. +// +// llama.cpp ggml_vec_dot_q4_0_q8_0: W4A8. The activation is pre-quantized to Q8_0 +// (per-32-block int8 ``aq`` + fp32 scale ``asb``); each block unpacks its 16 bytes to +// 32 int8 weights in [-8,7] (bytes_from_nibbles_32: low nibbles -> elems 0..15, high +// -> 16..31) and runs an integer block dot -- VPDPBUSD (AVX-VNNI) or VPMADDUBSW+VPMADDWD +// (AVX2) with the ggml sign trick |w|*(sign(w)*a)=w*a, or a scalar int loop -- then +// scales the block sum by wd*xd in fp32. No fp weight dequant / shuffle chain. The GPU +// offload path (ggml_moe_a8_vec / MMVQ) is also W4A8, so cpu and hybrid stay close. +using q4dot_fn = float (*)(const uint8_t*, const int8_t*, const float*, int); + +float q4_0_dot_i8_scalar(const uint8_t* w, const int8_t* aq, const float* asb, int K) { + float acc = 0.0f; + const int nb = K / 32; + for (int b = 0; b < nb; ++b) { + const uint8_t* blk = w + (size_t)b * 18; + uint16_t dh; + std::memcpy(&dh, blk, sizeof(dh)); + const uint8_t* q = blk + 2; // 16 nibble bytes + const int8_t* a = aq + (size_t)b * 32; + int isum = 0; + for (int j = 0; j < 16; ++j) { + isum += ((int)(q[j] & 0x0F) - 8) * (int)a[j]; // elem j + isum += ((int)(q[j] >> 4) - 8) * (int)a[16 + j]; // elem 16+j + } + acc += fp16_to_f32(dh) * asb[b] * (float)isum; + } + return acc; +} + +#if CPU_MOE_X86 +// fp16 block scale -> fp32 via HW F16C (single value in lane 0). +CPU_MOE_TARGET("f16c") +static inline float q4_scale(uint16_t h) { + return _mm_cvtss_f32(_mm_cvtph_ps(_mm_cvtsi32_si128((int)h))); +} + +// Unpack one Q4_0 block's 16 bytes -> 32 int8 weights in [-8,7] (elems 0..15 = low +// nibbles, 16..31 = high nibbles). ``eight`` = _mm256_set1_epi8(8). +CPU_MOE_TARGET("avx2") +static inline __m256i q4_unpack32(const uint8_t* blk, __m128i mask, __m256i eight) { + const __m128i qb = _mm_loadu_si128(reinterpret_cast(blk + 2)); + const __m128i lo = _mm_and_si128(qb, mask); + const __m128i hi = _mm_and_si128(_mm_srli_epi16(qb, 4), mask); + return _mm256_sub_epi8(_mm256_set_m128i(hi, lo), eight); +} + +// AVX2 W4A8 (llama.cpp non-VNNI mul_sum_i8_pairs): integer block dot via VPMADDUBSW + +// VPMADDWD (sign trick), scaled by wd*xd. |aw*sa| pair sums <= 8*127*2 < 32767 -> no +// int16 saturation. This is the fast path on AVX2 CPUs without AVX-VNNI (and the +// avx512-tier fallback, since the block dot is 256-bit either way). +CPU_MOE_TARGET("avx2,fma,f16c") +float q4_0_dot_i8_avx2(const uint8_t* w, const int8_t* aq, const float* asb, int K) { + const __m128i mask = _mm_set1_epi8(0x0F); + const __m256i eight = _mm256_set1_epi8(8); + const __m256i ones16 = _mm256_set1_epi16(1); + __m256 accF = _mm256_setzero_ps(); + const int nb = K / 32; + for (int b = 0; b < nb; ++b) { + const uint8_t* blk = w + (size_t)b * 18; + _mm_prefetch(reinterpret_cast(blk) + 512, _MM_HINT_T0); + uint16_t dh; + std::memcpy(&dh, blk, sizeof(dh)); + __m256i wq = q4_unpack32(blk, mask, eight); + __m256i a = _mm256_loadu_si256(reinterpret_cast(aq + (size_t)b * 32)); + __m256i aw = _mm256_sign_epi8(wq, wq); // |wq| (unsigned operand) + __m256i sa = _mm256_sign_epi8(a, wq); // sign(wq) * a (signed operand) + __m256i d32 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, sa), ones16); // 8 int32 + accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d32), _mm256_set1_ps(q4_scale(dh) * asb[b]), accF); + } + return hsum256(accF); +} + +// AVX-VNNI W4A8: one VPDPBUSD per block (the fast path on modern CPUs). +CPU_MOE_TARGET("avx2,avxvnni,fma,f16c") +float q4_0_dot_i8_vnni(const uint8_t* w, const int8_t* aq, const float* asb, int K) { + const __m128i mask = _mm_set1_epi8(0x0F); + const __m256i eight = _mm256_set1_epi8(8); + __m256 accF = _mm256_setzero_ps(); + const int nb = K / 32; + for (int b = 0; b < nb; ++b) { + const uint8_t* blk = w + (size_t)b * 18; + _mm_prefetch(reinterpret_cast(blk) + 512, _MM_HINT_T0); + uint16_t dh; + std::memcpy(&dh, blk, sizeof(dh)); + __m256i wq = q4_unpack32(blk, mask, eight); + __m256i a = _mm256_loadu_si256(reinterpret_cast(aq + (size_t)b * 32)); + __m256i aw = _mm256_sign_epi8(wq, wq); // |wq| (unsigned operand) + __m256i sa = _mm256_sign_epi8(a, wq); // sign(wq) * a (signed operand) + __m256i di = _mm256_dpbusd_avx_epi32(_mm256_setzero_si256(), aw, sa); + // All 32 elems of the block share wd*xd; distribute over di's 8 partial sums and + // reduce at the end (equivalent to scale * block_total). + accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(di), _mm256_set1_ps(q4_scale(dh) * asb[b]), accF); + } + return hsum256(accF); +} +#endif // CPU_MOE_X86 + +// All tiers are W4A8 (int8 activations pre-quantized to Q8_0). AVX-VNNI is orthogonal to +// the ISA tier (gated by cpu_has_avxvnni() / FREETOKEN_CPU_MOE_NO_VNNI), so it wins when +// present; otherwise the 256-bit VPMADDUBSW kernel covers both the avx2 and avx512 tiers. +q4dot_fn select_q4dot() { + const IsaTier t = pick_isa(); +#if CPU_MOE_X86 + if (cpu_has_avxvnni()) return q4_0_dot_i8_vnni; + if (t >= ISA_AVX2) return q4_0_dot_i8_avx2; +#endif + (void)t; + return q4_0_dot_i8_scalar; +} + +// Quantize a bf16 activation row to Q8_0 (llama.cpp): per-32-block symmetric int8 in +// natural order + one fp32 scale (amax/127) per block. Done once per token/route, +// amortized over every output row that the q4_0 W4A8 GEMV reads. K % 32 == 0. +void quant_q8_0(const bf16_t* x, int K, int8_t* aq, float* asb) { + const int nb = K / 32; + for (int b = 0; b < nb; ++b) { + const bf16_t* xb = x + (size_t)b * 32; + float xf[32], amax = 0.0f; + for (int j = 0; j < 32; ++j) { + xf[j] = bf16_to_f32(xb[j]); + amax = std::max(amax, std::fabs(xf[j])); + } + const float d = amax > 0.0f ? amax / 127.0f : 1.0f; + asb[b] = d; + const float inv = amax > 0.0f ? 1.0f / d : 0.0f; + int8_t* o = aq + (size_t)b * 32; + for (int j = 0; j < 32; ++j) + o[j] = (int8_t)std::max(-127, std::min(127, (int)std::lround(xf[j] * inv))); + } +} + +class Q4_0Format final : public MoeFormat { + public: + Q4_0Format(const MoeParams& p, const FormatArgs& a) : MoeFormat(p), dot_(select_q4dot()) { + gate_up_ = a.bank("gate_up"); + down_ = a.bank("down"); + const size_t tokens = p.max_tokens, routes = (size_t)p.max_tokens * p.top_k; + xi8_.resize(tokens * p.H); + xas_.resize(tokens * (p.H / 32)); + gi8_.resize(routes * p.I); + gas_.resize(routes * (p.I / 32)); + if (p.H % 32 != 0 || p.I % 32 != 0) + throw std::runtime_error("Q4_0 CPU MoE requires H and I to be multiples of 32"); + gu_row_bytes_ = (p.H / 32) * 18; // K = H (gate_up rows) + dn_row_bytes_ = (p.I / 32) * 18; // K = I (down rows) + // select_q4dot picks VPDPBUSD / VPMADDUBSW / scalar for the tier, so the tag reflects which of those it resolved to. + isa_ = std::string(tier_name(pick_isa())) + (cpu_has_avxvnni() ? "+vnni(q4_0-w4a8)" : "+q4_0-w4a8"); + } + + std::string isa_name() const override { return isa_; } + + + void prep_input(int tok, const bf16_t* x) override { + quant_q8_0(x, p.H, xi8_.data() + (size_t)tok * p.H, xas_.data() + (size_t)tok * (p.H / 32)); + } + + bool preps_inter() const override { return true; } + + void prep_inter(int64_t r, bf16_t* g) override { + quant_q8_0(g, p.I, gi8_.data() + (size_t)r * p.I, gas_.data() + (size_t)r * (p.I / 32)); + } + + void gate_up(int layer, int e, int tok, int i0, int i1, const bf16_t* /*x*/, float w_in, + bf16_t* g) const override { + const uint8_t* packed = static_cast(tbl_at(gate_up_, layer)); + const int8_t* xi8 = xi8_.data() + (size_t)tok * p.H; + const float* xas = xas_.data() + (size_t)tok * (p.H / 32); + gated_rows(p, i0, i1, w_in, g, [&](int row) { + return dot_(packed + ((size_t)e * (2 * p.I) + row) * (size_t)gu_row_bytes_, xi8, xas, p.H); + }); + } + + void down(int layer, int h0, int h1, const DownRoute* routes, int n, float* out) const override { + const uint8_t* packed = static_cast(tbl_at(down_, layer)); + for (int h = h0; h < h1; ++h) { + for (int j = 0; j < n; ++j) { + const size_t r = (size_t)routes[j].r; + out[j * kRowBlock + h - h0] = dot_(packed + ((size_t)routes[j].e * p.H + h) * (size_t)dn_row_bytes_, + gi8_.data() + r * p.I, gas_.data() + r * (p.I / 32), p.I); + } + } + } + + private: + q4dot_fn dot_; + // GGUF Q4_0 rows of 32-element blocks (fp16 scale + 16 nibble bytes): gate_up [E,2I,H/32*18], down [E,H,I/32*18]. + const uint64_t* gate_up_ = nullptr; + const uint64_t* down_ = nullptr; + int gu_row_bytes_ = 0, dn_row_bytes_ = 0; + std::string isa_; + // Q8_0 activations: per-32-block int8 + fp32 scale, for the input and the intermediate. + AlignedVector xi8_, gi8_; + AlignedVector xas_, gas_; +}; + +} // namespace + +std::unique_ptr make_q4_0_format(const MoeParams& p, const FormatArgs& a) { + return std::make_unique(p, a); +} + +} // namespace freetoken::cpu_moe diff --git a/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h index 1898b41189..8abcd1a098 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h +++ b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h @@ -44,6 +44,18 @@ #define cudaGetLastError hipGetLastError #endif +#ifndef cudaMalloc +#define cudaMalloc hipMalloc +#endif + +#ifndef cudaFree +#define cudaFree hipFree +#endif + +#ifndef cudaMemset +#define cudaMemset hipMemset +#endif + #ifndef cudaMallocHost #define cudaMallocHost hipHostMalloc #endif diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index 4918a032fe..82f09a9bb8 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -346,8 +346,9 @@ def _decode_hybrid( alphas=cache.alphas_for_slots(self.layer_id), is_prefill=False, ) - cpu_routed = cpu_routed_early if not _HYBRID_OVERLAP else executor.decode_sync(pending) - return gpu_routed + cpu_routed + if not _HYBRID_OVERLAP: + return gpu_routed + cpu_routed_early + return executor.decode_sync(pending, add_to=gpu_routed) def _prefill_routed( self, diff --git a/python/freetoken/moe/benchbw.py b/python/freetoken/moe/benchbw.py index e70cad31d3..9430ab2a84 100644 --- a/python/freetoken/moe/benchbw.py +++ b/python/freetoken/moe/benchbw.py @@ -490,17 +490,16 @@ def _cpu_moe_step_fns(ex, wl: Workload, E: int): disjoint experts (working set >> LLC) -- the DRAM-bound regime hybrid decode actually runs in; ``run()`` executes one step (releases the GIL inside the extension).""" bs, top_k = 1, wl.top_k - io = ex._io_for(bs) + io = ex._io io["x"].copy_(torch.randn(bs, wl.hidden, dtype=torch.bfloat16) * 0.1) io["w"].copy_(torch.rand(bs, top_k, dtype=torch.float32)) - task = ex._task_for(0, bs) base = torch.arange(top_k, dtype=torch.int32) def set_ids(i: int) -> None: io["ids"].copy_((((i * top_k) + base) % E).view(bs, top_k)) def run() -> None: - ex._ext.run_task(task) + ex._ext.run(0, bs) return set_ids, run @@ -566,7 +565,7 @@ def measure_overlap_bw(fmt: str, wl: Workload, device: torch.device, backend its bandwidth-matched fetch split: fetched : cpu-computed = pcie_ov : cpu_ov. Both sides hammer their own pinned banks flat out for ``seconds`` from a shared - barrier: a worker thread loops bs=1 CPU decode steps (``run_task`` releases the GIL) + barrier: a worker thread loops bs=1 CPU decode steps (``run`` releases the GIL) while the main thread loops full-layer ``copy_missing`` gathers (synchronized per copy so the DMA is really in flight, not just enqueued). Each side reports bytes / its own elapsed; the two windows differ by at most one CPU step + one gather. @@ -663,7 +662,7 @@ def _bench_format(fmt: str, wl: Workload, device: torch.device, threshold: float entry["isa_sweep"] = c["isa_sweep"] # The tier label comes from the bf16 dot chooser, so it can over-report tiers # the actual per-format kernel doesn't distinguish. Disclose the real story. - if fmt == "nvfp4": # WF_NVFP4 is the only use_vnni format + if fmt == "nvfp4": # nvfp4 is the only use_vnni format _note(entry, "isa tier labels are nominal: this W4A8 kernel rides AVX-VNNI " "regardless of the forced tier") elif fmt in ("mxfp4_triton", "ds_fp4", "fp8_block"): diff --git a/python/freetoken/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py index f519687ed6..bb9ee56385 100644 --- a/python/freetoken/moe/cpu_executor.py +++ b/python/freetoken/moe/cpu_executor.py @@ -1,18 +1,15 @@ """Python wrapper around the ``_cpu_moe`` C++ executor (--moe-strategy cpu). -Owns the persistent CPU worker pool, the per-batch-size pinned IO buffers, and -the per-(layer, batch-size) host-func task descriptors. ``decode`` issues the -whole CUDA-graph-capturable sequence on the current stream: +Owns the persistent CPU worker pool and one pinned IO block of ``max_tokens`` rows that +every layer and batch size shares (one task is in flight at a time). ``decode`` issues +the whole CUDA-graph-capturable sequence on the current stream: - D2H (hidden, topk_ids, topk_weights -> pinned) - -> submit host node (cudaLaunchHostFunc: enqueue MoE task to the pool) - -> sync host node (cudaLaunchHostFunc: spin until the pool drains) + D2H (hidden, topk_ids, topk_weights -> pinned IO rows [:bs]) + -> submit (doorbell naming (layer, bs), or a cudaLaunchHostFunc host node) + -> sync (wait for the CPU pool to finish) -> H2D (pinned expert output -> GPU) -Buffers and tasks are allocated lazily per batch size. GraphRunner runs an eager -``model.forward()`` at each batch size immediately before capturing it, so the -first (eager) call materializes the stable pinned buffers + task pointers that -the subsequent capture embeds in its host/memcpy nodes. +With the doorbell on CUDA, a publish kernel does the D2H and the doorbell, and skips the step when no route goes to the CPU; a combine kernel does the H2D (csrc/cpu_moe/gpu_sync_io.h). """ from __future__ import annotations @@ -35,28 +32,27 @@ # on a 75-layer model). Instead the GPU raises a mapped-pinned "ready" flag at submit; a # persistent CPU coordinator (in _cpu_moe) polls it, runs the layer, and sets a "done" # flag the GPU waits on at sync -- no host-func round-trip. Both GPU-side operations are -# STREAM MEMORY OPERATIONS (cuStreamWriteValue64 / cuStreamWaitValue64, resolved from the -# driver at runtime): they execute on the GPU front end with no SM-resident kernel, so +# STREAM MEMORY OPERATIONS (cuStreamWriteValue64 / cuStreamWaitValue64 on CUDA, see +# gpu_sync.h): they execute on the GPU front end with no SM-resident kernel, so # GPU utilization stays truthful during the CPU compute window. (The first cut used a # spin-wait kernel; that pinned reported utilization at 99% and laptop CPU/GPU dynamic # power schedulers responded by clamping the CPU frequency -- a net decode regression on -# power-coupled edge devices.) Each (layer, decode batch size) pair gets its own flag -# slot, so every captured decode graph rides the handshake. Where memops are unavailable -# (Windows WDDM, vGPU, old drivers -- functionally probed at startup) or the slot -# capacity is exceeded, decode keeps the host-func path (functional, slower). A Python -# watchdog thread turns a wedged coordinator into a loud RuntimeError (via err[] + -# raise_if_unhealthy) instead of an indefinite stream stall. +# power-coupled edge devices.) The doorbell word names the (layer, decode batch size) +# task, so one ready/done pair serves every layer and every captured graph (gpu_sync.h). +# Where memops are unavailable (some driver models, vGPU, old drivers -- functionally probed at +# startup), decode keeps the host-func path (functional, slower). ROCm 7.14 cannot replay +# that path from a graph, so capture fails loudly there instead. A Python watchdog thread +# turns a wedged coordinator into a loud RuntimeError (via raise_if_unhealthy) instead of +# an indefinite stream stall. # Caveat: the coordinator busy-polls one core while decode traffic flows (idle backoff # otherwise); FREETOKEN_CPU_MOE_FLAG_SYNC=0 opts out entirely. _FLAG_SYNC = os.getenv("FREETOKEN_CPU_MOE_FLAG_SYNC", "1") != "0" -# Flag slots per MoE layer: covers this many distinct decode batch sizes (captured graph -# sizes plus any eager padded sizes); more than that is unheard of, and the overflow -# just keeps the host-func path for the extra combos. -_FLAG_SLOTS_PER_LAYER = 16 # Pool threads spin up to 50 ms before they park. FREETOKEN_CPU_MOE_SPIN=0 opts out. _SPIN_WAIT = os.getenv("FREETOKEN_CPU_MOE_SPIN", "1") != "0" +# ds_fp4 runs its input FP8 round-trip on the GPU when it can; tests turn this off to compare against the CPU round-trip. +_DSFP4_GPU_PREQUANT = True -# Activation ids must match ActKind in csrc/cpu_moe/cpu_moe_ext.cpp. Id 3 is the +# Activation ids must match ActKind in csrc/cpu_moe/common.h. Id 3 is the # clamped (up + 1) swiglu: "swigluoai" runs it in the generic GEMV epilogue, # "gpt_oss_swiglu" is the same math fused inside the mxfp4 kernel. _ACT_IDS = { @@ -70,10 +66,6 @@ "swiglu_clamp": 4, } -# Weight-format ids must match WFmt in csrc/cpu_moe/cpu_moe_ext.cpp. -_WFMT_IDS = {"bf16": 0, "nvfp4": 1, "mxfp4_triton": 2, "ds_fp4": 3, "q4_0": 4, - "fp8_block": 5} - def compiled_extension_supports(activation: str) -> bool: """Whether the compiled ``_cpu_moe`` extension can serve ``activation`` @@ -157,7 +149,7 @@ def resolve_threads_and_affinity(requested: int) -> tuple[int, list[int]]: class CpuMoeExecutor: """Decode-time CPU expert compute over an ``OffloadMoeCache``'s host banks - (bf16, nvfp4, mxfp4_triton, ds_fp4, q4_0 or fp8_block — see ``_WFMT_IDS`` / ``_resolve_banks``).""" + (the formats ``_cpu_moe.weight_formats()`` lists — see ``_resolve_banks``).""" def __init__( self, @@ -178,26 +170,20 @@ def __init__( from freetoken.moe.legacy_format import canonical_role fmt = fmt or cache.quant_format - if fmt not in _WFMT_IDS: + if fmt not in _cpu_moe.weight_formats(): raise NotImplementedError( f"--moe-strategy cpu/hybrid computes experts on the CPU and supports " - f"{sorted(_WFMT_IDS)} formats, but this checkpoint's experts are " + f"{sorted(_cpu_moe.weight_formats())} formats, but this checkpoint's experts are " f"{fmt!r}; use --moe-strategy offload (GPU-side dequant) instead." ) if activation not in _ACT_IDS: raise NotImplementedError(f"CPU MoE backend: unsupported activation {activation!r}") + if apply_router_weight_on_input and fmt in ("mxfp4_triton", "ds_fp4"): + raise NotImplementedError(f"CPU MoE backend: {fmt} experts apply the router weight on the output only") # ABI probe: a stale prebuilt _cpu_moe.so accepts newer act ids without # error and silently computes the wrong activation in the generic # epilogue -- fail loudly with the rebuild instruction instead. (mxfp4 # handles its act inside the kernel and predates the marker.) - # A prebuilt .so that predates this format would accept the id and then take - # the wrong dequant branch -- silently wrong output rather than a crash. - if _WFMT_IDS[fmt] > getattr(_cpu_moe, "max_weight_format_id", lambda: 4)(): - raise RuntimeError( - f"the compiled _cpu_moe extension predates the {fmt!r} expert format; " - "rebuild it with `python setup.py build_ext --inplace` (or reinstall " - "the wheel) before serving this checkpoint on the cpu/hybrid backend." - ) if _ACT_IDS[activation] >= 3 and fmt != "mxfp4_triton": supported = getattr(_cpu_moe, "max_generic_act_id", lambda: 2)() if _ACT_IDS[activation] > supported: @@ -218,28 +204,42 @@ def __init__( # The per-layer tensors and their pointer tables must outlive the executor # (C++ holds raw addresses into both). self._banks: list[torch.Tensor] = [] - ptrs, (self.H, self.I) = self._resolve_banks( + banks, options, (self.H, self.I) = self._resolve_banks( {canonical_role(name): per_layer for name, per_layer in cache.bank_sources.items()}, fmt ) + # ds_fp4: the reference FP8 activation round-trip is a scalar per-element chain + # that the C++ side runs single-threaded on the CUDA host-callback thread -- + # straight on the decode critical path (~0.3ms/layer at H=4096, every worker and + # the GPU waiting on it). When a GPU is present we run the numerically identical + # round-trip as a captured GPU kernel BEFORE the D2H (see decode_submit) and tell + # the C++ side to skip its own. Measured on DeepSeek-V4-Flash bs=1 decode: + # 12.85 -> 15.65 tok/s, output bit-identical (tests/moe/test_dsfp4_prequant.py). + self._gpu_prequant = fmt == "ds_fp4" and device.type == "cuda" and _DSFP4_GPU_PREQUANT + # the W4A8 activation round-trip block follows the checkpoint's fp8 block (128 on V4, 32 on V4.1) + self._act_block = act_block + if fmt == "ds_fp4": + assert act_block is not None, "ds_fp4 experts need act_block, the checkpoint's fp8 activation quant block" + options.update(act_block=int(act_block), input_prequant=int(self._gpu_prequant)) + if self._gpu_prequant: + logger.info_rank0( + "ds_fp4: input FP8 round-trip moved to the GPU " + "(bit-identical grid; the CPU-side scalar round-trip is skipped)" + ) + # Decide the flag handshake up front (env + device + a functional stream-memop # probe): its coordinator needs a core of its own, which the auto thread sizing # below reserves (a coordinator time-slicing against the GEMV workers measurably # destabilizes throughput on fully-subscribed boxes). self._flag_sync = _FLAG_SYNC and device.type == "cuda" - self._cpu_moe = _cpu_moe # module ref for the decode-path memop calls - if self._flag_sync: - probe_scratch = alloc_pinned_tensor(1, dtype=torch.int64) - probe_scratch.zero_() - if not _cpu_moe.memops_probe( - torch.cuda.current_stream().cuda_stream, probe_scratch.data_ptr() - ): - logger.info_rank0( - "cpu-moe flag handshake unavailable: CUDA stream memory operations " - "are not supported here (Windows WDDM / vGPU / old driver); using " - "the cudaLaunchHostFunc sync" - ) - self._flag_sync = False + if self._flag_sync and not _cpu_moe.memops_probe(torch.cuda.current_stream().cuda_stream): + logger.info_rank0( + "cpu-moe flag handshake unavailable: stream memory operations failed the " + "startup probe; using the cudaLaunchHostFunc sync" + ) + self._flag_sync = False + # ROCm 7.14 returns wrong results when it replays host-function nodes, so there only the doorbell is graph-safe. + self._graph_capture_safe = self._flag_sync or _cpu_moe.host_func_capture_safe() nthreads, core_ids = resolve_threads_and_affinity(num_threads) coord_core = -1 @@ -250,6 +250,17 @@ def __init__( nthreads -= 1 core_ids = core_ids[:-1] self._coord_core = coord_core + alloc_io = ( + (lambda *shape, dtype: torch.empty(*shape, dtype=dtype)) + if device.type == "cpu" + else alloc_pinned_tensor + ) + self._io = { + "x": alloc_io(self.max_tokens, self.H, dtype=torch.bfloat16), + "ids": alloc_io(self.max_tokens, self.top_k, dtype=torch.int32), + "w": alloc_io(self.max_tokens, self.top_k, dtype=torch.float32), + "y": alloc_io(self.max_tokens, self.H, dtype=torch.bfloat16), + } self._ext = _cpu_moe.CpuMoeExecutor( num_threads=nthreads, num_layers=self.num_layers, @@ -260,11 +271,16 @@ def __init__( max_tokens=self.max_tokens, activation_id=_ACT_IDS[activation], apply_router_weight_on_input=1 if apply_router_weight_on_input else 0, - weight_format=_WFMT_IDS[fmt], + weight_format=fmt, + banks=banks, + options=options, swiglu_alpha=float(swiglu_alpha), swiglu_limit=float(swiglu_limit) if swiglu_limit is not None else float("inf"), core_ids=core_ids, - **ptrs, + x_ptr=self._io["x"].data_ptr(), + ids_ptr=self._io["ids"].data_ptr(), + w_ptr=self._io["w"].data_ptr(), + y_ptr=self._io["y"].data_ptr(), ) self.num_threads = nthreads self.core_ids = core_ids @@ -281,34 +297,24 @@ def __init__( ) torch.set_num_threads(clamp) - self._io: dict[int, dict[str, torch.Tensor]] = {} - self._tasks: dict[tuple[int, int], int] = {} - - # Flag-based handshake: mapped-pinned ready/done/err int64 arrays (one slot per - # (MoE layer, decode batch size) pair, allocated as tasks are created) + a - # persistent CPU coordinator that polls ready[], runs the slot's task on the - # pool, and sets done[]. Binary per-step protocol (GPU memops: done=0, ready=1; - # coordinator: consume ready, run, done=1; GPU waits done>=1 -- the WAIT - # immediate is constant, so CUDA-graph replays are safe). err[] is raised by the - # WATCHDOG thread when a ready flag stays unanswered (dead coordinator): it - # poisons done to unblock the stream and raise_if_unhealthy() turns the step - # into a loud error instead of silent stale output. Buffers are kept alive on - # self so the coordinator's pinned pointers stay valid for the executor's - # lifetime (flag_sync itself was decided above, before thread sizing). - self._ready = self._done = self._err = None - self._flag_slots: dict[tuple[int, int], int] = {} # (layer_id, bs) -> slot - self._flag_capacity = self.num_layers * _FLAG_SLOTS_PER_LAYER + # Flag-based handshake: the extension owns the ready/done words (gpu_sync.h) and a + # persistent CPU coordinator that polls the doorbell, runs the task it names on the + # pool, and sets done for the GPU. The WATCHDOG thread marks the executor unhealthy + # when a doorbell stays unanswered (dead coordinator): it poisons done to unblock the + # stream and raise_if_unhealthy() turns the step into a loud error instead of silent + # stale output (flag_sync itself was decided above, before thread sizing). + self._unhealthy = False + self._kernel_io = False if self._flag_sync: - self._ready = alloc_pinned_tensor(self._flag_capacity, dtype=torch.int64) - self._done = alloc_pinned_tensor(self._flag_capacity, dtype=torch.int64) - self._err = alloc_pinned_tensor(self._flag_capacity, dtype=torch.int64) - self._ready.zero_() - self._done.zero_() - self._err.zero_() - self._ext.start_flag_coordinator( - self._ready.data_ptr(), self._done.data_ptr(), self._flag_capacity, - self._coord_core, - ) + self._ext.start_flag_coordinator(self._coord_core) + # ROCm keeps the copies until the kernel IO has been checked on AMD GPUs. + if torch.version.hip is None: + reason = _cpu_moe.kernel_io_probe(torch.cuda.current_stream().cuda_stream) + if reason: + logger.warning_rank0(f"cpu-moe kernel IO unavailable ({reason}); using D2H/H2D copies around the doorbell") + else: + self._ext.enable_kernel_io() + self._kernel_io = True self._watchdog_stop = False # The thread target holds a WEAKREF and re-derefs it each tick: a bound # method would strong-reference the executor forever (the loop never ends @@ -325,32 +331,14 @@ def __init__( ) self._watchdog.start() - # ds_fp4: the reference FP8 activation round-trip is a scalar per-element chain - # that the C++ side runs single-threaded on the CUDA host-callback thread -- - # straight on the decode critical path (~0.3ms/layer at H=4096, every worker and - # the GPU waiting on it). When a GPU is present we run the numerically identical - # round-trip as a captured GPU kernel BEFORE the D2H (see decode_submit) and tell - # the C++ side to skip its own. Measured on DeepSeek-V4-Flash bs=1 decode: - # 12.85 -> 15.65 tok/s, output bit-identical (tests/moe/test_dsfp4_prequant.py). - self._gpu_prequant = fmt == "ds_fp4" and device.type == "cuda" - # the W4A8 activation round-trip block follows the checkpoint's fp8 block (128 on V4, 32 on V4.1) - self._act_block = act_block - if fmt == "ds_fp4": - assert act_block is not None, "ds_fp4 experts need act_block, the checkpoint's fp8 activation quant block" - self._ext.set_act_block(act_block) - if self._gpu_prequant: - self._ext.set_input_prequant(True) - logger.info_rank0( - "ds_fp4: input FP8 round-trip moved to the GPU " - "(bit-identical grid; the CPU-side scalar round-trip is skipped)" - ) logger.info_rank0( f"CPU MoE executor ready: threads={nthreads} (pinned to cores " f"{core_ids[0]}..{core_ids[-1]}) isa={self.isa} fmt={fmt} " f"H={self.H} I={self.I} experts={self.num_experts} layers={self.num_layers} " f"top_k={self.top_k} act={activation} max_tokens={self.max_tokens} " - f"spin_wait={self.spin_wait}" + f"spin_wait={self.spin_wait} " + f"handshake={'kernel' if self._kernel_io else 'memop' if self._flag_sync else 'host-func'}" ) def _make_table(self, layers: list[torch.Tensor]) -> torch.Tensor: @@ -368,7 +356,11 @@ def _make_table(self, layers: list[torch.Tensor]) -> torch.Tensor: self._banks.extend(layers) return table - def _resolve_fp8_block_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: + def _tables(self, **roles: list[torch.Tensor]) -> dict[str, int]: + """Per-layer table addresses keyed by canonical bank role.""" + return {role: self._make_table(layers).data_ptr() for role, layers in roles.items()} + + def _resolve_fp8_block_banks(self, banks: dict) -> tuple[dict, dict, tuple[int, int]]: """Block-FP8 schema: row-major e4m3 + one bf16 scale per 128x128 block; the scale banks pad their last dim, so C++ gets the real row stride.""" gu, gus = banks["gate_up"], banks["gate_up_scale"] dn, dns = banks["down"], banks["down_scale"] @@ -380,28 +372,19 @@ def _resolve_fp8_block_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: assert tuple(dn[0].shape[1:]) == (H, I), (dn[0].shape, H, I) assert gus[0].shape[1] == (2 * I + 127) // 128 and gus[0].shape[2] >= (H + 127) // 128, (gus[0].shape, H, I) assert dns[0].shape[1] == (H + 127) // 128 and dns[0].shape[2] >= (I + 127) // 128, (dns[0].shape, H, I) - ptrs = dict( - gate_up_ptr=self._make_table(gu).data_ptr(), - down_ptr=self._make_table(dn).data_ptr(), - gate_up_scale_ptr=self._make_table(gus).data_ptr(), - gate_up_global_ptr=0, - down_scale_ptr=self._make_table(dns).data_ptr(), - down_global_ptr=0, - gate_up_bias_ptr=0, - down_bias_ptr=0, - fp8_block_gu_scale_stride=int(gus[0].stride(1)), - fp8_block_dn_scale_stride=int(dns[0].stride(1)), - ) - return ptrs, (H, I) + banks = self._tables(gate_up=gu, down=dn, gate_up_scale=gus, down_scale=dns) + options = dict(gu_scale_stride=int(gus[0].stride(1)), dn_scale_stride=int(dns[0].stride(1))) + return banks, options, (H, I) - def _resolve_banks(self, banks: dict, fmt: str) -> tuple[dict, tuple[int, int]]: - """Return (pointer kwargs for the C++ ctor, (H, I)) for the given format. + def _resolve_banks(self, banks: dict, fmt: str) -> tuple[dict, dict, tuple[int, int]]: + """Return (banks, options, (H, I)) for the C++ ctor: banks maps each canonical role + the format reads to its per-layer table's address, options holds its integer settings. ``banks[name]`` is a list of ``num_layers`` ``[num_experts, ...]`` tensors (the per-layer host bank contract); shapes are read from the first layer so - per-partition (TP) sizes are exact. Unused pointers are 0. Every pointer kwarg - is actually a per-layer table's address (see ``_make_table``), not a single - bank's -- the C++ ctor resolves ``tbl[layer_id]`` per task. + per-partition (TP) sizes are exact. Every bank address is a per-layer table + (see ``_make_table``), not a single bank's -- the C++ side resolves + ``tbl[layer_id]`` per task. """ if fmt == "bf16": gate_up = banks["gate_up"] @@ -414,17 +397,8 @@ def _resolve_banks(self, banks: dict, fmt: str) -> tuple[dict, tuple[int, int]]: I = int(gate_up[0].shape[1] // 2) assert gate_up[0].shape[1] == 2 * I assert tuple(down[0].shape[1:]) == (H, I), (down[0].shape, H, I) - ptrs = dict( - gate_up_ptr=self._make_table(gate_up).data_ptr(), - down_ptr=self._make_table(down).data_ptr(), - gate_up_scale_ptr=0, - gate_up_global_ptr=0, - down_scale_ptr=0, - down_global_ptr=0, - gate_up_bias_ptr=0, - down_bias_ptr=0, - ) - return ptrs, (H, I) + banks = self._tables(gate_up=gate_up, down=down) + return banks, {}, (H, I) if fmt == "q4_0": return self._resolve_q4_0_banks(banks) @@ -452,19 +426,10 @@ def _resolve_banks(self, banks: dict, fmt: str) -> tuple[dict, tuple[int, int]]: assert tuple(gus[0].shape[1:]) == (2 * I, H // 16), (gus[0].shape, I, H) assert tuple(dns[0].shape[1:]) == (H, I // 16), (dns[0].shape, H, I) assert tuple(gug[0].shape[1:]) == (2 * I,) and tuple(dng[0].shape[1:]) == (H,) - ptrs = dict( - gate_up_ptr=self._make_table(gup).data_ptr(), - down_ptr=self._make_table(dnp).data_ptr(), - gate_up_scale_ptr=self._make_table(gus).data_ptr(), - gate_up_global_ptr=self._make_table(gug).data_ptr(), - down_scale_ptr=self._make_table(dns).data_ptr(), - down_global_ptr=self._make_table(dng).data_ptr(), - gate_up_bias_ptr=0, - down_bias_ptr=0, - ) - return ptrs, (H, I) + banks = self._tables(gate_up=gup, down=dnp, gate_up_scale=gus, gate_up_global=gug, down_scale=dns, down_global=dng) + return banks, {}, (H, I) - def _resolve_q4_0_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: + def _resolve_q4_0_banks(self, banks: dict) -> tuple[dict, dict, tuple[int, int]]: """Native GGUF Q4_0 schema (gemma4 GGUF): per-32 blocks (fp16 scale + 16 nibble bytes), row-major over K -- the *same* packed banks the GPU offload path streams. gate_up is [S, 2I, H//32*18], down is [S, H, I//32*18]; the C++ W4A16 GEMV reads a @@ -479,19 +444,10 @@ def _resolve_q4_0_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: assert H % 32 == 0 and I % 32 == 0, (H, I) assert int(gate_up[0].shape[2]) == (H // 32) * 18, (gate_up[0].shape, H) assert int(down[0].shape[2]) == (I // 32) * 18, (down[0].shape, I) - ptrs = dict( - gate_up_ptr=self._make_table(gate_up).data_ptr(), - down_ptr=self._make_table(down).data_ptr(), - gate_up_scale_ptr=0, - gate_up_global_ptr=0, - down_scale_ptr=0, - down_global_ptr=0, - gate_up_bias_ptr=0, - down_bias_ptr=0, - ) - return ptrs, (H, I) + banks = self._tables(gate_up=gate_up, down=down) + return banks, {}, (H, I) - def _resolve_mxfp4_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: + def _resolve_mxfp4_banks(self, banks: dict) -> tuple[dict, dict, tuple[int, int]]: """gpt-oss mxfp4 ``mxfp4_triton`` schema: transposed split-K blocks/scales (N innermost) + per-output-row biases. The C++ kernel streams K and accumulates a contiguous N-block, so the GPU-tiled layout is read in place @@ -510,19 +466,10 @@ def _resolve_mxfp4_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: assert tuple(gus[0].shape[1:]) == (H // 32, 2 * I), (gus[0].shape, H, I) assert tuple(dns[0].shape[1:]) == (I // 32, H), (dns[0].shape, H, I) assert tuple(gob[0].shape[1:]) == (2 * I,) and tuple(dob[0].shape[1:]) == (H,) - ptrs = dict( - gate_up_ptr=self._make_table(gub).data_ptr(), - down_ptr=self._make_table(dnb).data_ptr(), - gate_up_scale_ptr=self._make_table(gus).data_ptr(), - gate_up_global_ptr=0, - down_scale_ptr=self._make_table(dns).data_ptr(), - down_global_ptr=0, - gate_up_bias_ptr=self._make_table(gob).data_ptr(), - down_bias_ptr=self._make_table(dob).data_ptr(), - ) - return ptrs, (H, I) + banks = self._tables(gate_up=gub, down=dnb, gate_up_scale=gus, down_scale=dns, gate_up_bias=gob, down_bias=dob) + return banks, {}, (H, I) - def _resolve_dsfp4_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: + def _resolve_dsfp4_banks(self, banks: dict) -> tuple[dict, dict, tuple[int, int]]: """DeepSeek-V4 ``ds_fp4`` schema: row-major e2m1 (2/byte) + e8m0 per-32 block scales, no global, no bias. Layout matches nvfp4 (K contiguous per output row), so the C++ GEMV reads it in place. The kernel additionally FP8-round-trips the @@ -538,57 +485,8 @@ def _resolve_dsfp4_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: assert tuple(dnp[0].shape[1:]) == (H, I // 2), (dnp[0].shape, H, I) assert tuple(gus[0].shape[1:]) == (2 * I, H // 32), (gus[0].shape, I, H) assert tuple(dns[0].shape[1:]) == (H, I // 32), (dns[0].shape, H, I) - ptrs = dict( - gate_up_ptr=self._make_table(gup).data_ptr(), - down_ptr=self._make_table(dnp).data_ptr(), - gate_up_scale_ptr=self._make_table(gus).data_ptr(), - gate_up_global_ptr=0, - down_scale_ptr=self._make_table(dns).data_ptr(), - down_global_ptr=0, - gate_up_bias_ptr=0, - down_bias_ptr=0, - ) - return ptrs, (H, I) - - def _io_for(self, bs: int) -> dict[str, torch.Tensor]: - io = self._io.get(bs) - if io is None: - alloc_io = ( - (lambda *shape, dtype: torch.empty(*shape, dtype=dtype)) - if self.device.type == "cpu" - else alloc_pinned_tensor - ) - io = { - "x": alloc_io(bs, self.H, dtype=torch.bfloat16), - "ids": alloc_io(bs, self.top_k, dtype=torch.int32), - "w": alloc_io(bs, self.top_k, dtype=torch.float32), - "y": alloc_io(bs, self.H, dtype=torch.bfloat16), - } - self._io[bs] = io - return io - - def _task_for(self, layer_id: int, bs: int) -> int: - key = (layer_id, bs) - task = self._tasks.get(key) - if task is None: - io = self._io_for(bs) - task = self._ext.create_task( - layer_id, - bs, - io["x"].data_ptr(), - io["ids"].data_ptr(), - io["w"].data_ptr(), - io["y"].data_ptr(), - ) - self._tasks[key] = task - # Allocate this (layer, bs) combo a flag slot and register its task with the - # coordinator. Combos past the slot capacity keep the host-func path. - if self._flag_sync and key not in self._flag_slots: - slot = len(self._flag_slots) - if slot < self._flag_capacity: - self._flag_slots[key] = slot - self._ext.register_flag_task(slot, task) - return task + banks = self._tables(gate_up=gup, down=dnp, gate_up_scale=gus, down_scale=dns) + return banks, {}, (H, I) def decode( self, @@ -612,7 +510,7 @@ def decode_submit( topk_weights: torch.Tensor, topk_ids: torch.Tensor, ) -> tuple: - """Issue the D2H copies + the CPU-pool submit host node, then return without + """Issue the D2H copies + the CPU-pool submit host node (or the publish kernel), then return without waiting. Lets a caller (the hybrid backend) enqueue GPU work between this and :meth:`decode_sync` so the CPU compute overlaps the GPU GEMM / PCIe fetch. @@ -621,7 +519,14 @@ def decode_submit( :meth:`decode_sync`. The output tensor is allocated here so it stays live (and distinct from the interleaved GPU work) across the overlap window.""" bs = hidden_states.shape[0] - io = self._io_for(bs) + if bs > self.max_tokens: + raise ValueError(f"CPU MoE decode batch {bs} exceeds the executor's max_tokens {self.max_tokens}") + if not self._graph_capture_safe and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "CPU MoE decode cannot be captured into a CUDA graph here: the stream-memop " + "doorbell is off or failed its startup probe, and the host-function sync is not " + "replay-safe on this runtime; run with --cuda-graph-max-bs 0" + ) if self._gpu_prequant: # DSV4: apply the reference FP8 round-trip on the GPU (the same kernel the @@ -631,91 +536,86 @@ def decode_submit( hidden_states = act_quant_fp8_roundtrip(hidden_states, block=self._act_block) + stream = torch.cuda.current_stream().cuda_stream + if self._kernel_io and hidden_states.dtype == torch.bfloat16: + x = hidden_states.contiguous() + ids = topk_ids.to(torch.int32).contiguous() + w = topk_weights.to(torch.float32).contiguous() + assert x.shape == (bs, self.H) and ids.shape == w.shape == (bs, self.top_k), (x.shape, ids.shape, w.shape) + self._ext.flag_publish(stream, x.data_ptr(), ids.data_ptr(), w.data_ptr(), layer_id, bs) + return (bs, torch.empty_like(x), True) + # D2H: ship this step's activations + routing to pinned host memory. - io["x"].copy_(hidden_states, non_blocking=True) - io["ids"].copy_(topk_ids.to(torch.int32), non_blocking=True) - io["w"].copy_(topk_weights.to(torch.float32), non_blocking=True) + io = self._io + io["x"][:bs].copy_(hidden_states, non_blocking=True) + io["ids"][:bs].copy_(topk_ids.to(torch.int32), non_blocking=True) + io["w"][:bs].copy_(topk_weights.to(torch.float32), non_blocking=True) - task = self._task_for(layer_id, bs) out = torch.empty_like(hidden_states) - slot = self._flag_slots.get((layer_id, bs)) if self._flag_sync else None - if slot is not None: - # Front-end memops: done[slot]=0 then ready[slot]=1 (the coordinator's - # doorbell). No kernel launched; no host-func round trip. - self._cpu_moe.memop_submit( - torch.cuda.current_stream().cuda_stream, - self._done.data_ptr(), self._ready.data_ptr(), slot, - ) + if self._flag_sync: + # Front-end memops: reset done, then ring the doorbell with (layer, bs). + # No kernel launched; no host-func round trip. + self._ext.flag_signal(stream, layer_id, bs) else: - stream = torch.cuda.current_stream().cuda_stream - self._ext.submit_with_cuda_stream(stream, task) - return (bs, task, out, slot) + self._ext.submit_with_cuda_stream(stream, layer_id, bs) + return (bs, out, False) - def decode_sync(self, pending: tuple) -> torch.Tensor: + def decode_sync(self, pending: tuple, add_to: torch.Tensor | None = None) -> torch.Tensor: """Issue the CPU-pool sync + the H2D result copy for a prior :meth:`decode_submit`, and return the GPU output tensor. With flag-sync the wait is a front-end stream - memop on done[slot] (set by the CPU coordinator); otherwise a cudaLaunchHostFunc.""" - bs, task, out, slot = pending - if slot is not None: - # Front-end WAIT(done[slot] >= 1): blocks this stream's later nodes without + memop on done (set by the CPU coordinator); otherwise a cudaLaunchHostFunc. + With ``add_to``, return ``add_to`` + the CPU output; the kernel path adds in place into ``add_to``.""" + bs, out, kernel_io = pending + stream = torch.cuda.current_stream().cuda_stream + if self._flag_sync: + # Front-end WAIT(done >= 1): blocks this stream's later nodes without # occupying an SM, so GPU utilization stays truthful during the CPU window. - self._cpu_moe.memop_sync( - torch.cuda.current_stream().cuda_stream, self._done.data_ptr(), slot, - ) + self._ext.flag_wait(stream) else: - stream = torch.cuda.current_stream().cuda_stream - self._ext.sync_with_cuda_stream(stream, task) - io = self._io[bs] - out.copy_(io["y"], non_blocking=True) - return out - - def _watchdog_tick(self, suspects: dict) -> None: + self._ext.sync_with_cuda_stream(stream) + if kernel_io: + if add_to is not None and add_to.dtype == torch.bfloat16 and add_to.is_contiguous() and add_to.shape == out.shape: + self._ext.flag_combine(stream, add_to.data_ptr(), bs, True) + return add_to + self._ext.flag_combine(stream, out.data_ptr(), bs, False) + return out if add_to is None else add_to + out + out.copy_(self._io["y"][:bs], non_blocking=True) + return out if add_to is None else add_to + out + + def _watchdog_tick(self, suspect: dict) -> None: """One watchdog sampling round (called every 2 s by ``_watchdog_main``). - A slot is only declared dead when THREE things hold across >=10 s: its doorbell - is still pending (ready==1 && done==0), it was already pending when first - suspected, and the coordinator has served NOTHING on it since (flag_served_count - unchanged). The served-count criterion kills the false-positive window: two - point samples can land on the same slot's (independent, us-scale) pending - windows under heavy external load, but a coordinator that made progress in - between is alive by definition. ``suspects`` maps slot -> (first_seen, - served_at_first_sight) and persists across ticks.""" - stuck = (self._ready == 1) & (self._done == 0) - if not bool(stuck.any()): - suspects.clear() + The doorbell is only declared dead when THREE things hold across >=10 s: it is + still pending (ready set, done==0), it was already pending when first suspected, and + the coordinator has served NOTHING since (flag_served unchanged). The served-count + criterion kills the false-positive window: two point samples can land on + (independent, us-scale) pending windows under heavy external load, but a + coordinator that made progress in between is alive by definition. ``suspect`` + holds (first_seen, served_at_first_sight) and persists across ticks.""" + if not self._ext.flag_pending(): + suspect.clear() return now = time.monotonic() - pending = set(stuck.nonzero().flatten().tolist()) - for slot in list(suspects): - if slot not in pending: - del suspects[slot] - dead = [] - for slot in pending: - served = self._ext.flag_served_count(slot) - first_seen, served_then = suspects.get(slot, (None, None)) - if first_seen is None or served != served_then: - suspects[slot] = (now, served) # new suspect, or alive-but-loaded: rearm - continue - if now - first_seen >= 10.0: - dead.append(slot) - if not dead: + served = self._ext.flag_served() + if not suspect or suspect["served"] != served: + suspect.update(first_seen=now, served=served) # new suspect, or alive-but-loaded: rearm + return + if now - suspect["first_seen"] < 10.0: return logger.error( - f"cpu-moe flag watchdog: slots {dead} unanswered for >10s with no coordinator " - "progress (wedged/dead); poisoning done[] and failing the next step" + "cpu-moe flag watchdog: doorbell unanswered for >10s with no coordinator " + "progress (wedged/dead); poisoning done and failing the next step" ) - for i in dead: - self._err[i] = 1 - for i in dead: - self._done[i] = 1 # after err: unblock the stream into a checked failure - suspects.pop(i, None) + self._unhealthy = True + self._ext.poison_flag() # after the mark: unblock the stream into a checked failure + suspect.clear() def raise_if_unhealthy(self) -> None: """Raise if the flag watchdog fired (a doorbell stayed unanswered because the coordinator never responded). Called by the engine once per forward -- a single - pinned read -- so a dead coordinator surfaces as a loud error on the next step + flag read -- so a dead coordinator surfaces as a loud error on the next step instead of silently shipping stale expert outputs.""" - if self._err is not None and bool((self._err != 0).any()): + if self._unhealthy: raise RuntimeError( "CPU MoE flag-handshake watchdog fired: a decode step's doorbell was " "never answered by the coordinator thread (its outputs cannot be " @@ -728,13 +628,13 @@ def raise_if_unhealthy(self) -> None: def _watchdog_main(executor_ref) -> None: """Watchdog daemon body: weakref-deref per tick so the thread never keeps a dead executor alive (see the start site in ``CpuMoeExecutor.__init__``).""" - suspects: dict = {} + suspect: dict = {} while True: time.sleep(2.0) executor = executor_ref() - if executor is None or executor._watchdog_stop or executor._ready is None: + if executor is None or executor._watchdog_stop: return try: - executor._watchdog_tick(suspects) + executor._watchdog_tick(suspect) finally: del executor # drop the strong ref before the next sleep diff --git a/setup.py b/setup.py index b8e4693c89..b968373151 100644 --- a/setup.py +++ b/setup.py @@ -77,15 +77,22 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: if IS_ROCM: runtime_include_dirs, runtime_library_dirs, runtime_lib = _rocm_paths() runtime_link_args = [f"-Wl,-rpath,{runtime_library_dirs[0]}"] - # These extensions contain host code only. BuildExtension supplies the ROCm + # Host sources. BuildExtension supplies the ROCm # platform defines to the C++ compiler; offload architecture flags belong on # HIP device sources and would be rejected by the host compiler here. extra_compile = ["-O3", "-std=c++20"] + # Device sources (.cu, built by hipcc); BuildExtension adds the offload architectures. + device_compile = ["-O3", "-std=c++20"] else: runtime_include_dirs, runtime_library_dirs = _cuda_runtime_paths() runtime_lib = "cudart" runtime_link_args = [] extra_compile = ["-O3", "-std=c++20"] + # Device sources (.cu, built by nvcc): SASS for each major arch, which also runs on the newer minors of that major, + # plus PTX for the oldest arch nvcc 13 supports, which the driver builds for any other GPU at load. + # Naming archs also stops torch from picking them, which fails on a build machine without a GPU. + device_compile = ["-O3", "-std=c++20"] + [f"-gencode=arch=compute_{a},code=sm_{a}" for a in (75, 80, 90, 100, 120)] + device_compile.append("-gencode=arch=compute_75,code=compute_75") _check_toolchain() @@ -112,11 +119,17 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: name="freetoken.kernel._cpu_moe", sources=[ "python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp", + f"python/freetoken/kernel/csrc/cpu_moe/gpu_sync_{'hip' if IS_ROCM else 'cuda'}.cpp", + "python/freetoken/kernel/csrc/cpu_moe/gpu_sync_io.cu", + *sorted( + str(p.relative_to(ROOT)) + for p in (ROOT / "python/freetoken/kernel/csrc/cpu_moe/kernels").glob("*.cpp") + ), ], include_dirs=[KERNEL_INCLUDE, *runtime_include_dirs], library_dirs=runtime_library_dirs, libraries=[runtime_lib], - extra_compile_args=extra_compile + ["-pthread"], + extra_compile_args={"cxx": extra_compile + ["-pthread"], "nvcc": device_compile}, extra_link_args=runtime_link_args, ), # disk-backed row store (PLE / Engram tables); Linux-only until the TableFile/BatchReader seams grow Windows bodies diff --git a/tests/moe/hip_probe_faults.cpp b/tests/moe/hip_probe_faults.cpp new file mode 100644 index 0000000000..4426bf0d14 --- /dev/null +++ b/tests/moe/hip_probe_faults.cpp @@ -0,0 +1,75 @@ +#include +#include + +#include +#include +#include + +namespace { +std::atomic begins{0}, ends{0}, injected{0}, removed_nodes{0}, last_end_rc{0}; + +bool mode(const char* value) { + const char* selected = std::getenv("FREETOKEN_TEST_HIP_PROBE_FAULT"); + return selected != nullptr && std::strcmp(selected, value) == 0; +} + +template +Function real(const char* name) { + return reinterpret_cast(dlsym(RTLD_NEXT, name)); +} +} // namespace + +extern "C" hipError_t hipStreamBeginCapture(hipStream_t stream, hipStreamCaptureMode capture_mode) { + const hipError_t rc = real("hipStreamBeginCapture")(stream, capture_mode); + if (rc == hipSuccess) ++begins; + return rc; +} + +extern "C" hipError_t hipStreamWriteValue64(hipStream_t stream, void* address, uint64_t value, unsigned flags) { + // Force the compatibility path on runtimes that normally record native writes. + if (mode("dependency") && begins == 1) { + ++injected; + return hipSuccess; + } + return real("hipStreamWriteValue64")(stream, address, value, flags); +} + +extern "C" hipError_t hipStreamUpdateCaptureDependencies( + hipStream_t stream, hipGraphNode_t* dependencies, size_t count, unsigned flags) { + if (mode("dependency")) { + ++injected; + return hipErrorNotSupported; + } + return real("hipStreamUpdateCaptureDependencies")( + stream, dependencies, count, flags); +} + +extern "C" hipError_t hipGraphDestroyNode(hipGraphNode_t node) { + const hipError_t rc = real("hipGraphDestroyNode")(node); + if (rc == hipSuccess) ++removed_nodes; + return rc; +} + +extern "C" hipError_t hipStreamEndCapture(hipStream_t stream, hipGraph_t* graph) { + ++ends; + hipError_t rc = real("hipStreamEndCapture")(stream, graph); + if (mode("end-capture") && ends == 2 && rc == hipSuccess) { + // A real invalid API call also sets HIP's thread-local last-error state. + rc = real("hipGraphGetNodes")(*graph, nullptr, nullptr); + ++injected; + } + last_end_rc = static_cast(rc); + return rc; +} + +extern "C" int probe_test_injected() { return injected; } +extern "C" int probe_test_begins() { return begins; } +extern "C" int probe_test_ends() { return ends; } +extern "C" int probe_test_removed_nodes() { return removed_nodes; } +extern "C" int probe_test_last_end_rc() { return last_end_rc; } +extern "C" int probe_test_peek_error() { + return static_cast(real("hipPeekAtLastError")()); +} +extern "C" int probe_test_make_error() { + return static_cast(real("hipGraphGetNodes")(nullptr, nullptr, nullptr)); +} diff --git a/tests/moe/test_cpu_moe.py b/tests/moe/test_cpu_moe.py index 547c94c60a..38e79b2743 100644 --- a/tests/moe/test_cpu_moe.py +++ b/tests/moe/test_cpu_moe.py @@ -440,6 +440,111 @@ def _reference(cache, layer, hidden, ids, w, act="silu", apply_in=False): return out +@pytest.fixture(scope="module") +def hip_probe_fault_library(tmp_path_factory): + import importlib.util + import os + from pathlib import Path + import shutil + import subprocess + + if not torch.version.hip: + pytest.skip("requires the HIP sync backend") + compiler = shutil.which("c++") + if compiler is None: + pytest.skip("requires a C++ compiler for HIP fault injection") + candidates = [Path(os.getenv("ROCM_HOME", "/opt/rocm")), Path("/opt/rocm")] + spec = importlib.util.find_spec("_rocm_sdk_core") + if spec and spec.submodule_search_locations: + candidates.append(Path(next(iter(spec.submodule_search_locations)))) + include = next((root / "include" for root in candidates if (root / "include/hip/hip_runtime.h").exists()), None) + if include is None: + pytest.skip("requires HIP SDK headers for fault injection") + library = tmp_path_factory.mktemp("hip_probe") / "faults.so" + command = [ + compiler, "-std=c++20", "-shared", "-fPIC", "-D__HIP_PLATFORM_AMD__=1", + f"-I{include}", str(Path(__file__).with_name("hip_probe_faults.cpp")), + "-ldl", "-o", str(library), + ] + result = subprocess.run(command, capture_output=True, text=True, timeout=60) + assert result.returncode == 0, result.stdout + result.stderr + return library + + +@pytest.mark.parametrize("fault", ["dependency", "end-capture", "pending"]) +def test_cpu_moe_hip_probe_error_recovery(hip_probe_fault_library, fault): + """Rejected probes leave eager usable and preserve errors predating the probe.""" + import os + from pathlib import Path + import subprocess + import sys + import textwrap + + script = textwrap.dedent(""" + import ctypes + import os + import runpy + import sys + + import pytest + import torch + from freetoken.kernel import _cpu_moe + from freetoken.moe.cpu_executor import CpuMoeExecutor + + shim = ctypes.CDLL(os.environ["LD_PRELOAD"]) + stream_ptr = torch.cuda.current_stream().cuda_stream + fault = os.environ["FREETOKEN_TEST_HIP_PROBE_FAULT"] + if fault == "pending": + assert shim.probe_test_make_error() != 0 + pending = shim.probe_test_peek_error() + with pytest.raises(RuntimeError, match="pending error before probe"): + _cpu_moe.memops_probe(stream_ptr) + assert shim.probe_test_peek_error() == pending + assert shim.probe_test_begins() == 0 + sys.exit(0) + + helpers = runpy.run_path(sys.argv[1]) + torch.manual_seed(17) + cache = helpers["_make_cache"](2, 8, 256, 128, scale=0.03) + ex = CpuMoeExecutor( + cache, top_k=4, activation="silu", apply_router_weight_on_input=False, + num_threads=4, max_tokens=3, device=torch.device("cuda"), + ) + assert not ex._flag_sync and not ex._graph_capture_safe + assert shim.probe_test_injected() > 0 + assert shim.probe_test_begins() == shim.probe_test_ends() + assert shim.probe_test_peek_error() == 0 + if fault == "dependency": + assert shim.probe_test_removed_nodes() == 1 + assert shim.probe_test_last_end_rc() == 0 + else: + assert shim.probe_test_last_end_rc() != 0 + + for index in range(2): + torch.manual_seed(23 + index) + hidden = (torch.randn(3, 256) * 0.3).bfloat16().cuda() + ids = torch.randint(0, 8, (3, 4), dtype=torch.int32).cuda() + weights = torch.rand(3, 4).cuda() + actual = ex.decode(1, hidden, weights, ids).float().cpu() + expected = helpers["_reference"](cache, 1, hidden, ids.cpu(), weights.cpu()) + relative = (actual - expected).abs().max() / (expected.abs().max() + 1e-6) + assert relative < 2e-2, relative + + stream = torch.cuda.Stream() + with pytest.raises(RuntimeError, match="--cuda-graph-max-bs 0"): + with torch.cuda.graph(torch.cuda.CUDAGraph(), stream=stream): + ex.decode(1, hidden, weights, ids) + assert torch.isfinite(ex.decode(1, hidden, weights, ids).cpu()).all() + """) + env = dict(os.environ, LD_PRELOAD=str(hip_probe_fault_library), + FREETOKEN_TEST_HIP_PROBE_FAULT=fault, FREETOKEN_CPU_MOE_FLAG_SYNC="1") + result = subprocess.run( + [sys.executable, "-c", script, str(Path(__file__).resolve())], + env=env, capture_output=True, text=True, timeout=30, + ) + assert result.returncode == 0, result.stdout + result.stderr + + def test_cpu_moe_decode_cuda_graph_replay(): from freetoken.moe.cpu_executor import CpuMoeExecutor @@ -467,7 +572,7 @@ def test_cpu_moe_decode_cuda_graph_replay(): ids = torch.randint(0, E, (bs, top_k), device=dev, dtype=torch.int32) w = torch.rand(bs, top_k, device=dev, dtype=torch.float32) - # Eager forward: also materializes the per-bs pinned buffers + per-(layer,bs) task. + # Eager forward. out_eager = ex.decode(layer, hidden, w, ids) torch.cuda.synchronize() ref_eager = _reference(cache, layer, hidden, ids.cpu(), w.cpu()) @@ -500,11 +605,9 @@ def test_cpu_moe_decode_cuda_graph_replay(): # must actually have exercised it -- guard against a silent fallback to the # host-func path making this test vacuous for the capture-embedded memops. if ex._flag_sync: - assert (layer, bs) in ex._flag_slots, "flag slot expected for the decode task" - slot = ex._flag_slots[(layer, bs)] - assert ex._ext.flag_served_count(slot) >= 4, "1 eager + 3 replay dispatches expected" - assert int(ex._done[slot]) == 1 and int(ex._ready[slot]) == 0, "handshake at rest" - assert int(ex._err.sum()) == 0, "watchdog must not fire in normal operation" + assert ex._ext.flag_served() >= 4, "1 eager + 3 replay dispatches expected" + assert not ex._ext.flag_pending(), "handshake at rest" + assert not ex._unhealthy, "watchdog must not fire in normal operation" ex.raise_if_unhealthy() print("cpu moe cuda graph replay OK") diff --git a/tests/moe/test_cpu_moe_fp8_block.py b/tests/moe/test_cpu_moe_fp8_block.py index 06f954d8d9..baecb86f9a 100644 --- a/tests/moe/test_cpu_moe_fp8_block.py +++ b/tests/moe/test_cpu_moe_fp8_block.py @@ -74,11 +74,11 @@ def test_cpu_fp8_block_accepts_qwen38_padded_scale_bank(): device=torch.device("cpu"), ) hidden = torch.randn(1, H, dtype=torch.bfloat16) - io = ex._io_for(1) + io = ex._io io["x"].copy_(hidden) io["ids"].zero_() io["w"].fill_(1.0) - ex._ext.run_task(ex._task_for(0, 1)) + ex._ext.run(0, 1) cpu_out = io["y"][0].float().to(torch.float64) banks = cache.bank_sources gate_up = _reference_fp8_block_gemv(banks["gate_up"][0], banks["gate_up_scale"][0], hidden[0]) @@ -140,11 +140,8 @@ def _reference_fp8_block_gemv(weights, scales, x): @pytest.mark.parametrize("codes", ["finite", "subnormal"]) -@pytest.mark.parametrize( - "isa, expected_isa", - [("scalar", "scalar"), ("avx2", "avx2"), ("avx512", "avx512f"), ("avx512bf16", "avx512f")], -) -def test_cpu_fp8_block_matches_float64_reference(monkeypatch, isa, expected_isa, codes): +@pytest.mark.parametrize("isa", ["scalar", "avx2", "avx512"]) +def test_cpu_fp8_block_matches_float64_reference(monkeypatch, isa, codes): from freetoken.moe.cpu_executor import CpuMoeExecutor H, I = 259, 257 # ragged K for both GEMVs; 257 down rows exercise many reductions @@ -159,17 +156,16 @@ def test_cpu_fp8_block_matches_float64_reference(monkeypatch, isa, expected_isa, max_tokens=1, device=torch.device("cpu"), ) - tiers = {"scalar": 0, "avx2": 1, "avx512": 2, "avx512bf16": 3} + tiers = {"scalar": 0, "avx2": 1, "avx512": 2} if tiers[isa] > ("scalar", "avx2", "avx512f", "avx512bf16").index(ex.isa): pytest.skip(f"requested {isa}, host/build selected {ex.isa}") - assert ex._ext.fp8_block_isa_name() == expected_isa hidden = torch.randn(1, H, dtype=torch.bfloat16) - io = ex._io_for(1) + io = ex._io io["x"].copy_(hidden) io["ids"].zero_() io["w"].fill_(1.0) - ex._ext.run_task(ex._task_for(0, 1)) + ex._ext.run(0, 1) cpu_out = io["y"][0].float().to(torch.float64) banks = cache.bank_sources @@ -233,9 +229,8 @@ def test_cpu_fp8_block_matches_gpu_decode_kernel(bs, isa, monkeypatch): max_tokens=bs, device=dev, ) - got = ex._ext.fp8_block_isa_name() - if got != _ISA_TIERS[isa]: - pytest.skip(f"asked for {isa}, kernel selected {got}: not supported here") + if ex.isa != _ISA_TIERS[isa]: + pytest.skip(f"asked for {isa}, host/build selected {ex.isa}: not supported here") hidden = torch.randn(bs, H, device=dev, dtype=torch.bfloat16) ids = torch.stack([torch.randperm(E, device=dev)[:top_k] for _ in range(bs)]).to(torch.int32) @@ -254,7 +249,7 @@ def test_cpu_fp8_block_matches_gpu_decode_kernel(bs, isa, monkeypatch): ).float() rel = (cpu_out - gpu_out).abs().max() / (gpu_out.abs().max() + 1e-6) - assert rel < 2e-2, f"bs={bs} isa={got} rel err {rel.item()}" + assert rel < 2e-2, f"bs={bs} isa={ex.isa} rel err {rel.item()}" def test_fp16_bits_carry_every_e4m3_byte(): diff --git a/tests/moe/test_cpu_moe_q4_0.py b/tests/moe/test_cpu_moe_q4_0.py index b5ec2f09b4..7babb5d4ce 100644 --- a/tests/moe/test_cpu_moe_q4_0.py +++ b/tests/moe/test_cpu_moe_q4_0.py @@ -1,7 +1,7 @@ """CPU MoE executor -- native GGUF Q4_0 experts (``--moe-backend cpu``/``hybrid`` for the gemma4 GGUF checkpoint). -The CPU W4A16 GEMV (``q4_0_dot`` in csrc/cpu_moe/cpu_moe_ext.cpp) reads the *same* +The CPU W4A16 GEMV (``q4_0_dot`` in csrc/cpu_moe/kernels/q4_0.cpp) reads the *same* packed Q4_0 banks the GPU offload path streams and dequantizes weights inside the K-loop. We check it against the reference dequant (models/gguf/dequant.py) + the production bf16 GPU decode kernel on byte-identical banks: both are W4A16, so the diff --git a/tests/moe/test_dsfp4_prequant.py b/tests/moe/test_dsfp4_prequant.py index 11519502e2..8187402d3f 100644 --- a/tests/moe/test_dsfp4_prequant.py +++ b/tests/moe/test_dsfp4_prequant.py @@ -17,8 +17,9 @@ @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") -def test_dsfp4_gpu_prequant_bit_parity(): +def test_dsfp4_gpu_prequant_bit_parity(monkeypatch): from freetoken.kernel.pinned import alloc_pinned_tensor + from freetoken.moe import cpu_executor from freetoken.moe.cpu_executor import CpuMoeExecutor torch.manual_seed(0) @@ -44,14 +45,13 @@ def test_dsfp4_gpu_prequant_bit_parity(): wts = torch.rand(2, TOPK, device=dev) def run(cpu_roundtrip: bool) -> torch.Tensor: + # control: the original CPU-side scalar round-trip + monkeypatch.setattr(cpu_executor, "_DSFP4_GPU_PREQUANT", not cpu_roundtrip) ex = CpuMoeExecutor( cache, top_k=TOPK, activation="silu", apply_router_weight_on_input=False, num_threads=4, max_tokens=2, device=dev, swiglu_limit=10.0, act_block=128, ) - assert ex._gpu_prequant, "ds_fp4 on CUDA should enable GPU prequant" - if cpu_roundtrip: # control: the original CPU-side scalar round-trip - ex._gpu_prequant = False - ex._ext.set_input_prequant(False) + assert ex._gpu_prequant != cpu_roundtrip, "ds_fp4 on CUDA should enable GPU prequant" out = ex.decode(0, x, wts, ids.clone()).clone() torch.cuda.synchronize() del ex From f02fbb16a213f5c81094e4017c988d41f0164141 Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Tue, 13 Oct 2026 20:30:34 -0700 Subject: [PATCH 03/94] perf(moe): split hybrid decode routes inside the ensure kernel (#685) * perf(moe): split hybrid decode routes inside the ensure kernel * refactor(moe): always split routes in the hybrid ensure kernel * test(moe): drop the fetch split test left without its kernel check --- python/freetoken/layers/moe.py | 18 +--- python/freetoken/moe/offload_cache.py | 56 ++++------ python/freetoken/moe/offload_kernels.py | 137 ++++++++++++------------ tests/moe/test_hybrid_fetch.py | 75 +------------ 4 files changed, 95 insertions(+), 191 deletions(-) diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index 82f09a9bb8..0595b1d451 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -307,7 +307,7 @@ def _decode_hybrid( The CPU pool is kicked off (``decode_submit``) before the GPU PCIe fetch + GEMM so the CPU overflow GEMV runs concurrently with the GPU work. Capture-safe: the - routing split is device-side elementwise and the CPU submit/sync are host nodes. + routing split runs inside the hybrid ensure kernel and the CPU submit/sync are host nodes. Each route is computed exactly once -- CPU-assigned routes reach the GPU as slot -1 with weight 0 (a kernel with ``supports_inactive_slots`` skips them without a read; any other reads slot 0, whose storage is a payload or the bank's zero fill) and the CPU ids are -1 for @@ -315,14 +315,10 @@ def _decode_hybrid( """ executor = cache.cpu_executor assert executor is not None, "CPU MoE executor was not initialized" - raw_ids = topk_ids.clone() - cache.ensure_experts_hybrid(self.layer_id, topk_ids) - on_gpu = topk_ids >= 0 - cpu_ids = torch.where(on_gpu, -1, raw_ids).contiguous() - gpu_w = torch.where(on_gpu, topk_weights, 0.0).contiguous() - gpu_slots = topk_ids - if cache.collect_stats: - cache.record_decode_stats_hybrid(self.layer_id) + # the kernel reads every route's slot: give the inactive ones slot 0, which holds a + # complete payload or the bank's zero fill (never uninitialized storage), times weight 0 + clamp = not (self.quant_method is not None and self.quant_method.kernel.supports_inactive_slots) + gpu_slots, gpu_w, cpu_ids = cache.ensure_experts_hybrid(self.layer_id, topk_ids, topk_weights, clamp=clamp) pending = executor.decode_submit(self.layer_id, hidden_states, topk_weights, cpu_ids) # Measurement knob: FREETOKEN_HYBRID_OVERLAP=0 syncs the CPU pool *before* the @@ -332,10 +328,6 @@ def _decode_hybrid( ) cache.copy_missing() - if not (self.quant_method is not None and self.quant_method.kernel.supports_inactive_slots): - # the kernel reads every route's slot: give the inactive ones slot 0, which holds a - # complete payload or the bank's zero fill (never uninitialized storage), times weight 0 - gpu_slots = gpu_slots.clamp_min(0) gpu_routed = self._expert_gemm( cache, hidden_states, diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 0969f73829..984b7124ad 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -3,7 +3,7 @@ import math import os from dataclasses import dataclass -from typing import Iterator +from typing import TYPE_CHECKING, Iterator import torch from flashlib.kernels.slot_cache import N_STATS, Stat @@ -28,6 +28,9 @@ from freetoken.kernel.fast_index_copy import default_blocks_per_bank from freetoken.utils import init_logger +if TYPE_CHECKING: + from freetoken.moe.offload_kernels import HybridRoutes + logger = init_logger(__name__) # quant_format -> bank names, in registration order: the single place a format's bank @@ -224,24 +227,23 @@ def __post_init__(self) -> None: # is a one-off warm-up increment at capture time (<0.1% over a session). self.collect_stats = False # [num_layers, N_STATS] -- ensure_experts passes lru_stats[layer_id] straight to - # the kernel, which accumulates in the same launch. The stat_* tensors below stay - # for the hybrid path, whose kernel is still ours. + # the kernel, which accumulates in the same launch. The stat_* tensors below are + # the hybrid path's: views of one buffer its kernel accumulates into the same way. self.lru_stats = torch.zeros( (self.num_layers, N_STATS), dtype=torch.int64, device=self.device ) - self.stat_missing = torch.zeros((), dtype=torch.int64, device=self.device) - self.stat_active = torch.zeros((), dtype=torch.int64, device=self.device) - self.stat_calls = torch.zeros((), dtype=torch.int64, device=self.device) - # hybrid only: experts actually fetched over PCIe (<= stat_missing). The CPU - # computes stat_missing - stat_fetched of them. - self.stat_fetched = torch.zeros((), dtype=torch.int64, device=self.device) + # (missing, fetched, active, calls), then the same per layer + stat_buf = torch.zeros(4 * (self.num_layers + 1), dtype=torch.int64, device=self.device) + self.hybrid_stats = (stat_buf[:4], stat_buf[4:]) + # hybrid only: stat_fetched counts the experts actually fetched over PCIe (<= stat_missing). + # The CPU computes stat_missing - stat_fetched of them. + self.stat_missing, self.stat_fetched, self.stat_active, self.stat_calls = stat_buf[:4].unbind() # Per-layer counterparts of the scalars above (indexed by MoE-layer id). Same # device-side accumulation (graph-safe: layer_id is a static index per graph node), # so one req's per-layer miss rate is readable via decode_miss_stats_per_layer(). - self.stat_missing_layer = torch.zeros(self.num_layers, dtype=torch.int64, device=self.device) - self.stat_active_layer = torch.zeros(self.num_layers, dtype=torch.int64, device=self.device) - self.stat_fetched_layer = torch.zeros(self.num_layers, dtype=torch.int64, device=self.device) - self.stat_steps_layer = torch.zeros(self.num_layers, dtype=torch.int64, device=self.device) + per_layer = stat_buf[4:].view(self.num_layers, 4) + self.stat_missing_layer, self.stat_fetched_layer = per_layer[:, 0], per_layer[:, 1] + self.stat_active_layer, self.stat_steps_layer = per_layer[:, 2], per_layer[:, 3] # Opt-in decode routing histogram (per layer, per expert) for cache-skew # analysis. Accumulated in ``ensure_experts`` from the raw expert ids before the # kernel rewrites them to slots. Only accurate with CUDA graphs disabled (the @@ -855,7 +857,9 @@ def ensure_experts(self, layer_id: int, expert_ids: torch.Tensor) -> None: self._pending_whole_layer = False ensure_experts(self, layer_id, expert_ids) - def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None: + def ensure_experts_hybrid( + self, layer_id: int, expert_ids: torch.Tensor, weights: torch.Tensor, clamp: bool = False + ) -> HybridRoutes: """Capped-fetch LRU for the hybrid backend. Like :meth:`ensure_experts` but assigns slots to (and schedules copies for) at @@ -864,7 +868,8 @@ def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None left non-resident and ``expert_ids`` is rewritten to their cache slot (hit or freshly fetched) or ``-1`` (overflow -> compute on the CPU). ``num_indices`` holds the capped fetch count (for ``copy_missing``); ``num_missing_full`` the pre-cap - miss count (for stats). All device-side / fixed-shape, so it is CUDA-graph safe.""" + miss count (for stats). All device-side / fixed-shape, so it is CUDA-graph safe. + Returns the routes split between GPU and CPU (:class:`freetoken.moe.offload_kernels.HybridRoutes`).""" from freetoken.moe.offload_kernels import ensure_experts_hybrid if self.collect_decode_freq: @@ -872,8 +877,8 @@ def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None self.decode_freq[layer_id].scatter_add_(0, ids, torch.ones_like(ids)) self._pending_src_layer = layer_id self._pending_whole_layer = False - ensure_experts_hybrid( - self, layer_id, expert_ids, self.hybrid_max_fetch, self.hybrid_fetch_fraction + return ensure_experts_hybrid( + self, layer_id, expert_ids, weights, self.hybrid_max_fetch, self.hybrid_fetch_fraction, clamp ) def materialize_layer(self, layer_id: int) -> None: @@ -911,23 +916,6 @@ def record_decode_stats(self, layer_id: int) -> None: was eight torch ops per layer per step, all captured into the decode graph. """ - def record_decode_stats_hybrid(self, layer_id: int) -> None: - """Hybrid stats: full miss count (pre-cap), the PCIe-fetched count (capped), and - the active count. The CPU computes (missing - fetched) experts. Device-side; - accumulates both the scalar totals and the per-layer breakdown.""" - assert 0 <= layer_id < self.num_layers, f"layer_id {layer_id} out of range [0, {self.num_layers})" - missing = self.num_missing_full.sum() - fetched = self.num_indices.sum() - active = self.active_mask.sum() - self.stat_missing += missing - self.stat_fetched += fetched - self.stat_active += active - self.stat_calls += 1 - self.stat_missing_layer[layer_id] += missing - self.stat_fetched_layer[layer_id] += fetched - self.stat_active_layer[layer_id] += active - self.stat_steps_layer[layer_id] += 1 - def decode_miss_stats(self) -> dict: if self.decode_target == "hybrid": active = int(self.stat_active.item()) diff --git a/python/freetoken/moe/offload_kernels.py b/python/freetoken/moe/offload_kernels.py index cf513f52d1..ea12b6fb17 100644 --- a/python/freetoken/moe/offload_kernels.py +++ b/python/freetoken/moe/offload_kernels.py @@ -1,6 +1,7 @@ from __future__ import annotations import os +from typing import NamedTuple import torch import triton @@ -40,9 +41,23 @@ def ensure_experts(cache, layer_id: int, expert_ids: torch.Tensor) -> None: ) +class HybridRoutes(NamedTuple): + """One layer's routes split between the GPU and the CPU, each [bs, top_k].""" + + gpu_slots: torch.Tensor # cache slot of each route; a CPU route is -1, or 0 with clamp + gpu_w: torch.Tensor # router weights, 0 for a CPU route + cpu_ids: torch.Tensor # expert ids, -1 for a GPU route + + def ensure_experts_hybrid( - cache, layer_id: int, expert_ids: torch.Tensor, max_fetch: int, fetch_fraction: float = 0.0 -) -> None: + cache, + layer_id: int, + expert_ids: torch.Tensor, + weights: torch.Tensor, + max_fetch: int, + fetch_fraction: float = 0.0, + clamp: bool = False, +) -> HybridRoutes: """Capped-fetch variant of ``ensure_experts`` (hybrid backend). Identical LRU bookkeeping, but only the first ``max_fetch`` of this step's missing @@ -52,12 +67,17 @@ def ensure_experts_hybrid( split (fraction = pcie_bw / cpu_bw): fetch ~fraction of the step's misses, rounded to the integer that makes the PCIe fetch and the CPU overflow compute finish closest to together. ``num_indices`` = capped fetch count (copy_missing); ``num_missing_full`` = - pre-cap miss count (stats).""" - # Q16 fixed point so the GPU kernel and the CPU reference cap identically (no float). + pre-cap miss count (stats). + + The same launch splits the routes: ``expert_ids`` becomes the GPU slots in place, and + ``clamp`` gives CPU routes slot 0 instead of -1 for GEMM kernels that read every slot.""" + assert expert_ids.is_cuda, "the hybrid ensure kernel runs on the GPU" + # Q16 fixed point: the kernel caps in integers. frac_q16 = min(1 << 16, max(0, round(fetch_fraction * (1 << 16)))) - if not expert_ids.is_cuda: - return _ensure_experts_hybrid_cpu(cache, layer_id, expert_ids, max_fetch, frac_q16) - _ensure_experts_hybrid_gpu(cache, layer_id, expert_ids, max_fetch, frac_q16) + weights = weights.contiguous() + routes = HybridRoutes(expert_ids, torch.empty_like(weights), torch.empty_like(expert_ids)) + _ensure_experts_hybrid_gpu(cache, layer_id, expert_ids, weights, routes, clamp, max_fetch, frac_q16) + return routes def prefill_hit_compact(cache, layer_id: int, buffer_id: int) -> None: @@ -95,7 +115,14 @@ def reset_cache(cache) -> None: def _ensure_experts_hybrid_gpu( - cache, layer_id: int, expert_ids: torch.Tensor, max_fetch: int, frac_q16: int + cache, + layer_id: int, + expert_ids: torch.Tensor, + weights: torch.Tensor, + routes: HybridRoutes, + clamp: bool, + max_fetch: int, + frac_q16: int, ) -> None: block_e = triton.next_power_of_2(cache.num_experts) block_c = triton.next_power_of_2(cache.cache_size) @@ -112,6 +139,10 @@ def _ensure_experts_hybrid_gpu( cache.num_indices, cache.num_missing_full, cache.expert_recency, + weights, + routes.cpu_ids, + routes.gpu_w, + *(cache.hybrid_stats if cache.collect_stats else (expert_ids,) * 2), layer_id, expert_ids.numel(), int(max_fetch), @@ -121,73 +152,12 @@ def _ensure_experts_hybrid_gpu( BLOCK_E=block_e, BLOCK_C=block_c, BY_RECENCY=_HYBRID_FETCH_BY_RECENCY, + CLAMP=clamp, + STATS=cache.collect_stats, num_warps=num_warps, ) -def _ensure_experts_hybrid_cpu( - cache, layer_id: int, expert_ids: torch.Tensor, max_fetch: int, frac_q16: int -) -> None: - """CPU reference mirror of the hybrid kernel (eviction/fetch decisions bit-identical to - the GPU path; see tests/test_offload_lru_kernels.py). Fetches at most ``max_fetch`` (or - the bandwidth-matched ``~frac_q16/2^16 * misses`` when ``frac_q16`` > 0) of the missing - experts; overflow misses are rewritten to -1. With ``BY_RECENCY`` the fetch set is the - most-recently-active misses (ties -> lower id); else the lowest ids.""" - seen = [] - for expert in expert_ids.view(-1).tolist(): - if expert not in seen: - seen.append(expert) - - cache.active_mask.zero_() - step = int(cache.step.item()) + 1 - cache.step.fill_(step) - for expert in seen: - cache.active_mask[expert] = 1 - - for expert in seen: - slot = int(cache.slot_for_id[layer_id, expert].item()) - if slot != -1: - cache.usage[slot] = step - - missing = [e for e in seen if int(cache.slot_for_id[layer_id, e].item()) == -1] - if _HYBRID_FETCH_BY_RECENCY: - rec = cache.expert_recency[layer_id].tolist() - missing.sort(key=lambda e: (-rec[e], e)) - else: - missing.sort() - if frac_q16 > 0: - m, q = len(missing), 1 << 16 - lo = (m * frac_q16) >> 16 - cost = lambda f: max(f * (q - frac_q16), (m - f) * frac_q16) # noqa: E731 - max_fetch = lo if cost(lo) <= cost(lo + 1) else lo + 1 - num_fetch = min(len(missing), int(max_fetch)) - cache.num_missing_full.fill_(len(missing)) - cache.num_indices.fill_(num_fetch) - - usage = cache.usage.tolist() - for idx in range(num_fetch): - expert = missing[idx] - victim = min(range(cache.cache_size), key=lambda s: (usage[s], s)) - old_id = int(cache.id_of_slot[victim].item()) - if old_id >= 0: - cache.slot_for_id.view(-1)[old_id] = -1 - cache.id_of_slot[victim] = layer_id * cache.num_experts + expert - cache.slot_for_id[layer_id, expert] = victim - cache.usage[victim] = step - usage[victim] = step - cache.evict_slots[idx] = victim - cache.src_indices[idx] = expert # layer-local row - - if _HYBRID_FETCH_BY_RECENCY: - for expert in seen: - cache.expert_recency[layer_id, expert] = step - - # Overflow misses keep slot_for_id == -1, so the rewrite below yields -1 for them. - flat = expert_ids.view(-1) - for i in range(flat.numel()): - flat[i] = int(cache.slot_for_id[layer_id, int(flat[i].item())].item()) - - def _materialize_layer_gpu(cache, layer_id: int) -> None: block = triton.next_power_of_2(max(cache.num_experts, cache.cache_size)) _materialize_layer_kernel[(1,)]( @@ -300,6 +270,11 @@ def _ensure_experts_hybrid_kernel( num_indices_ptr, num_missing_full_ptr, expert_recency_ptr, + weights_ptr, + cpu_ids_ptr, + gpu_w_ptr, + stats_ptr, + stats_layer_ptr, layer_id, num_active, max_fetch, @@ -309,6 +284,8 @@ def _ensure_experts_hybrid_kernel( BLOCK_E: tl.constexpr, BLOCK_C: tl.constexpr, BY_RECENCY: tl.constexpr, + CLAMP: tl.constexpr, + STATS: tl.constexpr, ): """Capped-fetch timestamp-LRU (hybrid backend). @@ -401,8 +378,26 @@ def _ensure_experts_hybrid_kernel( for i in tl.range(num_active): e = tl.load(expert_ids_ptr + i) s = tl.load(slot_for_id_ptr + base + e) + on_gpu = s >= 0 + tl.store(cpu_ids_ptr + i, tl.where(on_gpu, -1, e)) + tl.store(gpu_w_ptr + i, tl.where(on_gpu, tl.load(weights_ptr + i), 0.0)) + if CLAMP: + s = tl.maximum(s, 0) tl.store(expert_ids_ptr + i, s) + # stats_ptr: (missing, fetched, active, calls); stats_layer_ptr: the same per layer, [num_layers, 4] + if STATS: + n_active = tl.sum((is_active & e_mask).to(tl.int64)) + tl.store(stats_ptr + 0, tl.load(stats_ptr + 0) + num_missing) + tl.store(stats_ptr + 1, tl.load(stats_ptr + 1) + num_fetch) + tl.store(stats_ptr + 2, tl.load(stats_ptr + 2) + n_active) + tl.store(stats_ptr + 3, tl.load(stats_ptr + 3) + 1) + row = stats_layer_ptr + layer_id * 4 + tl.store(row + 0, tl.load(row + 0) + num_missing) + tl.store(row + 1, tl.load(row + 1) + num_fetch) + tl.store(row + 2, tl.load(row + 2) + n_active) + tl.store(row + 3, tl.load(row + 3) + 1) + # Bump every active expert's recency to this step (LRU on the expert): an overflow miss # computed on the CPU now ranks high if it recurs, so it gets fetched next time. if BY_RECENCY: diff --git a/tests/moe/test_hybrid_fetch.py b/tests/moe/test_hybrid_fetch.py index a26630b894..fad4fa7843 100644 --- a/tests/moe/test_hybrid_fetch.py +++ b/tests/moe/test_hybrid_fetch.py @@ -1,43 +1,15 @@ """Hybrid decode's bandwidth-matched fetch split. -Covers the two halves of --moe-hybrid-max-fetch auto: the profile reader that turns -`ft bench bw` kernel bandwidths into a fetch fraction, and the ensure kernel's -per-step integer split (GPU kernel vs CPU reference mirror, and the balance rule). +Covers the profile reader of --moe-hybrid-max-fetch auto, which turns `ft bench bw` +kernel bandwidths into a fetch fraction. """ import json import os import pytest -import torch from freetoken.moe.bench_profile import PROFILE_VERSION, default_profile_path, load_backend_recommendation, load_hybrid_fetch_fraction -from freetoken.moe.offload_cache import OffloadMoeCache - -Q = 1 << 16 - - -def _balanced_fetch(num_missing: int, frac_q16: int) -> int: - """Reference split: F ~ frac * misses, rounded to whichever integer neighbor - minimizes the slower overlapped side (fetch ~ F*(1-frac), CPU ~ (M-F)*frac).""" - lo = (num_missing * frac_q16) >> 16 - cost = lambda f: max(f * (Q - frac_q16), (num_missing - f) * frac_q16) # noqa: E731 - return min(num_missing, lo if cost(lo) <= cost(lo + 1) else lo + 1) - - -def test_balanced_fetch_tracks_fraction(): - # The split follows fetched : cpu = pcie : (cpu - pcie) up to integer rounding, and - # never over/under-shoots by more than one expert. - for frac in (0.1, 0.415, 0.454, 0.7, 1.0): - q = round(frac * Q) - for m in range(0, 65): - f = _balanced_fetch(m, q) - assert 0 <= f <= m - assert abs(f - frac * m) <= 1.0 - # ceil would over-fetch here (the regression this rule fixed): 41.5% of 3 misses is - # 1.24 -> fetching 2 makes the PCIe side ~1.6x slower than balance; keep it at 1. - assert _balanced_fetch(3, round(0.415 * Q)) == 1 - assert _balanced_fetch(4, round(0.415 * Q)) == 2 def test_load_hybrid_fetch_fraction(tmp_path): @@ -87,46 +59,3 @@ def write(path, name, verdict): # this card's own file wins over the legacy one write(default_profile_path(uuid), "FAKE GPU", "offload") assert load_backend_recommendation("bf16", gpu_name="FAKE GPU", gpu_uuid=uuid) == "offload" - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") -def test_hybrid_fraction_gpu_matches_cpu_reference(): - torch.manual_seed(0) - num_experts, cache_size, top_k, frac = 32, 40, 8, 0.415 - - def make(): - return OffloadMoeCache( - num_layers=2, num_experts=num_experts, cache_size=cache_size, - device=torch.device("cuda"), quant_format="bf16", decode_target="hybrid", - hybrid_max_fetch=num_experts, hybrid_fetch_fraction=frac, - ) - - gpu, ref = make(), make() - frac_q16 = round(frac * Q) - for step in range(64): - ids = torch.randperm(num_experts)[:top_k].to(torch.int32) - g, c = ids.clone().cuda(), ids.clone() # a CPU ids tensor drives the reference path - gpu.ensure_experts_hybrid(0, g) - ref.ensure_experts_hybrid(0, c) - missing = int(gpu.num_missing_full.item()) - fetched = int(gpu.num_indices.item()) - assert missing == int(ref.num_missing_full.item()) - assert fetched == int(ref.num_indices.item()) == _balanced_fetch(missing, frac_q16) - # slot rewrites (hit/fetched -> slot, overflow -> -1) and LRU state stay identical - assert torch.equal(g.cpu(), c) - assert torch.equal(gpu.slot_for_id.cpu(), ref.slot_for_id.cpu()) - assert torch.equal(gpu.id_of_slot.cpu(), ref.id_of_slot.cpu()) - assert (g >= 0).sum().item() == len(set(ids.tolist())) - (missing - fetched) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") -def test_hybrid_fixed_cap_unchanged(): - # fraction 0 (no profile / explicit --moe-hybrid-max-fetch) keeps the fixed cap. - cache = OffloadMoeCache( - num_layers=1, num_experts=32, cache_size=40, device=torch.device("cuda"), - quant_format="bf16", decode_target="hybrid", hybrid_max_fetch=1, - ) - ids = torch.arange(8, dtype=torch.int32).cuda() - cache.ensure_experts_hybrid(0, ids) - assert int(cache.num_missing_full.item()) == 8 - assert int(cache.num_indices.item()) == 1 From b416f85c40970c033462b47773020872209978a3 Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Wed, 14 Oct 2026 19:44:55 -0700 Subject: [PATCH 04/94] feat(moe): split the offload expert cache into slot pools (#689) * feat(moe): split the offload expert cache into slot pools * perf(moe): stop recompiling the hybrid ensure kernel for each slot pool size * fix(engine): turn off prefill overlap when --moe-cache-mem leaves under two layers of slots * fix(engine): load the --moe-cache-split profile before the expert banks --- docs/cli.md | 4 +- python/freetoken/engine/cache_budget.py | 13 +- python/freetoken/engine/config.py | 6 + python/freetoken/engine/engine.py | 134 +++++-- python/freetoken/kvcache/cache_status.py | 3 + python/freetoken/layers/moe.py | 17 +- python/freetoken/models/deepseek_v4/moe.py | 5 +- python/freetoken/moe/benchbw.py | 9 +- python/freetoken/moe/expert_banks.py | 2 + python/freetoken/moe/offload_cache.py | 412 +++++++++++++++------ python/freetoken/moe/offload_kernels.py | 55 +-- python/freetoken/moe/pool_policy.py | 149 ++++++++ python/freetoken/server/args.py | 50 +++ tests/engine/test_cache_budget.py | 52 +++ tests/moe/test_nvfp4_backends.py | 12 +- tests/moe/test_offload.py | 186 +++++++++- 16 files changed, 921 insertions(+), 188 deletions(-) create mode 100644 python/freetoken/moe/pool_policy.py diff --git a/docs/cli.md b/docs/cli.md index 2166e9e86e..44b0e30ad9 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -84,7 +84,9 @@ See [models.md](models.md#moe-strategies) for what each strategy does. | `--moe-strategy` | auto | `fused`/`offload`/`cpu`/`hybrid`; auto → offload, or hybrid with a `ft bench bw` profile; fused on unified-memory GPUs (GB10), see [models.md](models.md#unified-memory-gpus-gb10--dgx-spark). `--moe-backend` is the deprecated old spelling | | `--quant-backend` | auto | Kernel per quantized layer type, `layer[.kind]=name` entries: `linear=marlin,moe=b12x` or `moe.nvfp4=triton`. A layer-level entry applies to every kind whose table lists the name | | `--nvfp4-backend` | — | Deprecated: stands in for `--quant-backend moe.nvfp4=` (`flashinfer` means b12x); cannot be combined with `--quant-backend` | -| `--moe-cache-size` / `--moe-cache-rate` / `--moe-cache-auto` | auto | GPU expert-cache size as slots / fraction of all experts / sized from free VRAM (mutually exclusive; auto is enabled by default for offload-family strategies) | +| `--moe-cache-size` / `--moe-cache-rate` / `--moe-cache-mem` / `--moe-cache-auto` | auto | GPU expert-cache size as slots / fraction of all experts / memory such as `12GiB` / sized from free VRAM (mutually exclusive; auto is enabled by default for offload-family strategies) | +| `--moe-cache-scope` | unified | `unified`: one LRU pool of slots per expert quant (quant type and bank layout), so all MoE layers share one pool when they share a quant format; `per-layer`: each layer has its own pool. `per-layer` does not combine with `--moe-strategy cpu` or `--moe-cpu-layers` | +| `--moe-cache-split` | equal | How the pools share the slots: `equal` (in proportion to their layer counts), or `profile:`: every pool keeps `top_k` slots and the rest go to the most used experts in the profile's per-layer expert counts. With `offload`, a pool smaller than the expert count lowers `--max-running-requests` to its slots / top_k | | `--kv-reserve-tokens` | 8192 | KV token floor reserved before `--moe-cache-auto` fills experts | | `--moe-cpu-threads` | physical cores | CPU worker threads for the cpu/hybrid executor | | `--moe-cpu-layers` | all on GPU | With `offload`: which MoE layers decode on CPU (`3,7,11`, a count, a fraction, or `auto`). `auto` is for Windows/WSL only, where CUDA pinned memory is capped; every value needs an expert format the CPU executor serves (bf16, fp8_block, nvfp4, mxfp4) | diff --git a/python/freetoken/engine/cache_budget.py b/python/freetoken/engine/cache_budget.py index b5f07ee8ab..891ac2ac11 100644 --- a/python/freetoken/engine/cache_budget.py +++ b/python/freetoken/engine/cache_budget.py @@ -14,8 +14,8 @@ import torch -def expert_bytes_per_slot(sources: dict[str, "list[torch.Tensor]"]) -> int: - """Bytes one expert slot occupies on GPU: summed row bytes over all banks. +def layer_slot_bytes(sources: dict[str, "list[torch.Tensor]"]) -> list[int]: + """Bytes one expert slot of each layer occupies on GPU: its row bytes summed over all banks. Each bank source is per-layer ``[num_experts, *row_shape]`` tensors and is already TP-sharded upstream, so the per-row byte count is the per-rank slot @@ -25,7 +25,14 @@ def expert_bytes_per_slot(sources: dict[str, "list[torch.Tensor]"]) -> int: # with cache_size), so they are intentionally excluded from the per-slot growth term. # tensor[0].numel() is the per-row element count (one expert slot); see the matching # slot-byte idiom in kvcache/linear_state_pool.py and kvcache/dsv4/v4_pool.py. - return sum(t[0][0].numel() * t[0].element_size() for t in sources.values()) + num_layers = len(next(iter(sources.values()))) + return [sum(t[l][0].numel() * t[l].element_size() for t in sources.values()) for l in range(num_layers)] + + +def expert_bytes_per_slot(sources: dict[str, "list[torch.Tensor]"]) -> int: + """Bytes one expert slot occupies on GPU, rounded up to the mean over layers when their bank layouts differ.""" + per_layer = layer_slot_bytes(sources) + return -(-sum(per_layer) // len(per_layer)) def net_cache_budget_bytes( diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index b6fa8b8949..ae2f1fd629 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -45,6 +45,12 @@ class EngineConfig: moe_cache_size: int = 0 moe_cache_rate: float | None = None moe_cache_auto: bool = False + # --moe-cache-mem: slot-cache budget in bytes, turned into a slot count once the expert banks are loaded + moe_cache_mem: int | None = None + # --moe-cache-scope: unified (one LRU pool per expert quant) or per-layer (one pool per layer) + moe_cache_scope: str = "unified" + # --moe-cache-split: how pools share the slots; equal or profile: (by expert use from a profile file) + moe_cache_split: str = "equal" kv_reserve_tokens: int = 8192 # KV floor for --moe-cache-auto; small by design (MoE-priority) moe_cache_policy: str = "lru" moe_prefill_overlap: bool = True diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index feb33fbb94..2757b10dea 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -23,7 +23,8 @@ from freetoken.moe import is_offload_moe_strategy from freetoken.moe.expert_banks import attach_resident_banks, load_expert_banks from freetoken.moe.host_banks import PinFailed -from freetoken.moe.offload_cache import OffloadMoeCache, attach_offload_moe_cache +from freetoken.moe.offload_cache import OffloadMoeCache, attach_offload_moe_cache, bank_layouts +from freetoken.moe.pool_policy import group_layers, make_pools, pool_of_layer, resolve_profile, split_slots from freetoken.utils import align_ceil, init_logger, is_sm90_family, is_sm100_family, mem_GB, torch_dtype from .config import EngineConfig @@ -38,6 +39,25 @@ logger = init_logger(__name__) +def _clamp_decode_batch_to_pools(config: EngineConfig, pool_sizes: list[int]) -> None: + """Clamp the decode batch so every GPU-decoded layer's routed experts fit in its pool.""" + num_experts = config.model_config.num_experts + small = [s for s in pool_sizes if s < num_experts] + if not small: + return + limit = min(small) // config.model_config.num_experts_per_tok + if config.max_running_req > limit: + logger.warning_rank0( + f"the smallest MoE slot pool has {min(small)} slots; lowering max_running_req " + f"{config.max_running_req} -> {limit} so a decode step fits in it" + ) + object.__setattr__(config, "max_running_req", limit) + if config.cuda_graph_max_bs is None or config.cuda_graph_max_bs > limit: + object.__setattr__(config, "cuda_graph_max_bs", limit) + if config.cuda_graph_bs is not None: + object.__setattr__(config, "cuda_graph_bs", [bs for bs in config.cuda_graph_bs if bs <= limit]) + + def _require_offload_cache_size(cache_size: int, num_experts: int) -> None: """The offload MoE cache needs at least one slot per expert per layer. A too-small size (e.g. a bare offload run with moe_cache_size unset and auto disabled) must fail loudly.""" @@ -623,7 +643,7 @@ def _run_mm_encoder(self, batch: Batch) -> None: if parts: batch.mm_embeds = torch.cat([p.to(self.dtype) for p in parts], dim=0) - def _resolve_auto_moe_cache_size(self, config: EngineConfig, banks, method=None) -> tuple[int, int, bool]: + def _resolve_auto_moe_cache_size(self, config: EngineConfig, banks, method=None, max_slots=None) -> tuple[int, int, bool]: """Resolve --moe-cache-auto into (moe_cache_size, num_pages, prefill_overlap). Pure glue over the Phase-1 budget policy; isolated here so it is unit-testable @@ -647,7 +667,7 @@ def _resolve_auto_moe_cache_size(self, config: EngineConfig, banks, method=None) prefill_overlap=config.moe_prefill_overlap, kv_reserve_tokens=max(config.kv_reserve_tokens, min_reserve), page_size=page_tokens, - max_slots=method.slot_limit() if method is not None else None, + max_slots=max_slots if max_slots is not None or method is None else method.slot_limit(), ) def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: @@ -703,6 +723,11 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: else HostResidency.PINNED.value for i in range(config.model_config.num_moe_layers) ] + # read before the expert load, so a bad profile fails in seconds + expert_counts = resolve_profile( + config.moe_cache_split, config.model_path, + config.model_config.num_moe_layers, config.model_config.num_experts, + ) try: with _weight_load_context(): banks = load_expert_banks( @@ -718,8 +743,27 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: ) except PinFailed as exc: raise RuntimeError(f"{exc}; {_pin_hint(self._host_tables_bytes)}") from exc + from freetoken.engine.cache_budget import expert_bytes_per_slot, layer_slot_bytes + + self._moe_layer_slot_bytes = layer_slot_bytes(banks.sources) + self._moe_pool_groups = group_layers(config.moe_cache_scope, bank_layouts(banks.sources, banks.layer_tags), self._moe_layer_slot_bytes) + if len(self._moe_pool_groups) > 1 and cpu_layer_ids: + raise ValueError( + "several MoE slot pools (--moe-cache-scope per-layer, or experts in several quant formats) need every MoE " + "layer to decode from the GPU slot cache; they do not combine with --moe-strategy cpu or --moe-cpu-layers" + ) + self._moe_expert_counts = expert_counts if len(self._moe_pool_groups) > 1 else None + if len(self._moe_pool_groups) == 1 and config.moe_cache_split != "equal": + logger.warning_rank0( + f"--moe-cache-split {config.moe_cache_split} is ignored: all MoE layers share one slot pool " + f"(--moe-cache-scope per-layer gives each layer its own)" + ) + self._moe_slot_limit = method.slot_limit() if method is not None else None + # with GPU decode, a pool below num_experts slots bounds the batch; the startup split keeps room for one token per step + pools = self._moe_pools(1) + size_cap = sum(p.max_num_slots for p in pools) if config.moe_cache_auto: - size, pages, overlap = self._resolve_auto_moe_cache_size(config, banks, method) + size, pages, overlap = self._resolve_auto_moe_cache_size(config, banks, method, max_slots=size_cap) object.__setattr__(config, "moe_cache_size", size) object.__setattr__(config, "moe_prefill_overlap", overlap) if config.num_page_override is None: @@ -735,7 +779,28 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: f"--moe-cache-auto resolved moe_cache_size={size} " f"num_pages={pages} (prefill_overlap={overlap})" ) + elif config.moe_cache_mem is not None: + per_slot = expert_bytes_per_slot(banks.sources) + size = min(config.moe_cache_mem // per_slot, size_cap) + # per_slot is a mean over layers, so with pools of different slot sizes the split can exceed the budget: shrink until it fits + while size > 0 and self._moe_pool_bytes(split_slots(size, pools, self._moe_expert_counts)) > config.moe_cache_mem: + size -= 1 + overlap = config.moe_prefill_overlap and size >= 2 * config.model_config.num_experts + object.__setattr__(config, "moe_cache_size", size) + object.__setattr__(config, "moe_prefill_overlap", overlap) + logger.info_rank0( + f"--moe-cache-mem {mem_GB(config.moe_cache_mem)} resolved moe_cache_size={size} " + f"({per_slot / 2**20:.2f} MiB per slot, prefill_overlap={overlap})" + ) _require_offload_cache_size(config.moe_cache_size, config.model_config.num_experts) + pool_sizes = split_slots(config.moe_cache_size, pools, self._moe_expert_counts) + if len(pools) > 1: + logger.info_rank0( + f"MoE slot pools: {config.moe_cache_scope}, {len(pools)} pools of {min(pool_sizes)}-{max(pool_sizes)} " + f"slots ({config.moe_cache_split})" + ) + if decode_target == "gpu": + _clamp_decode_batch_to_pools(config, pool_sizes) layout = max_slots = None if method is not None: if banks.kind is not None and (banks.kind, banks.kernel) != (method.kind, method.kernel.name): @@ -760,10 +825,12 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: hybrid_max_fetch=config.moe_hybrid_max_fetch, layout=layout, max_slots=max_slots, + pool_of_layer=pool_of_layer(pools, config.model_config.num_moe_layers), + pool_sizes=pool_sizes, ) # before set_bank_sources: the residency validation and the copy plan's skip of non-pinned layers key on the CPU-layer set cache.cpu_layer_ids = cpu_layer_ids - cache.set_bank_sources(banks.sources, layer_residency=banks.layer_residency) + cache.set_bank_sources(banks.sources, layer_residency=banks.layer_residency, layer_tags=banks.layer_tags) cache.set_alphas(banks.gate_up_alpha, banks.down_alpha) if decode_target == "hybrid": self._resolve_hybrid_fetch(config, cache) @@ -772,12 +839,23 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: cache.collect_stats = config.moe_collect_stats layers = attach_offload_moe_cache(self.model, cache) assert len(layers) == config.model_config.num_moe_layers + cache.zero_inactive_route_slots = cache.decode_target == "hybrid" and not all(layer.skips_inactive_slots for layer in layers) if cache.decode_target in ("cpu", "hybrid"): self._init_cpu_moe_executor(config, cache, layers) self.ctx.moe_offload_cache = cache self.moe_offload_cache = cache return cache + def _moe_pool_bytes(self, pool_sizes: list[int]) -> int: + return sum(s * self._moe_layer_slot_bytes[g[0]] for g, s in zip(self._moe_pool_groups, pool_sizes)) + + def _moe_pools(self, batch: int) -> list: + """The slot pools, each with room for the routed experts of ``batch`` decode tokens.""" + model_config = self.config.model_config + num_experts = model_config.num_experts + min_num_slots = min(num_experts, model_config.num_experts_per_tok * batch) + return make_pools(self._moe_pool_groups, num_experts, min_num_slots, self._moe_slot_limit) + def _resolve_hybrid_fetch(self, config: EngineConfig, cache) -> None: """Resolve --moe-hybrid-max-fetch -1 (auto) into a bandwidth-matched fetch fraction. @@ -867,19 +945,18 @@ def _sync_get_memory(self) -> Tuple[int, int]: return min_free_memory, max_free_memory - def _target_moe_and_expert_bytes(self, moe_cache_size: int | None) -> tuple[int, int]: - from freetoken.engine.cache_budget import expert_bytes_per_slot - - target_moe = ( - moe_cache_size - if moe_cache_size is not None - else (self.moe_offload_cache.cache_size if self.moe_offload_cache else 0) - ) - per_expert_bytes = ( - expert_bytes_per_slot(self.moe_offload_cache.bank_sources) - if self.moe_offload_cache is not None else 0 - ) - return target_moe, per_expert_bytes + def _target_moe_and_expert_bytes(self, moe_cache_size: int | None, pool_sizes: list[int] | None) -> tuple[int, int]: + cache = self.moe_offload_cache + if cache is None: + return moe_cache_size or 0, 0 + if moe_cache_size is None: + target_moe = cache.cache_size + nbytes = sum(c.numel() * c.element_size() for _, c in cache.banks) + else: + target_moe = moe_cache_size + nbytes = self._moe_pool_bytes(pool_sizes) + # pools of several bank layouts have different slot sizes; round up so target * bytes covers them + return target_moe, -(-nbytes // target_moe) def _resize_kv_pool(self, config, num_pages: int, num_swa_pages: int | None) -> None: # IN-PLACE, identity-preserving: the CacheManager's swa_pool reference, ctx.kv_cache and @@ -938,7 +1015,10 @@ def rebuild_runtime_cache( "moe_cache_size requested but this model has no MoE offload cache" ) try: - self.moe_offload_cache.validate_rebuild(moe_cache_size) + # a rebuild does not lower max_running_req again, so with GPU decode every pool must hold a full batch; hybrid caps its fetches instead + batch = config.max_running_req if self.moe_offload_cache.decode_target == "gpu" else 1 + moe_pool_sizes = split_slots(moe_cache_size, self._moe_pools(batch), self._moe_expert_counts) + self.moe_offload_cache.validate_rebuild(moe_cache_size, moe_pool_sizes) except ValueError as e: raise CacheRebuildRejected(str(e)) from e if num_pages is not None and num_pages <= 0: @@ -975,7 +1055,9 @@ def rebuild_runtime_cache( # must reject (recoverable) so the old caches stay intact and serving continues, # rather than freeing and then OOMing into permanent failure. The engine supplies # the memory account; the pool answers whether its target geometry fits. - target_moe, per_expert_bytes = self._target_moe_and_expert_bytes(moe_cache_size) + target_moe, per_expert_bytes = self._target_moe_and_expert_bytes( + moe_cache_size, moe_pool_sizes if moe_cache_size is not None else None + ) # Price the sibling GDN state pool at ITS target (physical slots = usable + padding # sink) and hand the bytes in -- the KV pool only budgets its own tiers. target_mamba = ( @@ -1019,7 +1101,7 @@ def rebuild_runtime_cache( object.__setattr__(config, "swa_num_pages_override", num_swa_pages) if moe_cache_size is not None: assert self.moe_offload_cache is not None, "no MoE offload cache to resize" - self.moe_offload_cache.rebuild(moe_cache_size) + self.moe_offload_cache.rebuild(moe_cache_size, moe_pool_sizes) if num_pages is not None: # sets self.num_pages (rebuilds KV + window) self._resize_kv_pool(config, num_pages, num_swa_pages) @@ -1455,6 +1537,9 @@ def _auto_cpu_layers(config: EngineConfig, num_moe_layers: int, *, reserved: int "moe_cache_size": 0, "moe_cache_rate": None, "moe_cache_auto": False, + "moe_cache_mem": None, + "moe_cache_scope": "unified", + "moe_cache_split": "equal", "moe_cpu_layers": None, "moe_cpu_threads": 0, "moe_hybrid_max_fetch": -1, @@ -1750,6 +1835,7 @@ def override(attr: str, value: Any): # this is dangerous, use with caution is_offload_moe_strategy(config.moe_strategy) and config.moe_cache_size <= 0 and config.moe_cache_rate is None + and config.moe_cache_mem is None and not getattr(config, "moe_cache_auto", False) ): # args.py's "no sizing flag -> default --moe-cache-auto" only fires when the @@ -1783,11 +1869,13 @@ def override(attr: str, value: Any): # this is dangerous, use with caution # An explicit 'fused' keeps the experts resident, so there is no slot cache to size. The # sizing flags no longer redirect the backend, so ignore them here and say so -- the # geometry the user asked for is what runs. Report the flag actually passed: --moe-cache- - # rate was already folded into moe_cache_size above, and the three are mutually exclusive. + # rate was already folded into moe_cache_size above, and the sizing flags are mutually exclusive. if config.moe_cache_rate: inert = f"--moe-cache-rate={config.moe_cache_rate}" elif config.moe_cache_size: inert = f"--moe-cache-size={config.moe_cache_size}" + elif config.moe_cache_mem is not None: + inert = f"--moe-cache-mem={config.moe_cache_mem}" elif getattr(config, "moe_cache_auto", False): inert = "--moe-cache-auto" else: @@ -1799,6 +1887,7 @@ def override(attr: str, value: Any): # this is dangerous, use with caution ) override("moe_cache_size", 0) override("moe_cache_rate", None) + override("moe_cache_mem", None) override("moe_cache_auto", False) if is_moe and config.moe_strategy == "cpu": @@ -1809,6 +1898,7 @@ def override(attr: str, value: Any): # this is dangerous, use with caution num_experts = config.model_config.num_experts if getattr(config, "moe_cache_auto", False): override("moe_cache_auto", False) + override("moe_cache_mem", None) override("moe_cache_size", 2 * num_experts) override("moe_prefill_overlap", True) logger.info_rank0( diff --git a/python/freetoken/kvcache/cache_status.py b/python/freetoken/kvcache/cache_status.py index 5b531ce265..8c29dccab5 100644 --- a/python/freetoken/kvcache/cache_status.py +++ b/python/freetoken/kvcache/cache_status.py @@ -40,6 +40,9 @@ def _moe() -> int: banks = getattr(engine.moe_offload_cache, "bank_caches", None) if not banks: return 0 + if getattr(engine.moe_offload_cache, "_raw_banks", False): + # pools of several bank layouts share byte segments: report the mean slot size + return int(sum(t.numel() for t in banks.values()) // engine.moe_offload_cache.cache_size) # Each bank cache is (cache_size, *row_shape); one slot's bytes = row bytes summed over # the format's banks (== cache_budget.expert_bytes_per_slot on the source rows). return int(sum(t[0].numel() * t.element_size() for t in banks.values())) diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index 0595b1d451..79ad9e3a5f 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -189,6 +189,10 @@ def __init__( ) self.offload_cache: OffloadMoeCache | None = None + @property + def skips_inactive_slots(self) -> bool: + return self.quant_method is not None and self.quant_method.kernel.supports_inactive_slots + def forward( self, hidden_states: torch.Tensor, @@ -289,7 +293,7 @@ def _decode_routed( hidden_states, topk_weights, topk_ids, - views=cache.bank_views(), + views=cache.pool_bank_views(self.layer_id), n=None, alphas=cache.alphas_for_slots(self.layer_id), is_prefill=False, @@ -317,7 +321,8 @@ def _decode_hybrid( assert executor is not None, "CPU MoE executor was not initialized" # the kernel reads every route's slot: give the inactive ones slot 0, which holds a # complete payload or the bank's zero fill (never uninitialized storage), times weight 0 - clamp = not (self.quant_method is not None and self.quant_method.kernel.supports_inactive_slots) + # with several quants this holds only because the cache zeroes staged slot 0 rows (zero_inactive_route_slots) + clamp = not self.skips_inactive_slots gpu_slots, gpu_w, cpu_ids = cache.ensure_experts_hybrid(self.layer_id, topk_ids, topk_weights, clamp=clamp) pending = executor.decode_submit(self.layer_id, hidden_states, topk_weights, cpu_ids) @@ -333,7 +338,7 @@ def _decode_hybrid( hidden_states, gpu_w, gpu_slots, - views=cache.bank_views(), + views=cache.pool_bank_views(self.layer_id), n=None, alphas=cache.alphas_for_slots(self.layer_id), is_prefill=False, @@ -370,16 +375,18 @@ def _prefill_routed( return out cache.materialize_layer(self.layer_id) cache.copy_missing() - return self._expert_gemm( + out = self._expert_gemm( cache, hidden_states, topk_weights, topk_ids, - views=cache.bank_views(self.num_experts), + views=cache.stage_views(self.layer_id, 0), n=self.num_experts, alphas=cache.alphas_for_layer(self.layer_id), is_prefill=True, ) + cache.release_prefill_layer(self.layer_id) + return out def _wait_prefill_overlap(self, cache: OffloadMoeCache) -> tuple[torch.Tensor, ...]: """Double-buffer choreography for this layer's overlap prefill: kick off the diff --git a/python/freetoken/models/deepseek_v4/moe.py b/python/freetoken/models/deepseek_v4/moe.py index 35a316f103..aadcdc0481 100644 --- a/python/freetoken/models/deepseek_v4/moe.py +++ b/python/freetoken/models/deepseek_v4/moe.py @@ -120,8 +120,9 @@ def _prefill_routed( cache = self.offload_cache assert cache is not None # unpinned (LOCKED) layers must take the base materialize path: their copy_missing is the whole-layer pageable branch with position == expert id, which ensure_experts's LRU slot remap would contradict (the GEMM would gather other experts' weights) + # one threshold for all layers (the smallest pool): with overlap prefill, the layers of a chunk must all stream or all take the slot path if ( - hidden_states.shape[0] * self.top_k >= self.num_experts + hidden_states.shape[0] * self.top_k >= min(self.num_experts, *cache.pool_sizes) or cache.is_unpinned_layer(self.layer_id) ): return super()._prefill_routed(hidden_states, topk_weights, topk_ids) @@ -134,7 +135,7 @@ def _prefill_routed( hidden_states, topk_weights, topk_ids, - views=cache.bank_views(), + views=cache.pool_bank_views(self.layer_id), n=None, alphas=cache.alphas_for_slots(self.layer_id), is_prefill=True, diff --git a/python/freetoken/moe/benchbw.py b/python/freetoken/moe/benchbw.py index 9430ab2a84..cd260cb2f6 100644 --- a/python/freetoken/moe/benchbw.py +++ b/python/freetoken/moe/benchbw.py @@ -413,14 +413,11 @@ def _build_gather_rig(fmt: str, wl: Workload, device: torch.device): specs = _offload_bank_specs(fmt, H, I) cache = OffloadMoeCache(num_layers=1, num_experts=E, cache_size=E, device=device, quant_format=fmt) total_bytes = 0 + sources = {} for name, (elems, dtype) in specs.items(): - src = alloc_pinned_tensor(E, elems, dtype=dtype) # cudaHostAlloc (mapped) like the loaders - dst = torch.empty(E, elems, dtype=dtype, device=device) - cache.bank_sources[name] = [src] - cache.bank_caches[name] = dst - cache.banks.append(([src], dst)) + sources[name] = [alloc_pinned_tensor(E, elems, dtype=dtype)] # cudaHostAlloc (mapped) like the loaders total_bytes += E * elems * dtype.itemsize - cache._build_copy_plan() # enable the fused path (production default) when aligned + cache.set_bank_sources(sources) # builds the fused copy plan (production default) when aligned cache.evict_slots[:E].copy_(torch.arange(E, dtype=torch.int32, device=device)) cache.src_indices[:E].copy_(torch.randperm(E, device=device).to(torch.int32)) diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py index 7500174ae4..5521c34e1f 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/python/freetoken/moe/expert_banks.py @@ -48,6 +48,8 @@ class ExpertBanks: kind: QuantKind | None = None kernel: str | None = None layout: dict | None = None + # per-layer quant tag (e.g. the gate_up and down GGUF types), for layers whose banks match in dtype and row shape but not in quant; None = one quant + layer_tags: list | None = None def _dummy_fill(role: str, tensor: torch.Tensor) -> None: diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 984b7124ad..701524e0a1 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -77,7 +77,7 @@ "down_bias", ), # DeepSeek-V4 FP4: packed e2m1 codes + e8m0 per-32 block scales, no global scale - # (4 banks). Read by DeepSeek-V4's own DS-FP4 grouped GEMV kernels via bank_views(). + # (4 banks). Read by DeepSeek-V4's own DS-FP4 grouped GEMV kernels via pool_bank_views(). "ds_fp4": ("gate_up_packed", "gate_up_scale", "down_packed", "down_scale"), } @@ -104,6 +104,14 @@ MARLIN_MAX_CACHE_SIZE = 992 +def bank_layouts(sources: dict[str, list[torch.Tensor]], layer_tags: list | None = None) -> list[tuple]: + """Each layer's layout key: its quant tag and every bank's (dtype, row shape). A slot pool holds layers of one key.""" + num_layers = len(next(iter(sources.values()))) + tags = layer_tags if layer_tags is not None else [None] * num_layers + assert len(tags) == num_layers, (len(tags), num_layers) + return [(tags[l], tuple((t[l].dtype, tuple(t[l].shape[1:])) for t in sources.values())) for l in range(num_layers)] + + @dataclass class OffloadMoeCache: num_layers: int @@ -148,6 +156,9 @@ class OffloadMoeCache: # bank layout from the expert kernel (a BankSpec per role); when given it replaces the _BANK_SCHEMAS lookup and the slot cap comes from max_slots layout: dict | None = None max_slots: int | None = None + # Each layer belongs to one pool: a contiguous slot range with its own LRU scope; pool p owns slots [sum(pool_sizes[:p]), +pool_sizes[p]). Slot ids in slot_for_id are local to the pool. None = one pool for all layers. + pool_of_layer: list[int] | None = None + pool_sizes: list[int] | None = None def __post_init__(self) -> None: policy_ids = {"lru": 0} @@ -155,6 +166,14 @@ def __post_init__(self) -> None: assert self.decode_target in ("gpu", "cpu", "hybrid"), self.decode_target if self.layout is None: assert self.quant_format in _BANK_SCHEMAS, f"unknown quant_format {self.quant_format!r}" + if self.pool_sizes is None: + self.pool_of_layer = [0] * self.num_layers + self.pool_sizes = [self.cache_size] + self.pool_of_layer = list(self.pool_of_layer) + self.pool_sizes = list(self.pool_sizes) + assert len(self.pool_of_layer) == self.num_layers, (len(self.pool_of_layer), self.num_layers) + assert sorted(set(self.pool_of_layer)) == list(range(len(self.pool_sizes))), self.pool_of_layer + assert sum(self.pool_sizes) == self.cache_size and min(self.pool_sizes) > 0, (self.pool_sizes, self.cache_size) # Attached by the engine for decode_target == "cpu" (CpuMoeExecutor); None # for the GPU decode path. self.cpu_executor = None @@ -162,8 +181,16 @@ def __post_init__(self) -> None: # offload/PCIe path. Set by the engine after construction (empty = all-GPU, # all layers = the plain --moe-strategy cpu case). self.cpu_layer_ids: frozenset = frozenset() + # Bank row layouts, set by set_bank_sources: layouts[g][k] is (dtype, row shape) of bank k (layouts of two quant tags may be equal), layer_layout[l] is layer l's g. + # With one (dtype, row shape) per bank a bank cache is typed [cache_size, *row]; with several it is a byte segment each pool views with its own layout. + self.layouts: list[tuple[tuple[torch.dtype, tuple[int, ...]], ...]] = [] + self.layer_layout: list[int] = [0] * self.num_layers + self._raw_banks = False + # Set by the engine when hybrid decode hands CPU routes to a kernel that reads slot 0 with weight 0 (no supports_inactive_slots). + # With several quants a staging buffer can leave another quant's bytes in a pool's slot 0, which may decode to NaN, so the last prefill layer zeroes those rows. + self.zero_inactive_route_slots = False # num_experts floor + nvfp4_marlin slot cap, shared with the runtime-rebuild path. - self.validate_rebuild(self.cache_size) + self.validate_rebuild(self.cache_size, self.pool_sizes) assert not self.prefill_overlap or self.cache_size >= 2 * self.num_experts, ( "Prefill overlap borrows two full expert-layer buffers from the unified MoE " "cache, so cache_size must be at least 2 * num_experts " @@ -254,25 +281,28 @@ def __post_init__(self) -> None: ) # (per-layer sources, cache) per bank, in schema order. Every piece of cache # machinery that moves bank bytes (copy_missing, the prefill double buffers, - # bank_views) iterates this list, so the slot cache is bank-count agnostic. + # the pool and staging views) iterates this list, so the slot cache is bank-count agnostic. self.banks: list[tuple[list[torch.Tensor], torch.Tensor]] = [] # Fused multi-bank copy descriptor (built by set_bank_sources/_build_copy_plan). # Source pointers are per layer (_copy_src_ptrs[layer_id] -> [num_banks] device - # tensor); dst/feat are layer-invariant. + # tensor); dst pointers are per pool and row bytes per layout. self._copy_fused_ok = False self._copy_blocks_per_bank = 8 self._copy_dst_ptrs: torch.Tensor | None = None + self._copy_dst_ptrs_pool: list[torch.Tensor] | None = None self._copy_src_ptrs: list[torch.Tensor] | None = None - self._copy_feat_bytes: torch.Tensor | None = None + self._copy_feat_bytes: list[torch.Tensor] | None = None # The layer whose misses ensure_experts/materialize_layer staged last; consumed # by copy_missing to pick the per-layer source (part of the same pending-copy # state as evict_slots/src_indices/num_indices). # _pending_whole_layer records WHICH staged it: the pageable branch is only sound after materialize_layer self._pending_src_layer: int | None = None self._pending_whole_layer = False - # Per-bank [2, num_experts, ...] double-buffer views over the slot cache's - # first 2 * num_experts slots (set up when prefill_overlap is enabled). - self.prefill_bank_buffers: list[torch.Tensor] = [] + # The two prefill staging buffers are the first 2 * _stage_bytes[k] bytes of every bank, on top of the first pools. + # _stage_views[g][b] are buffer b's per-bank [num_experts, ...] views for layout g; _inval[b] are the global slots it overwrites. + self._stage_bytes: list[int] = [] + self._stage_views: list[list[tuple[torch.Tensor, ...] | None]] = [] + self._inval: list[slice | torch.Tensor] = [] self.prefill_copy_stream: torch.cuda.Stream | None = None self.prefill_begin_event: torch.cuda.Event | None = None self.prefill_ready_events: list[torch.cuda.Event] = [] @@ -282,7 +312,7 @@ def __post_init__(self) -> None: self._prefill_buffer_has_release_event: list[bool] = [False, False] # hit-D2D split state: pinned begin-of-chunk snapshot of slot_for_id (the # classification input; frozen for the chunk -- no decode runs inside one, - # and buffer invalidation only clears slot < 2E entries, which classify as + # and buffer invalidation only clears slots under the staging buffers, which classify as # miss regardless), the lazily resolved batch-memcpy entry point (False = # unavailable), and row counters for cache reports. self._prefill_slot_snapshot: torch.Tensor | None = None @@ -292,11 +322,116 @@ def __post_init__(self) -> None: self._batch_memcpy = None self.prefill_hit_rows = 0 self.prefill_total_rows = 0 + self._slice_pools() + + def _pool_layout(self, pool: int) -> int: + return self.layer_layout[self.pool_of_layer.index(pool)] + + def _bank_rows(self, k: int, layout: int, offset: int, nrows: int) -> torch.Tensor: + """``nrows`` rows of bank ``k`` in ``layout``, starting ``offset`` bytes into the bank cache.""" + cache = self.banks[k][1] + row = self._row_bytes[layout][k] + if not self._raw_banks: + first = offset // row + return cache if first == 0 and nrows == cache.shape[0] else cache[first : first + nrows] + dtype, shape = self.layouts[layout][k] + return cache[offset : offset + nrows * row].view(dtype).view(nrows, *shape) + + def _stage_fits(self, buffers: int) -> bool: + return all(c.numel() * c.element_size() >= buffers * s for (_, c), s in zip(self.banks, self._stage_bytes)) + + def _slice_pools(self) -> None: + E = self.num_experts + bases = [0] + for size in self.pool_sizes[:-1]: + bases.append(bases[-1] + size) + self.pool_bases = bases + spans = list(zip(bases, self.pool_sizes)) + + def cut(t: torch.Tensor, b: int, s: int) -> torch.Tensor: + return t if s == t.shape[0] else t[b : b + s] + + self._pool_id_of_slot = [cut(self.id_of_slot, b, s) for b, s in spans] + self._pool_usage = [cut(self.usage, b, s) for b, s in spans] + # rows[k][p]: bytes of one slot of pool p in bank k; before the banks exist, one unit per slot with a stage of E + if self.banks: + rows = [[self._row_bytes[self._pool_layout(p)][k] for p in range(len(spans))] for k in range(len(self.banks))] + stage = self._stage_bytes + else: + rows, stage = [[1] * len(spans)], [E] + offsets = [] + for row in rows: + offsets.append([0]) + for s, r in zip(self.pool_sizes[:-1], row): + offsets[-1].append(offsets[-1][-1] + s * r) + self._pool_offsets = offsets + # local slots below this index overlap the prefill staging area in some bank + self._pool_stable_from = [ + max(min(max(-(-(2 * stage[k] - offsets[k][p]) // rows[k][p]), 0), s) for k in range(len(rows))) + for p, s in enumerate(self.pool_sizes) + ] + self._inval = [] + masks = [] + for b in range(2): + mask = torch.zeros(self.cache_size, dtype=torch.bool) + masks.append(mask) + for k in range(len(rows)): + lo, hi = b * stage[k], (b + 1) * stage[k] + for p, (base, s) in enumerate(spans): + first = max((lo - offsets[k][p]) // rows[k][p], 0) + end = min(-(-(hi - offsets[k][p]) // rows[k][p]), s) + mask[base + first : base + max(first, end)] = True + idx = mask.nonzero().flatten() + if idx.numel() == 0 or int(idx[-1] - idx[0]) + 1 == idx.numel(): + self._inval.append(slice(int(idx[0]), int(idx[-1]) + 1) if idx.numel() else slice(0, 0)) + else: + self._inval.append(idx.to(self.device)) + self._stage_dirty = [False, False] + self._stage_slot0_rows = [[], []] + if not self.banks: + self._pool_bank_views, self._stage_views = [], [] + return + self._pool_bank_views = [ + tuple(self._bank_rows(k, self._pool_layout(p), offsets[k][p], s) for k in range(len(self.banks))) + for p, s in enumerate(self.pool_sizes) + ] + self._stage_views = [ + [ + tuple(self._bank_rows(k, g, b * stage[k], E) for k in range(len(self.banks))) if self._stage_fits(b + 1) else None + for b in range(2) + ] + for g in range(len(self.layouts)) + ] + if len(self.layouts) > 1: + self._stage_slot0_rows = [ + [self._bank_rows(k, self._pool_layout(p), offsets[k][p], 1) for p, base in enumerate(bases) if masks[b][base] for k in range(len(self.banks))] + for b in range(2) + ] + + def stage_views(self, layer_id: int, buffer_id: int) -> tuple[torch.Tensor, ...]: + """Per-bank ``[num_experts, ...]`` views of prefill staging buffer ``buffer_id`` in ``layer_id``'s layout; row == expert id.""" + views = self._stage_views[self.layer_layout[layer_id]][buffer_id] + assert views is not None, f"the slot cache is too small for staging buffer {buffer_id}" + return views + + def pool_state(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: + """``(id_of_slot, usage)`` views of the pool that serves ``layer_id``.""" + p = self.pool_of_layer[layer_id] + return self._pool_id_of_slot[p], self._pool_usage[p] + + def pool_slots(self, layer_id: int) -> int: + return self.pool_sizes[self.pool_of_layer[layer_id]] + + def pool_bank_views(self, layer_id: int) -> tuple[torch.Tensor, ...]: + """Per-bank slot views of the pool that serves ``layer_id``, indexed by its local slot ids.""" + assert self.banks, "set_bank_sources must register the banks first" + return self._pool_bank_views[self.pool_of_layer[layer_id]] def set_bank_sources( self, sources: dict[str, list[torch.Tensor]], layer_residency: list[str] | None = None, + layer_tags: list | None = None, ) -> None: """Attach the host (CPU pinned) expert source banks and allocate a GPU slot cache per bank, following the format's bank schema. @@ -310,6 +445,8 @@ def set_bank_sources( ``layer_residency`` labels each layer with a ``HostResidency`` value (default: all pinned). Non-pinned (LOCKED/PAGEABLE) layers have no device address: they must already be routed to the CPU executor (``cpu_layer_ids``, set BEFORE this call), the copy plan skips their rows, and their only movement is ``copy_missing``'s whole-layer pageable prefill branch -- which is why prefill overlap is incompatible with them. + + ``layer_tags`` names each layer's quant (see :func:`bank_layouts`); layers of different tags never share a pool. """ from freetoken.moe.legacy_format import canonical_role from freetoken.moe.host_banks import HostResidency @@ -343,43 +480,62 @@ def set_bank_sources( for name in self.bank_schema: per_layer = sources[name] assert len(per_layer) == self.num_layers, (name, len(per_layer)) - head = per_layer[0] - if self.layout is not None: - spec = self.layout[name] - if tuple(head.shape[1:]) != tuple(spec.shape) or head.dtype != spec.dtype: - raise ValueError( - f"bank {name!r} rows are {tuple(head.shape[1:])} {head.dtype} but the expert kernel's layout " - f"wants {tuple(spec.shape)} {spec.dtype}; the banks were packed for another kernel" - ) for layer_id, source in enumerate(per_layer): assert source.is_contiguous(), f"bank {name!r} layer {layer_id} must be contiguous" assert source.size(0) == self.num_experts, (name, layer_id, source.shape) - assert source.shape == head.shape and source.dtype == head.dtype, ( - name, layer_id, source.shape, source.dtype, - ) + if self.layout is not None: + spec = self.layout[name] + if tuple(source.shape[1:]) != tuple(spec.shape) or source.dtype != spec.dtype: + raise ValueError( + f"bank {name!r} rows are {tuple(source.shape[1:])} {source.dtype} but the expert kernel's layout " + f"wants {tuple(spec.shape)} {spec.dtype}; the banks were packed for another kernel" + ) self.bank_sources[name] = list(per_layer) - self.bank_caches[name] = torch.zeros( - (self.cache_size, *head.shape[1:]), - dtype=head.dtype, - device=self.device, - ) - self.banks = [(self.bank_sources[n], self.bank_caches[n]) for n in self.bank_schema] + keys = bank_layouts(sources, layer_tags) + unique = list(dict.fromkeys(keys)) + self.layouts = [banks for _, banks in unique] + self.layer_layout = [unique.index(key) for key in keys] + for p in range(len(self.pool_sizes)): + mixed = {self.layer_layout[l] for l, q in enumerate(self.pool_of_layer) if q == p} + assert len(mixed) == 1, f"slot pool {p} holds layers of {len(mixed)} layouts; a pool needs one" + self._raw_banks = len(set(self.layouts)) > 1 + self._row_bytes = [[math.prod(shape) * dtype.itemsize for dtype, shape in layout] for layout in self.layouts] + if self._raw_banks: + # every pool view and staging buffer starts on a row boundary of its own layout + assert all(r % 16 == 0 for rows in self._row_bytes for r in rows), self._row_bytes + self._stage_bytes = [self.num_experts * max(rows[k] for rows in self._row_bytes) for k in range(len(self.bank_schema))] + self._alloc_banks() + if not self._stage_fits(1): + raise ValueError(f"MoE cache of {self.cache_size} slots cannot stage one whole layer of experts") + if self.prefill_overlap and not self._stage_fits(2): + logger.warning("Disabling MoE prefill overlap: the slot cache cannot hold two whole layers of experts") + self.prefill_overlap = False + self._slice_pools() self._build_copy_plan() if self.prefill_overlap: self._init_prefill_overlap_buffers() + def _alloc_banks(self) -> None: + for k, name in enumerate(self.bank_schema): + if self._raw_banks: + nbytes = sum(s * self._row_bytes[self._pool_layout(p)][k] for p, s in enumerate(self.pool_sizes)) + self.bank_caches[name] = torch.zeros(nbytes, dtype=torch.uint8, device=self.device) + else: + dtype, shape = self.layouts[0][k] + self.bank_caches[name] = torch.zeros((self.cache_size, *shape), dtype=dtype, device=self.device) + self.banks = [(self.bank_sources[n], self.bank_caches[n]) for n in self.bank_schema] + def _build_copy_plan(self) -> None: self._build_fused_copy_plan() if self._copy_fused_ok or self.device.type != "cuda" or not self.banks: return - for name in self.bank_schema: - cache = self.bank_caches[name] - feat = math.prod(cache.shape[1:]) * cache.element_size() - if feat % 128: - raise RuntimeError( - f"MoE bank {name!r} rows are {feat} bytes (not a multiple of 128): " - f"only the fused multi-bank copy can move them, but it is disabled" - ) + for k, name in enumerate(self.bank_schema): + for rows in self._row_bytes: + if rows[k] % 128: + raise RuntimeError( + f"MoE bank {name!r} rows are {rows[k]} bytes (not a multiple of 128): " + f"only the fused multi-bank copy can move them, but it is disabled" + ) def _build_fused_copy_plan(self) -> None: """Precompute the fused multi-bank copy descriptor (base addrs + per-row bytes). @@ -391,24 +547,25 @@ def _build_fused_copy_plan(self) -> None: """ self._copy_fused_ok = False self._copy_dst_ptrs = None + self._copy_dst_ptrs_pool = None self._copy_src_ptrs = None self._copy_feat_bytes = None self._copy_dst_ptrs_host: list[int] = [] self._copy_src_ptrs_host: list[list[int]] = [] - self._copy_feat_bytes_host: list[int] = [] - self._gather_bank_ids: list[int] = [] - self._gather_dst_ptrs: torch.Tensor | None = None - self._gather_feat_bytes: torch.Tensor | None = None + self._gather_bank_ids: list[list[int]] = [] + self._gather_feat_bytes: list[torch.Tensor | None] = [] + self._gather_stage_ptrs: list[list[torch.Tensor] | None] = [] + self._gather_pool_ptrs: list[torch.Tensor | None] = [] if not _FUSED_COPY or self.device.type != "cuda" or not self.banks: return from freetoken.kernel.pinned import device_ptr - dst_ptrs, feats = [], [] + aligned = [self._stage_bytes] + self._row_bytes + self._pool_offsets + if any(n % 16 for ns in aligned for n in ns) or any(c.data_ptr() % 16 for _, c in self.banks): + return # leave fused disabled; copy_missing uses the per-bank path + dst_ptrs = [cache.data_ptr() for _, cache in self.banks] layer_src_ptrs = [[] for _ in range(self.num_layers)] - for per_layer, cache in self.banks: - feat = math.prod(per_layer[0].shape[1:]) * per_layer[0].element_size() - if feat % 16 != 0 or cache.data_ptr() % 16 != 0: - return # leave fused disabled; copy_missing uses the per-bank path + for per_layer, _ in self.banks: for layer_id, source in enumerate(per_layer): if layer_id in self._unpinned_layers: # unregistered layer: no device alias exists, and the row is never consumed (CPU decode; pageable prefill) @@ -422,52 +579,65 @@ def _build_fused_copy_plan(self) -> None: if src_dev % 16 != 0: return layer_src_ptrs[layer_id].append(src_dev) - dst_ptrs.append(cache.data_ptr()) - feats.append(feat) - self._copy_dst_ptrs = torch.tensor(dst_ptrs, dtype=torch.int64, device=self.device) - self._copy_src_ptrs = [ - torch.tensor(ptrs, dtype=torch.int64, device=self.device) - for ptrs in layer_src_ptrs + + def ptrs(values: list[int]) -> torch.Tensor: + return torch.tensor(values, dtype=torch.int64, device=self.device) + + self._copy_dst_ptrs = ptrs(dst_ptrs) + self._copy_dst_ptrs_pool = [ + self._copy_dst_ptrs if all(o[p] == 0 for o in self._pool_offsets) + else ptrs([d + o[p] for d, o in zip(dst_ptrs, self._pool_offsets)]) + for p in range(len(self.pool_sizes)) ] - self._copy_feat_bytes = torch.tensor(feats, dtype=torch.int64, device=self.device) + self._copy_src_ptrs = [ptrs(values) for values in layer_src_ptrs] + self._copy_feat_bytes = [ptrs(rows) for rows in self._row_bytes] self._copy_dst_ptrs_host = dst_ptrs self._copy_src_ptrs_host = layer_src_ptrs - self._copy_feat_bytes_host = feats # hit-D2D gather serves only the big banks; small banks are whole-layer # H2D entries (see _SMALL_BANK_FEAT_BYTES), so their rows never need D2D. - self._gather_bank_ids = [i for i, f in enumerate(feats) if f >= _SMALL_BANK_FEAT_BYTES] - if len(self._gather_bank_ids) == len(feats): - self._gather_dst_ptrs = self._copy_dst_ptrs - self._gather_feat_bytes = self._copy_feat_bytes - elif self._gather_bank_ids: - self._gather_dst_ptrs = self._copy_dst_ptrs[self._gather_bank_ids].contiguous() - self._gather_feat_bytes = self._copy_feat_bytes[self._gather_bank_ids].contiguous() + # The gather copies cache rows of the layer's pool into its staging buffer, row == expert id. + self._gather_bank_ids = [[k for k, r in enumerate(rows) if r >= _SMALL_BANK_FEAT_BYTES] for rows in self._row_bytes] + for g, ids in enumerate(self._gather_bank_ids): + self._gather_feat_bytes.append(ptrs([self._row_bytes[g][k] for k in ids]) if ids else None) + self._gather_stage_ptrs.append( + [ptrs([dst_ptrs[k] + b * self._stage_bytes[k] for k in ids]) for b in range(2)] if ids else None + ) + for p in range(len(self.pool_sizes)): + ids = self._gather_bank_ids[self._pool_layout(p)] + self._gather_pool_ptrs.append(ptrs([dst_ptrs[k] + self._pool_offsets[k][p] for k in ids]) if ids else None) self._copy_blocks_per_bank = default_blocks_per_bank(self.device) self._copy_fused_ok = True - def validate_rebuild(self, cache_size: int) -> None: + def validate_rebuild(self, cache_size: int, pool_sizes: list[int] | None = None) -> None: """Pure geometry validation of a rebuild target (no GPU side effects). Raises ``ValueError`` if ``cache_size`` is below the ``num_experts`` floor or - above the marlin slot cap. Called by :meth:`rebuild` and by the engine's + a pool is above the expert kernel's slot cap, or the bank byte segments cannot stage one layer. Called by :meth:`rebuild` and by the engine's pre-teardown check, so an invalid target rejects with the old cache intact (no destructive free first). """ if cache_size < self.num_experts: raise ValueError(f"cache_size {cache_size} < num_experts {self.num_experts}") - if self.max_slots is not None and cache_size > self.max_slots: + # decode kernels see one pool's slots, so the cap applies per pool + largest = cache_size if pool_sizes is None else max(pool_sizes) + what = f"moe_cache_size={cache_size}" if largest == cache_size else f"a MoE slot pool of {largest} slots" + if self.max_slots is not None and largest > self.max_slots: raise ValueError( - f"moe_cache_size={cache_size} exceeds the expert kernel's slot limit of {self.max_slots}; " + f"{what} exceeds the expert kernel's slot limit of {self.max_slots}; " f"pass --moe-cache-size {self.max_slots} or less, or let the default kernel serve the experts" ) - if self.layout is None and self.quant_format == "nvfp4_marlin" and cache_size > MARLIN_MAX_CACHE_SIZE: + if self.layout is None and self.quant_format == "nvfp4_marlin" and largest > MARLIN_MAX_CACHE_SIZE: raise ValueError( - f"moe_cache_size={cache_size} exceeds the marlin backend's slot limit of " + f"{what} exceeds the marlin backend's slot limit of " f"{MARLIN_MAX_CACHE_SIZE} (vLLM moe_align_block_size caps padded experts at " "1024); reduce moe_cache_size or force --quant-backend moe.nvfp4=triton" ) + if self._raw_banks and pool_sizes is not None: + for k, stage in enumerate(self._stage_bytes): + if sum(s * self._row_bytes[self._pool_layout(p)][k] for p, s in enumerate(pool_sizes)) < stage: + raise ValueError(f"moe_cache_size={cache_size} cannot stage one whole layer of experts") - def rebuild(self, cache_size: int) -> None: + def rebuild(self, cache_size: int, pool_sizes: list[int] | None = None) -> None: """Resize the GPU slot cache + bookkeeping to ``cache_size`` IN PLACE. Keeps the CPU/pinned ``bank_sources`` and the GPU-resident alphas; never @@ -477,9 +647,10 @@ def rebuild(self, cache_size: int) -> None: ``ctx.moe_offload_cache`` stay valid. """ assert self.bank_sources, "set_bank_sources must run before rebuild" - self.validate_rebuild(cache_size) + sizes = [cache_size] if pool_sizes is None else list(pool_sizes) + assert len(sizes) == len(self.pool_sizes) and sum(sizes) == cache_size and min(sizes) > 0, (sizes, cache_size) + self.validate_rebuild(cache_size, sizes) # 1. Tear down prefill-overlap (its buffer views alias the old bank_caches). - self.prefill_bank_buffers = [] self.prefill_copy_stream = None self.prefill_begin_event = None self.prefill_ready_events = [] @@ -490,22 +661,22 @@ def rebuild(self, cache_size: int) -> None: # 2. Drop old GPU tensors (free-before-alloc). self.banks = [] self.bank_caches = {} + self._pool_bank_views = [] + self._stage_views = [] + self._stage_slot0_rows = [[], []] self.cache_size = cache_size + self.pool_sizes = sizes if self.device.type == "cuda": torch.cuda.synchronize(self.device) torch.cuda.empty_cache() # 3. Reallocate the slot cache from the retained host sources. - for name in self.bank_schema: - head = self.bank_sources[name][0] - self.bank_caches[name] = torch.zeros( - (cache_size, *head.shape[1:]), dtype=head.dtype, device=self.device - ) - self.banks = [(self.bank_sources[n], self.bank_caches[n]) for n in self.bank_schema] - self._build_copy_plan() # slot caches were reallocated -> refresh fused-copy addrs + self._alloc_banks() # 4. Reallocate cache_size-shaped bookkeeping; reset the slot map (cold start). self.slot_for_id.fill_(-1) self.id_of_slot = torch.full((cache_size,), -1, dtype=torch.int32, device=self.device) self.usage = torch.zeros((cache_size,), dtype=torch.int64, device=self.device) + self._slice_pools() + self._build_copy_plan() # slot caches were reallocated -> refresh fused-copy addrs plan_slots = max(self.num_experts, cache_size) self.evict_slots = torch.empty((plan_slots,), dtype=torch.int32, device=self.device) self.src_indices = torch.empty((plan_slots,), dtype=torch.int32, device=self.device) @@ -529,10 +700,10 @@ def rebuild(self, cache_size: int) -> None: self.prefill_total_rows = 0 self._hit_d2d_fallback_logged = False # geometry changed; re-log if still unusable # 5. Re-evaluate prefill overlap against the new size. - if self.prefill_overlap and cache_size < 2 * self.num_experts: + if self.prefill_overlap and not self._stage_fits(2): logger.warning( f"Disabling MoE prefill overlap on rebuild: cache_size {cache_size} " - f"< 2*num_experts {2 * self.num_experts}." + f"cannot hold two whole layers of experts (2*num_experts {2 * self.num_experts} slots of one layout)." ) self.prefill_overlap = False if self.prefill_overlap: @@ -585,8 +756,9 @@ def alphas_for_slots(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor] | ``layer_id`` -- are ever read by the grouped GEMM.""" if self.gate_up_alpha is None: return None + id_of_slot, _ = self.pool_state(layer_id) idx = layer_id * self.num_experts + ( - self.id_of_slot.clamp(min=0).long() % self.num_experts + id_of_slot.clamp(min=0).long() % self.num_experts ) return self.gate_up_alpha[idx], self.down_alpha[idx] @@ -600,25 +772,12 @@ def alphas_for_layer(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor] | hi = lo + self.num_experts return self.gate_up_alpha[lo:hi], self.down_alpha[lo:hi] - def bank_views(self, n: int | None = None) -> tuple[torch.Tensor, ...]: - """Per-bank cache views in registration order: the full ``[S]`` slot cache - (decode), or its first ``n`` slots (materialized layer).""" - assert self.banks, "set_bank_sources must register the banks first" - if n is None: - return tuple(cache for _, cache in self.banks) - return tuple(cache[:n] for _, cache in self.banks) - def _init_prefill_overlap_buffers(self) -> None: assert self.banks, "set_bank_sources must register the banks first" self._prefill_buffer_layer = [None, None] self._prefill_buffer_released = [True, True] self._prefill_buffer_has_release_event = [False, False] - # The double buffers borrow the slot cache's first 2 * num_experts slots - # (one full expert layer per buffer), one view per registered bank. - self.prefill_bank_buffers = [ - cache[: 2 * self.num_experts].view(2, self.num_experts, *cache.shape[1:]) - for _, cache in self.banks - ] + # the double buffers are the staging views (_slice_pools), on top of the first pools' slots if self.device.type == "cuda": self.prefill_copy_stream = torch.cuda.Stream(device=self.device) self.prefill_ready_events = [torch.cuda.Event() for _ in range(2)] @@ -638,14 +797,14 @@ def _init_prefill_overlap_buffers(self) -> None: self._prefill_hit_num = torch.zeros((1,), dtype=torch.int64, device=self.device) def _invalidate_prefill_buffer(self, buffer_id: int) -> None: - slot_start = buffer_id * self.num_experts - slot_end = slot_start + self.num_experts - old_ids = self.id_of_slot[slot_start:slot_end] + slots = self._inval[buffer_id] + old_ids = self.id_of_slot[slots] self.slot_for_id.view(-1)[old_ids[old_ids >= 0].long()] = -1 - old_ids.fill_(-1) + self.id_of_slot[slots] = -1 # usage=0 makes these slots the oldest, so the argmin(usage) victim selection in # ensure_experts evicts them first. - self.usage[slot_start:slot_end].zero_() + self.usage[slots] = 0 + self._stage_dirty[buffer_id] = True def begin_prefill(self) -> None: if not self.prefill_overlap: @@ -676,7 +835,7 @@ def prefetch_prefill_layer(self, layer_id: int) -> None: if layer_id < 0: raise ValueError(f"Invalid prefill layer id: {layer_id}") - assert self.banks and self.prefill_bank_buffers + assert self.banks and self._stage_views buffer_id = layer_id % 2 if self._prefill_buffer_layer[buffer_id] == layer_id: @@ -688,8 +847,8 @@ def prefetch_prefill_layer(self, layer_id: int) -> None: def copy() -> None: self._invalidate_prefill_buffer(buffer_id) - for (per_layer, _), buffer in zip(self.banks, self.prefill_bank_buffers): - buffer[buffer_id].copy_(per_layer[layer_id], non_blocking=True) + for (per_layer, _), view in zip(self.banks, self.stage_views(layer_id, buffer_id)): + view.copy_(per_layer[layer_id], non_blocking=True) if self._prefill_hit_d2d_active: self._prefetch_split(layer_id, buffer_id) @@ -720,10 +879,10 @@ def _hit_d2d_usable(self) -> bool: reason = "FREETOKEN_SKIP_FAST_INDEX_COPY is set (the hit gather would be a no-op)" elif not self._copy_fused_ok: reason = "the fused copy plan is unavailable (bank alignment or FREETOKEN_FUSED_COPY=0)" - elif self.cache_size <= 2 * self.num_experts: + elif all(stable >= s for stable, s in zip(self._pool_stable_from, self.pool_sizes)): reason = ( f"cache_size {self.cache_size} leaves no hit region " - f"(needs > {2 * self.num_experts} slots)" + f"(every slot lies under the prefill staging buffers of {2 * self.num_experts} experts)" ) elif not self._resolve_batch_memcpy(): reason = "cudaMemcpyBatchAsync is unavailable" # resolve logged the specifics @@ -761,7 +920,7 @@ def _prefetch_split(self, layer_id: int, buffer_id: int) -> None: run list comes from the begin-of-chunk snapshot because the batch API takes HOST pointer arrays. Live-vs-snapshot cannot disagree: the only chunk-internal writer (buffer invalidation) rewrites slots already below - the 2E threshold, and slots < 2E (including -1) are misses on both sides + the pool's staging threshold, and those slots (and -1) are misses on both sides -- the buffers own those slots, so their bytes are volatile within the chunk. Hit and miss row sets are disjoint, so the streams need no ordering against each other. @@ -772,18 +931,20 @@ def _prefetch_split(self, layer_id: int, buffer_id: int) -> None: from freetoken.moe.offload_kernels import prefill_hit_compact E = self.num_experts + layout = self.layer_layout[layer_id] + pool = self.pool_of_layer[layer_id] snap = self._prefill_snapshot_np[layer_id] - hit_mask = snap >= 2 * E + hit_mask = snap >= self._pool_stable_from[pool] self.prefill_hit_rows += int(hit_mask.sum()) self.prefill_total_rows += E - if self._gather_dst_ptrs is not None: - prefill_hit_compact(self, layer_id, buffer_id) + if self._gather_feat_bytes[layout] is not None: + prefill_hit_compact(self, layer_id) # blocks_per_bank=64 vs the PCIe-tuned default of 8: HBM D2D needs the # wider grid (~22 GB/s per 1024-thread block on H100). fast_index_copy_multi_jit( - self._gather_dst_ptrs, - self._gather_dst_ptrs, - self._gather_feat_bytes, + self._gather_stage_ptrs[layout][buffer_id], + self._gather_pool_ptrs[pool], + self._gather_feat_bytes[layout], self._prefill_hit_dst, self._prefill_hit_src, self._prefill_hit_num, @@ -799,16 +960,17 @@ def _prefetch_split(self, layer_id: int, buffer_id: int) -> None: starts = miss[run_starts] lengths = np.diff(np.concatenate((run_starts, [miss.size]))) dst, src, nbytes = [], [], [] - for b, feat in enumerate(self._copy_feat_bytes_host): + for b, feat in enumerate(self._row_bytes[layout]): + stage = self._copy_dst_ptrs_host[b] + buffer_id * self._stage_bytes[b] if feat < _SMALL_BANK_FEAT_BYTES: # Whole layer as one entry, EVEN with zero misses: it keeps every # batch entry above the driver's async floor and covers the hit # rows the gather skips for these banks. - dst.append(self._copy_dst_ptrs_host[b] + buffer_id * E * feat) + dst.append(stage) src.append(self._copy_src_ptrs_host[layer_id][b]) nbytes.append(E * feat) elif miss.size: - dst.extend(self._copy_dst_ptrs_host[b] + (buffer_id * E + starts) * feat) + dst.extend(stage + starts * feat) src.extend(self._copy_src_ptrs_host[layer_id][b] + starts * feat) nbytes.extend(lengths * feat) if dst: @@ -826,15 +988,22 @@ def wait_prefill_layer(self, layer_id: int) -> tuple[torch.Tensor, ...]: marlin/b12x ``(gate_up_packed, gate_up_scale, down_packed, down_scale)``; nvfp4 native adds the two global banks after each scale bank.""" assert self.prefill_overlap - assert self.prefill_bank_buffers + assert self._stage_views self.prefetch_prefill_layer(layer_id) buffer_id = layer_id % 2 assert self._prefill_buffer_layer[buffer_id] == layer_id if self.prefill_ready_events: torch.cuda.current_stream(self.device).wait_event(self.prefill_ready_events[buffer_id]) - return tuple(buffer[buffer_id] for buffer in self.prefill_bank_buffers) + return self.stage_views(layer_id, buffer_id) def release_prefill_layer(self, layer_id: int) -> None: + if self.zero_inactive_route_slots and layer_id == self.num_layers - 1: + # the staged slots were invalidated and nothing registers them within the forward, so their bytes are free to clear + for b in range(2): + if self._stage_dirty[b]: + for row in self._stage_slot0_rows[b]: + row.zero_() + self._stage_dirty[b] = False if not self.prefill_overlap: return buffer_id = layer_id % 2 @@ -959,6 +1128,8 @@ def decode_miss_stats_per_layer(self) -> dict: s, m, a, f = steps[L], missing[L], active[L], fetched[L] per_layer.append({ "layer": L, + "pool": self.pool_of_layer[L], + "pool_slots": self.pool_slots(L), "steps": s, "active_per_step": (a / s) if s else 0.0, "missing_per_step": (m / s) if s else 0.0, @@ -1010,11 +1181,13 @@ def copy_missing(self) -> None: f"pageable materialize (position == expert id); ensure_experts's " f"LRU slot remap cannot be honored without a device alias" ) - # the only copy a non-pinned layer ever needs is the non-overlap prefill materialize, which schedules the whole layer into slots [0, num_experts) with position == expert id -- a plain synchronous pageable H2D copy + # the only copy a non-pinned layer ever needs is the non-overlap prefill materialize, which schedules the whole layer into staging buffer 0 with position == expert id -- a plain synchronous pageable H2D copy # never CUDA-graph captured: prefill is not captured, and decode never reaches this branch (it routes to the CPU executor) - for per_layer, cache in self.banks: - cache[: self.num_experts].copy_(per_layer[layer_id]) + for (per_layer, _), view in zip(self.banks, self.stage_views(layer_id, 0)): + view.copy_(per_layer[layer_id]) return + # materialize stages rows of staging buffer 0; ensure stages slots local to the layer's pool + whole = self._pending_whole_layer if self._copy_fused_ok: from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit @@ -1023,9 +1196,9 @@ def copy_missing(self) -> None: # src_indices holds layer-local expert rows, resolved against this layer's # source pointers (layer_id is a static int per captured graph node). fast_index_copy_multi_jit( - self._copy_dst_ptrs, + self._copy_dst_ptrs if whole else self._copy_dst_ptrs_pool[self.pool_of_layer[layer_id]], self._copy_src_ptrs[layer_id], - self._copy_feat_bytes, + self._copy_feat_bytes[self.layer_layout[layer_id]], self.evict_slots, self.src_indices, self.num_indices, @@ -1035,7 +1208,8 @@ def copy_missing(self) -> None: from freetoken.kernel import fast_index_copy_jit - for per_layer, cache in self.banks: + views = self.stage_views(layer_id, 0) if whole else self.pool_bank_views(layer_id) + for (per_layer, _), cache in zip(self.banks, views): fast_index_copy_jit( cache, self.evict_slots, diff --git a/python/freetoken/moe/offload_kernels.py b/python/freetoken/moe/offload_kernels.py index ea12b6fb17..1d863a2f54 100644 --- a/python/freetoken/moe/offload_kernels.py +++ b/python/freetoken/moe/offload_kernels.py @@ -26,11 +26,12 @@ def ensure_experts(cache, layer_id: int, expert_ids: torch.Tensor) -> None: tensor. ``out_indices`` aliases the input, preserving the in-place rewrite every downstream GEMM depends on. """ + id_of_slot, usage = cache.pool_state(layer_id) lru_ensure( expert_ids, cache.slot_for_id.view(-1), - cache.id_of_slot, - cache.usage, + id_of_slot, + usage, cache.step, expert_ids, cache.src_indices, @@ -80,13 +81,13 @@ def ensure_experts_hybrid( return routes -def prefill_hit_compact(cache, layer_id: int, buffer_id: int) -> None: +def prefill_hit_compact(cache, layer_id: int) -> None: """Compact this layer's cache-resident experts into gather indices, device-side. - hit = slot_for_id[layer_id][e] >= 2 * num_experts (the double buffer owns the + hit = slot_for_id[layer_id][e] >= the pool's staging threshold (the double buffer owns the slots below, so those bytes are volatile within a prefill chunk and classify as miss). Writes fixed-shape ``_prefill_hit_dst``/``_prefill_hit_src`` (buffer - row / cache slot) and the count into ``_prefill_hit_num``; one launch on the + row / pool-local slot) and the count into ``_prefill_hit_num``; one launch on the current stream, no host sync. Safe against the concurrent buffer invalidation on the copy stream: that only rewrites entries already below the threshold.""" num_experts = cache.num_experts @@ -95,15 +96,23 @@ def prefill_hit_compact(cache, layer_id: int, buffer_id: int) -> None: cache._prefill_hit_dst, cache._prefill_hit_src, cache._prefill_hit_num, - buffer_id * num_experts, - 2 * num_experts, + cache._pool_stable_from[cache.pool_of_layer[layer_id]], num_experts, BLOCK=triton.next_power_of_2(num_experts), ) def materialize_layer(cache, layer_id: int) -> None: - _materialize_layer_gpu(cache, layer_id) + if len(cache.pool_sizes) == 1: + _materialize_layer_gpu(cache, layer_id) + return + # staging buffer 0 may overlap several pools, so it only holds the layer for this forward and registers nothing + E = cache.num_experts + cache._invalidate_prefill_buffer(0) + rows = torch.arange(E, dtype=torch.int32, device=cache.device) + cache.evict_slots[:E].copy_(rows) + cache.src_indices[:E].copy_(rows) + cache.num_indices.fill_(E) def reset_cache(cache) -> None: @@ -124,14 +133,16 @@ def _ensure_experts_hybrid_gpu( max_fetch: int, frac_q16: int, ) -> None: + id_of_slot, usage = cache.pool_state(layer_id) + num_slots = id_of_slot.numel() block_e = triton.next_power_of_2(cache.num_experts) - block_c = triton.next_power_of_2(cache.cache_size) + block_c = triton.next_power_of_2(num_slots) num_warps = 8 if block_c >= 2048 else 4 _ensure_experts_hybrid_kernel[(1,)]( expert_ids, cache.slot_for_id, - cache.id_of_slot, - cache.usage, + id_of_slot, + usage, cache.step, cache.active_mask, cache.evict_slots, @@ -148,7 +159,7 @@ def _ensure_experts_hybrid_gpu( int(max_fetch), int(frac_q16), cache.num_experts, - cache.cache_size, + num_slots, BLOCK_E=block_e, BLOCK_C=block_c, BY_RECENCY=_HYBRID_FETCH_BY_RECENCY, @@ -257,7 +268,7 @@ def _materialize_layer_kernel( -@triton.jit(do_not_specialize=["layer_id", "num_active", "max_fetch", "fetch_frac_q16"]) +@triton.jit(do_not_specialize=["layer_id", "num_active", "max_fetch", "fetch_frac_q16", "cache_size", "id_of_slot_ptr", "usage_ptr"]) def _ensure_experts_hybrid_kernel( expert_ids_ptr, slot_for_id_ptr, @@ -280,7 +291,7 @@ def _ensure_experts_hybrid_kernel( max_fetch, fetch_frac_q16, num_experts: tl.constexpr, - cache_size: tl.constexpr, + cache_size, BLOCK_E: tl.constexpr, BLOCK_C: tl.constexpr, BY_RECENCY: tl.constexpr, @@ -329,10 +340,11 @@ def _ensure_experts_hybrid_kernel( (lo + 1) * ((1 << 16) - fetch_frac_q16), (num_missing - lo - 1) * fetch_frac_q16 ) max_fetch = tl.where(cost_lo <= cost_hi, lo, lo + 1) - num_fetch = tl.minimum(num_missing, max_fetch) + is_hit = is_active & (slot >= 0) + # a pool smaller than num_experts may hold fewer free slots than misses + num_fetch = tl.minimum(tl.minimum(num_missing, max_fetch), cache_size - tl.sum(is_hit.to(tl.int32))) tl.store(num_missing_full_ptr, num_missing.to(tl.int64)) tl.store(num_indices_ptr, num_fetch.to(tl.int64)) - is_hit = is_active & (slot >= 0) tl.store(usage_ptr + slot, step, mask=is_hit) # Fetch-selection priority: encode (recency desc, id asc) into one strictly-ordered @@ -405,14 +417,13 @@ def _ensure_experts_hybrid_kernel( tl.store(expert_recency_ptr + base + off_e, step_vec, mask=is_active & e_mask) -@triton.jit(do_not_specialize=["buffer_base"]) +@triton.jit(do_not_specialize=["threshold"]) def _prefill_hit_compact_kernel( slot_ptr, # [num_experts] int32: this layer's slot_for_id row - dst_ptr, # [num_experts] int32 out: buffer rows, compacted - src_ptr, # [num_experts] int32 out: cache slots, compacted + dst_ptr, # [num_experts] int32 out: buffer rows (expert ids), compacted + src_ptr, # [num_experts] int32 out: pool-local slots, compacted num_ptr, # [1] int64 out: hit count - buffer_base, # buffer_id * num_experts - threshold, # 2 * num_experts + threshold, # first pool-local slot outside the staging buffers num_experts, BLOCK: tl.constexpr, ): @@ -421,6 +432,6 @@ def _prefill_hit_compact_kernel( slots = tl.load(slot_ptr + offs, mask=lane, other=-1) is_hit = lane & (slots >= threshold) pos = tl.cumsum(is_hit.to(tl.int32)) - 1 - tl.store(dst_ptr + pos, (buffer_base + offs).to(tl.int32), mask=is_hit) + tl.store(dst_ptr + pos, offs.to(tl.int32), mask=is_hit) tl.store(src_ptr + pos, slots, mask=is_hit) tl.store(num_ptr, tl.sum(is_hit.to(tl.int64))) diff --git a/python/freetoken/moe/pool_policy.py b/python/freetoken/moe/pool_policy.py new file mode 100644 index 0000000000..990c18ade2 --- /dev/null +++ b/python/freetoken/moe/pool_policy.py @@ -0,0 +1,149 @@ +"""Slot pool policy for the offload MoE cache: which layers share a pool, and how many slots each pool gets. + +Pure functions. The engine supplies the constraints and hands the result to ``OffloadMoeCache``: + + groups = group_layers(scope, layouts, slot_bytes) # --moe-cache-scope, once per load + pools = make_pools(groups, num_experts, min_num_slots, slot_limit) # the minimum follows the decode batch + sizes = split_slots(total, pools, counts) # --moe-cache-split, counts from resolve_profile + OffloadMoeCache(pool_of_layer=pool_of_layer(pools, num_layers), pool_sizes=sizes, ...) +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path + +import torch + +from freetoken.utils import init_logger + +logger = init_logger(__name__) + +SCOPES = ("unified", "per-layer") + + +@dataclass(frozen=True) +class PoolSpec: + layers: tuple[int, ...] + min_num_slots: int + max_num_slots: int + + +def group_layers(scope: str, layouts: list, slot_bytes: list[int]) -> list[list[int]]: + """The layers of each pool, in placement order. + + ``unified``: one pool per layout key (see ``bank_layouts``); ``per-layer``: one pool per layer. + Pools with wider slots come first: the prefill staging buffers sit at the start of every bank, so they then cover fewer slots. + """ + if scope == "unified": + groups: dict = {} + for layer, layout in enumerate(layouts): + groups.setdefault(layout, []).append(layer) + pools = list(groups.values()) + elif scope == "per-layer": + pools = [[layer] for layer in range(len(layouts))] + else: + raise ValueError(f"unknown --moe-cache-scope {scope!r}; want one of {SCOPES}") + return sorted(pools, key=lambda layers: (-slot_bytes[layers[0]], layers[0])) + + +def make_pools(groups: list[list[int]], num_experts: int, min_num_slots: int, slot_limit: int | None) -> list[PoolSpec]: + """Each group's slot range: at least ``min_num_slots``, at most one slot per expert of its layers and the kernel's ``slot_limit``.""" + pools = [] + for layers in groups: + max_num_slots = len(layers) * num_experts + if slot_limit is not None: + max_num_slots = min(max_num_slots, slot_limit) + pools.append(PoolSpec(tuple(layers), min(min_num_slots, max_num_slots), max_num_slots)) + return pools + + +def pool_of_layer(pools: list[PoolSpec], num_layers: int) -> list[int]: + """Each layer's pool index, the form ``OffloadMoeCache`` takes.""" + owner = [-1] * num_layers + for p, pool in enumerate(pools): + for layer in pool.layers: + owner[layer] = p + assert min(owner) >= 0, owner + return owner + + +def split_slots(total: int, pools: list[PoolSpec], counts: torch.Tensor | None = None) -> list[int]: + """Slots per pool, each in [min_num_slots, max_num_slots], summing to ``total``. + + With ``counts`` (per-layer expert use from ``resolve_profile``) the slots follow expert use, else the pools' layer counts. + """ + if len(pools) == 1: + return [total] + _check_budget(total, pools) + return _split_equal(total, pools) if counts is None else _split_hot(total, pools, counts) + + +def _check_budget(total: int, pools: list[PoolSpec]) -> None: + need = sum(p.min_num_slots for p in pools) + room = sum(p.max_num_slots for p in pools) + if total < need: + raise ValueError(f"MoE cache of {total} slots is below what the pools need ({need} slots)") + if total > room: + raise ValueError(f"MoE cache of {total} slots exceeds what the pools can hold ({room} slots)") + + +def _split_equal(total: int, pools: list[PoolSpec]) -> list[int]: + sizes = [p.min_num_slots for p in pools] + weights = [len(p.layers) for p in pools] + left = total - sum(sizes) + while left > 0: + open_ = [i for i, p in enumerate(pools) if sizes[i] < p.max_num_slots] + # give the next slot to the pool with the fewest slots per layer + i = min(open_, key=lambda j: (sizes[j] / weights[j], j)) + sizes[i] += 1 + left -= 1 + return sizes + + +def _split_hot(total: int, pools: list[PoolSpec], counts: torch.Tensor) -> list[int]: + sizes = [p.min_num_slots for p in pools] + # a pool's minimum covers its hottest experts; the rest, up to its maximum, competes across pools + cand = [] + for p, pool in enumerate(pools): + hot = counts[list(pool.layers)].flatten().sort(descending=True).values + cand.extend((float(c), p) for c in hot[pool.min_num_slots : pool.max_num_slots]) + cand.sort(key=lambda c: -c[0]) + for _, p in cand[: total - sum(sizes)]: + sizes[p] += 1 + return sizes + + +def resolve_profile(spec: str, model_path: str, num_layers: int, num_experts: int) -> torch.Tensor | None: + """Expert use for ``--moe-cache-split``: None for ``equal``, else the profile at ``profile:``.""" + if spec == "equal": + return None + path = Path(spec.removeprefix("profile:")) + logger.info_rank0(f"MoE pool profile: {path}") + return load_profile(path, num_layers, num_experts, model_path) + + +def load_profile(path: str | Path, num_layers: int, num_experts: int, model_path: str = "") -> torch.Tensor: + """Per-layer expert use ``[num_layers, num_experts]`` from a profile, each row scaled to sum to 1.""" + data = json.loads(Path(path).read_text()) + model = str(data.get("model", "")) + if model and model_path and not any(n.lower().startswith(model.lower()) for n in _model_names(model_path)): + logger.warning_rank0(f"MoE pool profile {path} was taken on {model!r}, not on {model_path}") + counts = torch.tensor(data["expert_counts"], dtype=torch.float64) + if tuple(counts.shape) != (num_layers, num_experts): + raise ValueError( + f"profile {path} has {tuple(counts.shape)} expert counts; this model needs ({num_layers}, {num_experts})" + ) + # every layer decodes the same tokens, so equal row sums keep the layers comparable + return counts / counts.sum(dim=1, keepdim=True).clamp(min=1e-12) + + +def _model_names(model_path: str) -> list[str]: + # every path component, so an HF cache snapshot (models--Org--Name/snapshots/) still finds its name + names = [] + for part in reversed(Path(model_path).parts): + name = part.split("--")[-1].removesuffix(".gguf") + if name: + names.append(name) + return names diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 7c0ed9b047..be244a117c 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -134,6 +134,7 @@ def parse_args( from freetoken.attention import validate_attn_backend from freetoken.kvcache import SUPPORTED_CACHE_MANAGER from freetoken.moe import MOE_STRATEGIES + from freetoken.moe.pool_policy import SCOPES def _parse_quant_backend(value: str) -> str: from freetoken.layers.quantization import QuantBackend @@ -153,6 +154,27 @@ def _parse_moe_cache_rate(value: str) -> float: raise argparse.ArgumentTypeError("must be in [0, 1]") return rate + def _parse_mem_size(value: str) -> int: + units = { + "b": 1, "k": 1 << 10, "kb": 1 << 10, "kib": 1 << 10, "m": 1 << 20, "mb": 1 << 20, "mib": 1 << 20, + "g": 1 << 30, "gb": 1 << 30, "gib": 1 << 30, + } + text = value.strip().lower() + number = text.rstrip("abgikm") + unit = text[len(number):] + try: + size = int(float(number) * units[unit]) + except (KeyError, ValueError, OverflowError) as exc: + raise argparse.ArgumentTypeError("must be a size with a unit, e.g. 12GiB or 800MiB") from exc + if size <= 0: + raise argparse.ArgumentTypeError("must be positive") + return size + + def _parse_cache_split(value: str) -> str: + if value == "equal" or (value.startswith("profile:") and len(value) > len("profile:")): + return value + raise argparse.ArgumentTypeError("must be equal or profile:") + def _positive_int(value: str) -> int: try: n = int(value) @@ -705,6 +727,12 @@ def _infer_reasoning_parser(model_path: str) -> str | None: default=ServerArgs.moe_cache_rate, help="The fraction of all MoE experts to keep in GPU cache.", ) + moe_cache_group.add_argument( + "--moe-cache-mem", + type=_parse_mem_size, + default=ServerArgs.moe_cache_mem, + help="GPU memory for the MoE expert slot cache, e.g. 12GiB or 800MiB (G, GB and GiB all mean GiB); the slot count is this over the bytes per slot.", + ) moe_cache_group.add_argument( "--moe-cache-auto", action="store_true", @@ -722,6 +750,27 @@ def _infer_reasoning_parser(model_path: str) -> str | None: help="KV-cache token floor reserved before --moe-cache-auto fills experts.", ) + parser.add_argument( + "--moe-cache-scope", + choices=SCOPES, + default=ServerArgs.moe_cache_scope, + help=( + "How MoE layers share the GPU expert slot cache: unified = one LRU pool per expert quant " + "(all layers when they share one quant format); per-layer = each layer has its own pool. " + "Pools are sized by --moe-cache-split." + ), + ) + parser.add_argument( + "--moe-cache-split", + type=_parse_cache_split, + default=ServerArgs.moe_cache_split, + help=( + "How pools share the slots: equal (by layer count), or profile: = by expert use from a " + "profile file. With --moe-strategy offload, a pool smaller than the expert count limits the " + "decode batch to its slots / top_k." + ), + ) + parser.add_argument( "--moe-cache-policy", default=ServerArgs.moe_cache_policy, @@ -892,6 +941,7 @@ def _infer_reasoning_parser(model_path: str) -> str | None: _no_cache_flag = ( kwargs["moe_cache_size"] == 0 and not kwargs["moe_cache_auto"] + and kwargs["moe_cache_mem"] is None and (kwargs["moe_cache_rate"] is None or kwargs["moe_cache_rate"] == 0) ) if is_offload_moe_strategy(kwargs["moe_strategy"]) and _no_cache_flag: diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py index a77409f013..7b37af9e0e 100644 --- a/tests/engine/test_cache_budget.py +++ b/tests/engine/test_cache_budget.py @@ -523,3 +523,55 @@ def test_uncapped_platform_stays_uncapped(monkeypatch): if hasattr(os, "uname") and "microsoft" in os.uname().release.lower(): pytest.skip("WSL caps pinning") assert _pin_budget_bytes(reserved=2**30) is None + + +def test_moe_slot_pool_split(tmp_path): + import json + + from freetoken.moe.pool_policy import group_layers, load_profile, make_pools, split_slots + + pools = make_pools(group_layers("per-layer", [0] * 4, [1] * 4), num_experts=16, min_num_slots=8, slot_limit=None) + assert split_slots(41, pools) == [11, 10, 10, 10] + with pytest.raises(ValueError, match="need"): + split_slots(31, pools) + with pytest.raises(ValueError, match="hold"): + split_slots(65, pools) + # one pool takes the whole cache, as before pools existed + assert split_slots(100, make_pools(group_layers("unified", [0] * 4, [1] * 4), 16, 8, None)) == [100] + # unified pools one bank layout each, widest slots first + assert group_layers("unified", ["a", "b", "a", "b"], [1, 3, 1, 3]) == [[1, 3], [0, 2]] + + # layer 0 uses two experts, layer 1 has a hot half and a cold half, layers 2 and 3 use all 16 evenly + counts = [[5] * 2 + [0] * 14, [3] * 8 + [1] * 8, [1] * 16, [1] * 16] + path = tmp_path / "profile.json" + path.write_text(json.dumps({"model": "toy", "expert_counts": counts})) + assert split_slots(48, pools, load_profile(path, 4, 16)) == [8, 8, 16, 16] + with pytest.raises(ValueError, match="needs"): + load_profile(path, 4, 32) + + + +def test_moe_pool_flags_and_decode_batch_clamp(): + from unittest.mock import patch + + from freetoken.engine.engine import _clamp_decode_batch_to_pools + from freetoken.server.args import parse_args + + hf = SimpleNamespace(to_dict=lambda: {"architectures": ["Qwen3MoeForCausalLM"], "torch_dtype": "bfloat16"}) + with patch("freetoken.utils.cached_load_hf_config", lambda _path: hf): + args, _ = parse_args([ + "--model", "/models/anon", "--moe-cache-mem", "1.5GiB", + "--moe-cache-scope", "per-layer", "--moe-cache-split", "profile:/p.json", + ]) + assert (args.moe_cache_mem, args.moe_cache_scope, args.moe_cache_split) == (3 << 29, "per-layer", "profile:/p.json") + for bad in (["--moe-cache-mem", "12"], ["--moe-cache-mem", "1GiB", "--moe-cache-size", "512"]): + with pytest.raises(SystemExit): + parse_args(["--model", "/models/anon", *bad]) + + # GPU decode needs batch * top_k distinct experts to fit the smallest pool below num_experts + config = SimpleNamespace( + model_config=SimpleNamespace(num_experts=16, num_experts_per_tok=4), + max_running_req=4, cuda_graph_max_bs=None, cuda_graph_bs=[1, 2, 4], + ) + _clamp_decode_batch_to_pools(config, [12, 16, 9]) + assert (config.max_running_req, config.cuda_graph_max_bs, config.cuda_graph_bs) == (2, 2, [1, 2]) diff --git a/tests/moe/test_nvfp4_backends.py b/tests/moe/test_nvfp4_backends.py index 0546fbcd6f..dbf06bcc6c 100644 --- a/tests/moe/test_nvfp4_backends.py +++ b/tests/moe/test_nvfp4_backends.py @@ -187,7 +187,7 @@ def test_marlin_prefill_matches_dequant_reference(): cache.materialize_layer(0) cache.copy_missing() g1, g2 = cache.alphas_for_layer(0) - gu_p, gu_s, dn_p, dn_s = cache.bank_views(E) + gu_p, gu_s, dn_p, dn_s = cache.stage_views(0, 0) out = marlin_fused_experts( hidden, gu_p, gu_s, g1, dn_p, dn_s, g2, topk_weights, topk_ids, "silu", False, @@ -214,7 +214,7 @@ def decode(layer_id, experts): cache.ensure_experts(layer_id, ids) # rewrites ids -> slots in place cache.copy_missing() g1, g2 = cache.alphas_for_slots(layer_id) - gu_p, gu_s, dn_p, dn_s = cache.bank_views() + gu_p, gu_s, dn_p, dn_s = cache.pool_bank_views(layer_id) out = marlin_fused_experts( hidden, gu_p, gu_s, g1, dn_p, dn_s, g2, topk_weights, ids, "silu", False, @@ -278,7 +278,7 @@ def test_marlin_overlap_prefill_matches_dequant_reference(): assert int(cache.num_indices.item()) == 2, "stale slot mappings survived the prefetch" cache.copy_missing() g1, g2 = cache.alphas_for_slots(0) - gu_p, gu_s, dn_p, dn_s = cache.bank_views() + gu_p, gu_s, dn_p, dn_s = cache.pool_bank_views(0) out = marlin_fused_experts( dec_hidden, gu_p, gu_s, g1, dn_p, dn_s, g2, dec_weights, ids, "silu", False, @@ -403,7 +403,7 @@ def decode(layer_id, experts): ref = _ref_moe(ref_sources, layer_id, hidden, topk_weights, ids) cache.ensure_experts(layer_id, ids) # rewrites ids -> slots in place cache.copy_missing() - gu_p, gu_s, gu_g, dn_p, dn_s, dn_g = cache.bank_views() + gu_p, gu_s, gu_g, dn_p, dn_s, dn_g = cache.pool_bank_views(layer_id) out = fused_experts_decode_nvfp4_marlin( hidden, gu_p, gu_s, gu_g, dn_p, dn_s, dn_g, topk_weights, ids, "silu", False ) @@ -433,7 +433,7 @@ def test_triton_decode_marlin_matches_baseline_kernel(): ids = torch.tensor([[1, 6]], dtype=torch.int32, device=device) cache.ensure_experts(0, ids) cache.copy_missing() - banks = cache.bank_views() + banks = cache.pool_bank_views(0) marlin = fused_experts_decode_nvfp4_marlin(hidden, *banks, topk_weights, ids, "silu", False) base = fused_experts_decode_nvfp4_serial(hidden, *banks, topk_weights, ids, "silu", False) torch.testing.assert_close(marlin.float(), base.float(), rtol=2e-3, atol=2e-3) @@ -461,7 +461,7 @@ def test_b12x_decode_matches_dequant_reference(): cache.ensure_experts(0, ids) cache.copy_missing() g1, g2 = cache.alphas_for_slots(0) - gu_p, gu_s, dn_p, dn_s = cache.bank_views() + gu_p, gu_s, dn_p, dn_s = cache.pool_bank_views(0) out = b12x_fused_experts( hidden, gu_p, gu_s, g1, dn_p, dn_s, g2, topk_weights, ids, "silu", False ) diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index cd5f5dc773..d3c8d9d843 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -209,7 +209,7 @@ def fake_fused( assert fused_calls[0]["w2_ptr"] == fused_calls[2]["w2_ptr"] assert fused_calls[0]["w1_ptr"] != fused_calls[1]["w1_ptr"] assert fused_calls[0]["w2_ptr"] != fused_calls[1]["w2_ptr"] - prefill_gate_up_buffer, prefill_down_buffer = cache.prefill_bank_buffers + prefill_gate_up_buffer, prefill_down_buffer = cache.stage_views(0, 0) assert prefill_gate_up_buffer.data_ptr() == cache.bank_caches["gate_up"].data_ptr() assert prefill_down_buffer.data_ptr() == cache.bank_caches["down"].data_ptr() @@ -639,6 +639,136 @@ def test_offload_cache_rebuild_resizes_and_preserves_sources(): assert torch.all(cache.id_of_slot == -1) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_slot_pools_of_two_bank_layouts_keep_their_rows_through_prefill(): + from freetoken.moe.offload_cache import OffloadMoeCache + + _init_tp() + torch.manual_seed(0) + L, E, top_k = 4, 8, 4 + # layers 1 and 3 keep wider down rows (uint8, like Q6_K next to Q4_K); their pools come first + wide = (1, 3) + sources = { + "gate_up": [torch.randn(E, 1024, 256, dtype=torch.bfloat16).pin_memory() for _ in range(L)], + "down": [ + torch.randint(0, 255, (E, 512, 768), dtype=torch.uint8).pin_memory() if layer in wide + else torch.randn(E, 512, 256, dtype=torch.bfloat16).pin_memory() + for layer in range(L) + ], + } + cache = OffloadMoeCache( + num_layers=L, num_experts=E, cache_size=33, device=torch.device("cuda"), + prefill_overlap=True, prefill_hit_d2d=True, pool_of_layer=[2, 0, 3, 1], pool_sizes=[6, 5, 12, 10], + ) + cache.set_bank_sources(sources) + assert [v[1].dtype for v in map(cache.pool_bank_views, range(L))] == [torch.bfloat16, torch.uint8] * 2 + # staging buffers take 2 x 8 rows of every bank: all of the wide pools, the first 8 slots of layer 0's pool + assert cache._pool_stable_from == [6, 5, 8, 0] + + def check_resident() -> None: + torch.cuda.synchronize() + for layer in range(L): + for e, s in enumerate(cache.slot_for_id[layer].tolist()): + if s >= 0: + for (src, _), view in zip(cache.banks, cache.pool_bank_views(layer)): + assert torch.equal(view[s].cpu(), src[layer][e]), (layer, e, s) + + def decode(steps: int) -> None: + for _ in range(steps): + for layer in range(L): + cache.ensure_experts(layer, torch.randperm(E)[:top_k].to(torch.int32).view(1, top_k).cuda()) + cache.copy_missing() + check_resident() + + decode(20) + cache.zero_inactive_route_slots = True + cache.begin_prefill() + for layer in range(L): + cache.prefetch_prefill_layer(layer) + cache.prefetch_prefill_layer(layer + 1) + views = cache.wait_prefill_layer(layer) + torch.cuda.synchronize() + for (src, _), view in zip(cache.banks, views): + assert torch.equal(view.cpu(), src[layer]) + cache.release_prefill_layer(layer) + assert cache.prefill_hit_rows > 0 + for p, stable in enumerate(cache._pool_stable_from): + if stable: + assert not any(view[0].any() for view in cache.pool_bank_views(cache.pool_of_layer.index(p))) + check_resident() + decode(20) + + cache.materialize_layer(3) + cache.copy_missing() + for (src, _), view in zip(cache.banks, cache.stage_views(3, 0)): + assert torch.equal(view.cpu(), src[3]) + check_resident() + + cache.rebuild(32, [5, 5, 10, 12]) + decode(20) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_slot_pools_split_layers_of_one_bank_layout_by_quant_tag(): + from freetoken.moe.offload_cache import OffloadMoeCache, bank_layouts + from freetoken.moe.pool_policy import group_layers + + _init_tp() + L, E = 4, 8 + # uint8 rows of one width, as GGUF Q4_K and IQ4_NL rows are + sources = { + "gate_up": [torch.randint(1, 255, (E, 64, 144), dtype=torch.uint8).pin_memory() for _ in range(L)], + "down": [torch.randint(1, 255, (E, 32, 144), dtype=torch.uint8).pin_memory() for _ in range(L)], + } + tags = [("ggml", "Q4_K", "Q4_K"), ("ggml", "IQ4_NL", "IQ4_NL")] * 2 + assert group_layers("unified", bank_layouts(sources, tags), [1] * L) == [[0, 2], [1, 3]] + + def make(pool_of_layer, pool_sizes): + return OffloadMoeCache( + num_layers=L, num_experts=E, cache_size=sum(pool_sizes), device=torch.device("cuda"), + prefill_overlap=True, pool_of_layer=pool_of_layer, pool_sizes=pool_sizes, + ) + + with pytest.raises(AssertionError, match="layouts"): + make([0] * L, [20]).set_bank_sources(sources, layer_tags=tags) + cache = make([0, 1, 0, 1], [10, 10]) + cache.set_bank_sources(sources, layer_tags=tags) + assert cache.bank_caches["down"].shape == (20, 32, 144) + # the staging buffers cover pool 0 and slots [0, 6) of pool 1, so both slot 0 rows end up with another quant's bytes + cache.zero_inactive_route_slots = True + cache.begin_prefill() + for layer in range(L): + cache.prefetch_prefill_layer(layer) + cache.prefetch_prefill_layer(layer + 1) + cache.wait_prefill_layer(layer) + cache.release_prefill_layer(layer) + torch.cuda.synchronize() + for layer in (0, 1): + for view in cache.pool_bank_views(layer): + assert not view[0].any() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_offload_cache_rebuild_frees_the_old_slots_first(): + from freetoken.moe.offload_cache import OffloadMoeCache + + _init_tp() + cache = OffloadMoeCache( + num_layers=2, num_experts=4, cache_size=32, device=torch.device("cuda"), + pool_of_layer=[0, 1], pool_sizes=[16, 16], + ) + cache.set_bank_sources({ + "gate_up": [torch.zeros(4, 1024, 1024, dtype=torch.bfloat16).pin_memory() for _ in range(2)], + "down": [torch.zeros(4, 512, 1024, dtype=torch.bfloat16).pin_memory() for _ in range(2)], + }) + torch.cuda.synchronize() + before = torch.cuda.memory_allocated() + torch.cuda.reset_peak_memory_stats() + cache.rebuild(32, [8, 24]) + # 96 MiB of slots: holding the old ones while allocating the new ones would double it + assert torch.cuda.max_memory_allocated() - before < 8 << 20 + + def test_offload_cache_rebuild_disables_prefill_overlap_when_too_small(): from freetoken.moe.offload_cache import OffloadMoeCache @@ -654,7 +784,7 @@ def test_offload_cache_rebuild_disables_prefill_overlap_when_too_small(): assert cache.cache_size == 5 assert cache.prefill_overlap is False - assert cache.prefill_bank_buffers == [] + assert cache._stage_views[0][1] is None # no room for the second staging buffer def test_offload_cache_rebuild_keeps_overlap_at_boundary(): @@ -690,6 +820,58 @@ def test_offload_cache_validate_rebuild_enforces_marlin_cap_and_floor(): bf16.validate_rebuild(3) # below the num_experts floor +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +@pytest.mark.parametrize("target", ["gpu", "hybrid"]) +def test_slot_pools_keep_each_layer_in_its_own_slots(target): + from freetoken.moe.offload_cache import OffloadMoeCache + + _init_tp() + torch.manual_seed(0) + L, E, sizes = 3, 16, [5, 7, 20] + # hybrid routes more experts than the smallest pool holds; GPU decode must fit them + top_k = 2 if target == "gpu" else 8 + sources = { + "gate_up": [torch.randn(E, 64, 32, dtype=torch.bfloat16).pin_memory() for _ in range(L)], + "down": [torch.randn(E, 32, 32, dtype=torch.bfloat16).pin_memory() for _ in range(L)], + } + + cache = OffloadMoeCache( + num_layers=L, num_experts=E, cache_size=sum(sizes), device=torch.device("cuda"), + decode_target=target, hybrid_max_fetch=E, pool_of_layer=[0, 1, 2], pool_sizes=sizes, + ) + cache.set_bank_sources(sources) + for _ in range(50): + for layer in range(L): + experts = torch.randperm(E)[:top_k].to(torch.int32) + slots = experts.cuda() + if target == "gpu": + cache.ensure_experts(layer, slots) + else: + cache.ensure_experts_hybrid(layer, slots, torch.rand(top_k, device="cuda")) + cache.copy_missing() + resident = [(e, s) for e, s in zip(experts.tolist(), slots.tolist()) if s >= 0] + assert len({s for _, s in resident}) == len(resident) + for e, s in resident: + assert s < sizes[layer] + for (src, _), view in zip(cache.banks, cache.pool_bank_views(layer)): + assert torch.equal(view[s].cpu(), src[layer][e]) + + # with several pools a whole-layer prefill stages slots [0, E) without registering them in any pool + cache.materialize_layer(1) + cache.copy_missing() + for (src, _), view in zip(cache.banks, cache.stage_views(1, 0)): + assert torch.equal(view.cpu(), src[1]) + slot_for_id, id_of_slot = cache.slot_for_id.cpu(), cache.id_of_slot.cpu() + for layer in range(L): + for e, s in enumerate(slot_for_id[layer].tolist()): + if s >= 0: + assert cache.pool_bases[layer] + s >= E + assert id_of_slot[cache.pool_bases[layer] + s] == layer * E + e + + cache.rebuild(40, [8, 8, 24]) + assert [v[0].shape[0] for v in map(cache.pool_bank_views, range(L))] == [8, 8, 24] + + def _make_split_cache(num_layers=2, locked=(1,), prefill_overlap=False, device="cpu"): """A [gate_up, down] bf16 cache with the given layers LOCKED (rest pinned).""" from freetoken.moe.host_banks import HostResidency From 0bca88f70c4909624f7f3c3e3a4fe6b32d183d56 Mon Sep 17 00:00:00 2001 From: LZ_QAQ <32613092+LZ-QWQ@users.noreply.github.com> Date: Thu, 15 Oct 2026 17:21:32 +0800 Subject: [PATCH 05/94] fix(rocm): reuse staged moe alignment on Triton 3.8 (#691) * fix(rocm): reuse staged moe alignment below triton 3.9 * fix(rocm): restrict moe alignment workaround to Triton 3.8 --- python/freetoken/moe/fused.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/python/freetoken/moe/fused.py b/python/freetoken/moe/fused.py index a90aa4c792..5c34fb33ca 100644 --- a/python/freetoken/moe/fused.py +++ b/python/freetoken/moe/fused.py @@ -84,7 +84,18 @@ def moe_align_block_size( - The padding ensures that the total number of tokens is now divisible by block_size for proper block matrix operations. """ - from freetoken.kernel.backend import is_sgl_kernel_installed + from freetoken.kernel.backend import is_rocm, is_sgl_kernel_installed + + if is_rocm(): + import triton + from packaging.version import Version + + # FIXME: remove once Triton 3.8 is unsupported. Its AMD backend folds compares on tl.histogram results to false. + # https://github.com/triton-lang/triton/pull/11246 + if Version(triton.__version__).release[:2] == (3, 8): + from freetoken.kernel import moe_align_block_size_triton + + return moe_align_block_size_triton(topk_ids, block_size, num_experts) if not is_sgl_kernel_installed(): from freetoken.kernel.triton.moe_align import ( From b43c5ccc2d36542cce11f355e78f37819cbfb80b Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Mon, 14 Sep 2026 18:23:04 -0400 Subject: [PATCH 06/94] test(kvcache): give the qsa fp8 pool fakes the mrope flag main now reads Assisted-by: Claude Opus 5 (cherry picked from commit 150c95d6f39b83e0516f5c9c42fef54c679eae09) (cherry picked from commit 4db8b858375b01ff7edf37031dc81671b0181853) (cherry picked from commit 495770e33dcb5c3aa6979c6515b61dfade5e29b3) --- tests/kvcache/test_qsa_pool_fp8.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/kvcache/test_qsa_pool_fp8.py b/tests/kvcache/test_qsa_pool_fp8.py index c0a91f862a..02b525f679 100644 --- a/tests/kvcache/test_qsa_pool_fp8.py +++ b/tests/kvcache/test_qsa_pool_fp8.py @@ -74,7 +74,7 @@ def _spec(): def _config(kv_quant, *, page_size=PAGE_SIZE, max_running_req=3): - mc = SimpleNamespace(num_layers=8, has_swa_attention=False, has_linear_attention=True) + mc = SimpleNamespace(num_layers=8, has_swa_attention=False, has_linear_attention=True, model_is_mrope=False) mc.kv_cache_group_specs = lambda: (_spec(),) return SimpleNamespace( model_config=mc, @@ -200,7 +200,7 @@ def test_factory_threads_kv_quant_into_the_qsa_pool(): mc = SimpleNamespace( num_layers=8, has_swa_attention=False, has_linear_attention=True, - num_kv_heads=HEADS, head_dim=DIM, dsv4_args=None, + num_kv_heads=HEADS, head_dim=DIM, dsv4_args=None, model_is_mrope=False, ) mc.kv_cache_group_specs = lambda: (_spec(),) pool = create_kvcache_pool( From dd652489e517ee861c89636f19fd8e8cb3c2e7ca Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Mon, 14 Sep 2026 23:02:28 -0400 Subject: [PATCH 07/94] fix(models): restore the os import and truthy set the fp8 env helpers read Assisted-by: Claude Opus 5 (cherry picked from commit 82a21ae2432fb891175c8c881f52cb4b710241bc) (cherry picked from commit 0369ed94d7d085ece40ffdcc0adf93a5c26a6b4e) (cherry picked from commit b5c7812fceb68b4b7835a80ab7c533e0385d9577) --- python/freetoken/models/config.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py index bfa3b6cbbb..0d4faae420 100644 --- a/python/freetoken/models/config.py +++ b/python/freetoken/models/config.py @@ -1,4 +1,6 @@ from __future__ import annotations + +import os from dataclasses import dataclass from typing import Any, ClassVar, Dict, List, Literal, Tuple, TypeAlias @@ -7,12 +9,14 @@ # State-dict key prefixes of the vision stack; load_weight drops them when the engine serves text-only. VISION_KEY_PREFIXES = ("vision_tower.", "embed_vision.", "vision_embedder.", "visual.") +_ENV_TRUE = {"1", "true", "yes", "on"} + def fp8_dense_enabled() -> bool: """Load-time FP8 for a model's bf16 attention / GDN projections (opt-in, default OFF): per-tensor e4m3 weights, per-tensor dynamic activation scale, cuBLASLt W8A8 GEMMs (``torch._scaled_mm``, sm_89+). ``FREETOKEN_FP8_DENSE=1``.""" - return os.getenv("FREETOKEN_FP8_DENSE", "0").strip().lower() in _VISION_TRUE + return os.getenv("FREETOKEN_FP8_DENSE", "0").strip().lower() in _ENV_TRUE def fp8_lmhead_enabled() -> bool: @@ -20,7 +24,7 @@ def fp8_lmhead_enabled() -> bool: for the rest). Separate from :func:`fp8_dense_enabled` because this one moves the logits: a per-tensor e4m3 vocab matrix changes every sampled token's score, so it carries its own quality gate. ``FREETOKEN_FP8_LMHEAD=1``.""" - return os.getenv("FREETOKEN_FP8_LMHEAD", "0").strip().lower() in _VISION_TRUE + return os.getenv("FREETOKEN_FP8_LMHEAD", "0").strip().lower() in _ENV_TRUE def detect_expert_quant(hf_config: Any) -> str: From 1dc6b7d56e7e1266770bef2c140f652dd3ac4ac3 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Mon, 14 Sep 2026 23:00:44 -0400 Subject: [PATCH 08/94] test(qwen4_exp): restore the rope device and clear get_rope after the cpu build test Assisted-by: Claude Opus 5 (cherry picked from commit 27ad92e0a554be016d6e533683c40a2eb67c6164) (cherry picked from commit e152f665cc1f35f073fbb82daedece77d02f91e1) (cherry picked from commit 333f7f9f0501116694ede067202070353faf341b) --- tests/models/qwen4_exp/test_build_cpu.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/tests/models/qwen4_exp/test_build_cpu.py b/tests/models/qwen4_exp/test_build_cpu.py index 0cf2027d14..58ad0c4dc3 100644 --- a/tests/models/qwen4_exp/test_build_cpu.py +++ b/tests/models/qwen4_exp/test_build_cpu.py @@ -10,7 +10,7 @@ import torch from freetoken.distributed import set_tp_info, try_get_tp_info -from freetoken.layers import set_rope_device +from freetoken.layers import rotary, set_rope_device from freetoken.models.qwen4_exp.config import parse_config from freetoken.models.qwen4_exp.model import Qwen4ExpForCausalLM @@ -20,9 +20,15 @@ def _build(): if try_get_tp_info() is None: set_tp_info(rank=0, size=1) + previous = rotary._ROPE_DEVICE set_rope_device(torch.device("cpu")) - with torch.device("meta"): - return Qwen4ExpForCausalLM(parse_config(toy_hf_config())) + try: + with torch.device("meta"): + return Qwen4ExpForCausalLM(parse_config(toy_hf_config())) + finally: + # get_rope is cached per geometry: a CPU cos/sin table left behind feeds later CUDA tests + rotary.get_rope.cache_clear() + set_rope_device(previous) def test_model_constructs_on_meta_without_cuda(): From c5010c97c50e06b56ff7e98232526e8b523eb6cd Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Mon, 14 Sep 2026 23:19:07 -0400 Subject: [PATCH 09/94] test(kernels): give the block_ends backend fake the kv scale accessors the quantized-kv backend reads Assisted-by: Claude Opus 5 (cherry picked from commit 600819fd7ce18c813cd4b8400cf534bff38f6f7f) (cherry picked from commit dd2ebdfda3e1c53e5269bd9adfbee7291cc8ffc2) (cherry picked from commit 12754d41c5fe646ccc6081ce26f05fdf7b0a57b9) --- tests/kernels/test_triton_attention.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/kernels/test_triton_attention.py b/tests/kernels/test_triton_attention.py index a5abb3d7b3..3176ad71ce 100644 --- a/tests/kernels/test_triton_attention.py +++ b/tests/kernels/test_triton_attention.py @@ -804,6 +804,12 @@ def k_cache(self, layer_id): def v_cache(self, layer_id): return self.v + def k_scale(self, layer_id): # 16-bit pool: no fp8 scales + return None + + def v_scale(self, layer_id): + return None + device = torch.device("cuda") head_dim = 256 kv_cache = FakeKVCache(device, head_dim) From 5e441ba089e979ccc5efc320a2522e478f2b9016 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 23 Sep 2026 19:17:58 -0400 Subject: [PATCH 10/94] merge fork #18: fix(moe): report the routing top-k figure as a static bound, not an oracle --- python/freetoken/engine/engine.py | 19 +++++++---------- python/freetoken/moe/offload_cache.py | 12 +++++------ tests/moe/test_moe_collect_stats.py | 30 +++++++++++++++++++++++---- 3 files changed, 40 insertions(+), 21 deletions(-) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index ed60d8d340..27fdb7d716 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -900,9 +900,8 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: cache.collect_stats = config.moe_collect_stats # The routing histogram rides the same switch: on its own the miss rate says how # often we fetch, but not whether a smarter policy could have avoided the fetch. - # decode_routing_stats turns it into an oracle hit rate -- the ceiling any policy - # holding this many slots could reach on the observed routing -- which is the number - # worth having before anyone rewrites eviction. + # decode_routing_stats turns it into routing-skew numbers, including the hit rate of + # the best fixed expert set per layer (not a ceiling: LRU's temporal locality can beat it). cache.collect_decode_freq = config.moe_collect_stats # attach_offload_moe_cache walks for OffloadMoELayers, or defers to a model's # _iter_offload_moe_layers() hook when its MoE blocks are bespoke nn.Modules (DSV4). @@ -1272,8 +1271,8 @@ def _emit_moe_stats(self) -> None: Accumulation is device-side and captured into the decode graph, so it is free to leave running; reading it is not (the counters have to come back to the host), which is why this only fires every MOE_STATS_INTERVAL decode steps. The routing histogram - is deliberately *not* reset -- the oracle bound wants the whole run's distribution, - not one window's. + is deliberately *not* reset -- it describes the whole run's distribution, not one + window's, so the static top-k figure flattens as a session covers more topics. """ cache = self.moe_offload_cache agg = cache.decode_miss_stats() @@ -1299,14 +1298,12 @@ def _emit_moe_stats(self) -> None: ) routing = cache.decode_routing_stats() if routing: - # oracle_hit_at_slots is the upper bound on hit rate for *any* policy with this - # many slots per layer. If it sits near the realized hit rate, the cache is - # already doing as well as the routing allows and the win has to come from - # somewhere else (more slots, more bandwidth); if it sits far above, eviction - # policy is leaving something on the table. + # static_topk_hit is what pinning each layer's most frequent experts would score over + # the whole run. A realized hit rate above it means the cache is winning on temporal + # locality; it says nothing about how close to optimal eviction is. logger.info_rank0( "MoE routing: " - f"oracle_hit={routing['oracle_hit_at_slots']:.3f} " + f"static_topk_hit={routing['static_topk_hit_at_slots']:.3f} " f"(realized {1.0 - agg['miss_rate']:.3f}), " f"slots/layer={routing['slots_per_layer']:.1f}, " f"working_set={routing['working_set_mean']:.1f}" diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 701524e0a1..82e3cc23a2 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -1141,10 +1141,10 @@ def decode_miss_stats_per_layer(self) -> dict: def decode_routing_stats(self) -> dict: """Per-layer decode routing concentration, for cache-skew analysis. - Uses the histogram from ``collect_decode_freq``. The ``oracle_hit`` is the best a - per-layer LRU holding ``cache_size/num_layers`` slots could achieve on the observed - (stationary) routing distribution -- i.e. an upper bound on hit rate that depends - purely on how skewed routing is, independent of any LRU/LFU dynamics. + Uses the histogram from ``collect_decode_freq``. ``static_topk_hit_at_slots`` is the + hit rate of the best *fixed* expert set: each layer's ``cache_size/num_layers`` (rounded) + most frequent experts over the whole histogram. It measures routing skew only. It is + not an upper bound on a dynamic cache: LRU exploits temporal locality and can beat it. """ freq = self.decode_freq.float() total = freq.sum(dim=1) @@ -1154,7 +1154,7 @@ def decode_routing_stats(self) -> dict: slots_per_layer = self.cache_size / self.num_layers C = max(1, int(round(slots_per_layer))) sorted_f, _ = torch.sort(freq, dim=1, descending=True) - oracle_hit = (sorted_f[:, :C].sum(dim=1)[valid] / total[valid]).mean().item() + static_hit = (sorted_f[:, :C].sum(dim=1)[valid] / total[valid]).mean().item() ws = (freq > 0).sum(dim=1).float() cdf = torch.cumsum(sorted_f, dim=1) / total.clamp(min=1).unsqueeze(1) cover90 = ((cdf < 0.9).sum(dim=1).float() + 1)[valid] @@ -1166,7 +1166,7 @@ def decode_routing_stats(self) -> dict: "working_set_mean": ws[valid].mean().item(), "working_set_max": int(ws[valid].max().item()), "experts_for_90pct": cover90.mean().item(), - "oracle_hit_at_slots": oracle_hit, + "static_topk_hit_at_slots": static_hit, "norm_entropy": norm_ent, } diff --git a/tests/moe/test_moe_collect_stats.py b/tests/moe/test_moe_collect_stats.py index f960abbdc0..bd1c4e9b3f 100644 --- a/tests/moe/test_moe_collect_stats.py +++ b/tests/moe/test_moe_collect_stats.py @@ -10,7 +10,10 @@ import io from types import SimpleNamespace +import torch + from freetoken.engine.engine import MOE_STATS_INTERVAL, Engine +from freetoken.moe.offload_cache import OffloadMoeCache from freetoken.server.args import ServerArgs, parse_args @@ -63,7 +66,7 @@ def decode_routing_stats(self): "working_set_mean": 173.1, "working_set_max": 243, "experts_for_90pct": 92.3, - "oracle_hit_at_slots": 0.764, + "static_topk_hit_at_slots": 0.764, "norm_entropy": 0.813, } @@ -82,9 +85,9 @@ def test_emit_reports_and_resets_the_window(caplog): cache = _StubCache() out = _emit(cache, caplog) assert "miss_rate=0.250" in out - # The oracle bound is the whole point of the report: it says how much room a different - # eviction policy could possibly have. - assert "oracle_hit=0.764" in out + # Labelled as the fixed-set figure it is; "oracle" implied an upper bound it is not. + assert "static_topk_hit=0.764" in out + assert "oracle" not in out assert "(realized 0.750)" in out # Ranked worst-first, and the layer that never ran is left out entirely. assert "L0=0.500, L1=0.100" in out @@ -107,3 +110,22 @@ def test_idle_window_emits_nothing_and_keeps_counters(caplog): def test_interval_is_a_sane_window(): assert MOE_STATS_INTERVAL >= 1 + + +def test_static_topk_hit_is_not_an_upper_bound_on_lru(): + """One slot, four experts each routed in a run of four tokens: A A A A B B B B ... + + The best fixed expert catches 4 of 16 tokens; an LRU holding one slot misses only on + each run's first token. The figure is routing skew, not a ceiling on a dynamic cache. + """ + trace = [e for e in range(4) for _ in range(4)] + cached, lru_hits = None, 0 + for e in trace: + lru_hits += e == cached + cached = e + freq = torch.bincount(torch.tensor(trace), minlength=4).unsqueeze(0) + cache = SimpleNamespace(decode_freq=freq, cache_size=1, num_layers=1, num_experts=4) + stats = OffloadMoeCache.decode_routing_stats(cache) + assert stats["static_topk_hit_at_slots"] == 0.25 + assert "oracle_hit_at_slots" not in stats + assert lru_hits / len(trace) == 0.75 From c9f71582e67f9440f732dd76df05fb4309ef2e2d Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 23 Sep 2026 19:17:58 -0400 Subject: [PATCH 11/94] merge fork #19: fix(qwen4_exp): refuse an ftw checkpoint whose fp8 split differs from the fp8 env flags --- python/freetoken/models/qwen4_exp/__init__.py | 2 + python/freetoken/models/qwen4_exp/weight.py | 26 +++++++ python/freetoken/models/weight.py | 10 ++- tests/models/qwen4_exp/test_fp8_dense.py | 69 +++++++++++++++++++ 4 files changed, 106 insertions(+), 1 deletion(-) diff --git a/python/freetoken/models/qwen4_exp/__init__.py b/python/freetoken/models/qwen4_exp/__init__.py index 0fe4a4ffa8..84154b26f4 100644 --- a/python/freetoken/models/qwen4_exp/__init__.py +++ b/python/freetoken/models/qwen4_exp/__init__.py @@ -12,6 +12,7 @@ from .config import parse_config from .model import Qwen4ExpForCausalLM, Qwen4ExpForConditionalGeneration from .weight import ( + check_ftw_weights, ftw_side_files, nvfp4_expert_spec, iter_vision_weights, @@ -24,6 +25,7 @@ from freetoken.models.qwen3_5_moe.weight import iter_expert_pieces __all__ = [ + "check_ftw_weights", "ftw_side_files", "nvfp4_expert_spec", "Qwen4ExpForCausalLM", diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index 30a3d1f797..91c8c1bd17 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -417,6 +417,32 @@ def _ple_table_files(folder: str) -> list[str]: return sorted(os.path.join(folder, shard) for shard in files) +def check_ftw_weights(config, model_path: str, names: set[str]) -> None: + """Refuse an FTW checkpoint whose stored FP8 split disagrees with this process's FP8 flags. + + The FTW holds the tensors ``iter_weights`` emitted at conversion, so a mismatch otherwise + surfaces only after the dense read, as a KeyError on the first fused GDN projection.""" + from .config import use_fp8_lmhead + + stored = { + "FREETOKEN_FP8_DENSE": any(n.endswith(".in_proj_qkvz.weight_scale") for n in names), + "FREETOKEN_FP8_LMHEAD": "lm_head.weight_scale" in names, + } + wanted = { + "FREETOKEN_FP8_DENSE": config.attn_quant == "fp8_dynamic", + "FREETOKEN_FP8_LMHEAD": use_fp8_lmhead(config), + } + wrong = [var for var in stored if stored[var] != wanted[var]] + if wrong: + was = ", ".join(f"{var}={int(stored[var])}" for var in wrong) + now = ", ".join(f"{var}={int(wanted[var])}" for var in wrong) + raise ValueError( + f"FTW checkpoint {model_path} stores its dense weights as converted with {was}, " + f"but this process has {now}; serve it with the conversion-time setting, or " + "reconvert with ft checkpoint under the setting you want" + ) + + def ftw_side_files(model_path: str, out_dir: str) -> list[str]: """Write the PLE n-gram table tensors, and only those, into ``ple-table-*.safetensors`` next to an FTW checkpoint. diff --git a/python/freetoken/models/weight.py b/python/freetoken/models/weight.py index e3382b1b59..9016b9477e 100644 --- a/python/freetoken/models/weight.py +++ b/python/freetoken/models/weight.py @@ -218,12 +218,20 @@ def load_weight( # decided at conversion (offload -> experts live in banks, not here); a backend mismatch # fails loudly in load_state_dict (strict missing/unexpected expert keys), so the reader # just yields the stored weight tensors regardless of the include_moe_experts flag. - from freetoken.checkpoint.ftw import is_ftw_checkpoint, iter_ftw_weights + from freetoken.checkpoint.ftw import FTWReader, is_ftw_checkpoint, iter_ftw_weights from freetoken.models.config import VISION_KEY_PREFIXES # a text-only engine never built the tower, so its tensors are not even read keep = None if include_vision else (lambda name: not name.startswith(VISION_KEY_PREFIXES)) if is_ftw_checkpoint(model_path): + spec = get_model_spec(cached_load_hf_config(model_path).architectures[0]) + # a family whose stored tensors depend on load-time settings checks them against the index before the read + check = _model_override(spec, "check_ftw_weights") + if check is not None: + reader = FTWReader(model_path) + names = {e["name"] for e in reader.entries("weight")} + reader.close() + check(_spec_for_model_path(model_path)[0], model_path, names) weights = iter_ftw_weights(model_path, keep=keep) else: _config, spec = _spec_for_model_path(model_path) diff --git a/tests/models/qwen4_exp/test_fp8_dense.py b/tests/models/qwen4_exp/test_fp8_dense.py index 058515b183..74da973341 100644 --- a/tests/models/qwen4_exp/test_fp8_dense.py +++ b/tests/models/qwen4_exp/test_fp8_dense.py @@ -151,3 +151,72 @@ def test_reader_gates_lm_head_behind_its_own_flag(): passthrough = dict(_fp8_dense("x.self_attn.qkv_proj.weight", t, cfg, 1, dense=False, lm_head=True)) assert list(passthrough) == ["x.self_attn.qkv_proj.weight"] + + +def _ftw_with(tmp_path, names): + from freetoken.checkpoint.ftw import FTWWriter + + writer = FTWWriter(str(tmp_path)) + for name in names: + writer.add_tensor(name, torch.zeros(4, dtype=torch.bfloat16)) + writer.finalize({}) + return str(tmp_path) + + +_BF16_NAMES = ["model.layers.0.linear_attn.in_proj.weight", "lm_head.weight"] +_FP8_NAMES = [ + "model.layers.0.linear_attn.in_proj_qkvz.weight", + "model.layers.0.linear_attn.in_proj_qkvz.weight_scale", + "model.layers.0.linear_attn.in_proj_ba.weight", + "lm_head.weight", + "lm_head.weight_scale", +] + + +@pytest.mark.parametrize( + "names, env, bad", + [ + (_BF16_NAMES, {"FREETOKEN_FP8_DENSE": "1", "FREETOKEN_FP8_LMHEAD": "1"}, + "FREETOKEN_FP8_DENSE=0, FREETOKEN_FP8_LMHEAD=0"), + (_FP8_NAMES, {}, "FREETOKEN_FP8_DENSE=1, FREETOKEN_FP8_LMHEAD=1"), + (_FP8_NAMES, {"FREETOKEN_FP8_DENSE": "1"}, "converted with FREETOKEN_FP8_LMHEAD=1, but"), + ], +) +def test_ftw_load_refuses_fp8_flags_that_differ_from_conversion(tmp_path, monkeypatch, names, env, bad): + """#8: the FTW replays conversion-time tensors, so the FP8 flags must match before any byte is read.""" + from freetoken.models import weight as model_weight + + for var in ("FREETOKEN_FP8_DENSE", "FREETOKEN_FP8_LMHEAD"): + monkeypatch.delenv(var, raising=False) + for var, value in env.items(): + monkeypatch.setenv(var, value) + path = _ftw_with(tmp_path, names) + spec = SimpleNamespace(module="freetoken.models.qwen4_exp") + config = SimpleNamespace( + attn_quant="fp8_dynamic" if env.get("FREETOKEN_FP8_DENSE") else "none", + tie_word_embeddings=False, + quant=None, + ) + monkeypatch.setattr(model_weight, "cached_load_hf_config", lambda _: SimpleNamespace(architectures=["x"])) + monkeypatch.setattr(model_weight, "get_model_spec", lambda _: spec) + monkeypatch.setattr(model_weight, "_spec_for_model_path", lambda _: (config, spec)) + with pytest.raises(ValueError, match=bad): + next(model_weight.load_weight(path, torch.device("cpu"))) + + +@pytest.mark.parametrize("names, fp8", [(_BF16_NAMES, False), (_FP8_NAMES, True)]) +def test_ftw_load_accepts_matching_fp8_flags(tmp_path, monkeypatch, names, fp8): + from freetoken.models import weight as model_weight + + for var in ("FREETOKEN_FP8_DENSE", "FREETOKEN_FP8_LMHEAD"): + if fp8: + monkeypatch.setenv(var, "1") + else: + monkeypatch.delenv(var, raising=False) + path = _ftw_with(tmp_path, names) + spec = SimpleNamespace(module="freetoken.models.qwen4_exp") + config = SimpleNamespace(attn_quant="fp8_dynamic" if fp8 else "none", tie_word_embeddings=False, quant=None) + monkeypatch.setattr(model_weight, "cached_load_hf_config", lambda _: SimpleNamespace(architectures=["x"])) + monkeypatch.setattr(model_weight, "get_model_spec", lambda _: spec) + monkeypatch.setattr(model_weight, "_spec_for_model_path", lambda _: (config, spec)) + assert [name for name, _ in model_weight.load_weight(path, torch.device("cpu"))] == names From fbd2004a4b154fad6bf5eeafb439258633d96ae1 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 23 Sep 2026 19:17:58 -0400 Subject: [PATCH 12/94] merge fork #20: fix(engine): let moe-cache-auto drop prefill overlap when its slot floor does not fit --- python/freetoken/engine/cache_budget.py | 16 +++++++++-- python/freetoken/engine/engine.py | 5 ++++ tests/engine/test_cache_budget.py | 37 +++++++++++++++++++++++++ 3 files changed, 55 insertions(+), 3 deletions(-) diff --git a/python/freetoken/engine/cache_budget.py b/python/freetoken/engine/cache_budget.py index 891ac2ac11..00a50b86ad 100644 --- a/python/freetoken/engine/cache_budget.py +++ b/python/freetoken/engine/cache_budget.py @@ -61,6 +61,7 @@ def plan_cache_budget( prefill_overlap: bool, kv_reserve_pages: int, max_slots: int, + page_size: int = 1, ) -> tuple[int, int, bool]: """Split ``budget_bytes`` MoE-first into (moe_cache_size, num_pages, prefill_overlap). @@ -68,7 +69,8 @@ def plan_cache_budget( weights + fixed_cache_size; the (1-memory_ratio) remainder is the graph headroom). Experts greedily fill the budget after reserving ``kv_reserve_pages`` for KV, clamped to ``[floor, min(total_experts, max_slots)]`` (floor is ``2*num_experts`` when prefill - overlap is feasible else ``num_experts``); KV pages take whatever remains. + overlap is feasible, in slots and in bytes, else ``num_experts``); KV pages take whatever + remains. ``page_size`` (tokens per page) only feeds the error message. """ assert per_expert_bytes > 0, "per_expert_bytes must be positive" assert cache_per_page > 0, "cache_per_page must be positive (owned-KV models unsupported here)" @@ -81,6 +83,10 @@ def plan_cache_budget( assert hi >= lo, f"slot cap {hi} below the minimum {lo} slots" kv_reserve_bytes = kv_reserve_pages * cache_per_page + # Overlap is a prefill speedup, not a requirement: when its 2*num_experts floor leaves no + # room for the KV reserve (and the 2 pages asserted below), plan without it. + if overlap and 2 * num_experts * per_expert_bytes + max(kv_reserve_pages, 2) * cache_per_page > budget_bytes: + overlap, lo = False, num_experts # MoE-priority: reserve KV first, then experts greedily take the remaining budget. raw = (budget_bytes - kv_reserve_bytes) // per_expert_bytes moe_cache_size = max(lo, min(raw, hi)) @@ -93,10 +99,13 @@ def plan_cache_budget( # the reserve (or negative), yielding a plan that exceeds budget_bytes. Reject here so # --moe-cache-auto fails in arithmetic instead of OOMing in a later CUDA allocation. total = moe_cache_size * per_expert_bytes + num_pages * cache_per_page + fit_pages = max(0, (budget_bytes - moe_cache_size * per_expert_bytes) // cache_per_page) assert total <= budget_bytes, ( f"cache budget too small: minimum plan (moe={moe_cache_size} slots, " - f"kv={num_pages} pages) needs {total} B > budget {budget_bytes} B " - "(raise memory_ratio, lower kv_reserve_tokens, or free GPU memory)" + f"kv={num_pages} pages) needs {total / 2**30:.2f} GiB > budget {budget_bytes / 2**30:.2f} GiB; " + f"beside {moe_cache_size} slots at most {fit_pages} KV pages ({fit_pages * page_size} tokens) fit. " + "Lower --kv-reserve-tokens, use a smaller --kv-cache-dtype, raise --memory-ratio, " + "or set --moe-cache-size and --max-seq-len-override explicitly" ) assert num_pages > 1, "not enough memory for KV cache after MoE allocation" return moe_cache_size, num_pages, overlap @@ -137,4 +146,5 @@ def resolve_moe_cache_auto( prefill_overlap=prefill_overlap, kv_reserve_pages=kv_reserve_pages, max_slots=max_slots, + page_size=page_size, ) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 27fdb7d716..4ff3e6f144 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -825,6 +825,11 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: size_cap = sum(p.max_num_slots for p in pools) if config.moe_cache_auto: size, pages, overlap = self._resolve_auto_moe_cache_size(config, banks, method, max_slots=size_cap) + if config.moe_prefill_overlap and not overlap: + logger.info_rank0( + f"--moe-cache-auto: prefill overlap disabled, its {2 * config.model_config.num_experts} " + "slot floor does not fit beside the KV reserve" + ) object.__setattr__(config, "moe_cache_size", size) object.__setattr__(config, "moe_prefill_overlap", overlap) if config.num_page_override is None: diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py index 7b37af9e0e..66d0a8df7f 100644 --- a/tests/engine/test_cache_budget.py +++ b/tests/engine/test_cache_budget.py @@ -80,6 +80,43 @@ def test_budget_too_small_for_min_moe_plus_reserve_raises(): ) # min moe = 4 slots (400 B) + reserve (10 pages = 100 B) = 500 B > 300 B budget +def test_overlap_floor_that_does_not_fit_falls_back_to_num_experts(): + # #4: 2*num_experts slots (800 B) + the KV reserve (100 B) exceed the 700 B budget, but + # num_experts slots do fit. Plan without overlap instead of refusing to start. + size, pages, overlap = plan_cache_budget( + budget_bytes=700, per_expert_bytes=100, cache_per_page=10, + num_experts=4, total_experts=50, prefill_overlap=True, + kv_reserve_pages=10, max_slots=50, + ) + assert overlap is False + assert size == 6 # (700 - 100) // 100, above the num_experts floor + assert pages == 10 + assert size * 100 + pages * 10 <= 700 + + +def test_overlap_kept_when_its_floor_fits_the_budget(): + size, pages, overlap = plan_cache_budget( + budget_bytes=900, per_expert_bytes=100, cache_per_page=10, + num_experts=4, total_experts=50, prefill_overlap=True, + kv_reserve_pages=10, max_slots=50, + ) + assert (size, pages, overlap) == (8, 10, True) + + +def test_budget_too_small_message_names_what_fits_and_the_flags(): + # num_experts slots (400 B) + 10 reserved pages (100 B) > 450 B: 5 pages fit beside them. + with pytest.raises(AssertionError) as err: + plan_cache_budget( + budget_bytes=450, per_expert_bytes=100, cache_per_page=10, + num_experts=4, total_experts=50, prefill_overlap=True, + kv_reserve_pages=10, max_slots=50, page_size=64, + ) + msg = str(err.value) + assert "beside 4 slots at most 5 KV pages (320 tokens) fit" in msg + for flag in ("--kv-reserve-tokens", "--kv-cache-dtype", "--memory-ratio", "--moe-cache-size"): + assert flag in msg + + def test_prefill_overlap_false_is_honored(): # Even when the cache could fit 2*num_experts, an explicit False stays False. size, pages, overlap = plan_cache_budget( From 570456b6753f906fb59bcffcfe0c549c6adc5924 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 23 Sep 2026 19:17:59 -0400 Subject: [PATCH 13/94] merge fork #21: fix(triton): keep the float glm dsa sparse launches two-stage to fit sm_120 shared memory --- python/freetoken/kernel/triton/glm_dsa_sparse.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/python/freetoken/kernel/triton/glm_dsa_sparse.py b/python/freetoken/kernel/triton/glm_dsa_sparse.py index d73ece8cf9..bac2cd7eed 100644 --- a/python/freetoken/kernel/triton/glm_dsa_sparse.py +++ b/python/freetoken/kernel/triton/glm_dsa_sparse.py @@ -498,9 +498,9 @@ def glm_dsa_sparse_attn( D_V=d_v, D_R=d_r, BLOCK_H=BLOCK_H, BLOCK_T=block_t, HAS_COUNTS=has_counts, HAS_ROPE=d_r > 0, HAS_FP8=has_fp8, HAS_NVFP4=has_nvfp4, NUM_SPLITS=n_splits, - # The 512-wide latent accumulator exceeds the 99 KiB shared-memory - # limit of consumer Blackwell GPUs with a two-stage pipeline. - num_warps=4, num_stages=1, + # Consumer Blackwell caps kernel shared memory at 99 KiB. Single-stage, the float path + # needs 100 KiB; two stages (its setting before quantized KV) fit. + num_warps=4, num_stages=1 if has_fp8 or has_nvfp4 else 2, ) grid2 = (m, b, h) _glm_dsa_merge_kernel[grid2]( @@ -527,7 +527,8 @@ def glm_dsa_sparse_attn( D_V=d_v, D_R=d_r, BLOCK_H=BLOCK_H, BLOCK_T=block_t, HAS_COUNTS=has_counts, HAS_ROPE=d_r > 0, HAS_FP8=has_fp8, HAS_NVFP4=has_nvfp4, - num_warps=4, num_stages=1, + # same shared-memory split as the split-k launch above + num_warps=4, num_stages=1 if has_fp8 or has_nvfp4 else 2, ) return o From 935e121f7e1a6ae20c7309320d8ebe97e8d29b78 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 23 Sep 2026 19:17:59 -0400 Subject: [PATCH 14/94] merge fork #23 (upstream FlashML-org/FreeToken#17 rebased): fix(server): resolve Hugging Face repo ids to a local checkpoint dir # Conflicts: # python/freetoken/utils/hf.py --- python/freetoken/server/args.py | 20 +++-- python/freetoken/utils/__init__.py | 2 + python/freetoken/utils/hf.py | 45 ++++++++++- tests/server/test_model_path_resolution.py | 87 ++++++++++++++++++++++ 4 files changed, 148 insertions(+), 6 deletions(-) create mode 100644 tests/server/test_model_path_resolution.py diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index b3125074ad..62c33ad853 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -988,16 +988,26 @@ def _infer_reasoning_parser(model_path: str) -> str | None: if is_offload_moe_strategy(kwargs["moe_strategy"]) and _no_cache_flag: kwargs["moe_cache_auto"] = True - if kwargs["model_source"] == "modelscope": - model_path = kwargs["model_path"] - if not os.path.isdir(model_path): + # Resolve a hub repo id ("org/model") to a local checkpoint directory. This has to + # happen here, after served_model_name and the parser cascade have read the repo id + # (they want "DeepSeek-V4-Flash-0731", not a snapshot hash) and before anything opens + # the checkpoint: config parsing reads model-specific files straight off disk, so a + # bare repo id reaches `open()` as a relative path and dies with a FileNotFoundError. + model_path = kwargs["model_path"] + if not os.path.isdir(model_path): + if kwargs["model_source"] == "modelscope": from modelscope import snapshot_download ignore_patterns = [] if kwargs["use_dummy_weight"]: ignore_patterns = ["*.bin", "*.safetensors", "*.pt", "*.ckpt"] - model_path = snapshot_download(model_path, ignore_patterns=ignore_patterns) - kwargs["model_path"] = model_path + kwargs["model_path"] = snapshot_download(model_path, ignore_patterns=ignore_patterns) + else: + from freetoken.utils import download_hf_checkpoint + + kwargs["model_path"] = download_hf_checkpoint( + model_path, dummy_weight=kwargs["use_dummy_weight"] + ) del kwargs["model_source"] # "auto" (or an unspecified dtype) resolves to the checkpoint's dtype. Multimodal / diff --git a/python/freetoken/utils/__init__.py b/python/freetoken/utils/__init__.py index af68bf5816..a1dd10b7c3 100644 --- a/python/freetoken/utils/__init__.py +++ b/python/freetoken/utils/__init__.py @@ -9,6 +9,7 @@ ) from .hf import ( cached_load_hf_config, + download_hf_checkpoint, download_hf_weight, load_eos_token_ids, load_generation_sampling, @@ -30,6 +31,7 @@ __all__ = [ "cached_load_hf_config", + "download_hf_checkpoint", "download_hf_weight", "load_eos_token_ids", "load_generation_sampling", diff --git a/python/freetoken/utils/hf.py b/python/freetoken/utils/hf.py index 489585e8c2..6c6a5f7a00 100644 --- a/python/freetoken/utils/hf.py +++ b/python/freetoken/utils/hf.py @@ -2,6 +2,7 @@ import functools import json import os +import re from typing import Any, FrozenSet, Mapping from huggingface_hub import hf_hub_download, snapshot_download @@ -296,4 +297,46 @@ def download_hf_weight(model_path: str) -> str: except Exception as e: raise ValueError( f"Model path '{model_path}' is neither a local directory nor a valid model ID: {e}" - ) \ No newline at end of file + ) + + +# Weight formats FreeToken never loads -- it reads safetensors (plus GGUF for Gemma-4). +# Plenty of repos still ship a duplicate .bin/.pth copy of the same tensors, so skipping +# them can halve a first-time pull. +_UNLOADABLE_WEIGHTS = ("*.bin", "*.pt", "*.pth", "*.h5", "*.msgpack", "*.onnx") + + +# A hub repo id is "org/name" and nothing else. Anything with a leading separator, a +# "~", or more than one "/" is a filesystem path the user got wrong, and turning that +# into a hub lookup replaces "no such directory" with a confusing download error. +_HUB_REPO_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*/[A-Za-z0-9][A-Za-z0-9._-]*$") + + +def looks_like_hub_repo_id(model_path: str) -> bool: + """Whether ``model_path`` should be resolved against the Hub rather than the disk.""" + return bool(_HUB_REPO_ID.match(model_path)) + + +def download_hf_checkpoint(model_path: str, *, dummy_weight: bool = False) -> str: + """Resolve a Hugging Face repo id to a complete local checkpoint directory. + + ``download_hf_weight`` fetches only ``*.safetensors``, which is all the weight + loaders need once the engine is running. Config parsing happens earlier and wants + more than that: the tokenizer, and the model-specific subdirs some architectures + read directly off disk -- DeepSeek-V4 fails without ``inference/config.json`` + (see ``models/deepseek_v4/args.py``). So this pulls the whole snapshot. + + A path that is already a local directory is returned untouched. + """ + if os.path.isdir(model_path) or not looks_like_hub_repo_id(model_path): + return model_path + ignore_patterns = list(_UNLOADABLE_WEIGHTS) + if dummy_weight: + # Config and tokenizer are still needed; the tensors are not. + ignore_patterns += ["*.safetensors", "*.gguf"] + try: + return snapshot_download(model_path, ignore_patterns=ignore_patterns) + except Exception as e: + raise ValueError( + f"Model path '{model_path}' is neither a local directory nor a valid model ID: {e}" + ) diff --git a/tests/server/test_model_path_resolution.py b/tests/server/test_model_path_resolution.py new file mode 100644 index 0000000000..76016e6869 --- /dev/null +++ b/tests/server/test_model_path_resolution.py @@ -0,0 +1,87 @@ +"""``--model org/name``: a hub repo id resolves to a local checkpoint directory. + +``ft serve --model deepseek-ai/DeepSeek-V4-Flash-0731`` used to pass the repo id through +untouched as a filesystem path. Config parsing then opened it relative to cwd: + + FileNotFoundError: No DeepSeek-V4 ModelArgs JSON found under + deepseek-ai/DeepSeek-V4-Flash-0731 (looked for inference/config.json) + +Only the modelscope branch ever downloaded, and ``download_hf_weight`` (used later, by the +weight loaders) fetches ``*.safetensors`` alone -- so it would never have produced +``inference/config.json`` either. + +These pin the resolve for both sources, and pin the two things that must keep reading the +*repo id* rather than the resolved snapshot path: the served model name and the parser +cascade. Resolving too early renames the served model to a snapshot hash. +""" + +from __future__ import annotations + +import sys +import types +from unittest.mock import patch + +import pytest + +from freetoken.server.args import parse_args + +REPO_ID = "deepseek-ai/DeepSeek-V4-Flash-0731" +SNAPSHOT = "/cache/hub/models--deepseek-ai--DeepSeek-V4-Flash-0731/snapshots/7872f01" + +# Explicit dtype and parsers keep parse_args off the network: all three otherwise read the +# checkpoint config, which is exactly what cannot be loaded until the path resolves. +BASE_ARGS = ["--dtype", "bfloat16", "--tool-call-parser", "llama3", "--reasoning-parser", "off"] + + +def _parse(model_path: str, *extra: str): + return parse_args([*BASE_ARGS, "--model", model_path, *extra])[0] + + +@pytest.fixture +def hf_download(): + with patch("freetoken.utils.hf.snapshot_download", return_value=SNAPSHOT) as m: + yield m + + +def test_repo_id_resolves_to_local_snapshot(hf_download): + assert _parse(REPO_ID).model_path == SNAPSHOT + assert hf_download.call_args.args[0] == REPO_ID + + +def test_served_model_name_keeps_the_repo_id(hf_download): + """Not the snapshot hash -- this is the id clients send in the `model` field.""" + assert _parse(REPO_ID).served_model_name == "DeepSeek-V4-Flash-0731" + + +def test_local_dir_is_never_downloaded(tmp_path, hf_download): + assert _parse(str(tmp_path)).model_path == str(tmp_path) + hf_download.assert_not_called() + + +def test_formats_freetoken_cannot_load_are_skipped(hf_download): + """Many repos ship a duplicate .bin copy; the loaders only read safetensors/gguf.""" + _parse(REPO_ID) + ignored = hf_download.call_args.kwargs["ignore_patterns"] + assert "*.bin" in ignored + assert "*.safetensors" not in ignored + + +def test_dummy_weight_skips_the_tensors(hf_download): + """--dummy-weight still needs config + tokenizer, but not 156 GiB of weights.""" + _parse(REPO_ID, "--dummy-weight") + ignored = hf_download.call_args.kwargs["ignore_patterns"] + assert "*.safetensors" in ignored + assert "*.gguf" in ignored + + +def test_modelscope_source_still_uses_modelscope(monkeypatch): + """The isdir check moved out of the modelscope branch; it must still short-circuit.""" + fake = types.ModuleType("modelscope") + calls = [] + fake.snapshot_download = lambda repo, **kw: calls.append((repo, kw)) or SNAPSHOT + monkeypatch.setitem(sys.modules, "modelscope", fake) + + with patch("freetoken.utils.hf.snapshot_download") as hf: + assert _parse(REPO_ID, "--model-source", "modelscope").model_path == SNAPSHOT + hf.assert_not_called() + assert calls[0][0] == REPO_ID From 5829f75e9ee03610f78bcc72a2db27214d1b70c5 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 23 Sep 2026 19:17:59 -0400 Subject: [PATCH 15/94] merge fork #47 (upstream FlashML-org/FreeToken#393 rebased): OpenAI protocol coverage --- docs/openai_api.md | 101 +++++ python/freetoken/core.py | 33 +- python/freetoken/engine/sample.py | 204 ++++++++-- python/freetoken/message/frontend.py | 3 + python/freetoken/message/tokenizer.py | 4 + python/freetoken/scheduler/scheduler.py | 9 +- python/freetoken/server/api_models.py | 58 +++ python/freetoken/server/api_server.py | 10 +- python/freetoken/server/control_api.py | 18 + python/freetoken/server/generation.py | 58 ++- python/freetoken/server/openai_api.py | 471 ++++++++++++++++------ python/freetoken/server/stats.py | 41 ++ python/freetoken/tokenizer/detokenize.py | 54 ++- python/freetoken/tokenizer/server.py | 24 +- python/freetoken/tokenizer/tokenize.py | 8 +- tests/engine/test_logits_processors.py | 155 +++++++ tests/server/test_openai_extras.py | 383 ++++++++++++++++++ tests/tokenizer/test_detokenize_extras.py | 81 ++++ 18 files changed, 1538 insertions(+), 177 deletions(-) create mode 100644 docs/openai_api.md create mode 100644 tests/engine/test_logits_processors.py create mode 100644 tests/server/test_openai_extras.py create mode 100644 tests/tokenizer/test_detokenize_extras.py diff --git a/docs/openai_api.md b/docs/openai_api.md new file mode 100644 index 0000000000..0b1cc9d93c --- /dev/null +++ b/docs/openai_api.md @@ -0,0 +1,101 @@ +# OpenAI-compatible API: endpoints and parameters + +What `ft serve` accepts on its OpenAI-shaped routes, what it ignores, and what it refuses, +next to the same surface in vLLM, SGLang and llama.cpp's server. "honoured" means the value +changes the generation; "accepted" means the request is not rejected but the value has no +effect; "400" means the server refuses the request with a clear message. + +## Endpoints + +| route | FreeToken | vLLM | SGLang | llama.cpp | +|---|---|---|---|---| +| `POST /v1/chat/completions` | yes | yes | yes | yes | +| `POST /v1/completions` | yes | yes | yes | yes | +| `GET /v1/models`, `GET /v1/models/{id}` | yes | yes | yes | yes | +| `POST /v1/responses` (+ get, cancel) | yes | yes | yes | yes | +| `POST /v1/messages` (Anthropic), `count_tokens` | yes | yes | no | yes | +| `POST /tokenize`, `POST /detokenize` (also under `/v1/`) | yes | yes | yes | yes (`/tokenize`, `/detokenize`) | +| `GET /metrics` (Prometheus) | yes | yes | yes | yes | +| `GET /version` | yes | yes | yes | no (`/props`) | +| `GET /health` | yes | yes | yes | yes | +| `GET /v1/stats`, `/v1/cache/*`, `/v1/requests`, `/v1/admin/*` | yes (FreeToken control plane) | `/metrics` only | `/get_server_info`, `/flush_cache` | `/slots`, `/props`, `/metrics` | +| `POST /v1/embeddings`, `/rerank`, `/score`, `/pooling` | no (generative models only) | embedding / reranker models | embedding / reranker models | yes, any model with pooling | +| `POST /v1/audio/*`, `/v1/files`, `/v1/batches` | no | transcriptions, batch runner (offline) | transcriptions, files, batches | no | +| `POST /infill`, `/apply-template`, `/completion` (non-OpenAI) | `suffix` on `/v1/completions`; `/tokenize` with `messages` renders the template | no | no | yes | + +## Chat completions: request + +| parameter | FreeToken | vLLM | SGLang | llama.cpp | +|---|---|---|---|---| +| `model`, `messages` (text + `image_url` parts) | honoured | honoured | honoured | honoured | +| `max_tokens`, `max_completion_tokens` | honoured (default 32k) | honoured | honoured | `max_tokens` | +| `temperature`, `top_p`, `top_k` | honoured | honoured | honoured | honoured | +| `min_p` | honoured (prob. below `min_p` x top prob dropped) | honoured | honoured | honoured | +| `presence_penalty`, `frequency_penalty` | honoured, over generated tokens | honoured | honoured | honoured | +| `repetition_penalty` | honoured, over prompt + generated tokens (HF semantics) | honoured | honoured | `repeat_penalty` | +| `logit_bias` | honoured (token id -> bias, clamped to [-100, 100]) | honoured | honoured | honoured | +| `min_tokens` | honoured: no EOS / stop token before this many generated tokens | honoured | `min_new_tokens` | no | +| `stop` (strings), `stop_token_ids` | honoured | honoured | honoured | `stop` only | +| `include_stop_str_in_output` | honoured | honoured | `no_stop_trim` | no | +| `skip_special_tokens` | honoured; default off here, the reasoning and tool parsers read the tags | honoured (default on) | honoured (default on) | n/a | +| `n` | honoured, 1..16, one generation per choice (the prefix cache shares the prompt) | honoured | honoured | `n_cmpl` | +| `stream`, `stream_options.include_usage` | honoured | honoured | honoured | honoured | +| `logprobs`, `top_logprobs` | 400 on this route (see #224); `/v1/completions` `logprobs` 400 | honoured | honoured | `n_probs` | +| `tools`, `tool_choice` (`auto` / `none` / `required` / named), `parallel_tool_calls` | honoured | honoured | honoured | honoured | +| `reasoning_effort`, `thinking`, `chat_template_kwargs` | honoured | `chat_template_kwargs` | honoured | honoured | +| `continue_final_message` | honoured (the last assistant message is continued, no generation prompt) | honoured | honoured | no | +| `response_format` `json_object` / `json_schema`, guided decoding | 400 (no constrained decoding) | honoured (xgrammar / guidance) | honoured (xgrammar / outlines / llguidance) | `grammar`, `json_schema` | +| `seed` | accepted, not honoured (one batched sampling kernel, no per-request generator) | honoured | `sampling_seed` | honoured | +| `user`, `metadata`, `store`, `service_tier` | accepted, ignored | accepted, ignored | accepted, ignored | accepted, ignored | +| `request_id` | honoured: becomes the response id | honoured | `rid` | `id_slot` (different meaning) | +| `function_call` (legacy) | 400 | 400 | 400 | no | +| `best_of`, `use_beam_search`, `length_penalty`, `prompt_logprobs`, `echo` (chat), `truncate_prompt_tokens`, `priority`, `cache_salt`, `lora` | no | honoured | partly | partly | +| `ignore_eos` | honoured (benchmarking) | honoured | honoured | honoured | + +## Legacy completions: request + +`prompt` as a string or a list of strings (token-id prompts are refused); the same sampling +parameters as above; plus: + +| parameter | FreeToken | vLLM | SGLang | llama.cpp | +|---|---|---|---|---| +| `echo` | honoured (the prompt leads the text; not with `logprobs`) | honoured | honoured | no | +| `suffix` | honoured on models with fill-in-the-middle tokens (`<|fim_prefix|>` ... `<|fim_middle|>`): the prompt becomes prefix-suffix-middle | 400 | 400 | `/infill` | +| `n` | honoured, prompt-major choice order (`prompt i`, sample `k` -> `index i*n+k`) | honoured | honoured | no | +| `logprobs` | 400 (see #224 for the sampled-token logprobs) | honoured | honoured | `n_probs` | + +## Responses + +| field | FreeToken | notes | +|---|---|---| +| `id` | `chatcmpl-` / `cmpl-`, or `-` when the client sent one | | +| `created`, `model`, `object` | yes | | +| `system_fingerprint` | `null` | vLLM and SGLang also send `null` | +| `choices[].message.reasoning_content` | yes, when the model reasoned | vLLM `reasoning_content`, llama.cpp `reasoning_content` | +| `choices[].message.tool_calls` | yes | | +| `choices[].finish_reason` | `stop`, `length`, `tool_calls` | | +| `usage.prompt_tokens`, `completion_tokens`, `total_tokens` | yes | for `n > 1` the prompt is counted once, completions summed | +| `usage.prompt_tokens_details.cached_tokens` | with `--enable-cache-report` | vLLM, SGLang: always | +| `usage.completion_tokens_details.reasoning_tokens` | yes, when the model emitted a reasoning end tag | vLLM, SGLang | +| streaming: one `[DONE]`, an optional final usage chunk | yes | | +| llama.cpp `timings` | no; `/v1/stats` and `/metrics` carry the server-side rates | | + +## Tokenize / detokenize + +``` +POST /tokenize {"prompt": "text", "add_special_tokens": true, "return_token_strs": false} +POST /tokenize {"messages": [...], "add_generation_prompt": true, "chat_template_kwargs": {}} + -> {"count": N, "max_model_len": L, "tokens": [...], "token_strs": [...]?} +POST /detokenize {"tokens": [...], "skip_special_tokens": false} -> {"prompt": "text"} +``` + +`messages` are rendered through the model's chat template exactly as a generation would +render them, so `count` equals the `prompt_tokens` that generation reports. + +## Not implemented, on purpose + +- Constrained decoding (`response_format` json / json_schema, regex, grammar): needs a + grammar engine on the sampling path. +- Prompt logprobs and `echo` with `logprobs`: need the prefill logits. +- Per-request `seed`: the batched sampling kernel draws from one generator. +- Embeddings, reranking, scoring, audio: not generation. diff --git a/python/freetoken/core.py b/python/freetoken/core.py index 812f9a4349..a160beaf81 100644 --- a/python/freetoken/core.py +++ b/python/freetoken/core.py @@ -24,13 +24,45 @@ class SamplingParams: # Stop strings (OpenAI `stop` / Anthropic `stop_sequences`). Generation finishes when one # appears in the decoded output; the matched substring (and anything after) is trimmed. stop_strs: list[str] = field(default_factory=list) + # ---- logits processors (engine/sample.py), all off by default ---- + # min_p: drop tokens whose probability is below min_p x the top probability. + min_p: float = 0.0 + # OpenAI presence/frequency penalties over the generated tokens; repetition_penalty is the + # HF/vLLM multiplicative penalty over prompt + generated tokens (1.0 = off). presence_penalty: float = 0.0 frequency_penalty: float = 0.0 + repetition_penalty: float = 1.0 + # OpenAI logit_bias as [[token id, bias], ...] (clamped to [-100, 100] by the API). A list, + # not a dict: the messages cross process boundaries as msgpack, which refuses int map keys. + logit_bias: list[list[float]] | None = None + # No EOS / stop token before this many generated tokens (vLLM min_tokens, SGLang + # min_new_tokens). The scheduler fills min_tokens_stop_ids with the ids to mask. + min_tokens: int = 0 + min_tokens_stop_ids: list[int] = field(default_factory=list) + # Extra token ids that end generation with finish_reason "stop" (vLLM stop_token_ids). + stop_token_ids: list[int] = field(default_factory=list) + # Keep the matched stop string in the output instead of trimming it (vLLM + # include_stop_str_in_output, SGLang no_stop_trim). + include_stop_str_in_output: bool = False + # Decode with skip_special_tokens. Off by default here: the reasoning and tool parsers + # consume the /tool tags from the decoded text. + skip_special_tokens: bool = False @property def is_greedy(self) -> bool: return self.temperature <= 0.0 or self.top_k == 1 + @property + def needs_logits_processing(self) -> bool: + return bool( + self.min_p > 0.0 + or self.presence_penalty != 0.0 + or self.frequency_penalty != 0.0 + or self.repetition_penalty != 1.0 + or self.logit_bias + or self.min_tokens > 0 + ) + @dataclass(eq=False) class Req: @@ -68,7 +100,6 @@ class Req: # handler must not free resources under an in-flight forward; it sets this flag and # _process_last_data frees the request when the batch drains (after copy_done.synchronize). aborted: bool = False - output_token_counts: torch.Tensor | None = field(default=None, init=False, repr=False) def __post_init__(self) -> None: assert self.input_ids.is_cpu diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py index c6fb2f40ee..71dd79b844 100644 --- a/python/freetoken/engine/sample.py +++ b/python/freetoken/engine/sample.py @@ -1,6 +1,6 @@ from __future__ import annotations -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import TYPE_CHECKING, List import torch @@ -10,19 +10,114 @@ from freetoken.core import Batch +@dataclass +class LogitsPlan: + """Inputs of the logits processors for the rows of one batch that asked for any + (presence / frequency / repetition penalties, logit_bias, min_tokens, min_p). + + Built on the host in ``Sampler.prepare`` from the requests' SamplingParams and + token histories, applied on the device by ``apply_logits_processors`` before the + sampling kernel. ``None`` when no request in the batch uses a processor, so the + default path stays exactly as it was. Row indices inside the plan are LOCAL (0..r-1, + in the order of ``rows``); ``rows`` maps them back to the batch. + """ + + rows: torch.Tensor # [r] int64 batch rows + temps: torch.Tensor # [r, 1] float32 temperatures (min_p works on probabilities) + presence: torch.Tensor | None = None # [r, 1] float32 + frequency: torch.Tensor | None = None # [r, 1] float32 + repetition: torch.Tensor | None = None # [r, 1] float32 (1.0 = off) + # Generated ids per row, right-padded with vocab_size (a scratch column that is + # dropped), for the presence / frequency counts and the repetition mask. + hist_ids: torch.Tensor | None = None # [r, L] int64 + # Prompt ids per row, same padding: repetition_penalty also covers the prompt. + prompt_ids: torch.Tensor | None = None # [r, P] int64 + bias_rows: torch.Tensor | None = None # [b] int64 local rows + bias_ids: torch.Tensor | None = None # [b] int64 + bias_vals: torch.Tensor | None = None # [b] float32 + # Rows still under min_tokens and the ids (EOS + stop_token_ids, padded) they must + # not sample yet. + min_rows: torch.Tensor | None = None # [m] int64 local rows + min_ids: torch.Tensor | None = None # [m, S] int64 + min_p: torch.Tensor | None = None # [r, 1] float32 (0 = off for that row) + + @dataclass class BatchSamplingArgs: temperatures: torch.Tensor | None top_k: torch.Tensor | None = None top_p: torch.Tensor | None = None greedy_mask: torch.Tensor | None = None - penalties: list[tuple[int, torch.Tensor, float, float]] = field(default_factory=list) + plan: LogitsPlan | None = None def make_device_tensor(data: List, dtype: torch.dtype, device: torch.device) -> torch.Tensor: + if device.type != "cuda": # the unit tests build plans on the CPU + return torch.tensor(data, dtype=dtype, device=device) return torch.tensor(data, dtype=dtype, pin_memory=True).to(device, non_blocking=True) +def _padded_rows(rows: list[list[int]], pad: int, device: torch.device) -> torch.Tensor: + """[len(rows), max_len] int64 with ``pad`` on the right. A synchronous copy: the + histories can be long (prompt + output) and a pinned buffer must not be freed under + an in-flight non_blocking copy.""" + width = max(1, max((len(r) for r in rows), default=1)) + out = torch.full((len(rows), width), pad, dtype=torch.int64) + for i, r in enumerate(rows): + if r: + out[i, : len(r)] = torch.as_tensor(r, dtype=torch.int64) + return out.to(device) + + +def apply_logits_processors( + logits: torch.Tensor, plan: LogitsPlan, vocab_size: int +) -> torch.Tensor: + """A float32 copy of ``logits`` with the plan applied to its rows, in vLLM's order: + repetition -> presence -> frequency penalties, logit_bias, the min_tokens mask, then + min_p. Rows outside the plan are copied unchanged.""" + out = logits.to(torch.float32, copy=True) + sub = out[plan.rows] # [r, V] (advanced indexing copies) + r = sub.shape[0] + device = sub.device + neg_inf = float("-inf") + + if plan.hist_ids is not None: + counts = torch.zeros((r, vocab_size + 1), dtype=torch.float32, device=device) + counts.scatter_add_(1, plan.hist_ids, torch.ones_like(plan.hist_ids, dtype=torch.float32)) + counts = counts[:, :vocab_size] + if plan.repetition is not None: + seen = counts > 0 + if plan.prompt_ids is not None: + in_prompt = torch.zeros((r, vocab_size + 1), dtype=torch.bool, device=device) + in_prompt.scatter_(1, plan.prompt_ids, torch.ones_like(plan.prompt_ids, dtype=torch.bool)) + seen = seen | in_prompt[:, :vocab_size] + # HF / vLLM semantics: divide positive logits, multiply negative ones. + penalized = torch.where(sub > 0, sub / plan.repetition, sub * plan.repetition) + sub = torch.where(seen, penalized, sub) + if plan.presence is not None: + sub = sub - plan.presence * (counts > 0).to(sub.dtype) + if plan.frequency is not None: + sub = sub - plan.frequency * counts + + if plan.bias_ids is not None: + sub.index_put_((plan.bias_rows, plan.bias_ids), plan.bias_vals, accumulate=True) + + if plan.min_ids is not None: + block = torch.zeros((plan.min_ids.shape[0], vocab_size + 1), dtype=torch.bool, device=device) + block.scatter_(1, plan.min_ids, torch.ones_like(plan.min_ids, dtype=torch.bool)) + rows = sub[plan.min_rows] + sub[plan.min_rows] = torch.where(block[:, :vocab_size], torch.full_like(rows, neg_inf), rows) + + if plan.min_p is not None: + probs = torch.softmax(sub / plan.temps, dim=-1) + keep = probs >= plan.min_p * probs.amax(dim=-1, keepdim=True) + keep = keep | (plan.min_p <= 0.0) # rows without min_p stay untouched + sub = torch.where(keep, sub, torch.full_like(sub, neg_inf)) + + out[plan.rows] = sub + return out + + def sample_impl( logits: torch.Tensor, temperatures: torch.Tensor, @@ -57,23 +152,88 @@ class Sampler: device: torch.device vocab_size: int + def _plan(self, batch: Batch, params) -> LogitsPlan | None: + """The LogitsPlan for this batch, or None when no request asks for a processor. + Token histories come from the host-side ``req.input_ids``; under overlap + scheduling the previous step's token may not be appended yet, so a penalty sees + the history one token late. That is the accepted cost of keeping the counts off + the device.""" + MIN_T = 1e-6 + rows = [i for i, p in enumerate(params) if p.needs_logits_processing] + if not rows: + return None + picked = [(batch.reqs[i], params[i]) for i in rows] + temps = [[max(0.0 if p.is_greedy else p.temperature, MIN_T)] for _, p in picked] + plan = LogitsPlan( + rows=make_device_tensor(rows, torch.int64, self.device), + temps=make_device_tensor(temps, torch.float32, self.device), + ) + pad = self.vocab_size + + want_hist = any( + p.presence_penalty != 0.0 or p.frequency_penalty != 0.0 or p.repetition_penalty != 1.0 + for _, p in picked + ) + if want_hist: + hist: list[list[int]] = [] + prompts: list[list[int]] = [] + want_prompt = False + for req, p in picked: + prompt_len = req.max_device_len - req.output_len + ids = req.input_ids + penalized = p.presence_penalty != 0.0 or p.frequency_penalty != 0.0 or p.repetition_penalty != 1.0 + hist.append(ids[prompt_len:].tolist() if penalized else []) + if p.repetition_penalty != 1.0: + want_prompt = True + prompts.append(ids[:prompt_len].tolist()) + else: + prompts.append([]) + plan.hist_ids = _padded_rows(hist, pad, self.device) + if want_prompt: + plan.prompt_ids = _padded_rows(prompts, pad, self.device) + if any(p.presence_penalty != 0.0 for _, p in picked): + plan.presence = make_device_tensor([[p.presence_penalty] for _, p in picked], torch.float32, self.device) + if any(p.frequency_penalty != 0.0 for _, p in picked): + plan.frequency = make_device_tensor([[p.frequency_penalty] for _, p in picked], torch.float32, self.device) + if any(p.repetition_penalty != 1.0 for _, p in picked): + plan.repetition = make_device_tensor([[p.repetition_penalty] for _, p in picked], torch.float32, self.device) + + bias_rows: list[int] = [] + bias_ids: list[int] = [] + bias_vals: list[float] = [] + for local, (_, p) in enumerate(picked): + for tid, val in p.logit_bias or (): + if 0 <= int(tid) < self.vocab_size: + bias_rows.append(local) + bias_ids.append(int(tid)) + bias_vals.append(float(val)) + if bias_ids: + plan.bias_rows = make_device_tensor(bias_rows, torch.int64, self.device) + plan.bias_ids = make_device_tensor(bias_ids, torch.int64, self.device) + plan.bias_vals = make_device_tensor(bias_vals, torch.float32, self.device) + + min_rows: list[int] = [] + min_ids: list[list[int]] = [] + for local, (req, p) in enumerate(picked): + if p.min_tokens > 0 and p.min_tokens_stop_ids: + generated = len(req.input_ids) - (req.max_device_len - req.output_len) + if generated < p.min_tokens: + min_rows.append(local) + min_ids.append([t for t in p.min_tokens_stop_ids if 0 <= t < self.vocab_size]) + if min_rows: + plan.min_rows = make_device_tensor(min_rows, torch.int64, self.device) + plan.min_ids = _padded_rows(min_ids, pad, self.device) + + if any(p.min_p > 0.0 for _, p in picked): + plan.min_p = make_device_tensor([[p.min_p] for _, p in picked], torch.float32, self.device) + return plan + def prepare(self, batch: Batch) -> BatchSamplingArgs: params = [r.sampling_params for r in batch.reqs] is_greedy = [p.is_greedy for p in params] - penalties = [] - for row, req in enumerate(batch.reqs): - p = req.sampling_params - if not (p.presence_penalty or p.frequency_penalty) or not req.can_decode: - continue - if req.output_token_counts is None: - req.output_token_counts = torch.zeros( - self.vocab_size, dtype=torch.int32, device=self.device - ) - penalties.append( - (row, req.output_token_counts, p.presence_penalty, p.frequency_penalty) - ) + plan = self._plan(batch, params) if all(is_greedy): - return BatchSamplingArgs(temperatures=None, penalties=penalties) + return BatchSamplingArgs(temperatures=None, plan=plan) MIN_P = MIN_T = 1e-6 # Greedy outputs are selected explicitly in sample(); use neutral sampling @@ -97,16 +257,14 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs: make_device_tensor(is_greedy, torch.bool, self.device) if any(is_greedy) else None ) return BatchSamplingArgs( - temperatures, top_k=top_k, top_p=top_p, greedy_mask=greedy_mask, penalties=penalties + temperatures, top_k=top_k, top_p=top_p, greedy_mask=greedy_mask, plan=plan ) @nvtx_annotate("Sampler") def sample(self, logits: torch.Tensor, args: BatchSamplingArgs) -> torch.Tensor: with torch.cuda.nvtx.range("Sampler"): - if args.penalties: - logits = logits.float().clone() - for row, counts, presence, frequency in args.penalties: - logits[row] -= frequency * counts + presence * (counts > 0) + if args.plan is not None: + logits = apply_logits_processors(logits, args.plan, self.vocab_size) if args.temperatures is None: # greedy sampling tokens = torch.argmax(logits, dim=-1) else: @@ -116,10 +274,4 @@ def sample(self, logits: torch.Tensor, args: BatchSamplingArgs) -> torch.Tensor: # greedy rows must follow argmax's deterministic tie-breaking. greedy_tokens = torch.argmax(logits, dim=-1).to(tokens.dtype) tokens = torch.where(args.greedy_mask, greedy_tokens, tokens) - # Update on the sampling stream: overlapped scheduling can prepare the next - # batch before the previous token reaches Req.input_ids on the CPU. - for row, counts, _, _ in args.penalties: - counts.scatter_add_( - 0, tokens[row : row + 1].long(), counts.new_ones(1) - ) return tokens diff --git a/python/freetoken/message/frontend.py b/python/freetoken/message/frontend.py index 24725567b9..f3200650d8 100644 --- a/python/freetoken/message/frontend.py +++ b/python/freetoken/message/frontend.py @@ -56,6 +56,9 @@ class UserReply(BaseFrontendMsg): finish_reason: str | None = None # The stop string that ended generation (Anthropic reports it as stop_reason='stop_sequence'). matched_stop: str | None = None + # On the finished reply: tokens generated up to and including the reasoning end tag + # (0 when the model emitted none) -> usage.completion_tokens_details.reasoning_tokens. + reasoning_tokens: int = 0 @dataclass diff --git a/python/freetoken/message/tokenizer.py b/python/freetoken/message/tokenizer.py index d52d658e63..0d2116038f 100644 --- a/python/freetoken/message/tokenizer.py +++ b/python/freetoken/message/tokenizer.py @@ -36,6 +36,10 @@ class DetokenizeMsg(BaseTokenizerMsg): # The request's stop strings (None when it has none), so the detokenizer can hold back # a trailing partial-stop prefix instead of streaming it and then needing to retract. stop_strs: list[str] | None = None + # Keep the matched stop string in the output (SamplingParams.include_stop_str_in_output). + keep_stop_str: bool = False + # Decode this request with skip_special_tokens (SamplingParams.skip_special_tokens). + skip_special_tokens: bool = False # KV page-pool usage snapshot at this step (not-evictable used/total), passed # through to the frontend for the shell status bar. 0/0 for owned-KV models. kv_used_pages: int = 0 diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index d057d3b908..f508a9dc89 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -353,7 +353,7 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: hit_length = req.input_ids.numel() >= req.max_device_len hit_eos = ( not req.sampling_params.ignore_eos and next_token in self.eos_token_ids - ) + ) or next_token in req.sampling_params.stop_token_ids matched_stop = ( self._match_stop_str(req) if not hit_eos and req.sampling_params.stop_strs @@ -379,6 +379,8 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: finish_reason=finish_reason, matched_stop=matched_stop, stop_strs=req.sampling_params.stop_strs or None, + keep_stop_str=req.sampling_params.include_stop_str_in_output, + skip_special_tokens=req.sampling_params.skip_special_tokens, ) ) @@ -542,6 +544,11 @@ def _process_one_msg(self, msg: BaseBackendMsg) -> None: logger.warning_rank0( f"Adjust max_tokens to {max_output_len} for request {msg.uid}." ) + sp = msg.sampling_params + if sp.min_tokens > 0: + # min_tokens: the sampler masks these ids while the output is shorter. + sp.min_tokens = min(sp.min_tokens, sp.max_tokens) + sp.min_tokens_stop_ids = sorted(set(self.eos_token_ids) | set(sp.stop_token_ids)) self.prefill_manager.add_one_req(msg) elif isinstance(msg, AbortBackendMsg): logger.debug_rank0("Aborting request %d", msg.uid) diff --git a/python/freetoken/server/api_models.py b/python/freetoken/server/api_models.py index 6c3ade0924..474af6b961 100644 --- a/python/freetoken/server/api_models.py +++ b/python/freetoken/server/api_models.py @@ -90,11 +90,30 @@ class ChatCompletionRequest(BaseModel): function_call: Any | None = None logit_bias: dict[str, float] | None = None response_format: dict[str, Any] | None = None + # The last message is an assistant prefix the model continues (no generation prompt). + continue_final_message: bool = False + # ---- sampling extras shared with vLLM / SGLang / llama.cpp (all optional) ---- + min_p: float | None = None + repetition_penalty: float | None = None + min_tokens: int = 0 + stop_token_ids: list[int] | None = None + include_stop_str_in_output: bool = False + skip_special_tokens: bool | None = None + # Accepted for OpenAI compatibility, not honoured: the batched sampling kernel has no + # per-request generator (seed); user / metadata / store are opaque to the server. + seed: int | None = None + user: str | None = None + # A client-chosen id echoed as the response id (vLLM request_id / SGLang rid). + request_id: str | None = None @model_validator(mode="after") def _sync_max_completion_tokens(self) -> "ChatCompletionRequest": if self.max_completion_tokens is not None: self.max_tokens = self.max_completion_tokens + if self.continue_final_message and ( + not self.messages or self.messages[-1].role != "assistant" + ): + raise ValueError("continue_final_message needs a final assistant message") return self @@ -120,6 +139,19 @@ class CompletionRequest(BaseModel): suffix: str | None = None logit_bias: dict[str, float] | None = None response_format: dict[str, Any] | None = None + # ---- sampling extras shared with vLLM / SGLang / llama.cpp (all optional) ---- + min_p: float | None = None + repetition_penalty: float | None = None + min_tokens: int = 0 + stop_token_ids: list[int] | None = None + include_stop_str_in_output: bool = False + skip_special_tokens: bool | None = None + # Accepted for OpenAI compatibility, not honoured: the batched sampling kernel has no + # per-request generator (seed); user / metadata / store are opaque to the server. + seed: int | None = None + user: str | None = None + # A client-chosen id echoed as the response id (vLLM request_id / SGLang rid). + request_id: str | None = None @model_validator(mode="after") def _sync_max_completion_tokens(self) -> "CompletionRequest": @@ -128,6 +160,32 @@ def _sync_max_completion_tokens(self) -> "CompletionRequest": return self +MAX_N = 16 # choices per request; each is a full generation, so keep the fan-out bounded + + +class TokenizeRequest(BaseModel): + """``POST /tokenize`` (the vLLM / SGLang shape): a raw prompt, or chat messages that + are rendered through the model's chat template first.""" + + model_config = ConfigDict(extra="allow") + + model: str | None = None + prompt: str | None = None + messages: list[Message] | None = None + add_special_tokens: bool = True + add_generation_prompt: bool = True + return_token_strs: bool = False + chat_template_kwargs: dict[str, Any] = Field(default_factory=dict) + + +class DetokenizeRequest(BaseModel): + model_config = ConfigDict(extra="allow") + + model: str | None = None + tokens: list[int] + skip_special_tokens: bool = False + + class ModelCard(BaseModel): id: str object: str = "model" diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index a9b16adfd7..f74750e691 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -361,16 +361,20 @@ async def stream_generate(self, uid: int): yield b"data: [DONE]\n\n" logger.debug("Finished streaming response for user %s", uid) - async def stream_with_cancellation(self, generator, request: Request, uid: int): + async def stream_with_cancellation(self, generator, request: Request, uid): + """``uid`` is one request id or a sequence of them (an ``n > 1`` fan-out streams + several generations into one response): a disconnect aborts every one.""" + uids = list(uid) if isinstance(uid, (list, tuple)) else [uid] try: async for chunk in generator: # detect if the client has disconnected if await request.is_disconnected(): - logger.info("Client disconnected for user %s", uid) + logger.info("Client disconnected for user %s", uids) raise asyncio.CancelledError yield chunk except asyncio.CancelledError: - asyncio.create_task(self.abort_user(uid)) + for one in uids: + asyncio.create_task(self.abort_user(one)) raise async def abort_user(self, uid: int): diff --git a/python/freetoken/server/control_api.py b/python/freetoken/server/control_api.py index 7158e4e5fa..cef632306e 100644 --- a/python/freetoken/server/control_api.py +++ b/python/freetoken/server/control_api.py @@ -68,6 +68,24 @@ async def list_requests(since: int = 0, limit: int = 100): from .stats import build_stats + @app.get("/metrics") + async def metrics(): + """Prometheus text exposition of /v1/stats (the llama.cpp / vLLM convention).""" + from fastapi.responses import PlainTextResponse + + from .stats import prometheus_text + + doc = build_stats( + get_state(), request_ring.requests_p95_ms(), request_ring.requests_ttft_mean_ms() + ) + return PlainTextResponse(prometheus_text(doc), media_type="text/plain; version=0.0.4") + + @app.get("/version") + async def version(): + from freetoken import __version__ + + return {"version": __version__} + @app.get("/v1/stats") async def stats(): doc = build_stats( diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index 2e342ff899..c7d658688c 100644 --- a/python/freetoken/server/generation.py +++ b/python/freetoken/server/generation.py @@ -115,6 +115,7 @@ class GenDone: completion_tokens: int matched_stop: str | None = None cached_tokens: int = 0 + reasoning_tokens: int = 0 GenEvent = ReasoningDelta | ContentDelta | ToolCallStart | ToolCallArgsDelta | ToolCallsDelta | GenDone @@ -130,6 +131,7 @@ class GenResult: completion_tokens: int matched_stop: str | None = None cached_tokens: int = 0 + reasoning_tokens: int = 0 @dataclass @@ -169,11 +171,19 @@ def resolve_sampling( model_sampling: dict[str, Any], stop: str | list[str] | None = None, default_max_tokens: int = DEFAULT_MAX_OUTPUT_TOKENS, + min_p: float | None = None, presence_penalty: float = 0.0, frequency_penalty: float = 0.0, + repetition_penalty: float | None = None, + logit_bias: dict[str, float] | dict[int, float] | None = None, + min_tokens: int = 0, + stop_token_ids: list[int] | None = None, + include_stop_str_in_output: bool = False, + skip_special_tokens: bool | None = None, ) -> SamplingParams: """Map a protocol's sampling fields onto the engine's neutral SamplingParams, - filling unspecified fields from the checkpoint's recommended defaults.""" + filling unspecified fields from the checkpoint's recommended defaults. Range errors + raise ValueError (the adapters answer 400).""" def pick(value, key, framework): return value if value is not None else model_sampling.get(key, framework) @@ -185,16 +195,50 @@ def pick(value, key, framework): # non-positive value is a client error. if max_tokens is not None and max_tokens < 1: raise ValueError(f"max_tokens must be at least 1, got {max_tokens}") - return SamplingParams( + resolved_min_p = float(pick(min_p, "min_p", 0.0)) + resolved_rep = float(pick(repetition_penalty, "repetition_penalty", 1.0)) + if not 0.0 <= resolved_min_p <= 1.0: + raise ValueError(f"min_p must be in [0, 1], got {resolved_min_p}") + for name, value in (("presence_penalty", presence_penalty), ("frequency_penalty", frequency_penalty)): + if not -2.0 <= float(value) <= 2.0: + raise ValueError(f"{name} must be in [-2, 2], got {value}") + if resolved_rep <= 0.0: + raise ValueError(f"repetition_penalty must be positive, got {resolved_rep}") + if min_tokens < 0: + raise ValueError(f"min_tokens must be >= 0, got {min_tokens}") + bias: list[list[float]] | None = None + if logit_bias: + bias = [] + for key, value in logit_bias.items(): + try: + tid = int(key) + except (TypeError, ValueError): + raise ValueError(f"logit_bias keys must be token ids, got {key!r}") from None + if tid < 0: + raise ValueError(f"logit_bias token id must be >= 0, got {tid}") + bias.append([tid, max(-100.0, min(100.0, float(value)))]) # the OpenAI range + ids = [int(t) for t in (stop_token_ids or [])] + if any(t < 0 for t in ids): + raise ValueError("stop_token_ids must be non-negative token ids") + params = SamplingParams( ignore_eos=ignore_eos, max_tokens=default_max_tokens if max_tokens is None else max_tokens, temperature=pick(temperature, "temperature", 0.0), top_k=pick(top_k, "top_k", -1), top_p=pick(top_p, "top_p", 1.0), stop_strs=[s for s in stop_list if s], # drop empty strings (would match everything) - presence_penalty=presence_penalty, - frequency_penalty=frequency_penalty, + min_p=resolved_min_p, + presence_penalty=float(presence_penalty), + frequency_penalty=float(frequency_penalty), + repetition_penalty=resolved_rep, + logit_bias=bias, + min_tokens=int(min_tokens), + stop_token_ids=ids, + include_stop_str_in_output=bool(include_stop_str_in_output), ) + if skip_special_tokens is not None: + params.skip_special_tokens = bool(skip_special_tokens) + return params def render_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: @@ -607,6 +651,7 @@ async def _generate_events_impl(uid: int, spec: GenSpec, state: Any) -> AsyncIte prompt_tokens = 0 completion_tokens = 0 cached_tokens = 0 + reasoning_tokens = 0 pending = "" parse_tools = spec.parse_tools reasoning_parser = _make_reasoning_parser(spec, state) @@ -702,6 +747,7 @@ def _route_tool_text(piece: str) -> list[GenEvent]: prompt_tokens += ack.prompt_tokens_delta completion_tokens += ack.completion_tokens_delta cached_tokens += ack.cached_tokens + reasoning_tokens = max(reasoning_tokens, getattr(ack, "reasoning_tokens", 0)) content_delta = ack.incremental_output if reasoning_parser is not None and content_delta: reasoning_delta, content_delta = reasoning_parser.parse_stream_chunk(content_delta) @@ -784,6 +830,7 @@ def _route_tool_text(piece: str) -> list[GenEvent]: yield GenDone( finish_reason, prompt_tokens, completion_tokens, matched_stop=engine_matched_stop, cached_tokens=cached_tokens, + reasoning_tokens=reasoning_tokens, ) @@ -794,6 +841,7 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any) -> GenResult: prompt_tokens = 0 completion_tokens = 0 cached_tokens = 0 + reasoning_tokens = 0 engine_finish_reason: str | None = None engine_matched_stop: str | None = None async for ack in state.wait_for_ack(uid): @@ -802,6 +850,7 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any) -> GenResult: prompt_tokens += ack.prompt_tokens_delta completion_tokens += ack.completion_tokens_delta cached_tokens += ack.cached_tokens + reasoning_tokens = max(reasoning_tokens, getattr(ack, "reasoning_tokens", 0)) full_content += ack.incremental_output if ack.finished: engine_finish_reason = getattr(ack, "finish_reason", None) @@ -828,4 +877,5 @@ async def _generate_full_impl(uid: int, spec: GenSpec, state: Any) -> GenResult: completion_tokens=completion_tokens, matched_stop=engine_matched_stop, cached_tokens=cached_tokens, + reasoning_tokens=reasoning_tokens, ) diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index 99a4283549..f1bb0888f0 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -14,6 +14,9 @@ from freetoken.tokenizer.effort import EFFORT_SCALE, KNOWN_REASONING_EFFORTS from .api_models import ( + MAX_N, + DetokenizeRequest, + TokenizeRequest, ChatCompletionRequest, CompletionRequest, ModelCard, @@ -68,20 +71,11 @@ def chat_request_to_genspec( thinking_type = _thinking_type(req) if req.reasoning_effort or thinking_type: ctk = effort_toggle_kwargs(req.reasoning_effort, ctk, thinking_type=thinking_type) + if req.continue_final_message: + ctk = {**ctk, "continue_final_message": True} return GenSpec( messages=render_messages([m.model_dump(exclude_none=True) for m in req.messages]), - sampling_params=resolve_sampling( - temperature=req.temperature, - top_k=req.top_k, - top_p=req.top_p, - max_tokens=req.max_tokens, - ignore_eos=req.ignore_eos, - model_sampling=model_sampling, - stop=req.stop, - default_max_tokens=default_max_tokens, - presence_penalty=req.presence_penalty, - frequency_penalty=req.frequency_penalty, - ), + sampling_params=_resolve_sampling(req, model_sampling, default_max_tokens=default_max_tokens), chat_template_kwargs=ctk, template_tools=_tools_for_template(req), parser_tools=(_all_tool_dicts(req.tools) if _should_parse_tools(req) else None), @@ -133,6 +127,16 @@ async def v1_completions(req: CompletionRequest, request: Request): return gate return await handle_completion(req, request, state, get_model_sampling()) + @app.post("/tokenize") + @app.post("/v1/tokenize") + async def tokenize(req: TokenizeRequest): + return await handle_tokenize(req, get_state()) + + @app.post("/detokenize") + @app.post("/v1/detokenize") + async def detokenize(req: DetokenizeRequest): + return await handle_detokenize(req, get_state()) + @app.get("/v1/models") async def v1_models(): state = get_state() @@ -157,15 +161,13 @@ async def handle_chat_completion( ): if req.function_call is not None: return create_error_response("function_call is not supported; use tools/tool_choice instead") - if req.logit_bias is not None: - return create_error_response("logit_bias is not supported") if _response_format_unsupported(req.response_format): return create_error_response( "response_format json_object/json_schema is not supported (no constrained decoding)", param="response_format", ) - if req.n != 1: - return create_error_response("Only n=1 is supported", param="n") + if not 1 <= req.n <= MAX_N: + return create_error_response(f"n must be between 1 and {MAX_N}", param="n") # Case/whitespace and the "off" disable synonym stay accepted here because # effort_toggle_kwargs normalizes and honors them downstream. effort = req.reasoning_effort.strip().lower() if isinstance(req.reasoning_effort, str) else None @@ -198,43 +200,65 @@ async def handle_chat_completion( if err is not None: return create_error_response(str(err), code=err.code) + # n > 1: one generation per choice, submitted together (the prefix cache serves the + # shared prompt after the first prefill). + uids: list[int] = [] try: - uid = await submit_generation(spec, state) + for _ in range(req.n): + uids.append(await submit_generation(spec, state)) except GenerationError as exc: + for u in uids: # choices already queued before a later one failed + await state.abort_user(u) return create_error_response(str(exc), code=exc.code) + uid = uids[0] if req.stream: - chunks = stream_chat_completion_chunks(uid, req, state, spec) + if req.n == 1: + chunks = stream_chat_completion_chunks(uid, req, state, spec) + else: + chunks = _merge_streams( + [ + stream_chat_completion_chunks(u, req, state, spec, index=i, terminal=False) + for i, u in enumerate(uids) + ], + lambda: _chat_chunk(req, uid, []), + include_usage=bool(req.stream_options and req.stream_options.include_usage), + state=state, + ) if request is not None: - chunks = state.stream_with_cancellation(chunks, request, uid) + chunks = state.stream_with_cancellation(chunks, request, uids) return StreamingResponse(chunks, media_type="text/event-stream") try: - result = await generate_full(uid, spec, state, source="/v1/chat/completions") + results = await asyncio.gather( + *(generate_full(u, spec, state, source="/v1/chat/completions") for u in uids) + ) except GenerationError as exc: + for u in uids: + await state.abort_user(u) return create_error_response(str(exc), code=exc.code) - message: dict[str, Any] = {"role": "assistant", "content": result.content} - if result.reasoning: - message["reasoning_content"] = result.reasoning - if result.tool_calls: - message["tool_calls"] = _tool_calls_to_openai(result.tool_calls) - + choices: list[dict[str, Any]] = [] + for index, result in enumerate(results): + message: dict[str, Any] = {"role": "assistant", "content": result.content} + if result.reasoning: + message["reasoning_content"] = result.reasoning + if result.tool_calls: + message["tool_calls"] = _tool_calls_to_openai(result.tool_calls) + choices.append({"index": index, "message": message, "finish_reason": result.finish_reason}) + + first = results[0] return { - "id": f"chatcmpl-{uid}", + "id": _response_id("chatcmpl", req, uid), "object": "chat.completion", "created": int(time.time()), "model": req.model, - "choices": [ - { - "index": 0, - "message": message, - "finish_reason": result.finish_reason, - } - ], + "system_fingerprint": None, + "choices": choices, "usage": _usage( - result.prompt_tokens, - result.completion_tokens, - _reported_cached(state, result.cached_tokens), + first.prompt_tokens, + sum(r.completion_tokens for r in results), + _reported_cached(state, max(r.cached_tokens for r in results)), + reasoning_tokens=sum(r.reasoning_tokens for r in results), ), } @@ -244,21 +268,28 @@ async def stream_chat_completion_chunks( req: ChatCompletionRequest, state: Any, spec: GenSpec | None = None, + index: int = 0, + terminal: bool = True, ) -> AsyncIterator[bytes]: - """Format generate_events() into the OpenAI chat.completion.chunk SSE stream.""" + """Format generate_events() into the OpenAI chat.completion.chunk SSE stream. + + ``index`` is the choice index (``n > 1`` runs one of these per choice); with + ``terminal=False`` the usage chunk and ``[DONE]`` are left to the merger, which reads + this stream's usage from the ``_StreamUsage`` yielded last.""" if spec is None: spec = chat_request_to_genspec(req, {}) yield _sse( _chat_chunk( req, uid, - [{"delta": {"role": "assistant", "content": ""}, "index": 0, "finish_reason": None}], + [{"delta": {"role": "assistant", "content": ""}, "index": index, "finish_reason": None}], ) ) prompt_tokens = 0 completion_tokens = 0 cached_tokens = 0 + reasoning_tokens = 0 tool_calls_sent = 0 open_tool: dict[str, Any] | None = None events = generate_events(uid, spec, state, source="/v1/chat/completions") @@ -279,7 +310,7 @@ async def stream_chat_completion_chunks( _chat_chunk( req, uid, - [{"delta": {"reasoning_content": ev.text}, "index": 0, "finish_reason": None}], + [{"delta": {"reasoning_content": ev.text}, "index": index, "finish_reason": None}], ) ) elif isinstance(ev, ContentDelta): @@ -287,7 +318,7 @@ async def stream_chat_completion_chunks( _chat_chunk( req, uid, - [{"delta": {"content": ev.text}, "index": 0, "finish_reason": None}], + [{"delta": {"content": ev.text}, "index": index, "finish_reason": None}], ) ) elif isinstance(ev, ToolCallStart): @@ -307,7 +338,7 @@ async def stream_chat_completion_chunks( "type": "function", "function": {"name": ev.name, "arguments": ""}, }]}, - "index": 0, "finish_reason": None, + "index": index, "finish_reason": None, }], ) ) @@ -324,7 +355,7 @@ async def stream_chat_completion_chunks( "index": open_tool["index"], "function": {"arguments": ev.fragment}, }]}, - "index": 0, "finish_reason": None, + "index": index, "finish_reason": None, }], ) ) @@ -348,7 +379,7 @@ async def stream_chat_completion_chunks( "index": open_tool["index"], "function": {"arguments": remainder}, }]}, - "index": 0, "finish_reason": None, + "index": index, "finish_reason": None, }], ) ) @@ -360,7 +391,7 @@ async def stream_chat_completion_chunks( yield _sse( _chat_chunk( req, uid, - [{"delta": {"tool_calls": [delta]}, "index": 0, "finish_reason": None}], + [{"delta": {"tool_calls": [delta]}, "index": index, "finish_reason": None}], ) ) tool_calls_sent += 1 @@ -368,21 +399,18 @@ async def stream_chat_completion_chunks( prompt_tokens = ev.prompt_tokens completion_tokens = ev.completion_tokens cached_tokens = ev.cached_tokens - yield _sse(_chat_chunk(req, uid, [{"delta": {}, "index": 0, "finish_reason": ev.finish_reason}])) + reasoning_tokens = ev.reasoning_tokens + yield _sse(_chat_chunk(req, uid, [{"delta": {}, "index": index, "finish_reason": ev.finish_reason}])) + usage = _usage( + prompt_tokens, completion_tokens, _reported_cached(state, cached_tokens), + reasoning_tokens=reasoning_tokens, + ) + if not terminal: + yield _StreamUsage(usage) + return if req.stream_options and req.stream_options.include_usage: - yield _sse( - { - "id": f"chatcmpl-{uid}", - "object": "chat.completion.chunk", - "created": int(time.time()), - "model": req.model, - "choices": [], - "usage": _usage( - prompt_tokens, completion_tokens, _reported_cached(state, cached_tokens) - ), - } - ) + yield _sse({**_chat_chunk(req, uid, []), "usage": usage}) yield b"data: [DONE]\n\n" @@ -396,47 +424,63 @@ async def handle_completion( unsupported = _completion_unsupported_reason(req) if unsupported is not None: return create_error_response(unsupported) - try: # surfaces an out-of-range max_tokens as a 400 rather than a 500 from the worker + if not 1 <= req.n <= MAX_N: + return create_error_response(f"n must be between 1 and {MAX_N}", param="n") + try: # surfaces an out-of-range value as a 400 rather than a 500 from the worker default_max_tokens = ( getattr(state.config, "max_output_tokens", None) or DEFAULT_MAX_OUTPUT_TOKENS ) - _resolve_sampling(req, model_sampling, default_max_tokens=default_max_tokens) + sampling = _resolve_sampling(req, model_sampling, default_max_tokens=default_max_tokens) except ValueError as exc: - return create_error_response(str(exc), param="max_tokens") + return create_error_response(str(exc)) prompts = [req.prompt] if isinstance(req.prompt, str) else req.prompt assert isinstance(prompts, list) + if req.suffix is not None: + markers = await _fim_markers(state) + if markers is None: + return create_error_response( + "suffix needs a model with fill-in-the-middle tokens " + "(<|fim_prefix|> / <|fim_suffix|> / <|fim_middle|>)", + param="suffix", + ) + prompts = [_fim_prompt(p, req.suffix, markers) for p in prompts] + echo_texts = list(prompts) if req.echo else [""] * len(prompts) + if req.suffix is not None and req.echo: + # echo returns what the client sent, not the FIM-wrapped prompt + echo_texts = [req.prompt] if isinstance(req.prompt, str) else list(req.prompt) + + # choices are ordered prompt-major: prompt i, sample k -> index i * n + k (OpenAI) + uids: list[int] = [] + for prompt in prompts: + for _ in range(req.n): + uid = state.new_user() + await state.send_one(TokenizeMsg(uid=uid, text=prompt, sampling_params=sampling)) + uids.append(uid) + if req.stream: - if len(prompts) != 1: - return create_error_response("Streaming completions only support a single text prompt") - uid = state.new_user() - await state.send_one( - TokenizeMsg(uid=uid, text=prompts[0], sampling_params=_resolve_sampling( - req, model_sampling, default_max_tokens=default_max_tokens - )) - ) - chunks = stream_completion_chunks(uid, req, state) + if len(uids) == 1: + chunks = stream_completion_chunks(uids[0], req, state, echo_text=echo_texts[0]) + else: + chunks = _merge_streams( + [ + stream_completion_chunks( + u, req, state, index=i, terminal=False, echo_text=echo_texts[i // req.n] + ) + for i, u in enumerate(uids) + ], + lambda: _completion_chunk(req, uids[0], []), + include_usage=bool(req.stream_options and req.stream_options.include_usage), + state=state, + ) if request is not None: - chunks = state.stream_with_cancellation(chunks, request, uid) + chunks = state.stream_with_cancellation(chunks, request, uids) return StreamingResponse(chunks, media_type="text/event-stream") - choices: list[dict[str, Any]] = [] - prompt_tokens = 0 - completion_tokens = 0 - cached_tokens = 0 - for index, prompt in enumerate(prompts): - uid = state.new_user() - await state.send_one( - TokenizeMsg( - uid=uid, - text=prompt, - sampling_params=_resolve_sampling( - req, model_sampling, default_max_tokens=default_max_tokens - ), - ) - ) + async def _drain_one(uid: int, index: int) -> dict[str, Any] | JSONResponse: text = "" finish_reason = "stop" + prompt_tokens = completion_tokens = cached_tokens = 0 async for ack in state.wait_for_ack(uid): if getattr(ack, "error", None): return create_error_response(ack.error) @@ -447,77 +491,225 @@ async def handle_completion( if ack.finished: finish_reason = getattr(ack, "finish_reason", None) or "stop" break - choices.append({"index": index, "text": text, "finish_reason": finish_reason, "logprobs": None}) + return { + "index": index, + "text": echo_texts[index // req.n] + text, + "finish_reason": finish_reason, + "logprobs": None, + "_usage": (prompt_tokens, completion_tokens, cached_tokens), + } + + drained = await asyncio.gather(*(_drain_one(u, i) for i, u in enumerate(uids))) + choices: list[dict[str, Any]] = [] + prompt_tokens = completion_tokens = cached_tokens = 0 + for choice in drained: + if isinstance(choice, JSONResponse): + return choice + p, c, cached = choice.pop("_usage") + # the prompt is counted once per prompt, not once per sample of it + if choice["index"] % req.n == 0: + prompt_tokens += p + completion_tokens += c + cached_tokens = max(cached_tokens, cached) + choices.append(choice) return { - "id": f"cmpl-{uuid.uuid4().hex}", + "id": _response_id("cmpl", req, uuid.uuid4().hex), "object": "text_completion", "created": int(time.time()), "model": req.model, + "system_fingerprint": None, "choices": choices, "usage": _usage(prompt_tokens, completion_tokens, _reported_cached(state, cached_tokens)), } -async def stream_completion_chunks(uid: int, req: CompletionRequest, state: Any) -> AsyncIterator[bytes]: +def _completion_chunk(req: CompletionRequest, uid: Any, choices: list[dict[str, Any]]) -> dict[str, Any]: + return { + "id": _response_id("cmpl", req, uid), + "object": "text_completion.chunk", + "created": int(time.time()), + "model": req.model, + "choices": choices, + } + + +async def stream_completion_chunks( + uid: int, + req: CompletionRequest, + state: Any, + index: int = 0, + terminal: bool = True, + echo_text: str = "", +) -> AsyncIterator[bytes]: prompt_tokens = 0 completion_tokens = 0 cached_tokens = 0 finish_reason = "stop" + if echo_text: + # echo: the prompt leads the stream (vLLM sends it in the first chunk) + yield _sse(_completion_chunk(req, uid, [{"text": echo_text, "index": index, "finish_reason": None, "logprobs": None}])) async for ack in state.wait_for_ack(uid): if getattr(ack, "error", None): yield _sse({"error": {"message": ack.error, "type": "invalid_request_error", "code": None}}) - yield b"data: [DONE]\n\n" + if terminal: + yield b"data: [DONE]\n\n" return prompt_tokens += ack.prompt_tokens_delta completion_tokens += ack.completion_tokens_delta cached_tokens += ack.cached_tokens if ack.incremental_output: yield _sse( - { - "id": f"cmpl-{uid}", - "object": "text_completion.chunk", - "created": int(time.time()), - "model": req.model, - "choices": [ - { - "text": ack.incremental_output, - "index": 0, - "finish_reason": None, - "logprobs": None, - } - ], - } + _completion_chunk( + req, uid, + [{"text": ack.incremental_output, "index": index, "finish_reason": None, "logprobs": None}], + ) ) if ack.finished: finish_reason = getattr(ack, "finish_reason", None) or "stop" break - yield _sse( - { - "id": f"cmpl-{uid}", - "object": "text_completion.chunk", - "created": int(time.time()), - "model": req.model, - "choices": [{"text": "", "index": 0, "finish_reason": finish_reason, "logprobs": None}], - } - ) + yield _sse(_completion_chunk(req, uid, [{"text": "", "index": index, "finish_reason": finish_reason, "logprobs": None}])) + usage = _usage(prompt_tokens, completion_tokens, _reported_cached(state, cached_tokens)) + if not terminal: + yield _StreamUsage(usage) + return if req.stream_options and req.stream_options.include_usage: - yield _sse( - { - "id": f"cmpl-{uid}", - "object": "text_completion.chunk", - "created": int(time.time()), - "model": req.model, - "choices": [], - "usage": _usage( - prompt_tokens, completion_tokens, _reported_cached(state, cached_tokens) - ), - } + yield _sse({**_completion_chunk(req, uid, []), "usage": usage}) + yield b"data: [DONE]\n\n" + + +class _StreamUsage: + """Yielded last by a non-terminal sub-stream: its usage dict for the merger.""" + + __slots__ = ("usage",) + + def __init__(self, usage: dict[str, Any]) -> None: + self.usage = usage + + +async def _merge_streams( + streams: list[AsyncIterator[Any]], + envelope: Callable[[], dict[str, Any]], + *, + include_usage: bool, + state: Any, +) -> AsyncIterator[bytes]: + """Interleave the SSE chunks of several choice streams into one response (``n > 1``), + then one usage chunk (prompt counted once, completions summed) and one ``[DONE]``.""" + queue: asyncio.Queue[tuple[int, Any]] = asyncio.Queue() + usages: dict[int, dict[str, Any]] = {} + + async def pump(slot: int, stream: AsyncIterator[Any]) -> None: + try: + async for item in stream: + if isinstance(item, _StreamUsage): + usages[slot] = item.usage + else: + await queue.put((slot, item)) + finally: + await queue.put((slot, None)) + + tasks = [asyncio.create_task(pump(i, st)) for i, st in enumerate(streams)] + try: + open_slots = len(streams) + while open_slots: + slot, item = await queue.get() + if item is None: + open_slots -= 1 + continue + yield item + finally: + for task in tasks: + task.cancel() + if include_usage and usages: + first = usages[min(usages)] + merged = _usage( + first["prompt_tokens"], + sum(u["completion_tokens"] for u in usages.values()), + max(u.get("prompt_tokens_details", {}).get("cached_tokens", 0) for u in usages.values()), + reasoning_tokens=sum( + u.get("completion_tokens_details", {}).get("reasoning_tokens", 0) for u in usages.values() + ), ) + yield _sse({**envelope(), "usage": merged}) yield b"data: [DONE]\n\n" +_FIM_TOKENS = ("<|fim_prefix|>", "<|fim_suffix|>", "<|fim_middle|>") + + +async def _fim_markers(state: Any) -> tuple[str, str, str] | None: + """The model's fill-in-the-middle marker strings when its vocabulary has them + (Qwen / DeepSeek-Coder style), else None.""" + build = getattr(state, "frontend_tokenizer", None) + if build is None: + return None + try: + manager = await asyncio.to_thread(build) + tok = manager.tokenizer + unk = getattr(tok, "unk_token_id", None) + for name in _FIM_TOKENS: + tid = tok.convert_tokens_to_ids(name) + if not isinstance(tid, int) or tid < 0 or tid == unk: + return None + except Exception: # noqa: BLE001 -- no tokenizer here means no FIM + return None + return _FIM_TOKENS + + +def _fim_prompt(prompt: str, suffix: str, markers: tuple[str, str, str]) -> str: + """The prefix-suffix-middle prompt of OpenAI's ``suffix`` (llama.cpp's /infill).""" + prefix_tok, suffix_tok, middle_tok = markers + return f"{prefix_tok}{prompt}{suffix_tok}{suffix}{middle_tok}" + + +async def handle_tokenize(req: TokenizeRequest, state: Any): + """``POST /tokenize``: the vLLM / SGLang shape. A raw ``prompt`` is encoded as is; chat + ``messages`` are rendered through the model's chat template first (the exact prompt a + generation would tokenize).""" + if (req.prompt is None) == (req.messages is None): + return create_error_response("pass exactly one of prompt or messages") + build = getattr(state, "frontend_tokenizer", None) + if build is None: + return create_error_response("no tokenizer on this server", status_code=503, err_type="api_error") + try: + manager = await asyncio.to_thread(build) + if req.prompt is not None: + ids = await asyncio.to_thread( + manager.tokenizer.encode, req.prompt, add_special_tokens=req.add_special_tokens + ) + else: + messages = render_messages([m.model_dump(exclude_none=True) for m in req.messages]) + ctk = dict(req.chat_template_kwargs) + if not req.add_generation_prompt: + ctk["continue_final_message"] = True + msg = TokenizeMsg(uid=-1, text=messages, sampling_params=SamplingParams(), chat_template_kwargs=ctk) + text = await asyncio.to_thread(manager.render_prompt, msg) + ids = await asyncio.to_thread(manager.tokenizer.encode, text, add_special_tokens=False) + except ValueError as exc: + return create_error_response(str(exc)) + body: dict[str, Any] = { + "count": len(ids), + "max_model_len": _model_context_length(state), + "tokens": [int(t) for t in ids], + } + if req.return_token_strs: + body["token_strs"] = manager.tokenizer.convert_ids_to_tokens(ids) + return body + + +async def handle_detokenize(req: DetokenizeRequest, state: Any): + build = getattr(state, "frontend_tokenizer", None) + if build is None: + return create_error_response("no tokenizer on this server", status_code=503, err_type="api_error") + manager = await asyncio.to_thread(build) + text = await asyncio.to_thread( + manager.tokenizer.decode, req.tokens, skip_special_tokens=req.skip_special_tokens + ) + return {"prompt": text} + + def create_error_response( message: str, status_code: int = 400, @@ -552,11 +744,25 @@ def _resolve_sampling( model_sampling=model_sampling, stop=req.stop, default_max_tokens=default_max_tokens, + min_p=req.min_p, presence_penalty=req.presence_penalty, frequency_penalty=req.frequency_penalty, + repetition_penalty=req.repetition_penalty, + logit_bias=req.logit_bias, + min_tokens=req.min_tokens, + stop_token_ids=req.stop_token_ids, + include_stop_str_in_output=req.include_stop_str_in_output, + skip_special_tokens=req.skip_special_tokens, ) +def _response_id(prefix: str, req: Any, fallback: Any) -> str: + """``request_id`` from the client when given (vLLM request_id / SGLang rid), else the + server's own id.""" + rid = getattr(req, "request_id", None) + return f"{prefix}-{rid}" if rid else f"{prefix}-{fallback}" + + def _tools_for_template(req: ChatCompletionRequest) -> list[dict[str, Any]] | None: if not _should_parse_tools(req): return None @@ -617,10 +823,11 @@ def _tool_call_id(name: str | None, index: int) -> str: def _chat_chunk(req: ChatCompletionRequest, uid: int, choices: list[dict[str, Any]]) -> dict[str, Any]: return { - "id": f"chatcmpl-{uid}", + "id": _response_id("chatcmpl", req, uid), "object": "chat.completion.chunk", "created": int(time.time()), "model": req.model, + "system_fingerprint": None, "choices": choices, } @@ -634,7 +841,9 @@ def _reported_cached(state: Any, cached_tokens: int) -> int: return cached_tokens if getattr(state.config, "enable_cache_report", False) else 0 -def _usage(prompt_tokens: int, completion_tokens: int, cached_tokens: int = 0) -> dict[str, Any]: +def _usage( + prompt_tokens: int, completion_tokens: int, cached_tokens: int = 0, reasoning_tokens: int = 0 +) -> dict[str, Any]: usage: dict[str, Any] = { "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, @@ -644,6 +853,10 @@ def _usage(prompt_tokens: int, completion_tokens: int, cached_tokens: int = 0) - # disabled report and a 0-token hit serialize identically. if cached_tokens > 0: usage["prompt_tokens_details"] = {"cached_tokens": cached_tokens} + # OpenAI's completion_tokens_details.reasoning_tokens: tokens up to and including the + # reasoning end tag, only when the model emitted one. + if reasoning_tokens > 0: + usage["completion_tokens_details"] = {"reasoning_tokens": reasoning_tokens} return usage @@ -657,12 +870,6 @@ def _completion_unsupported_reason(req: CompletionRequest) -> str | None: return "OpenAI token-id prompt inputs are not supported; pass text prompt strings instead" if req.logprobs is not None: return "logprobs is not supported" - if req.echo: - return "echo is not supported" - if req.suffix is not None: - return "suffix is not supported" - if req.logit_bias is not None: - return "logit_bias is not supported" if _response_format_unsupported(req.response_format): return "response_format json_object/json_schema is not supported (no constrained decoding)" return None diff --git a/python/freetoken/server/stats.py b/python/freetoken/server/stats.py index 3d48ee96a0..c9d0212c32 100644 --- a/python/freetoken/server/stats.py +++ b/python/freetoken/server/stats.py @@ -174,3 +174,44 @@ def build_stats(state: Any, p95_ms: int, ttft_mean_ms: int) -> dict: "completion_tokens_total": tr.completion_tokens_total, }, } + + + +def prometheus_text(doc: dict) -> str: + """The /v1/stats document as Prometheus text exposition (gauges and counters).""" + lines: list[str] = [] + + def gauge(name: str, value, help_text: str, kind: str = "gauge") -> None: + if value is None: + return + lines.append(f"# HELP {name} {help_text}") + lines.append(f"# TYPE {name} {kind}") + lines.append(f"{name} {value}") + + model = doc.get("model") or {} + model_id = str(model.get("id") or model.get("name") or "").replace('"', "'") + lines.append("# HELP freetoken_info Server identity.") + lines.append("# TYPE freetoken_info gauge") + lines.append(f'freetoken_info{{instance_id="{doc.get("instance_id")}",model="{model_id}"}} 1') + gauge("freetoken_uptime_seconds", doc.get("uptime_s"), "Seconds since the server became ready.") + req = doc.get("requests") or {} + gauge("freetoken_requests_active", req.get("active"), "Requests in flight.") + gauge("freetoken_requests_completed_total", req.get("completed"), "Requests completed.", "counter") + gauge("freetoken_request_latency_p95_milliseconds", req.get("p95_ms"), "p95 request latency over the recent window.") + gauge("freetoken_ttft_mean_milliseconds", req.get("ttft_mean_ms"), "Mean time to first token over the recent window.") + gauge("freetoken_prompt_tokens_total", req.get("prompt_tokens_total"), "Prompt tokens processed.", "counter") + gauge("freetoken_completion_tokens_total", req.get("completion_tokens_total"), "Tokens generated.", "counter") + tp = doc.get("throughput") or {} + gauge("freetoken_decode_tokens_per_second", tp.get("decode_tps"), "Decode throughput, sliding window.") + gauge("freetoken_prefill_tokens_per_second", tp.get("prefill_tps"), "Prefill throughput, sliding window.") + for pool in ("kv", "swa"): + p = doc.get(pool) + if p: + gauge(f"freetoken_{pool}_pages_used", p.get("used_pages"), f"{pool} pages in use.") + gauge(f"freetoken_{pool}_pages_total", p.get("total_pages"), f"{pool} pages in the pool.") + mamba = doc.get("mamba") + if mamba: + gauge("freetoken_mamba_slots_used", mamba.get("used_slots"), "GDN state slots in use.") + gauge("freetoken_mamba_slots_total", mamba.get("total_slots"), "GDN state slots in the pool.") + gauge("freetoken_vram_bytes", doc.get("vram_bytes"), "GPU memory in use by the engine.") + return "\n".join(lines) + "\n" diff --git a/python/freetoken/tokenizer/detokenize.py b/python/freetoken/tokenizer/detokenize.py index d138ac1a68..27dd3b8034 100644 --- a/python/freetoken/tokenizer/detokenize.py +++ b/python/freetoken/tokenizer/detokenize.py @@ -71,11 +71,16 @@ class DecodeStatus: read_offset: int # length of read ids surr_offset: int # length of surr ids sent_offset: int # length of sent out string + skip_special: bool = False # decode this request with skip_special_tokens + think_end_at: int | None = None # generated tokens up to and incl. the reasoning end tag class DetokenizeManager: def __init__( - self, tokenizer: PreTrainedTokenizerBase, eos_token_ids: FrozenSet[int] | None = None + self, + tokenizer: PreTrainedTokenizerBase, + eos_token_ids: FrozenSet[int] | None = None, + think_end_id: int | None = None, ) -> None: # uid -> DecodeStatus self.decode_map: Dict[int, DecodeStatus] = {} @@ -83,6 +88,9 @@ def __init__( self.eos_token_ids = ( eos_token_ids if eos_token_ids is not None else frozenset({tokenizer.eos_token_id}) ) + # The reasoning end tag's id (e.g. ````): tokens generated up to and + # including it are reported as usage.completion_tokens_details.reasoning_tokens. + self.think_end_id = think_end_id if think_end_id is None else int(think_end_id) def discard(self, uid: int) -> None: """Drop a uid's decode state without a finished DetokenizeMsg. An aborted or errored @@ -91,8 +99,26 @@ def discard(self, uid: int) -> None: self.decode_map.pop(uid, None) def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]: + return self.detokenize_with_meta(msgs)[0] + + def _batch_decode(self, ids: List[List[int]], skip: List[bool]) -> List[str]: + """batch_decode honouring a per-request skip_special_tokens flag (one call per + distinct flag value, results back in order).""" + out: List[str | None] = [None] * len(ids) + for flag in (False, True): + idx = [i for i, f in enumerate(skip) if f == flag] + if idx: + texts = self.tokenizer.batch_decode([ids[i] for i in idx], skip_special_tokens=flag) + for i, t in zip(idx, texts, strict=True): + out[i] = t + return [t if t is not None else "" for t in out] + + def detokenize_with_meta(self, msgs: List[DetokenizeMsg]) -> tuple[List[str], List[int]]: + """Incremental text per message plus, for a finished message, its reasoning token + count (0 otherwise or when the model emitted no reasoning end tag).""" read_ids: List[List[int]] = [] surr_ids: List[List[int]] = [] + skip: List[bool] = [] for msg in msgs: if msg.uid not in self.decode_map: self.decode_map[msg.uid] = DecodeStatus( @@ -101,17 +127,26 @@ def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]: read_offset=0, surr_offset=0, sent_offset=0, + skip_special=bool(getattr(msg, "skip_special_tokens", False)), ) s = self.decode_map[msg.uid] if not (msg.finished and msg.next_token in self.eos_token_ids): s.decoded_ids.append(msg.next_token) + if ( + self.think_end_id is not None + and s.think_end_at is None + and msg.next_token == self.think_end_id + ): + s.think_end_at = len(s.decoded_ids) read_ids.append(s.decoded_ids[s.surr_offset :]) surr_ids.append(s.decoded_ids[s.surr_offset : s.read_offset]) + skip.append(s.skip_special) - read_texts = self.tokenizer.batch_decode(read_ids) - surr_texts = self.tokenizer.batch_decode(surr_ids) + read_texts = self._batch_decode(read_ids, skip) + surr_texts = self._batch_decode(surr_ids, skip) incremental_strs: List[str] = [] + reasoning_tokens: List[int] = [] for msg, read_str, surr_str in zip(msgs, read_texts, surr_texts, strict=True): s = self.decode_map[msg.uid] new_text = read_str[len(surr_str) :] @@ -127,10 +162,16 @@ def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]: prev_sent = s.sent_offset if msg.finished: - # Generation is over: flush everything, trimming at the matched stop string. + # Generation is over: flush everything, trimming at the matched stop string + # (or keeping it, for include_stop_str_in_output). if msg.matched_stop: cut = output_str.find(msg.matched_stop) - emit_end = cut if cut >= 0 else len(output_str) + if cut < 0: + emit_end = len(output_str) + elif getattr(msg, "keep_stop_str", False): + emit_end = cut + len(msg.matched_stop) + else: + emit_end = cut else: emit_end = len(output_str) elif msg.stop_strs: @@ -141,7 +182,8 @@ def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]: incremental_output = output_str[prev_sent:emit_end] if emit_end > prev_sent else "" s.sent_offset = max(prev_sent, emit_end) incremental_strs.append(incremental_output) + reasoning_tokens.append((s.think_end_at or 0) if msg.finished else 0) if msg.finished: del self.decode_map[msg.uid] - return incremental_strs + return incremental_strs, reasoning_tokens diff --git a/python/freetoken/tokenizer/server.py b/python/freetoken/tokenizer/server.py index d3365bfecd..2050fea3f6 100644 --- a/python/freetoken/tokenizer/server.py +++ b/python/freetoken/tokenizer/server.py @@ -81,6 +81,19 @@ def _send_generation_replies( _put_user_replies(send_frontend, terminal_errors) +def _reasoning_end_id(tokenizer) -> int | None: + """The id of the reasoning end tag when the vocabulary has one (```` for the + Qwen / DeepSeek / GLM families), else None. Only used to count reasoning tokens.""" + for tag in ("", "<|end_of_thinking|>"): + try: + tid = tokenizer.convert_tokens_to_ids(tag) + except Exception: # noqa: BLE001 -- a tokenizer without the tag answers unk or raises + continue + if isinstance(tid, int) and tid >= 0 and tid != getattr(tokenizer, "unk_token_id", -1): + return tid + return None + + def _tokenize_requests( tokenize_manager: Any, messages: List[TokenizeMsg], @@ -151,7 +164,9 @@ def tokenize_worker( tokenize_manager = TokenizeManager(tokenizer, get_mm_processor(tokenizer_path, mm)) detokenize_manager = DetokenizeManager( - tokenizer, load_eos_token_ids(tokenizer_path, tokenizer) + tokenizer, + load_eos_token_ids(tokenizer_path, tokenizer), + think_end_id=_reasoning_end_id(tokenizer), ) if ack_queue is not None: @@ -209,7 +224,7 @@ def tokenize_worker( ) sampled_replies: List[UserReply] = [] if len(detokenize_msg) > 0: - replies = detokenize_manager.detokenize(detokenize_msg) + replies, reasoning_counts = detokenize_manager.detokenize_with_meta(detokenize_msg) sampled_replies = [ UserReply( uid=msg.uid, @@ -217,6 +232,7 @@ def tokenize_worker( finished=msg.finished, finish_reason=msg.finish_reason, matched_stop=msg.matched_stop, + reasoning_tokens=reasoning_tokens, completion_tokens_delta=1, kv_used_pages=msg.kv_used_pages, kv_total_pages=msg.kv_total_pages, @@ -226,7 +242,9 @@ def tokenize_worker( swa_total_tokens=msg.swa_total_tokens, gpu_mem_bytes=msg.gpu_mem_bytes, ) - for msg, reply in zip(detokenize_msg, replies, strict=True) + for msg, reply, reasoning_tokens in zip( + detokenize_msg, replies, reasoning_counts, strict=True + ) ] # An error reply and a client abort are both terminal for their uid, and neither diff --git a/python/freetoken/tokenizer/tokenize.py b/python/freetoken/tokenizer/tokenize.py index fc66391cf2..f512930dee 100644 --- a/python/freetoken/tokenizer/tokenize.py +++ b/python/freetoken/tokenizer/tokenize.py @@ -176,10 +176,16 @@ def _render( if tools is not None: chat_template_kwargs = {**chat_template_kwargs, "tools": tools} messages = _map_developer_role(messages, getattr(self.tokenizer, "chat_template", None)) + # continue_final_message (an OpenAI-compatible extra, as in vLLM / SGLang): the last + # message is an assistant prefix the model must continue, so no generation prompt. + chat_template_kwargs = dict(chat_template_kwargs) + continue_final = bool(chat_template_kwargs.pop("continue_final_message", False)) + if continue_final: + chat_template_kwargs["continue_final_message"] = True prompt = self.tokenizer.apply_chat_template( messages, tokenize=False, - add_generation_prompt=True, + add_generation_prompt=not continue_final, **chat_template_kwargs, ) assert isinstance(prompt, str) diff --git a/tests/engine/test_logits_processors.py b/tests/engine/test_logits_processors.py new file mode 100644 index 0000000000..f051dec51f --- /dev/null +++ b/tests/engine/test_logits_processors.py @@ -0,0 +1,155 @@ +"""The logits processors behind min_p, the penalties, logit_bias and min_tokens: pure +torch on the CPU, so the math is checked without a GPU or the sampling kernels.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import torch +from freetoken.core import SamplingParams +from freetoken.engine.sample import LogitsPlan, Sampler, apply_logits_processors + +V = 8 +CPU = torch.device("cpu") + + +def _plan(rows, **fields) -> LogitsPlan: + plan = LogitsPlan( + rows=torch.tensor(rows, dtype=torch.int64), + temps=torch.ones((len(rows), 1), dtype=torch.float32), + ) + for k, v in fields.items(): + setattr(plan, k, v) + return plan + + +def _ids(rows: list[list[int]]) -> torch.Tensor: + width = max(len(r) for r in rows) + out = torch.full((len(rows), width), V, dtype=torch.int64) + for i, r in enumerate(rows): + out[i, : len(r)] = torch.tensor(r) + return out + + +def test_presence_and_frequency_penalties_count_generated_tokens(): + logits = torch.zeros((2, V)) + plan = _plan( + [1], + hist_ids=_ids([[3, 3, 5]]), + presence=torch.tensor([[0.5]]), + frequency=torch.tensor([[0.25]]), + ) + out = apply_logits_processors(logits, plan, V) + assert torch.equal(out[0], logits[0]) # row outside the plan untouched + assert out[1, 3].item() == -(0.5 + 2 * 0.25) + assert out[1, 5].item() == -(0.5 + 0.25) + assert out[1, 0].item() == 0.0 + + +def test_repetition_penalty_covers_prompt_and_output_with_hf_semantics(): + logits = torch.tensor([[2.0, -2.0, 2.0, -2.0, 1.0, 1.0, 1.0, 1.0]]) + plan = _plan( + [0], + hist_ids=_ids([[0, 1]]), # generated: 0 (positive), 1 (negative) + prompt_ids=_ids([[2, 3]]), # prompt: 2 (positive), 3 (negative) + repetition=torch.tensor([[2.0]]), + ) + out = apply_logits_processors(logits, plan, V) + assert out[0, 0].item() == 1.0 and out[0, 2].item() == 1.0 # divided + assert out[0, 1].item() == -4.0 and out[0, 3].item() == -4.0 # multiplied + assert out[0, 4].item() == 1.0 # unseen + + +def test_logit_bias_adds_per_row(): + logits = torch.zeros((2, V)) + plan = _plan( + [0, 1], + bias_rows=torch.tensor([0, 1, 1]), + bias_ids=torch.tensor([2, 2, 6]), + bias_vals=torch.tensor([5.0, -100.0, 1.5]), + ) + out = apply_logits_processors(logits, plan, V) + assert ( + out[0, 2].item() == 5.0 + and out[1, 2].item() == -100.0 + and out[1, 6].item() == 1.5 + ) + + +def test_min_tokens_masks_the_stop_ids_of_the_row_only(): + logits = torch.zeros((2, V)) + plan = _plan([0, 1], min_rows=torch.tensor([1]), min_ids=_ids([[0, 7]])) + out = apply_logits_processors(logits, plan, V) + assert out[1, 0].item() == float("-inf") and out[1, 7].item() == float("-inf") + assert out[1, 3].item() == 0.0 + assert torch.isfinite(out[0]).all() + + +def test_min_p_drops_tokens_below_the_fraction_of_the_top_probability(): + logits = torch.log(torch.tensor([[0.5, 0.3, 0.1, 0.05, 0.03, 0.01, 0.005, 0.005]])) + plan = _plan( + [0], min_p=torch.tensor([[0.15]]) + ) # threshold 0.075, between 0.1 and 0.05 + out = apply_logits_processors(logits, plan, V) + kept = torch.isfinite(out[0]) + assert kept.tolist() == [True, True, True, False, False, False, False, False] + # a row with min_p 0 is left alone + plan0 = _plan([0], min_p=torch.tensor([[0.0]])) + assert torch.isfinite(apply_logits_processors(logits, plan0, V)[0]).all() + + +def _req( + prompt: list[int], generated: list[int], max_tokens: int, **sp +) -> SimpleNamespace: + ids = torch.tensor(prompt + generated, dtype=torch.int32) + return SimpleNamespace( + input_ids=ids, + output_len=max_tokens, + max_device_len=len(prompt) + max_tokens, + sampling_params=SamplingParams(**sp), + ) + + +def test_prepare_builds_no_plan_without_processors_and_a_plan_with_them(): + sampler = Sampler(CPU, V) + plain = SimpleNamespace(reqs=[_req([1, 2], [3], 4)]) + assert sampler.prepare(plain).plan is None + + # all greedy: the argmax path runs on the CPU, the sampling kernels need a GPU + batch = SimpleNamespace( + reqs=[ + _req([1, 2], [3], 4), + _req( + [1, 2], + [3, 3], + 4, + presence_penalty=1.0, + repetition_penalty=1.5, + logit_bias=[[5, 2.0]], + ), + _req([4], [], 4, min_tokens=2, min_tokens_stop_ids=[0, 7]), + _req( + [4], [6, 6], 4, min_tokens=2, min_tokens_stop_ids=[0] + ), # already past min_tokens + ] + ) + args = sampler.prepare(batch) + plan = args.plan + assert plan is not None + assert plan.rows.tolist() == [1, 2, 3] + assert plan.hist_ids[0].tolist()[:2] == [3, 3] # row 1's generated tokens + assert plan.prompt_ids[0].tolist()[:2] == [1, 2] + assert plan.bias_rows.tolist() == [0] and plan.bias_ids.tolist() == [5] + assert plan.min_rows.tolist() == [1] # local row of the third request only + assert plan.min_ids[0].tolist() == [0, 7] + + logits = torch.zeros((4, V)) + out = apply_logits_processors(logits, plan, V) + assert ( + out[1, 3].item() == -1.0 + ) # presence on the repeated 3 (repetition on 0 is a no-op) + assert out[1, 5].item() == 2.0 + assert out[2, 0].item() == float("-inf") and out[2, 7].item() == float("-inf") + assert torch.isfinite(out[3]).all() + # greedy path: argmax over the processed logits + assert sampler.sample(logits, args)[1].item() == 5 diff --git a/tests/server/test_openai_extras.py b/tests/server/test_openai_extras.py new file mode 100644 index 0000000000..115112b9d0 --- /dev/null +++ b/tests/server/test_openai_extras.py @@ -0,0 +1,383 @@ +"""OpenAI protocol coverage: n > 1 fan-out, echo and suffix (FIM) on completions, the +sampling extras reaching SamplingParams, request_id, usage details, and the tokenize / +detokenize routes.""" + +from __future__ import annotations + +import asyncio +import json +from types import SimpleNamespace + +import pytest +from freetoken.message import TokenizeMsg, UserReply +from freetoken.server.api_models import ( + ChatCompletionRequest, + CompletionRequest, + DetokenizeRequest, + TokenizeRequest, +) +from freetoken.server.openai_api import ( + _usage, + chat_request_to_genspec, + handle_chat_completion, + handle_completion, + handle_detokenize, + handle_tokenize, +) +from freetoken.server.stats import prometheus_text + + +def run(coro): + return asyncio.run(coro) + + +class FakeState: + """Distinct uids per generation; every uid gets the same two-ack reply.""" + + def __init__(self, text=("hel", "lo"), reasoning_tokens=0, tokenizer=None): + self.config = SimpleNamespace( + model_path="/models/unit-model", + served_model_name="unit-model", + tool_call_parser="llama3", + reasoning_parser=None, + max_seq_len=4096, + ) + self.text = text + self.reasoning_tokens = reasoning_tokens + self._next = 100 + self.sent: list[TokenizeMsg] = [] + self.aborted: list[int] = [] + if tokenizer is not None: + manager = SimpleNamespace( + tokenizer=tokenizer, + render_prompt=lambda msg: "RENDERED:" + json.dumps(msg.text), + ) + self.frontend_tokenizer = lambda: manager + + def new_user(self) -> int: + self._next += 1 + return self._next + + async def send_one(self, msg): + self.sent.append(msg) + + async def abort_user(self, uid): + self.aborted.append(uid) + + async def wait_for_ack(self, uid: int): + yield UserReply( + uid=uid, incremental_output="", finished=False, prompt_tokens_delta=5 + ) + for i, piece in enumerate(self.text): + last = i == len(self.text) - 1 + yield UserReply( + uid=uid, + incremental_output=piece, + finished=last, + completion_tokens_delta=1, + finish_reason="stop" if last else None, + reasoning_tokens=self.reasoning_tokens if last else 0, + ) + + +def chat(**kwargs) -> ChatCompletionRequest: + body = {"model": "unit-model", "messages": [{"role": "user", "content": "hi"}]} + body.update(kwargs) + return ChatCompletionRequest(**body) + + +def parse_sse(chunks: list[bytes]) -> list: + out = [] + for chunk in chunks: + for line in chunk.decode().split("\n"): + if not line.startswith("data: "): + continue + payload = line[len("data: ") :] + out.append("[DONE]" if payload == "[DONE]" else json.loads(payload)) + return out + + +async def _collect(generator): + return [item async for item in generator] + + +# ---- n > 1 ------------------------------------------------------------------------------ + + +def test_chat_n_returns_one_choice_per_generation_and_counts_the_prompt_once(): + state = FakeState(reasoning_tokens=2) + body = run(handle_chat_completion(chat(n=2), None, state, {})) + assert [c["index"] for c in body["choices"]] == [0, 1] + assert all(c["message"]["content"] == "hello" for c in body["choices"]) + assert len(state.sent) == 2 and len({m.uid for m in state.sent}) == 2 + usage = body["usage"] + assert ( + usage["prompt_tokens"] == 5 + and usage["completion_tokens"] == 4 + and usage["total_tokens"] == 9 + ) + assert usage["completion_tokens_details"] == {"reasoning_tokens": 4} + assert body["system_fingerprint"] is None + + +def test_chat_n_stream_interleaves_choice_indices_and_ends_once(): + state = FakeState() + response = run( + handle_chat_completion( + chat(n=2, stream=True, stream_options={"include_usage": True}), + None, + state, + {}, + ) + ) + events = parse_sse(run(_collect(response.body_iterator))) + assert events.count("[DONE]") == 1 and events[-1] == "[DONE]" + indices = {c["index"] for ev in events[:-1] for c in ev.get("choices", [])} + assert indices == {0, 1} + usage_chunks = [ev for ev in events[:-1] if ev.get("usage")] + assert len(usage_chunks) == 1 + assert usage_chunks[0]["usage"]["prompt_tokens"] == 5 + assert usage_chunks[0]["usage"]["completion_tokens"] == 4 + finals = [ + c for ev in events[:-1] for c in ev.get("choices", []) if c.get("finish_reason") + ] + assert len(finals) == 2 + + +def test_n_out_of_range_is_400(): + for n in (0, 17): + body = run(handle_chat_completion(chat(n=n), None, FakeState(), {})) + assert body.status_code == 400 + assert json.loads(body.body)["error"]["param"] == "n" + + +# ---- completions: echo, suffix, n --------------------------------------------------------- + + +def test_completion_echo_prepends_the_prompt(): + state = FakeState() + req = CompletionRequest(model="unit-model", prompt="Once upon", echo=True) + body = run(handle_completion(req, None, state, {})) + assert body["choices"][0]["text"] == "Once uponhello" + + +def test_completion_echo_streams_the_prompt_first(): + state = FakeState() + req = CompletionRequest( + model="unit-model", prompt="Once upon", echo=True, stream=True + ) + response = run(handle_completion(req, None, state, {})) + events = parse_sse(run(_collect(response.body_iterator))) + assert events[0]["choices"][0]["text"] == "Once upon" + assert events[-1] == "[DONE]" + + +class _FimTokenizer: + unk_token_id = -1 + + def convert_tokens_to_ids(self, name): + return {"<|fim_prefix|>": 10, "<|fim_suffix|>": 11, "<|fim_middle|>": 12}.get( + name, -1 + ) + + +class _NoFimTokenizer: + unk_token_id = 3 + + def convert_tokens_to_ids(self, name): + return 3 + + +def test_completion_suffix_builds_a_fill_in_the_middle_prompt(): + state = FakeState(tokenizer=_FimTokenizer()) + req = CompletionRequest(model="unit-model", prompt="def f(", suffix=" return 1") + run(handle_completion(req, None, state, {})) + assert ( + state.sent[0].text + == "<|fim_prefix|>def f(<|fim_suffix|> return 1<|fim_middle|>" + ) + + +def test_completion_suffix_needs_fim_tokens(): + state = FakeState(tokenizer=_NoFimTokenizer()) + req = CompletionRequest(model="unit-model", prompt="def f(", suffix="x") + body = run(handle_completion(req, None, state, {})) + assert ( + body.status_code == 400 and json.loads(body.body)["error"]["param"] == "suffix" + ) + + +def test_completion_n_orders_choices_prompt_major(): + state = FakeState() + req = CompletionRequest(model="unit-model", prompt=["a", "b"], n=2) + body = run(handle_completion(req, None, state, {})) + assert [c["index"] for c in body["choices"]] == [0, 1, 2, 3] + assert [m.text for m in state.sent] == ["a", "a", "b", "b"] + assert body["usage"]["prompt_tokens"] == 10 # two prompts, each counted once + assert body["usage"]["completion_tokens"] == 8 + + +# ---- sampling extras -------------------------------------------------------------------- + + +def test_sampling_extras_reach_sampling_params(): + spec = chat_request_to_genspec( + chat( + min_p=0.1, + repetition_penalty=1.2, + presence_penalty=0.5, + frequency_penalty=0.3, + logit_bias={"5": 150, "7": -3}, + min_tokens=3, + stop_token_ids=[9], + include_stop_str_in_output=True, + skip_special_tokens=True, + seed=7, + user="u1", + ), + {}, + ) + sp = spec.sampling_params + assert sp.min_p == 0.1 and sp.repetition_penalty == 1.2 + assert sp.presence_penalty == 0.5 and sp.frequency_penalty == 0.3 + assert sp.logit_bias == [[5, 100.0], [7, -3.0]] # clamped to the OpenAI range + assert sp.min_tokens == 3 and sp.stop_token_ids == [9] + assert sp.include_stop_str_in_output is True and sp.skip_special_tokens is True + assert sp.needs_logits_processing + + +@pytest.mark.parametrize( + "field", + [ + {"min_p": 1.5}, + {"presence_penalty": 3.0}, + {"frequency_penalty": -2.5}, + {"repetition_penalty": 0.0}, + {"logit_bias": {"x": 1.0}}, + {"min_tokens": -1}, + {"stop_token_ids": [-1]}, + ], +) +def test_bad_sampling_extras_are_400(field): + body = run(handle_chat_completion(chat(**field), None, FakeState(), {})) + assert body.status_code == 400, field + + +def test_continue_final_message_needs_an_assistant_tail_and_reaches_the_template(): + spec = chat_request_to_genspec( + chat( + messages=[ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "Sure,"}, + ], + continue_final_message=True, + ), + {}, + ) + assert spec.chat_template_kwargs["continue_final_message"] is True + with pytest.raises(ValueError): + chat(continue_final_message=True) + + +# ---- ids and usage ------------------------------------------------------------------------ + + +def test_request_id_becomes_the_response_id(): + state = FakeState() + body = run(handle_chat_completion(chat(request_id="abc-1"), None, state, {})) + assert body["id"] == "chatcmpl-abc-1" + response = run( + handle_chat_completion(chat(request_id="abc-2", stream=True), None, state, {}) + ) + events = parse_sse(run(_collect(response.body_iterator))) + assert events[0]["id"] == "chatcmpl-abc-2" + + +def test_usage_details_appear_only_when_nonzero(): + assert "completion_tokens_details" not in _usage(10, 5) + assert _usage(10, 5, reasoning_tokens=3)["completion_tokens_details"] == { + "reasoning_tokens": 3 + } + assert _usage(10, 5, cached_tokens=4)["prompt_tokens_details"] == { + "cached_tokens": 4 + } + + +# ---- tokenize / detokenize / metrics ----------------------------------------------------- + + +class _Tokenizer: + def encode(self, text, add_special_tokens=True): + ids = [ord(c) for c in text] + return ([1] + ids) if add_special_tokens else ids + + def decode(self, ids, skip_special_tokens=False): + return "".join(chr(i) for i in ids if not (skip_special_tokens and i == 1)) + + def convert_ids_to_tokens(self, ids): + return [chr(i) for i in ids] + + +def test_tokenize_prompt_and_messages_and_detokenize(): + state = FakeState(tokenizer=_Tokenizer()) + body = run(handle_tokenize(TokenizeRequest(prompt="ab"), state)) + assert body == {"count": 3, "max_model_len": 4096, "tokens": [1, 97, 98]} + body = run( + handle_tokenize( + TokenizeRequest( + prompt="ab", add_special_tokens=False, return_token_strs=True + ), + state, + ) + ) + assert body["tokens"] == [97, 98] and body["token_strs"] == ["a", "b"] + body = run( + handle_tokenize( + TokenizeRequest(messages=[{"role": "user", "content": "x"}]), state + ) + ) + assert body["count"] == len( + "RENDERED:" + json.dumps([{"role": "user", "content": "x"}]) + ) + bad = run( + handle_tokenize( + TokenizeRequest(prompt="a", messages=[{"role": "user", "content": "x"}]), + state, + ) + ) + assert bad.status_code == 400 + body = run( + handle_detokenize( + DetokenizeRequest(tokens=[1, 104, 105], skip_special_tokens=True), state + ) + ) + assert body == {"prompt": "hi"} + + +def test_metrics_exposition_lists_the_stats_document(): + doc = { + "instance_id": "abc", + "model": {"id": "unit-model"}, + "uptime_s": 12, + "kv": {"used_pages": 3, "total_pages": 10, "page_size": 64}, + "mamba": {"used_slots": 1, "total_slots": 4}, + "swa": None, + "vram_bytes": 123, + "throughput": {"decode_tps": 99.5, "prefill_tps": 1000.0}, + "requests": { + "active": 1, + "completed": 7, + "p95_ms": 50, + "ttft_mean_ms": 20, + "prompt_tokens_total": 100, + "completion_tokens_total": 40, + }, + } + text = prometheus_text(doc) + assert 'freetoken_info{instance_id="abc",model="unit-model"} 1' in text + assert "freetoken_requests_active 1" in text + assert "freetoken_requests_completed_total 7" in text + assert "freetoken_kv_pages_total 10" in text + assert "freetoken_mamba_slots_used 1" in text + assert "freetoken_decode_tokens_per_second 99.5" in text + assert "# TYPE freetoken_prompt_tokens_total counter" in text diff --git a/tests/tokenizer/test_detokenize_extras.py b/tests/tokenizer/test_detokenize_extras.py new file mode 100644 index 0000000000..5617a10537 --- /dev/null +++ b/tests/tokenizer/test_detokenize_extras.py @@ -0,0 +1,81 @@ +"""Per-request detokenizer options: include_stop_str_in_output, skip_special_tokens, and +the reasoning token count reported on the finished reply.""" + +from __future__ import annotations + +from freetoken.message import DetokenizeMsg +from freetoken.tokenizer.detokenize import DetokenizeManager + +EOS, SPECIAL, THINK_END = 0, 99, 7 + + +class _Tok: + """Token ids decode to letters; 99 is a special marker, 7 the reasoning end tag.""" + + eos_token_id = EOS + unk_token_id = -1 + + def decode(self, ids, skip_special_tokens=False): + out = [] + for t in ids: + if t == SPECIAL: + if not skip_special_tokens: + out.append("") + elif t == THINK_END: + out.append("") + else: + out.append(chr(ord("a") + t)) + return "".join(out) + + def batch_decode(self, ids_list, skip_special_tokens=False): + return [self.decode(ids, skip_special_tokens) for ids in ids_list] + + +def _run(manager, uid, tokens, *, finished_kwargs=None, **msg_kwargs): + text = "" + meta = 0 + for i, t in enumerate(tokens): + last = i == len(tokens) - 1 + kwargs = dict(msg_kwargs) + if last and finished_kwargs: + kwargs.update(finished_kwargs) + strs, counts = manager.detokenize_with_meta( + [DetokenizeMsg(uid=uid, next_token=t, finished=last, **kwargs)] + ) + text += strs[0] + meta = counts[0] + return text, meta + + +def test_stop_string_is_trimmed_by_default_and_kept_on_request(): + manager = DetokenizeManager(_Tok(), frozenset({EOS})) + trimmed, _ = _run( + manager, 1, [1, 2, 3], stop_strs=["cd"], finished_kwargs={"matched_stop": "cd"} + ) + assert trimmed == "b" + kept, _ = _run( + manager, + 2, + [1, 2, 3], + stop_strs=["cd"], + keep_stop_str=True, + finished_kwargs={"matched_stop": "cd"}, + ) + assert kept == "bcd" + + +def test_skip_special_tokens_is_per_request(): + manager = DetokenizeManager(_Tok(), frozenset({EOS})) + shown, _ = _run(manager, 1, [1, SPECIAL, 2, EOS]) + hidden, _ = _run(manager, 2, [1, SPECIAL, 2, EOS], skip_special_tokens=True) + assert shown == "bc" + assert hidden == "bc" + + +def test_reasoning_tokens_count_up_to_the_end_tag(): + manager = DetokenizeManager(_Tok(), frozenset({EOS}), think_end_id=THINK_END) + text, reasoning = _run(manager, 1, [1, 2, THINK_END, 3, 4, EOS]) + assert text == "bcde" + assert reasoning == 3 # b, c and the tag itself + _, none = _run(manager, 2, [1, 2, EOS]) + assert none == 0 From 9042a72b83ed41ef3394fe5a75d5b2ef22814bdc Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 17 Sep 2026 22:53:22 -0400 Subject: [PATCH 16/94] fix(models): skip the ftw layout check for a bare ftw without an hf config Fork #19 asks the model family to check the stored fp8 split against the env flags; a synthetic FTW written straight from tensors (upstream's test_ftw_weights) carries no config.json, so there is no family to ask and the read proceeds as before. Assisted-by: Claude Fable 5.1 (cherry picked from commit cab9be35b8e56fa3cf08a109017183691e2e69c5) (cherry picked from commit e59e5441e44dac96d530f4e2cac34d2c6273411f) --- python/freetoken/models/weight.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/python/freetoken/models/weight.py b/python/freetoken/models/weight.py index 9016b9477e..29cce4eefa 100644 --- a/python/freetoken/models/weight.py +++ b/python/freetoken/models/weight.py @@ -224,9 +224,13 @@ def load_weight( # a text-only engine never built the tower, so its tensors are not even read keep = None if include_vision else (lambda name: not name.startswith(VISION_KEY_PREFIXES)) if is_ftw_checkpoint(model_path): - spec = get_model_spec(cached_load_hf_config(model_path).architectures[0]) - # a family whose stored tensors depend on load-time settings checks them against the index before the read - check = _model_override(spec, "check_ftw_weights") + # a family whose stored tensors depend on load-time settings checks them against the index before the + # read; a bare FTW without a readable HF config (synthetic test checkpoints) has no family to ask + try: + spec = get_model_spec(cached_load_hf_config(model_path).architectures[0]) + except (OSError, ValueError): + spec = None + check = _model_override(spec, "check_ftw_weights") if spec is not None else None if check is not None: reader = FTWReader(model_path) names = {e["name"] for e in reader.entries("weight")} From a2d3bbb05b7164acaefef103f8c01a8529a62ef6 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 17 Sep 2026 22:53:25 -0400 Subject: [PATCH 17/94] test(scheduler): give the rejection test's engine fake the pool size the admission clamp reads Assisted-by: Claude Fable 5.1 (cherry picked from commit 4584957f3d63449a153097f380ffdac035a245dd) (cherry picked from commit 0ab1764343228845db270505185dd2be7090086d) --- tests/scheduler/test_cost_accounting_core.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/scheduler/test_cost_accounting_core.py b/tests/scheduler/test_cost_accounting_core.py index 045010c6b4..116d4b88c7 100644 --- a/tests/scheduler/test_cost_accounting_core.py +++ b/tests/scheduler/test_cost_accounting_core.py @@ -132,7 +132,8 @@ def fail_prepare(_batch): def test_scheduler_rejection_emits_error_but_no_admission(): scheduler = Scheduler.__new__(Scheduler) - scheduler.engine = SimpleNamespace(max_seq_len=4) + scheduler.engine = SimpleNamespace(max_seq_len=4, num_pages=4) + scheduler.config = SimpleNamespace(page_size=1) added = [] scheduler.prefill_manager = SimpleNamespace(add_one_req=added.append) sent = [] From 17166c124e66c71d2e1b9f03da8dbc7a0c271c05 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 17 Sep 2026 22:41:19 -0400 Subject: [PATCH 18/94] test(server): abort every sample of an n>1 non-streaming request on disconnect Guards the #222 resolution against #393's fan-out: the disconnect watcher takes the whole uid list, not the first sample's uid. Assisted-by: Claude Fable 5.1 (cherry picked from commit 93ff96349235ea3ca1504733a2abbd6e6bcccf94) (cherry picked from commit 69a5efa3300e93bd331b4394e54eea8cef82ad3f) (cherry picked from commit 40e994121fb79a21ba0d7ffff70e4394a4498899) --- tests/server/test_stream_cancellation.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/tests/server/test_stream_cancellation.py b/tests/server/test_stream_cancellation.py index e38fa56bfb..ba7da7c17b 100644 --- a/tests/server/test_stream_cancellation.py +++ b/tests/server/test_stream_cancellation.py @@ -294,3 +294,25 @@ def test_non_stream_connected_client_gets_result_without_abort(monkeypatch): assert result["choices"][0]["message"]["content"] == "Hi" assert state.aborted == [] + + +def test_chat_non_stream_disconnect_aborts_every_sample(monkeypatch): + # n > 1 fans out into one engine request per sample; a disconnect must stop all of them + monkeypatch.setattr(openai_api, "_DISCONNECT_POLL_SECONDS", 0.01) + + class _State(_ApiState): + def __init__(self): + super().__init__(acks=None) + self._next = 7 + + def new_user(self): + self._next += 1 + return self._next + + state = _State() + resp = asyncio.run( + openai_api.handle_chat_completion(_chat_req(n=2), _Request(disconnected=True), state, {}) + ) + + assert resp.status_code == 499 + assert sorted(state.aborted) == [8, 9] From 455ad07d9a40f579acd9b7e70a9aaa52844eaf95 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 17 Sep 2026 23:04:45 -0400 Subject: [PATCH 19/94] fix(server): record requests with a pure asgi middleware so handlers can see a client disconnect Starlette's @app.middleware("http") (BaseHTTPMiddleware) wraps the handler's receive channel and never yields the client's http.disconnect while a non-streaming request is still computing, so request.is_disconnected() stays False and the disconnect watcher from FlashML-org/FreeToken#222 could not stop an abandoned request. Streaming requests only worked because the socket write failed. The recorder is the one such middleware; as pure ASGI it records the same row at response start. Assisted-by: Claude Fable 5.1 (cherry picked from commit 8adde913c94e2f76b3962b8d10f10ded4b1cccc6) (cherry picked from commit 202d7fec30b29a9a4320370b0c62426422422339) --- python/freetoken/server/api_server.py | 79 ++++++++++++------- .../test_request_recorder_disconnect.py | 75 ++++++++++++++++++ 2 files changed, 124 insertions(+), 30 deletions(-) create mode 100644 tests/server/test_request_recorder_disconnect.py diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index 9d53df1bf0..0184e51230 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -506,38 +506,57 @@ def _served_model_name() -> str | None: return getattr(cfg, "served_model_name", None) -@app.middleware("http") -async def _record_request_middleware(request: Request, call_next): +class _RecordRequestMiddleware: """Time every generation request into the ring for /v1/requests + /v1/stats p95. Single- model server, so model = served_model_name; stream is inferred from the response media - type. Token counts are P3 (SSE usage arrives after the handler returns) — kept as None.""" - path = request.url.path - if path.startswith(_UNTRACKED_REQUEST_PREFIXES) or not path.startswith( - _TRACKED_REQUEST_PREFIXES - ): - return await call_next(request) - import time as _time - - start = _time.monotonic() - response = await call_next(request) - duration_ms = int((_time.monotonic() - start) * 1000) - ctype = response.headers.get("content-type", "") - request_ring.record_request( - request_ring.RequestRecord( - ts=_time.strftime("%Y-%m-%dT%H:%M:%SZ", _time.gmtime()), - method=request.method, - path=path, - status=response.status_code, - model=_served_model_name(), - duration_ms=duration_ms, - ttft_ms=None, - prompt_tokens=None, - completion_tokens=None, - stream=ctype.startswith("text/event-stream"), - error=None, - ) - ) - return response + type. Token counts are P3 (SSE usage arrives after the handler returns) — kept as None. + + Pure ASGI on purpose: Starlette's ``@app.middleware("http")`` (BaseHTTPMiddleware) hands the + handler a wrapped receive channel that never yields the client's ``http.disconnect`` while a + non-streaming request is still computing, so ``request.is_disconnected()`` stays False and an + abandoned request runs to max_tokens. The recorder is the only such middleware here.""" + + def __init__(self, app): + self.app = app + + async def __call__(self, scope, receive, send): + path = scope.get("path", "") if scope["type"] == "http" else "" + if not path.startswith(_TRACKED_REQUEST_PREFIXES) or path.startswith(_UNTRACKED_REQUEST_PREFIXES): + await self.app(scope, receive, send) + return + import time as _time + + start = _time.monotonic() + + async def send_recording(message): + # the base-class version recorded once the handler returned its response object, + # i.e. at response start, before a streaming body -- keep that timing + if message["type"] == "http.response.start": + ctype = "" + for name, value in message.get("headers", ()): + if name.lower() == b"content-type": + ctype = value.decode("latin-1") + request_ring.record_request( + request_ring.RequestRecord( + ts=_time.strftime("%Y-%m-%dT%H:%M:%SZ", _time.gmtime()), + method=scope["method"], + path=path, + status=message["status"], + model=_served_model_name(), + duration_ms=int((_time.monotonic() - start) * 1000), + ttft_ms=None, + prompt_tokens=None, + completion_tokens=None, + stream=ctype.startswith("text/event-stream"), + error=None, + ) + ) + await send(message) + + await self.app(scope, receive, send_recording) + + +app.add_middleware(_RecordRequestMiddleware) class CacheRebuildRequest(BaseModel): diff --git a/tests/server/test_request_recorder_disconnect.py b/tests/server/test_request_recorder_disconnect.py new file mode 100644 index 0000000000..06ff17f99b --- /dev/null +++ b/tests/server/test_request_recorder_disconnect.py @@ -0,0 +1,75 @@ +"""The request recorder must not hide a client disconnect from a non-streaming handler. + +Real sockets on purpose: the failure only shows through uvicorn's receive channel (a +BaseHTTPMiddleware wrapper never yields ``http.disconnect`` while the handler computes), which +TestClient does not model. +""" + +from __future__ import annotations + +import asyncio +import socket +import threading +import time + +from fastapi import FastAPI, Request + +from freetoken.server.api_server import _RecordRequestMiddleware + + +def _free_port() -> int: + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def _serve(app: FastAPI, port: int): + import uvicorn + + server = uvicorn.Server(uvicorn.Config(app, host="127.0.0.1", port=port, log_level="error")) + thread = threading.Thread(target=server.run, daemon=True) + thread.start() + deadline = time.monotonic() + 10 + while not server.started: + assert time.monotonic() < deadline, "uvicorn did not start" + time.sleep(0.02) + return server + + +def _post_and_hang_up(port: int) -> None: + body = b'{"prompt": "x"}' + with socket.create_connection(("127.0.0.1", port)) as sock: + sock.sendall( + b"POST /v1/completions HTTP/1.1\r\nHost: t\r\nContent-Type: application/json\r\n" + b"Content-Length: " + str(len(body)).encode() + b"\r\n\r\n" + body + ) + time.sleep(0.2) # the handler is now polling + + +def test_non_streaming_handler_sees_the_client_leave(): + app = FastAPI() + app.add_middleware(_RecordRequestMiddleware) + seen: dict[str, float | None] = {} + + @app.post("/v1/completions") + async def slow(request: Request): + t0 = time.monotonic() + for _ in range(60): + await asyncio.sleep(0.05) + if await request.is_disconnected(): + seen["after"] = time.monotonic() - t0 + return {"detected": True} + seen["after"] = None + return {"detected": False} + + port = _free_port() + server = _serve(app, port) + try: + _post_and_hang_up(port) + deadline = time.monotonic() + 5 + while "after" not in seen and time.monotonic() < deadline: + time.sleep(0.05) + finally: + server.should_exit = True + assert seen.get("after") is not None, "handler never saw the disconnect" + assert seen["after"] < 2.0 From 34a99887ec838236a7434827887c93143e313d82 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Sat, 19 Sep 2026 17:21:12 -0400 Subject: [PATCH 20/94] fix(kvcache): refuse the host kv tiers on an nvfp4 pool whose block scales they cannot mirror Upstream #499/#525 mirror KV codes plus the fp8 pool's per-token _scale_buffer; the nvfp4 pool (#408) keeps per-block scales in _block_scale_buffer, which neither the offloader nor the prefix tier copies, so --kv-host-pages or FT_PREFIX_HOST would rehydrate codes without scales. Both constructors now raise a clear error naming the knob; the GDN-snapshot-only tier (FT_GDN_HOST_TIER) never reads the KV pool and stays usable. Assisted-by: Claude Fable 5.1 (cherry picked from commit 4bb35fc34b0cb65b41584c65d7390b6c5f4f854a) --- python/freetoken/kvcache/host_prefix_tier.py | 3 ++ python/freetoken/kvcache/kv_host_offload.py | 14 +++++++ tests/kvcache/test_host_tier_nvfp4_guard.py | 44 ++++++++++++++++++++ 3 files changed, 61 insertions(+) create mode 100644 tests/kvcache/test_host_tier_nvfp4_guard.py diff --git a/python/freetoken/kvcache/host_prefix_tier.py b/python/freetoken/kvcache/host_prefix_tier.py index 0265b126ef..3b2d2d5dc3 100644 --- a/python/freetoken/kvcache/host_prefix_tier.py +++ b/python/freetoken/kvcache/host_prefix_tier.py @@ -79,6 +79,9 @@ def __init__(self, state_pool, kv_pool=None, *, gdn_slots_host: int = 64, self._kv_scales: HostBank | None = None self._kv_free: list[int] = [] if kv_pool is not None and kv_budget_bytes > 0: + from freetoken.kvcache.kv_host_offload import _reject_block_scaled_pool + + _reject_block_scaled_pool(kv_pool, "FT_PREFIX_HOST") codes = kv_pool._kv_buffer # [2, L, P, page, H, D] scales = kv_pool._scale_buffer # [2, L, P*page, H] | None (fp8) _, layers, _, page_size, heads, dim = codes.shape diff --git a/python/freetoken/kvcache/kv_host_offload.py b/python/freetoken/kvcache/kv_host_offload.py index dbb49f4814..1f8c358426 100644 --- a/python/freetoken/kvcache/kv_host_offload.py +++ b/python/freetoken/kvcache/kv_host_offload.py @@ -33,10 +33,24 @@ logger = init_logger(__name__) +def _reject_block_scaled_pool(pool, knob: str) -> None: + """The host mirrors carry codes plus the fp8 pool's per-token scales; an nvfp4 pool + (--kv-cache-dtype nvfp4) keeps per-block scales in _block_scale_buffer, which no mirror + copies, so a rehydrated page would be codes without scales. Refuse instead of decoding + garbage.""" + if getattr(pool, "_block_scale_buffer", None) is not None: + raise ValueError( + f"{knob} needs a KV pool with per-token scales (bf16 or fp8); the nvfp4 pool keeps " + "block scales the host mirror does not carry -- leave the tier off or serve with " + "--kv-cache-dtype fp8" + ) + + class KVHostOffloader: """GPU-slot LRU cache over the pinned host mirror of the logical KV page space.""" def __init__(self, pool, num_logical_pages: int, device: torch.device) -> None: + _reject_block_scaled_pool(pool, "--kv-host-pages") buf = pool._kv_buffer # [2, L, P_gpu + 1, page_size, kv_heads, head_dim] assert buf.shape[0] == 2 and buf.is_contiguous() _, num_layers, num_slots, page_size, kv_heads, head_dim = buf.shape diff --git a/tests/kvcache/test_host_tier_nvfp4_guard.py b/tests/kvcache/test_host_tier_nvfp4_guard.py new file mode 100644 index 0000000000..a793fb294a --- /dev/null +++ b/tests/kvcache/test_host_tier_nvfp4_guard.py @@ -0,0 +1,44 @@ +"""The host KV tiers (upstream #499 / #525) mirror codes plus the fp8 pool's per-token scales; +an nvfp4 pool keeps block scales they do not copy, so both must refuse it up front.""" + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.kvcache.host_prefix_tier import HostPrefixTier +from freetoken.kvcache.kv_host_offload import KVHostOffloader + + +def _kv_pool(block_scaled: bool): + return SimpleNamespace( + _kv_buffer=torch.zeros(2, 1, 3, 4, 1, 8, dtype=torch.uint8), # [2, L, P+1, page, H, D] + _scale_buffer=None, + _block_scale_buffer=torch.zeros(1) if block_scaled else None, + ) + + +def _state_pool(): + return SimpleNamespace( + conv_states=torch.zeros(2, 2, 4, 3), + recurrent_states=torch.zeros(2, 2, 2, 4), + slot_states={}, + ) + + +def test_kv_host_offloader_refuses_an_nvfp4_pool(): + with pytest.raises(ValueError, match="--kv-host-pages .* nvfp4"): + KVHostOffloader(_kv_pool(block_scaled=True), num_logical_pages=4, device=torch.device("cpu")) + + +def test_prefix_tier_refuses_kv_pages_on_an_nvfp4_pool(): + with pytest.raises(ValueError, match="FT_PREFIX_HOST .* nvfp4"): + HostPrefixTier(_state_pool(), kv_pool=_kv_pool(block_scaled=True), gdn_slots_host=2, + kv_budget_bytes=1 << 16, device="cpu") + + +def test_prefix_tier_gdn_snapshots_do_not_touch_the_kv_pool(): + # the snapshot-only tier (no KV budget) never reads the pool, nvfp4 or not + tier = HostPrefixTier(_state_pool(), kv_pool=_kv_pool(block_scaled=True), gdn_slots_host=2, + kv_budget_bytes=0, device="cpu") + assert tier.num_slots == 2 and not tier.has_kv From c7af50f9e996747c13246899c82ec93d8da0670f Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 23 Sep 2026 19:57:03 -0400 Subject: [PATCH 21/94] test(scheduler): give the scheduler stubs the prefill-start map the metrics timing reads Upstream #504 times the prefill span in Scheduler._prefill_start; the stubs from #222 and #505's tests build the scheduler without it. Assisted-by: Claude Opus 5.5 --- tests/scheduler/test_abort_inflight_prefill.py | 1 + tests/scheduler/test_hybrid_cache_manager.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/scheduler/test_abort_inflight_prefill.py b/tests/scheduler/test_abort_inflight_prefill.py index 85c217859f..f2cbd6e0d2 100644 --- a/tests/scheduler/test_abort_inflight_prefill.py +++ b/tests/scheduler/test_abort_inflight_prefill.py @@ -63,6 +63,7 @@ def _setup(): config=SimpleNamespace(page_size=1), status_reporter=SimpleNamespace(report_batch=lambda *_, **__: None), send_result=sent.extend, + _prefill_start={}, _kv_usage_pages=cm.page_usage, _mamba_slot_usage=lambda: None, _swa_token_usage=lambda: None, diff --git a/tests/scheduler/test_hybrid_cache_manager.py b/tests/scheduler/test_hybrid_cache_manager.py index 282ab2ca43..07385d7046 100644 --- a/tests/scheduler/test_hybrid_cache_manager.py +++ b/tests/scheduler/test_hybrid_cache_manager.py @@ -213,7 +213,7 @@ def test_chunked_prefill_retains_resumable_snapshot(monkeypatch, full_chunks, ta assert torch.all(pool.conv_states[:, match.mamba_value] == states[expected]) # Finish/abort may occur before the prefill commit; use the real idempotent cleanup path. - stub = SimpleNamespace(cache_manager=cm, table_manager=tm) + stub = SimpleNamespace(cache_manager=cm, table_manager=tm, _prefill_start={}) Scheduler._free_req_resources(stub, final) Scheduler._free_req_resources(stub, final) match = cm.match_req(_pend(prompt.tolist() + [999])) From 70e942e72d4ca01343ec6f1f5043b6f09a02c546 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 23 Sep 2026 19:59:17 -0400 Subject: [PATCH 22/94] test(engine): give the logprobs sampler fakes the logits-processing flag the sampler plan reads Fork #47's sampler plan asks each request whether it needs logits processing; the upstream #224 test fakes predate it. Assisted-by: Claude Opus 5.5 --- tests/engine/test_sample_logprobs.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/engine/test_sample_logprobs.py b/tests/engine/test_sample_logprobs.py index aabbdc5fe7..6bd204604d 100644 --- a/tests/engine/test_sample_logprobs.py +++ b/tests/engine/test_sample_logprobs.py @@ -137,8 +137,8 @@ def test_compute_logprobs_with_prepare_keeps_max_top_logprobs_clamped(monkeypatc batch = SimpleNamespace( reqs=[ - SimpleNamespace(sampling_params=SimpleNamespace(logprobs=True, top_logprobs=17, is_greedy=True)), - SimpleNamespace(sampling_params=SimpleNamespace(logprobs=False, top_logprobs=0, is_greedy=True)), + SimpleNamespace(sampling_params=SimpleNamespace(needs_logits_processing=False, logprobs=True, top_logprobs=17, is_greedy=True)), + SimpleNamespace(sampling_params=SimpleNamespace(needs_logits_processing=False, logprobs=False, top_logprobs=0, is_greedy=True)), ] ) args = sampler.prepare(batch) From 4d91b476d416ab4af93d3fd7c75e6cddcfeac0d6 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 23 Sep 2026 20:48:29 -0400 Subject: [PATCH 23/94] merge fork review fixes: 11 bugs found re-reviewing the 5080deploy stack Assisted-by: Claude Opus 5.5 # Conflicts: # python/freetoken/server/openai_api.py # tests/server/test_stream_cancellation.py --- python/freetoken/engine/sample.py | 4 +- python/freetoken/kvcache/host_prefix_tier.py | 53 +++++++++-------- python/freetoken/models/qwen4_exp/ple.py | 13 +++-- python/freetoken/scheduler/cache.py | 7 ++- python/freetoken/server/anthropic_api.py | 8 ++- python/freetoken/server/generation.py | 13 ++++- python/freetoken/server/openai_api.py | 15 +++-- python/freetoken/server/responses_api.py | 8 ++- python/freetoken/tokenizer/tokenize.py | 10 +++- python/freetoken/utils/mp.py | 4 +- tests/engine/test_logits_processors.py | 10 ++++ tests/kvcache/test_host_tier_live_pool.py | 53 +++++++++++++++++ .../kvcache/test_host_tier_rehydrate_exact.py | 39 +++++++++++++ tests/models/qwen4_exp/test_ple.py | 13 +++++ tests/server/test_openai_api.py | 20 +++++++ tests/server/test_sampling_validation.py | 23 ++++++++ tests/server/test_stream_cancellation.py | 58 +++++++++++++++++++ tests/tokenizer/test_tokenize.py | 18 ++++++ tests/utils/test_mp_large_frame.py | 13 +++++ 19 files changed, 338 insertions(+), 44 deletions(-) create mode 100644 tests/kvcache/test_host_tier_live_pool.py create mode 100644 tests/kvcache/test_host_tier_rehydrate_exact.py create mode 100644 tests/server/test_sampling_validation.py create mode 100644 tests/utils/test_mp_large_frame.py diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py index f1984f7576..bca85c56fa 100644 --- a/python/freetoken/engine/sample.py +++ b/python/freetoken/engine/sample.py @@ -187,7 +187,9 @@ def _plan(self, batch: Batch, params) -> LogitsPlan | None: hist.append(ids[prompt_len:].tolist() if penalized else []) if p.repetition_penalty != 1.0: want_prompt = True - prompts.append(ids[:prompt_len].tolist()) + # multimodal placeholders are pseudo-ids >= vocab_size (MM_PAD_SHIFT_VALUE); + # send them to the scratch column, not past the scatter target + prompts.append(ids[:prompt_len].clamp(max=pad).tolist()) else: prompts.append([]) plan.hist_ids = _padded_rows(hist, pad, self.device) diff --git a/python/freetoken/kvcache/host_prefix_tier.py b/python/freetoken/kvcache/host_prefix_tier.py index 3b2d2d5dc3..39c96dfd6d 100644 --- a/python/freetoken/kvcache/host_prefix_tier.py +++ b/python/freetoken/kvcache/host_prefix_tier.py @@ -59,13 +59,15 @@ def __init__(self, state_pool, kv_pool=None, *, gdn_slots_host: int = 64, tensors = {"conv": state_pool.conv_states, "rec": state_pool.recurrent_states} tensors.update(state_pool.slot_states) self._order = list(tensors) - self._banks: dict[str, tuple[torch.Tensor, HostBank]] = {} + # banks only: the pool tensors are re-read on every copy, since + # LinearStatePool.rebuild replaces them (a runtime cache rebuild) + self._banks: dict[str, HostBank] = {} for name, t in tensors.items(): # pool: [layers, slots, *rest] -> bank: [host_slots, layers, *rest] bank = HostBank((gdn_slots_host, t.shape[0], *t.shape[2:]), t.dtype) if self.device.type == "cuda" and torch.cuda.is_available(): bank.pin() - self._banks[name] = (t, bank) + self._banks[name] = bank self.num_slots = gdn_slots_host self._free = list(range(gdn_slots_host)) self._in_use: set[int] = set() @@ -102,40 +104,41 @@ def __init__(self, state_pool, kv_pool=None, *, gdn_slots_host: int = 64, torch.cuda.Stream(self.device) if self.device.type == "cuda" else None) # ---------------- GDN snapshot storage ---------------- + def _pairs(self): + """(live pool tensor [layers, slots, ...], host bank) per state tensor.""" + pool = self.pool + live = {"conv": pool.conv_states, "rec": pool.recurrent_states, **pool.slot_states} + return [(live[name], self._banks[name]) for name in self._order] + + def _copy(self, copies) -> None: + if self.copy_stream is None: + for dst, src in copies: + dst.copy_(src) + return + # Order the copies after the caller's stream: the scheduler stream (itself made to wait + # on the engine stream) or the engine stream. Overlap launches the next batch before + # the previous one is drained, so a just-freed slot may still be written by it. + self.copy_stream.wait_stream(torch.cuda.current_stream(self.device)) + with torch.cuda.stream(self.copy_stream): + for dst, src in copies: + dst.copy_(src, non_blocking=True) + self.copy_stream.synchronize() + def write_gdn(self, gpu_slot: int) -> int: """D2H do snapshot do gpu_slot pra um slot host livre. Retorna o host_slot. - Deve ser enfileirado ANTES de state_pool.free(gpu_slot): a ordem de stream do - engine garante leitura antes de qualquer reescrita do slot (ver spec).""" + Deve ser enfileirado ANTES de state_pool.free(gpu_slot).""" if not self._free: raise RuntimeError("HostPrefixTier: sem slots host GDN livres (chame evict_host)") hs = self._free.pop() self._in_use.add(hs) - if self.copy_stream is not None: - with torch.cuda.stream(self.copy_stream): - for name in self._order: - t, bank = self._banks[name] - bank.tensor[hs].copy_(t[:, gpu_slot], non_blocking=True) - self.copy_stream.synchronize() - else: - for name in self._order: - t, bank = self._banks[name] - bank.tensor[hs].copy_(t[:, gpu_slot]) + self._copy([(bank.tensor[hs], t[:, gpu_slot]) for t, bank in self._pairs()]) return hs def read_gdn(self, host_slot: int, dst_gpu_slot: int) -> None: """H2D host_slot -> dst_gpu_slot do pool (já alocado pelo caller).""" if host_slot not in self._in_use: raise KeyError(f"HostPrefixTier: slot host {host_slot} não está em uso") - if self.copy_stream is not None: - with torch.cuda.stream(self.copy_stream): - for name in self._order: - t, bank = self._banks[name] - t[:, dst_gpu_slot].copy_(bank.tensor[host_slot], non_blocking=True) - self.copy_stream.synchronize() - else: - for name in self._order: - t, bank = self._banks[name] - t[:, dst_gpu_slot].copy_(bank.tensor[host_slot]) + self._copy([(t[:, dst_gpu_slot], bank.tensor[host_slot]) for t, bank in self._pairs()]) def free_gdn(self, host_slot: int) -> None: self._in_use.discard(host_slot) @@ -198,6 +201,7 @@ def _go(): del staged_s if self.copy_stream is not None: + self.copy_stream.wait_stream(torch.cuda.current_stream(self.device)) with torch.cuda.stream(self.copy_stream): _go() self.copy_stream.synchronize() @@ -235,6 +239,7 @@ def _go(): del staged_s if self.copy_stream is not None: + self.copy_stream.wait_stream(torch.cuda.current_stream(self.device)) with torch.cuda.stream(self.copy_stream): _go() self.copy_stream.synchronize() diff --git a/python/freetoken/models/qwen4_exp/ple.py b/python/freetoken/models/qwen4_exp/ple.py index 284e023bdd..44973bc0a6 100644 --- a/python/freetoken/models/qwen4_exp/ple.py +++ b/python/freetoken/models/qwen4_exp/ple.py @@ -49,6 +49,10 @@ _FUSED_HASH_ENV = "FREETOKEN_PLE_FUSED_HASH" +def _capturing(device: torch.device) -> bool: + return device.type == "cuda" and torch.cuda.is_current_stream_capturing() + + def _fused_row_ids_enabled() -> bool: """The fused hash kernel, on unless ``FREETOKEN_PLE_FUSED_HASH=0`` takes it back to torch ops.""" return (os.getenv(_FUSED_HASH_ENV) or "1").strip() not in ("0", "false", "False") @@ -482,12 +486,13 @@ def _token_index(self, meta: PLEMetadata) -> Tuple[torch.Tensor, torch.Tensor]: The fused kernel addresses the hash window through these instead of materializing the ``[B, ctx+max_len]`` packed window. Memoized on the shape (which is all they depend on) - so a captured replay reads a stable address instead of re-running the build; a build - that happens DURING capture is not cached, since its buffers live in the graph pool. + for eager forwards only: a capture always builds its own, which the graph recomputes + into its pool on every replay. A cached pair (the eager warmup forward right before + capture leaves one) would be baked in by address, then evicted and reused. """ device = meta.input_ids.device num_tokens = meta.input_ids.numel() - capturing = device.type == "cuda" and torch.cuda.is_current_stream_capturing() + capturing = _capturing(device) # is_decode is part of the key, not just the shape: a decode of B requests and a # prefill of ONE B-token request are the same [T] and mean opposite things (one token # per request at offset 0 vs B offsets inside one request). @@ -496,7 +501,7 @@ def _token_index(self, meta: PLEMetadata) -> Tuple[torch.Tensor, torch.Tensor]: (num_tokens,) if meta.is_decode else tuple(meta.seq_lens), str(device), ) - cached = self._token_index_cache.get(key) + cached = None if capturing else self._token_index_cache.get(key) if cached is not None: return cached if meta.is_decode: # one token per request, each at offset 0 diff --git a/python/freetoken/scheduler/cache.py b/python/freetoken/scheduler/cache.py index 7e4aa8e695..4e6b91df88 100644 --- a/python/freetoken/scheduler/cache.py +++ b/python/freetoken/scheduler/cache.py @@ -624,8 +624,11 @@ def _rehydrate_snapshot(self, node) -> int | None: prefix, copy the snapshot back into a fresh pool slot and return the slot; None = keep walking (the match truncates as before).""" tier = self.host_tier - hit = tier.lookup(self.prefix_cache._collect_key(node), self.page_size) - if hit is None or hit.gdn_slot is None: + key = self.prefix_cache._collect_key(node) + hit = tier.lookup(key, self.page_size) + # lookup is longest-prefix; a shorter entry's state would resume this node's full + # depth with the tokens in between missing from the recurrence + if hit is None or hit.gdn_slot is None or hit.host_len != len(key): return None pool = self.linear_state_pool if pool.num_free_slots < 1: diff --git a/python/freetoken/server/anthropic_api.py b/python/freetoken/server/anthropic_api.py index 385c87f3f1..4cb4b34822 100644 --- a/python/freetoken/server/anthropic_api.py +++ b/python/freetoken/server/anthropic_api.py @@ -55,6 +55,7 @@ submit_generation, with_keepalive, ) +from .openai_api import _await_watching_disconnect from .request_logger import log_request # Emit a protocol-native `ping` event after this many seconds of stream silence, @@ -141,9 +142,14 @@ async def handle_anthropic_messages( return StreamingResponse(events, media_type="text/event-stream") try: - result = await generate_full(uid, spec, state, source="/v1/messages") + # an abandoned request must not keep decoding to max_tokens (#222's watcher) + result = await _await_watching_disconnect( + generate_full(uid, spec, state, source="/v1/messages"), request, state, [uid] + ) except GenerationError as exc: return _anthropic_error_response(400, "invalid_request_error", str(exc)) + if result is None: + return _anthropic_error_response(499, "api_error", "client disconnected before the response was ready") response = anthropic_full_response( result, req.model, uid, cache_report=cache_report, metrics=metrics ) diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index e7d2b970a7..3db3913909 100644 --- a/python/freetoken/server/generation.py +++ b/python/freetoken/server/generation.py @@ -14,6 +14,7 @@ from __future__ import annotations import asyncio +import math import json import time from collections.abc import AsyncIterator @@ -231,8 +232,9 @@ def pick(value, key, framework): for name, value in (("presence_penalty", presence_penalty), ("frequency_penalty", frequency_penalty)): if not -2.0 <= float(value) <= 2.0: raise ValueError(f"{name} must be in [-2, 2], got {value}") - if resolved_rep <= 0.0: - raise ValueError(f"repetition_penalty must be positive, got {resolved_rep}") + # not `<= 0`: NaN compares False and would poison every seen token's logit + if not (math.isfinite(resolved_rep) and resolved_rep > 0.0): + raise ValueError(f"repetition_penalty must be a positive finite number, got {resolved_rep}") if min_tokens < 0: raise ValueError(f"min_tokens must be >= 0, got {min_tokens}") bias: list[list[float]] | None = None @@ -487,7 +489,12 @@ def _make_reasoning_parser(spec: GenSpec, state: Any) -> ReasoningParser | None: # The qwen3 chat template opens an implicit (thinking on) unless # enable_thinking is explicitly false, so the model emits only the closing # . Mirror that default here, else the chain-of-thought leaks into content. - force_reasoning = (spec.chat_template_kwargs or {}).get("enable_thinking") is not False + # continue_final_message is the exception: the template renders the final assistant + # turn as ... + content, so the continuation starts in content. + ctk = spec.chat_template_kwargs or {} + force_reasoning = ctk.get("enable_thinking") is not False and not ctk.get( + "continue_final_message" + ) elif parser_name == "glm": # GLM's template honors enable_thinking (default on) even with tools; the # generic fallback would force thinking and mislabel disabled output as reasoning. diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index b67825bfaf..9ce63060ab 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -401,14 +401,17 @@ async def handle_chat_completion( chunks = state.stream_with_cancellation(chunks, request, uids) return StreamingResponse(chunks, media_type="text/event-stream") + tasks = [ + asyncio.ensure_future(generate_full(u, spec, state, source="/v1/chat/completions")) + for u in uids + ] try: - results = await _await_watching_disconnect( - asyncio.gather(*(generate_full(u, spec, state, source="/v1/chat/completions") for u in uids)), - request, - state, - uids, - ) + results = await _await_watching_disconnect(asyncio.gather(*tasks), request, state, uids) except GenerationError as exc: + # gather does not cancel the other choices; the aborts below drop the events they + # wait on, so without the cancel they would stay pending forever + for task in tasks: + task.cancel() for u in uids: await state.abort_user(u) return create_error_response(str(exc), code=exc.code) diff --git a/python/freetoken/server/responses_api.py b/python/freetoken/server/responses_api.py index eea41722bb..9fd063a992 100644 --- a/python/freetoken/server/responses_api.py +++ b/python/freetoken/server/responses_api.py @@ -79,6 +79,7 @@ submit_generation, with_keepalive, ) +from .openai_api import _await_watching_disconnect from .request_logger import log_request # Seconds of event silence before a keep-alive frame is emitted on the stream. @@ -177,9 +178,14 @@ async def handle_responses( return StreamingResponse(events, media_type="text/event-stream") try: - result = await generate_full(uid, spec, state, source="/v1/responses") + # an abandoned request must not keep decoding to max_tokens (#222's watcher) + result = await _await_watching_disconnect( + generate_full(uid, spec, state, source="/v1/responses"), request, state, [uid] + ) except GenerationError as exc: return _error_response(400, str(exc), exc.code) + if result is None: + return _error_response(499, "client disconnected before the response was ready") response = build_responses_response( result, req, response_id, created, cache_report=cache_report, metrics=metrics ) diff --git a/python/freetoken/tokenizer/tokenize.py b/python/freetoken/tokenizer/tokenize.py index f512930dee..9648974dab 100644 --- a/python/freetoken/tokenizer/tokenize.py +++ b/python/freetoken/tokenizer/tokenize.py @@ -67,7 +67,15 @@ def _map_developer_role( m.get("role") == "developer" for m in messages ): return messages - return [{**m, "role": "system"} if m.get("role") == "developer" else m for m in messages] + mapped = [{**m, "role": "system"} if m.get("role") == "developer" else m for m in messages] + # The frontend hoisted and merged system messages before this mapping ran (templates such + # as Qwen's raise on a system message that is not first), so do it again for the new ones. + system = [m for m in mapped if m.get("role") == "system"] + if len(system) == 1 and mapped[0] is system[0]: + return mapped + rest = [m for m in mapped if m.get("role") != "system"] + content = "\n\n".join(str(m.get("content") or "") for m in system if m.get("content")) + return [{"role": "system", "content": content}] + rest class TokenizeManager: diff --git a/python/freetoken/utils/mp.py b/python/freetoken/utils/mp.py index f91498fbb2..ee35c057c1 100644 --- a/python/freetoken/utils/mp.py +++ b/python/freetoken/utils/mp.py @@ -22,7 +22,9 @@ class _CoalescedUnpacker: """ def __init__(self) -> None: - self._unpacker = msgpack.Unpacker(raw=False) + # 0 = effectively uncapped (2**31-1), as unpackb had; the Unpacker default of 100 MiB would + # kill the worker on a large message (e.g. an image's pixel values) + self._unpacker = msgpack.Unpacker(raw=False, max_buffer_size=0) self._pending: list[Any] = [] def feed(self, frame: bytes) -> None: diff --git a/tests/engine/test_logits_processors.py b/tests/engine/test_logits_processors.py index f051dec51f..a9645f4210 100644 --- a/tests/engine/test_logits_processors.py +++ b/tests/engine/test_logits_processors.py @@ -153,3 +153,13 @@ def test_prepare_builds_no_plan_without_processors_and_a_plan_with_them(): assert torch.isfinite(out[3]).all() # greedy path: argmax over the processed logits assert sampler.sample(logits, args)[1].item() == 5 + + +def test_repetition_penalty_ignores_multimodal_placeholder_ids_in_the_prompt(): + # image tokens carry pseudo-ids >= MM_PAD_SHIFT_VALUE; scattering them into the + # [rows, V+1] prompt mask indexed out of bounds (a device-side assert on CUDA) + sampler = Sampler(CPU, V) + batch = SimpleNamespace(reqs=[_req([1, 1_000_000, 2], [3], 4, repetition_penalty=2.0)]) + args = sampler.prepare(batch) + out = apply_logits_processors(torch.ones(1, V), args.plan, V) + assert out[0, 1] == 0.5 and out[0, 2] == 0.5 and out[0, 4] == 1.0 diff --git a/tests/kvcache/test_host_tier_live_pool.py b/tests/kvcache/test_host_tier_live_pool.py new file mode 100644 index 0000000000..bd12836893 --- /dev/null +++ b/tests/kvcache/test_host_tier_live_pool.py @@ -0,0 +1,53 @@ +"""The GDN snapshot tier (upstream #525) must copy against the pool's CURRENT tensors, which +LinearStatePool.rebuild replaces, and must order its copies after the engine stream.""" + +import contextlib +from types import SimpleNamespace + +import torch + +import freetoken.kvcache.host_prefix_tier as tier_module +from freetoken.kvcache.host_prefix_tier import HostPrefixTier + + +def _pool(slots: int): + return SimpleNamespace( + conv_states=torch.zeros(2, slots, 4, 3), + recurrent_states=torch.zeros(2, slots, 2, 4), + slot_states={}, + ) + + +def test_read_after_a_pool_rebuild_lands_in_the_new_tensors(): + pool = _pool(3) + tier = HostPrefixTier(pool, gdn_slots_host=2, device="cpu") + pool.conv_states[:, 1] = 5.0 + pool.recurrent_states[:, 1] = 7.0 + host = tier.write_gdn(1) + # a runtime rebuild swaps in fresh (larger) state tensors + pool.conv_states = torch.zeros(2, 6, 4, 3) + pool.recurrent_states = torch.zeros(2, 6, 2, 4) + tier.read_gdn(host, 5) + assert torch.all(pool.conv_states[:, 5] == 5.0) + assert torch.all(pool.recurrent_states[:, 5] == 7.0) + + +def test_copies_wait_for_the_engine_stream(monkeypatch): + events = [] + engine = object() + + class FakeStream: + def wait_stream(self, other): + events.append(("wait", other)) + + def synchronize(self): + events.append(("sync",)) + + monkeypatch.setattr(tier_module.torch.cuda, "current_stream", lambda device=None: engine) + monkeypatch.setattr(tier_module.torch.cuda, "stream", lambda s: contextlib.nullcontext()) + pool = _pool(3) + tier = HostPrefixTier(pool, gdn_slots_host=2, device="cpu") + tier.copy_stream = FakeStream() + host = tier.write_gdn(1) + tier.read_gdn(host, 2) + assert events == [("wait", engine), ("sync",), ("wait", engine), ("sync",)] diff --git a/tests/kvcache/test_host_tier_rehydrate_exact.py b/tests/kvcache/test_host_tier_rehydrate_exact.py new file mode 100644 index 0000000000..f67402041f --- /dev/null +++ b/tests/kvcache/test_host_tier_rehydrate_exact.py @@ -0,0 +1,39 @@ +"""The GDN snapshot tier (upstream #525) may only revive a tombstone from a snapshot of +exactly that node's prefix: the host index lookup is longest-prefix, and a shorter +snapshot would resume the recurrence with the tokens in between silently missing.""" + +from types import SimpleNamespace + +import torch + +from freetoken.kvcache.host_prefix_tier import HostMatch +from freetoken.scheduler.cache import CacheManager + + +def _manager(host_len: int, key_len: int = 128): + reads = [] + tier = SimpleNamespace( + lookup=lambda key, page_size: HostMatch(host_len, None, 7), + read_gdn=lambda host_slot, gpu_slot: reads.append((host_slot, gpu_slot)), + ) + pool = SimpleNamespace(num_free_slots=4, alloc=lambda n: [3]) + manager = SimpleNamespace( + host_tier=tier, + prefix_cache=SimpleNamespace(_collect_key=lambda node: torch.arange(key_len)), + page_size=64, + linear_state_pool=pool, + ensure_mamba_slots=lambda n: None, + ) + return manager, reads + + +def test_a_shorter_host_snapshot_does_not_revive_a_deeper_node(): + manager, reads = _manager(host_len=64) + assert CacheManager._rehydrate_snapshot(manager, node=object()) is None + assert reads == [] + + +def test_an_exact_host_snapshot_revives_the_node(): + manager, reads = _manager(host_len=128) + assert CacheManager._rehydrate_snapshot(manager, node=object()) == 3 + assert reads == [(7, 3)] diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index 74d4fd614e..02abaee1f4 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -201,6 +201,19 @@ def test_token_index_cache_is_keyed_by_shape(): assert embedding._token_index(other)[0] is not first[0] +def test_token_index_capture_never_reuses_an_evictable_cache_entry(monkeypatch): + """The eager warmup forward caches the decode pair right before capture; a capture that + reused it would bake an address the FIFO later evicts while the graph still reads it.""" + embedding = _make_layer(_config()).ple_embedding + decode = _meta([[5]], [[EOS, EOS]], decode=True) + eager = embedding._token_index(decode) # the warmup forward + monkeypatch.setattr(ple_module, "_capturing", lambda device: True) + captured = embedding._token_index(decode) + assert captured[0] is not eager[0] and captured[1] is not eager[1] + assert all(v[0] is not captured[0] for v in embedding._token_index_cache.values()) + assert captured[0].tolist() == [0] and captured[1].tolist() == [0] + + def test_token_index_does_not_confuse_a_decode_with_a_one_request_prefill(): """Same [T], opposite meaning: B decode rows at offset 0 vs B offsets in one request.""" layer = _make_layer(_config()) diff --git a/tests/server/test_openai_api.py b/tests/server/test_openai_api.py index ce675a2508..7d5b3cd114 100644 --- a/tests/server/test_openai_api.py +++ b/tests/server/test_openai_api.py @@ -223,6 +223,26 @@ def test_glm_reasoning_parser_honors_disabled_thinking_with_tools(): assert parser is not None and parser.detector.force_reasoning is True +def test_qwen3_continuation_of_a_final_assistant_message_is_content_not_reasoning(): + # The Qwen template closes the think block before the assistant's content, so a + # continue_final_message continuation never emits ; starting the parser inside + # reasoning filed the whole continuation under reasoning_content. + from freetoken.server.generation import _make_reasoning_parser + + state = FakeState([], reasoning_parser="qwen3") + req = chat_request( + messages=[ + {"role": "user", "content": "Count to five."}, + {"role": "assistant", "content": "One, two,"}, + ], + continue_final_message=True, + ) + parser = _make_reasoning_parser(chat_request_to_genspec(req, {}), state) + assert parser is not None and parser.detector.force_reasoning is False + plain = _make_reasoning_parser(chat_request_to_genspec(chat_request(), {}), state) + assert plain.detector.force_reasoning is True + + def test_non_stream_chat_completion_returns_openai_tool_calls_and_sends_tools(): output = '[TOOL_CALLS] [{"name":"get_weather","arguments":{"city":"Paris"}}]' state = FakeState( diff --git a/tests/server/test_sampling_validation.py b/tests/server/test_sampling_validation.py new file mode 100644 index 0000000000..3b2e1d6473 --- /dev/null +++ b/tests/server/test_sampling_validation.py @@ -0,0 +1,23 @@ +"""resolve_sampling rejects non-finite penalties: NaN passes a `<= 0` check and would +turn every already-seen token's logit into NaN.""" + +import pytest + +from freetoken.server.generation import resolve_sampling + + +def _resolve(**extra): + return resolve_sampling( + temperature=None, top_k=None, top_p=None, max_tokens=8, ignore_eos=False, + model_sampling={}, **extra, + ) + + +@pytest.mark.parametrize("value", [float("nan"), float("inf"), 0.0, -1.0]) +def test_repetition_penalty_must_be_positive_and_finite(value): + with pytest.raises(ValueError, match="repetition_penalty"): + _resolve(repetition_penalty=value) + + +def test_a_normal_repetition_penalty_is_kept(): + assert _resolve(repetition_penalty=1.1).repetition_penalty == pytest.approx(1.1) diff --git a/tests/server/test_stream_cancellation.py b/tests/server/test_stream_cancellation.py index ba7da7c17b..71d5ac8b85 100644 --- a/tests/server/test_stream_cancellation.py +++ b/tests/server/test_stream_cancellation.py @@ -316,3 +316,61 @@ def new_user(self): assert resp.status_code == 499 assert sorted(state.aborted) == [8, 9] + + +def test_a_failed_choice_cancels_the_other_choices(monkeypatch): + # n > 1: one choice errors while another is still generating. gather passes the error up + # without cancelling the sibling, and the aborts drop the event it waits on. + monkeypatch.setattr(openai_api, "_DISCONNECT_POLL_SECONDS", 0.01) + + class _State(_ApiState): + def __init__(self): + super().__init__(acks=None) + self._next = 7 + + def new_user(self): + self._next += 1 + return self._next + + async def wait_for_ack(self, uid): + if uid == 8: + yield SimpleNamespace(**{**vars(_ack("")), "error": "boom"}) + return + await asyncio.sleep(3600) + yield _ack("never") + + async def scenario(): + state = _State() + resp = await openai_api.handle_chat_completion( + _chat_req(n=2), _Request(disconnected=False), state, {} + ) + await asyncio.sleep(0) + pending = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()] + return resp, state, pending + + resp, state, pending = asyncio.run(scenario()) + assert resp.status_code == 400 + assert sorted(state.aborted) == [8, 9] + assert pending == [] + + +def test_messages_and_responses_non_stream_disconnect_delivers_abort(monkeypatch): + # #222 only wrapped the OpenAI handlers; the Anthropic and Responses endpoints kept + # decoding an abandoned non-streaming request to max_tokens. + from freetoken.server import anthropic_api, responses_api + from freetoken.server.anthropic_models import AnthropicMessagesRequest + from freetoken.server.responses_api import ResponsesRequest + + monkeypatch.setattr(openai_api, "_DISCONNECT_POLL_SECONDS", 0.01) + + state = _ApiState(acks=None) + req = AnthropicMessagesRequest(model="m", max_tokens=8, messages=[{"role": "user", "content": "hi"}]) + handler = anthropic_api.handle_anthropic_messages(req, _Request(disconnected=True), state, {}) + resp = asyncio.run(asyncio.wait_for(handler, timeout=5)) # unwatched: hangs, then times out + assert resp.status_code == 499 and state.aborted == [7] + + state = _ApiState(acks=None) + req = ResponsesRequest(model="m", input="hi", max_output_tokens=8) + handler = responses_api.handle_responses(req, _Request(disconnected=True), state, {}) + resp = asyncio.run(asyncio.wait_for(handler, timeout=5)) + assert resp.status_code == 499 and state.aborted == [7] diff --git a/tests/tokenizer/test_tokenize.py b/tests/tokenizer/test_tokenize.py index c82a76861c..390fcd4cce 100644 --- a/tests/tokenizer/test_tokenize.py +++ b/tests/tokenizer/test_tokenize.py @@ -827,3 +827,21 @@ def test_developer_role_maps_to_system_unless_the_template_knows_it(): assert mapped[0]["content"] == "be brief" and msgs[0]["role"] == "developer" # input untouched assert _map_developer_role(msgs, "{% if message.role == 'developer' %}") is msgs assert _map_developer_role(msgs[1:], None) == msgs[1:] # nothing to map + + +def test_mapped_developer_messages_are_merged_into_one_leading_system_message(): + """The frontend hoists system messages before the tokenizer maps developer -> system; a + template that rejects a non-first system message must still get exactly one, first.""" + from freetoken.tokenizer.tokenize import _map_developer_role + + template = "{% if message.role == 'system' %}" + msgs = [ + {"role": "system", "content": "be brief"}, + {"role": "developer", "content": "answer in German"}, + {"role": "user", "content": "hi"}, + ] + mapped = _map_developer_role(msgs, template) + assert [m["role"] for m in mapped] == ["system", "user"] + assert mapped[0]["content"] == "be brief\n\nanswer in German" + late = [{"role": "user", "content": "hi"}, {"role": "developer", "content": "be brief"}] + assert [m["role"] for m in _map_developer_role(late, template)] == ["system", "user"] diff --git a/tests/utils/test_mp_large_frame.py b/tests/utils/test_mp_large_frame.py new file mode 100644 index 0000000000..f91ee5604d --- /dev/null +++ b/tests/utils/test_mp_large_frame.py @@ -0,0 +1,13 @@ +"""The coalesced-frame decoder (#466) must accept a message larger than msgpack's +Unpacker default buffer (100 MiB), which unpackb never limited.""" + +import msgpack + +from freetoken.utils.mp import _CoalescedUnpacker + + +def test_a_message_over_100_mib_decodes(): + payload = {"pixels": b"x" * (101 * 1024 * 1024)} + unpacker = _CoalescedUnpacker() + unpacker.feed(msgpack.packb(payload)) + assert len(unpacker) == 1 and len(unpacker.take()["pixels"]) == len(payload["pixels"]) From 4f2528fa4db421acd2d93f4d2853b87909074fe7 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 23 Sep 2026 23:52:44 -0400 Subject: [PATCH 24/94] merge fork fix (issue #65): run the PLE graph fill on its own thread, submitted before the launch # Conflicts: # python/freetoken/models/qwen4_exp/ple_disk.py --- python/freetoken/models/qwen4_exp/ple_disk.py | 126 ++++++++++++---- .../qwen4_exp/test_ple_disk_fill_order.py | 136 ++++++++++++++++++ 2 files changed, 237 insertions(+), 25 deletions(-) create mode 100644 tests/models/qwen4_exp/test_ple_disk_fill_order.py diff --git a/python/freetoken/models/qwen4_exp/ple_disk.py b/python/freetoken/models/qwen4_exp/ple_disk.py index 69ee9fa93a..d2eaa99df8 100644 --- a/python/freetoken/models/qwen4_exp/ple_disk.py +++ b/python/freetoken/models/qwen4_exp/ple_disk.py @@ -6,6 +6,9 @@ from __future__ import annotations import os +import queue +import threading +from concurrent.futures import Future from contextlib import contextmanager from dataclasses import dataclass from typing import Sequence @@ -100,6 +103,52 @@ def resolve_row_source(folder: str) -> PleRowSource: return source_from_safetensors(folder) +class _Filler: + """The one thread that drives the PleStore (engine-thread-only C++, no locks), FIFO. + + Graph fills run here, submitted before the launch: the fill releases the captured decode's + WAIT, and cuGraphLaunch may block until the GPU drains, so it must not wait for the launch + to return (a large graph deadlocked that way; fork issue #65).""" + + def __init__(self) -> None: + self._jobs: queue.SimpleQueue = queue.SimpleQueue() + self._thread = threading.Thread(target=self._run, name="ple-filler", daemon=True) + self._thread.start() + + def submit(self, fn) -> Future: + fut: Future = Future() + self._jobs.put((fn, fut)) + return fut + + def _run(self) -> None: + while (item := self._jobs.get()) is not None: + fn, fut = item + try: + fut.set_result(fn()) + except BaseException as exc: # noqa: BLE001 -- surfaced to the engine thread + fut.set_exception(exc) + + def close(self) -> None: + self._jobs.put(None) + self._thread.join(timeout=5) + + +class _PendingFill: + """A graph fill submitted before its launch; ``cancel`` after a failed launch.""" + + def __init__(self, table: "DiskRowTable", future: Future, cancelled: threading.Event) -> None: + self._table, self._future, self._cancelled = table, future, cancelled + + def cancel(self) -> None: + # no WAIT consumed the signal: clear it, or the next step's WAIT passes on stale rows + self._cancelled.set() + try: + self._future.result() + except BaseException: # noqa: BLE001 -- the launch error is the one to surface + pass + self._table._flag.zero_() + + class DiskRowTable: """``PLETableBackend`` whose rows are read from disk per fill (--ple-backend disk).""" @@ -150,9 +199,6 @@ def __init__( self._graph_dev = torch.empty( max_graph_rows * self._token_bytes, dtype=torch.uint8, device=self._device ) - # completes once the last graph launch that read ``_graph_pinned`` has run (recorded by the - # dispatch context after the step's fill: a record inside the capture would not be observable) - self._graph_consumed = torch.cuda.Event() eager_bytes = max_extend_tokens * self._token_bytes # the overlap scheduler fills batch k+1 while batch k's lookup copy may still be queued: the # eager staging alternates two pinned buffers, and each is rewritten only once the H2D copy @@ -177,10 +223,36 @@ def __init__( # ---------------- host side (engine thread, before the forward launches) ---------------- + def _filler(self) -> _Filler: + filler = self.__dict__.get("_filler_thread") + if filler is None: + filler = self._filler_thread = _Filler() + return filler + + def close(self) -> None: + filler = self.__dict__.pop("_filler_thread", None) + if filler is not None: + filler.close() + + def _current_stream(self): + return torch.cuda.current_stream(self._device) + + def _raise_failed_fill(self) -> None: + pending = self.__dict__.pop("_pending_fill", None) + if pending is not None and pending.done() and pending.exception() is not None: + raise pending.exception() + if pending is not None and not pending.done(): + self._pending_fill = pending + + def _fill_now(self, runs: Sequence[torch.Tensor], *, graph: bool) -> None: + """A fill that completes before its launch, on the store's thread (after any queued graph fill).""" + self._filler().submit(lambda: self.fill(runs, graph=graph)).result() + def fill(self, runs: Sequence[torch.Tensor], *, graph: bool) -> None: """Stage per-request token runs (two context ids, then the new tokens) in batch order.""" if graph: - self._graph_consumed.synchronize() + # no launch fence: the job already waited on the readback recorded after the last graph, + # and a fence recorded after this step's launch would wait on the WAIT this fill releases pinned = self._graph_pinned else: self._eager_slot ^= 1 @@ -202,22 +274,27 @@ def _ple_context(self, ids: torch.Tensor, position: int) -> list[int]: return _context(ids, position, self.eos_token_id) return [self.image_token_id if t >= MM_PAD_SHIFT_VALUE else t for t in _context(ids, position, self.eos_token_id)] - def host_fill_batch(self, batch: Batch, use_graph: bool): - """Stage this batch's rows; returns the post-dispatch fill callable under flag-sync, else None.""" - eos = self.eos_token_id + def host_fill_batch(self, batch: Batch, use_graph: bool) -> _PendingFill | None: + """Stage this batch's rows before the dispatch. Under flag-sync a graph fill is only submitted: + it completes on the filler thread, independently of the launch; returns its handle, else None.""" + self._raise_failed_fill() if batch.is_decode: reqs = list(batch.reqs) + # the context comes from the request as it is NOW: the engine mutates it once the launch returns + contexts = [self._ple_context(r.input_ids, r.device_len - 1) for r in reqs] if use_graph and self._wait_sync: bs = batch.padded_size self._token_readback[:bs].copy_(batch.input_ids, non_blocking=True) - self._readback_event.record(torch.cuda.current_stream(self._device)) + self._readback_event.record(self._current_stream()) + cancelled = threading.Event() - def _complete() -> None: + def _job() -> None: try: self._readback_event.synchronize() tokens = self._token_readback[:bs].to(torch.int64).tolist() - runs = [torch.tensor([*self._ple_context(r.input_ids, r.device_len - 1), t], dtype=torch.int64) - for r, t in zip(reqs, tokens)] + if cancelled.is_set(): + return + runs = [torch.tensor([*ctx, t], dtype=torch.int64) for ctx, t in zip(contexts, tokens)] self.fill(runs, graph=True) except BaseException: from freetoken.kernel.row_store import signal @@ -226,12 +303,12 @@ def _complete() -> None: signal(self._flag) raise - return _complete + self._pending_fill = self._filler().submit(_job) + return _PendingFill(self, self._pending_fill, cancelled) # launch-gating: this D2H is the step's readback and orders the fill after sampling tokens = batch.input_ids.to("cpu").to(torch.int64).tolist() - runs = [torch.tensor([*self._ple_context(r.input_ids, r.device_len - 1), t], dtype=torch.int64) - for r, t in zip(reqs, tokens)] - self.fill(runs, graph=use_graph) + runs = [torch.tensor([*ctx, t], dtype=torch.int64) for ctx, t in zip(contexts, tokens)] + self._fill_now(runs, graph=use_graph) return None runs = [ torch.cat(( @@ -240,20 +317,19 @@ def _complete() -> None: )) for req in batch.padded_reqs ] - self.fill(runs, graph=False) + self._fill_now(runs, graph=False) return None @contextmanager def forward_host_ctx(self, batch: Batch, use_graph: bool): - """Around one dispatch: stage on enter, run the deferred fill+signal on exit.""" - deferred = self.host_fill_batch(batch, use_graph) - yield - # no try/finally: a failed launch leaves no WAIT pending, so the fill must not run - if deferred is not None: - deferred() - if use_graph: - # only after this step's fill: the next fill waits on this launch, not the fill on it - self._graph_consumed.record(torch.cuda.current_stream(self._device)) + """Around one dispatch: stage (or submit the graph fill) on enter; cancel a submitted fill if the launch fails.""" + pending = self.host_fill_batch(batch, use_graph) + try: + yield + except BaseException: + if pending is not None: + pending.cancel() + raise # ---------------- device side (PLETableBackend protocol) ---------------- diff --git a/tests/models/qwen4_exp/test_ple_disk_fill_order.py b/tests/models/qwen4_exp/test_ple_disk_fill_order.py new file mode 100644 index 0000000000..e7fe58f7e0 --- /dev/null +++ b/tests/models/qwen4_exp/test_ple_disk_fill_order.py @@ -0,0 +1,136 @@ +"""The PLE disk backend's wait-sync fill must not depend on the graph launch returning. + +The captured decode graph WAITs on ``_flag``; ``cuGraphLaunch`` may block until the GPU drains +(a large graph fills the launch queue), and the GPU is parked on that WAIT. So the fill that sets +the flag has to run on its own thread, started before the launch (fork issue #65). +These tests stand in for the GPU with host fakes: the "launch" blocks until the flag is set. +""" + +from __future__ import annotations + +import threading +import time +from types import SimpleNamespace + +import pytest +import torch + +import freetoken.models.qwen4_exp.ple_disk as ple_disk +from freetoken.models.qwen4_exp.ple_disk import DiskRowTable + +EOS = 9 + + +class _Store: + """Stands in for the C++ PleStore: records calls and the thread that made them.""" + + def __init__(self, flag: torch.Tensor): + self.flag = flag + self.calls: list[tuple[str, list[int] | None, str]] = [] + + def stage(self, tokens_addr, n, staging_addr): + tokens = ctypes_read(tokens_addr, n + 2) + self.calls.append(("stage", tokens, threading.current_thread().name)) + + def flush(self, signal_addr): + self.calls.append(("flush", None, threading.current_thread().name)) + if signal_addr: + self.flag[0] = 1 + + +def ctypes_read(addr: int, n: int) -> list[int]: + import ctypes + + return list((ctypes.c_int64 * n).from_address(addr)) + + +class _Event: + def record(self, stream=None): + pass + + def synchronize(self): + pass + + +def _table(monkeypatch) -> DiskRowTable: + # host-only: no CUDA stream behind the (faked) readback event + monkeypatch.setattr(ple_disk.torch.cuda, "current_stream", lambda *a, **k: None) + t = DiskRowTable.__new__(DiskRowTable) + t._device = None + t._wait_sync = True + t._flag = torch.zeros(1, dtype=torch.int64) + t._store = _Store(t._flag) + t._token_readback = torch.zeros(8, dtype=torch.int32) + t._readback_event = _Event() + t._graph_pinned = torch.zeros(64, dtype=torch.uint8) + t._eager_pinned = [torch.zeros(64, dtype=torch.uint8), torch.zeros(64, dtype=torch.uint8)] + t._eager_read = [_Event(), _Event()] + t._eager_slot = 0 + t._token_bytes = 1 + t.eos_token_id = EOS + t.image_token_id = None + return t + + +def _close(t: DiskRowTable) -> None: + close = getattr(t, "close", None) + if close is not None: + close() + + +def _decode_batch(prev_tokens: list[int], new_token: int): + req = SimpleNamespace(input_ids=torch.tensor(prev_tokens + [new_token]), device_len=len(prev_tokens) + 1) + return SimpleNamespace(is_decode=True, reqs=[req], padded_size=1, + input_ids=torch.tensor([new_token], dtype=torch.int32)) + + +def _blocking_launch(flag: torch.Tensor, timeout: float = 5.0) -> bool: + """A cuGraphLaunch that cannot return until the GPU passes the graph's WAIT(flag).""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if int(flag[0]) == 1: + flag[0] = 0 # the graph's WAIT-then-RESET + return True + time.sleep(0.001) + return False + + +def test_graph_fill_does_not_wait_for_the_launch_to_return(monkeypatch): + t = _table(monkeypatch) + with t.forward_host_ctx(_decode_batch([1, 2, 3], 7), use_graph=True): + assert _blocking_launch(t._flag), "deadlock: the fill only runs after the launch returns" + _close(t) + + +def test_fills_run_in_step_order_on_one_thread_with_contexts_taken_before_the_launch(monkeypatch): + t = _table(monkeypatch) + for prev, new in (([1, 2, 3], 7), ([1, 2, 3, 7], 8), ([1, 2, 3, 7, 8], 4)): + batch = _decode_batch(prev, new) + with t.forward_host_ctx(batch, use_graph=True): + # the engine mutates the request as soon as the launch returns (overlap scheduling) + batch.reqs[0].input_ids = torch.tensor([0, 0, 0, 0, 0, 0]) + batch.reqs[0].device_len += 1 + assert _blocking_launch(t._flag) + _close(t) + staged = [c for c in t._store.calls if c[0] == "stage"] + assert [c[1] for c in staged] == [[2, 3, 7], [3, 7, 8], [7, 8, 4]] + assert len({c[2] for c in t._store.calls}) == 1 # the store stays single-threaded + + +def test_eager_fill_completes_before_the_launch(monkeypatch): + t = _table(monkeypatch) + req = SimpleNamespace(input_ids=torch.tensor([1, 2, 3, 4]), cached_len=1, device_len=4) + batch = SimpleNamespace(is_decode=False, padded_reqs=[req]) + with t.forward_host_ctx(batch, use_graph=False): + assert [c[0] for c in t._store.calls] == ["stage", "flush"] + _close(t) + + +def test_a_failed_launch_leaves_no_stale_signal(monkeypatch): + t = _table(monkeypatch) + with pytest.raises(RuntimeError, match="launch failed"): + with t.forward_host_ctx(_decode_batch([1, 2, 3], 7), use_graph=True): + raise RuntimeError("launch failed") + # no WAIT consumed it, so a leftover 1 would let the next step read stale rows + assert int(t._flag[0]) == 0 + _close(t) From c8709d9ba2c2f4ffc0c07f1b3f7748c7d359b8f3 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 23 Sep 2026 23:52:44 -0400 Subject: [PATCH 25/94] merge fork perf: load only the routed experts for short prefills (FREETOKEN_MOE_SMALL_PREFILL_TOKENS, off by default) --- python/freetoken/layers/moe.py | 19 +++++++ tests/moe/test_small_prefill_on_demand.py | 69 +++++++++++++++++++++++ 2 files changed, 88 insertions(+) create mode 100644 tests/moe/test_small_prefill_on_demand.py diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index 79ad9e3a5f..ee1c5c5eed 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -24,6 +24,9 @@ # default. Set FREETOKEN_HYBRID_OVERLAP=0 to force the serial path (CPU sync before the # GPU work) -- a measurement-only escape hatch to A/B the overlap benefit. _HYBRID_OVERLAP = os.getenv("FREETOKEN_HYBRID_OVERLAP", "1") != "0" +# Prefills of at most this many tokens load only their routed experts (the decode path's +# on-demand LRU) instead of streaming every expert of every layer; 0 = off. +_SMALL_PREFILL_TOKENS = int(os.getenv("FREETOKEN_MOE_SMALL_PREFILL_TOKENS", "0")) class MoELayer(BaseOP): @@ -359,6 +362,22 @@ def _prefill_routed( pass through unmapped.""" cache = self.offload_cache assert cache is not None + if 0 < hidden_states.shape[0] <= _SMALL_PREFILL_TOKENS and cache.decode_target in ("gpu", "hybrid"): + # A short extension (an agent turn over a cached prefix) touches a fraction of each + # layer's experts; fetching those beats streaming all of them (~68 GB here). The same + # token count holds in every layer, so the overlap double buffer is never begun. + cache.ensure_experts(self.layer_id, topk_ids) + cache.copy_missing() + return self._expert_gemm( + cache, + hidden_states, + topk_weights, + topk_ids, + views=cache.bank_views(), + n=None, + alphas=cache.alphas_for_slots(self.layer_id), + is_prefill=False, + ) if cache.prefill_overlap: views = self._wait_prefill_overlap(cache) out = self._expert_gemm( diff --git a/tests/moe/test_small_prefill_on_demand.py b/tests/moe/test_small_prefill_on_demand.py new file mode 100644 index 0000000000..5d984a0dec --- /dev/null +++ b/tests/moe/test_small_prefill_on_demand.py @@ -0,0 +1,69 @@ +"""Short prefills load only their routed experts (FREETOKEN_MOE_SMALL_PREFILL_TOKENS) instead of +streaming whole layers; above the threshold, or with it off, the whole-layer path is unchanged.""" + +import pytest +import torch + +import freetoken.layers.moe as moe_mod +from freetoken.distributed import set_tp_info, try_get_tp_info + + +def _layer_and_cache(device, num_experts=8, top_k=2, hidden=16, inter=32): + from freetoken.layers.moe import OffloadMoELayer + from freetoken.layers.quantization import NoQuantConfig + from freetoken.moe.offload_cache import OffloadMoeCache + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + layer = OffloadMoELayer(0, num_experts, top_k, hidden, inter, quant_config=NoQuantConfig(), + prefix="model.layers.0.mlp.experts") + cache = OffloadMoeCache(num_layers=1, num_experts=num_experts, cache_size=num_experts, device=device) + g = torch.Generator().manual_seed(0) + cache.set_bank_sources({ + "gate_up": [torch.randn(num_experts, 2 * inter, hidden, generator=g, dtype=torch.bfloat16) * 0.1], + "down": [torch.randn(num_experts, hidden, inter, generator=g, dtype=torch.bfloat16) * 0.1], + }) + layer.offload_cache = cache + return layer, cache + + +@pytest.mark.parametrize("tokens, threshold, on_demand", [ + (3, 8, True), (8, 8, True), (9, 8, False), (3, 0, False), +]) +def test_small_prefill_dispatch(monkeypatch, tokens, threshold, on_demand): + layer, cache = _layer_and_cache(torch.device("cpu")) + monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", threshold) + calls = [] + monkeypatch.setattr(cache, "ensure_experts", lambda layer_id, ids: calls.append("ensure")) + monkeypatch.setattr(cache, "materialize_layer", lambda layer_id: calls.append("materialize")) + monkeypatch.setattr(cache, "copy_missing", lambda: calls.append("copy")) + monkeypatch.setattr(cache, "begin_prefill", lambda: calls.append("begin_prefill"), raising=False) + got = {} + + def fake_gemm(cache_, hs, w, ids, *, views, n, alphas, is_prefill): + got.update(n=n, is_prefill=is_prefill) + return hs + + monkeypatch.setattr(layer, "_expert_gemm", fake_gemm) + hs = torch.randn(tokens, 16) + layer._prefill_routed(hs, torch.full((tokens, 2), 0.5), torch.zeros(tokens, 2, dtype=torch.int32)) + if on_demand: + assert calls == ["ensure", "copy"] and got == {"n": None, "is_prefill": False} + else: + assert "ensure" not in calls and "begin_prefill" not in calls[1:] and got["n"] == layer.num_experts + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_small_prefill_matches_the_whole_layer_path(monkeypatch): + from freetoken.core import get_global_ctx # noqa: F401 (layer reads the ctx only in forward()) + + dev = torch.device("cuda") + hs = torch.randn(5, 16, dtype=torch.bfloat16, device=dev) * 0.5 + ids = torch.tensor([[0, 3], [1, 3], [7, 2], [3, 4], [5, 0]], dtype=torch.int32, device=dev) + w = torch.tensor([[0.6, 0.4], [0.5, 0.5], [0.9, 0.1], [0.3, 0.7], [0.2, 0.8]], device=dev) + outs = {} + for threshold in (0, 8): + monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", threshold) + layer, cache = _layer_and_cache(dev) + outs[threshold] = layer._prefill_routed(hs.clone(), w.clone(), ids.clone()).float() + torch.testing.assert_close(outs[8], outs[0], rtol=2e-2, atol=2e-2) From c3ec1551f87116dd08a87fb52c60786a10452582 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 24 Sep 2026 00:27:13 -0400 Subject: [PATCH 26/94] merge fork fix: small-prefill on-demand path hands the LRU kernel unique ids under Triton's element limit --- python/freetoken/layers/moe.py | 49 +++++++++++++++++------ tests/moe/test_small_prefill_on_demand.py | 37 +++++++++++++++-- 2 files changed, 70 insertions(+), 16 deletions(-) diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index ee1c5c5eed..9d795df3cc 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -27,6 +27,11 @@ # Prefills of at most this many tokens load only their routed experts (the decode path's # on-demand LRU) instead of streaming every expert of every layer; 0 = off. _SMALL_PREFILL_TOKENS = int(os.getenv("FREETOKEN_MOE_SMALL_PREFILL_TOKENS", "0")) +_TRITON_MAX_NUMEL = 1 << 20 # Triton's per-tensor element limit (the LRU kernel's block) + + +def triton_next_power_of_2(n: int) -> int: + return 1 << max(0, int(n) - 1).bit_length() class MoELayer(BaseOP): @@ -366,18 +371,18 @@ def _prefill_routed( # A short extension (an agent turn over a cached prefix) touches a fraction of each # layer's experts; fetching those beats streaming all of them (~68 GB here). The same # token count holds in every layer, so the overlap double buffer is never begun. - cache.ensure_experts(self.layer_id, topk_ids) - cache.copy_missing() - return self._expert_gemm( - cache, - hidden_states, - topk_weights, - topk_ids, - views=cache.bank_views(), - n=None, - alphas=cache.alphas_for_slots(self.layer_id), - is_prefill=False, - ) + slots = self._ensure_unique(cache, topk_ids) + if slots is not None: + return self._expert_gemm( + cache, + hidden_states, + topk_weights, + slots, + views=cache.bank_views(), + n=None, + alphas=cache.alphas_for_slots(self.layer_id), + is_prefill=False, + ) if cache.prefill_overlap: views = self._wait_prefill_overlap(cache) out = self._expert_gemm( @@ -407,6 +412,26 @@ def _prefill_routed( cache.release_prefill_layer(self.layer_id) return out + def _ensure_unique(self, cache: OffloadMoeCache, topk_ids: torch.Tensor) -> torch.Tensor | None: + """Make this layer's routed experts resident and return ``topk_ids`` mapped to slots, or + None (take the whole-layer path) when they cannot all be resident at once. + + The LRU kernel is sized for decode: its block is next_pow2(ids) x next_pow2(slots), so a + prefill's tokens x top_k ids overflow Triton's element limit. It gets the unique ids + instead (at most num_experts), chunked under that limit; each chunk's misses are copied + before the next ensure reuses the copy plan.""" + uniq, inverse = torch.unique(topk_ids.reshape(-1), return_inverse=True) + n = uniq.numel() + if n > cache.cache_size // 2: # leave room so a chunk never evicts the one before it + return None + uniq = uniq.to(torch.int32) + per_call = max(1, _TRITON_MAX_NUMEL // triton_next_power_of_2(cache.cache_size)) + for start in range(0, n, per_call): + part = uniq[start : start + per_call] # a contiguous view: ensure rewrites it in place + cache.ensure_experts(self.layer_id, part) + cache.copy_missing() + return uniq[inverse].view_as(topk_ids) + def _wait_prefill_overlap(self, cache: OffloadMoeCache) -> tuple[torch.Tensor, ...]: """Double-buffer choreography for this layer's overlap prefill: kick off the next layer's full-layer H2D copy, then return this layer's bank views (in diff --git a/tests/moe/test_small_prefill_on_demand.py b/tests/moe/test_small_prefill_on_demand.py index 5d984a0dec..e7fb70721a 100644 --- a/tests/moe/test_small_prefill_on_demand.py +++ b/tests/moe/test_small_prefill_on_demand.py @@ -8,7 +8,7 @@ from freetoken.distributed import set_tp_info, try_get_tp_info -def _layer_and_cache(device, num_experts=8, top_k=2, hidden=16, inter=32): +def _layer_and_cache(device, num_experts=8, top_k=2, hidden=16, inter=32, cache_size=None): from freetoken.layers.moe import OffloadMoELayer from freetoken.layers.quantization import NoQuantConfig from freetoken.moe.offload_cache import OffloadMoeCache @@ -17,7 +17,7 @@ def _layer_and_cache(device, num_experts=8, top_k=2, hidden=16, inter=32): set_tp_info(rank=0, size=1) layer = OffloadMoELayer(0, num_experts, top_k, hidden, inter, quant_config=NoQuantConfig(), prefix="model.layers.0.mlp.experts") - cache = OffloadMoeCache(num_layers=1, num_experts=num_experts, cache_size=num_experts, device=device) + cache = OffloadMoeCache(num_layers=1, num_experts=num_experts, cache_size=cache_size or num_experts, device=device) g = torch.Generator().manual_seed(0) cache.set_bank_sources({ "gate_up": [torch.randn(num_experts, 2 * inter, hidden, generator=g, dtype=torch.bfloat16) * 0.1], @@ -64,6 +64,35 @@ def test_small_prefill_matches_the_whole_layer_path(monkeypatch): outs = {} for threshold in (0, 8): monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", threshold) - layer, cache = _layer_and_cache(dev) - outs[threshold] = layer._prefill_routed(hs.clone(), w.clone(), ids.clone()).float() + layer, cache = _layer_and_cache(dev, cache_size=32) + outs[threshold] = _run_and_check_path(layer, hs, w, ids, on_demand=threshold > 0) torch.testing.assert_close(outs[8], outs[0], rtol=2e-2, atol=2e-2) + + +def _run_and_check_path(layer, hs, w, ids, *, on_demand): + taken = [] + real = getattr(layer, "_ensure_unique", None) + if real is not None: + layer._ensure_unique = lambda cache, t: taken.append(r := real(cache, t)) or r + out = layer._prefill_routed(hs.clone(), w.clone(), ids.clone()).float() + assert (bool(taken) and taken[0] is not None) == on_demand + return out + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_a_real_sized_small_prefill_fits_the_lru_kernel(monkeypatch): + # 200 tokens x top-10 = 2000 routed ids against a cache that rounds up to 2048 slots: the + # LRU kernel's block is next_pow2(ids) x next_pow2(slots), so the raw ids overflow Triton's + # 1M-element limit; the path must hand it the (at most num_experts) unique ids instead. + dev = torch.device("cuda") + tokens, experts, top_k = 200, 512, 10 + g = torch.Generator(device=dev).manual_seed(1) + hs = torch.randn(tokens, 16, dtype=torch.bfloat16, device=dev, generator=g) * 0.5 + ids = torch.randint(0, experts, (tokens, top_k), dtype=torch.int32, device=dev, generator=g) + w = torch.softmax(torch.randn(tokens, top_k, device=dev, generator=g), dim=-1) + outs = {} + for threshold in (0, 1024): + monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", threshold) + layer, cache = _layer_and_cache(dev, num_experts=experts, top_k=top_k, cache_size=1100) + outs[threshold] = _run_and_check_path(layer, hs, w, ids, on_demand=threshold > 0) + torch.testing.assert_close(outs[1024], outs[0], rtol=2e-2, atol=2e-2) From 6593b585f2d8e1082d8181e566add37e2c89e149 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 24 Sep 2026 00:50:11 -0400 Subject: [PATCH 27/94] merge fork fix: small-prefill ensures its experts in decode-width chunks --- python/freetoken/layers/moe.py | 19 ++++++++----------- tests/moe/test_small_prefill_on_demand.py | 6 +++++- 2 files changed, 13 insertions(+), 12 deletions(-) diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index 9d795df3cc..c49eab87eb 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -27,11 +27,9 @@ # Prefills of at most this many tokens load only their routed experts (the decode path's # on-demand LRU) instead of streaming every expert of every layer; 0 = off. _SMALL_PREFILL_TOKENS = int(os.getenv("FREETOKEN_MOE_SMALL_PREFILL_TOKENS", "0")) -_TRITON_MAX_NUMEL = 1 << 20 # Triton's per-tensor element limit (the LRU kernel's block) - - -def triton_next_power_of_2(n: int) -> int: - return 1 << max(0, int(n) - 1).bit_length() +# ids per LRU ensure in that path: the kernel's block is next_pow2(ids) x next_pow2(slots), and a +# block much wider than decode's spills to local memory whose launch reservation can OOM +_ENSURE_CHUNK_IDS = 32 class MoELayer(BaseOP): @@ -417,17 +415,16 @@ def _ensure_unique(self, cache: OffloadMoeCache, topk_ids: torch.Tensor) -> torc None (take the whole-layer path) when they cannot all be resident at once. The LRU kernel is sized for decode: its block is next_pow2(ids) x next_pow2(slots), so a - prefill's tokens x top_k ids overflow Triton's element limit. It gets the unique ids - instead (at most num_experts), chunked under that limit; each chunk's misses are copied - before the next ensure reuses the copy plan.""" + prefill's tokens x top_k ids overflow Triton's element limit, and even a block under it + can fail to launch on a full GPU. It gets the unique ids (at most num_experts) in + decode-width chunks; each chunk's misses are copied before the next ensure reuses the plan.""" uniq, inverse = torch.unique(topk_ids.reshape(-1), return_inverse=True) n = uniq.numel() if n > cache.cache_size // 2: # leave room so a chunk never evicts the one before it return None uniq = uniq.to(torch.int32) - per_call = max(1, _TRITON_MAX_NUMEL // triton_next_power_of_2(cache.cache_size)) - for start in range(0, n, per_call): - part = uniq[start : start + per_call] # a contiguous view: ensure rewrites it in place + for start in range(0, n, _ENSURE_CHUNK_IDS): + part = uniq[start : start + _ENSURE_CHUNK_IDS] # a contiguous view: ensure rewrites it in place cache.ensure_experts(self.layer_id, part) cache.copy_missing() return uniq[inverse].view_as(topk_ids) diff --git a/tests/moe/test_small_prefill_on_demand.py b/tests/moe/test_small_prefill_on_demand.py index e7fb70721a..eea5842d23 100644 --- a/tests/moe/test_small_prefill_on_demand.py +++ b/tests/moe/test_small_prefill_on_demand.py @@ -83,7 +83,8 @@ def _run_and_check_path(layer, hs, w, ids, *, on_demand): def test_a_real_sized_small_prefill_fits_the_lru_kernel(monkeypatch): # 200 tokens x top-10 = 2000 routed ids against a cache that rounds up to 2048 slots: the # LRU kernel's block is next_pow2(ids) x next_pow2(slots), so the raw ids overflow Triton's - # 1M-element limit; the path must hand it the (at most num_experts) unique ids instead. + # 1M-element limit, and a 512-id block failed to launch on a full GPU (CUDA OOM). The path + # must hand it the unique ids in decode-width chunks. dev = torch.device("cuda") tokens, experts, top_k = 200, 512, 10 g = torch.Generator(device=dev).manual_seed(1) @@ -94,5 +95,8 @@ def test_a_real_sized_small_prefill_fits_the_lru_kernel(monkeypatch): for threshold in (0, 1024): monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", threshold) layer, cache = _layer_and_cache(dev, num_experts=experts, top_k=top_k, cache_size=1100) + widths, ensure = [], cache.ensure_experts + monkeypatch.setattr(cache, "ensure_experts", lambda lid, q: (widths.append(q.numel()), ensure(lid, q))[1]) outs[threshold] = _run_and_check_path(layer, hs, w, ids, on_demand=threshold > 0) + assert widths and max(widths) <= 32, f"LRU ensure widths {sorted(set(widths))}: wider than decode's" torch.testing.assert_close(outs[1024], outs[0], rtol=2e-2, atol=2e-2) From 9975dd08fe71b894f132b4468db2a35cac44b7f6 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Sat, 26 Sep 2026 16:14:49 -0400 Subject: [PATCH 28/94] merge fork perf: scored expert eviction (--moe-cache-policy), shared expert beside the miss copy (FREETOKEN_MOE_COPY_OVERLAP), opt-in slim copy --- benchmarks/README.md | 30 + benchmarks/bench_expert_copy_corun.py | 518 +++++++++++++++ benchmarks/bench_moe_copy_overlap.py | 233 +++++++ benchmarks/bench_scored_ensure.py | 158 +++++ python/freetoken/env.py | 2 + python/freetoken/kernel/aot.py | 19 + .../kernel/csrc/jit/fast_index_copy.cuh | 109 +++ python/freetoken/kernel/fast_index_copy.py | 42 ++ python/freetoken/layers/moe.py | 21 +- python/freetoken/models/deepseek_v4/moe.py | 2 +- python/freetoken/models/qwen4_exp/moe.py | 12 +- python/freetoken/moe/offload_cache.py | 170 ++++- python/freetoken/moe/offload_kernels.py | 45 ++ python/freetoken/moe/scored_ensure.py | 395 +++++++++++ python/freetoken/server/args.py | 9 +- tests/README.md | 1 + tests/models/qwen4_exp/test_skeleton.py | 124 ++++ tests/moe/ref_scored_cache.py | 159 +++++ tests/moe/test_fused_copy.py | 173 ++++- tests/moe/test_offload.py | 117 +++- tests/moe/test_scored_ensure.py | 622 ++++++++++++++++++ tests/moe/test_small_prefill_on_demand.py | 43 +- 22 files changed, 2971 insertions(+), 33 deletions(-) create mode 100644 benchmarks/bench_expert_copy_corun.py create mode 100644 benchmarks/bench_moe_copy_overlap.py create mode 100644 benchmarks/bench_scored_ensure.py create mode 100644 python/freetoken/moe/scored_ensure.py create mode 100644 tests/moe/ref_scored_cache.py create mode 100644 tests/moe/test_scored_ensure.py diff --git a/benchmarks/README.md b/benchmarks/README.md index 5cbb0b34e7..3f0df72f81 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -27,6 +27,36 @@ batch size x miss rate. python benchmarks/bench_offload_cache_copy.py ``` +**`bench_scored_ensure.py`** — synthetic (no checkpoint): us per decode `ensure` call at a +fixed miss count, flashlib `lru_ensure` vs the `--moe-cache-policy` kernels (lru, kd, kdfb, +rule), from a CUDA graph of two decode steps on a full cache. `--stats` times it with +`--moe-collect-stats` accumulation on. + +```bash +python benchmarks/bench_scored_ensure.py --slots 1650 --k 10,20 +``` + +**`bench_moe_copy_overlap.py`** — synthetic (no checkpoint): `FREETOKEN_MOE_COPY_OVERLAP` +off vs on for Qwen4ExpMoE decode layers at the Qwen3.8-Flash-Next expert geometry, as CUDA +graphs. Checks the outputs are bitwise identical, times replays ABAB, and profiles the copy's +start after the ensure and how much of the side-stream shared expert runs inside it. + +```bash +python benchmarks/bench_moe_copy_overlap.py --bs 1 --slots 160 +``` + +**`bench_expert_copy_corun.py`** — synthetic (no checkpoint): the fused expert copy kernel +alone and beside compute, on registered host banks with the qwen4_exp NVFP4 row sizes. Sweeps +`fast_index_copy_multi` (threads x blocks/bank) and `fast_index_copy_multi_slim` +(blocks x threads x unroll) over 1..10-expert plans (GB/s, "hot" back-to-back and "cold" +spacer-separated copies), then reports the stretch sigma of bf16/fp8 GEMVs, a small matmul and +the NVFP4 routed decode GEMV while a copy loop runs on another stream. `--json` keeps the raw +per-trial times. + +```bash +python benchmarks/bench_expert_copy_corun.py --trials 5 --json copy_corun.json +``` + For host RAM vs PCIe bandwidth and the offload/hybrid backend pick, use `ft bench bw` instead — it writes the JSON profile the engine reads. diff --git a/benchmarks/bench_expert_copy_corun.py b/benchmarks/bench_expert_copy_corun.py new file mode 100644 index 0000000000..b76b82ee6f --- /dev/null +++ b/benchmarks/bench_expert_copy_corun.py @@ -0,0 +1,518 @@ +"""Expert-copy kernel sweep and co-run stretch on real-size NVFP4 expert banks. + +Two measurements, no checkpoint needed: + +* sweep: GB/s of one fused expert copy (all 6 banks, 1..10 experts per plan) for each + ``fast_index_copy_multi`` (num_threads, blocks_per_bank) config, and for each + ``fast_index_copy_multi_slim`` (blocks, threads, unroll) config when that kernel exists; +* co-run: the stretch ``sigma = t_corun / t_alone - 1`` of decode compute kernels (bf16 and + fp8 GEMV [1,H]x[H,N], a small bf16 matmul, the NVFP4 routed decode GEMV) while a copy loop + runs on another stream, and the copy's GB/s while each compute loop runs. + +The host banks are registered anonymous mmaps (``HostBank``, like the serving banks) with +the qwen4_exp triton-NVFP4 row sizes. Source rows are not reused within a loop, the GEMV +weights are cycled over copies that together exceed L2, and every loop is a CUDA graph. +Alone and co-run trials are interleaved; each cell reports min and median over trials. + +Copy loops run in two modes. "hot": copy nodes back to back. "cold": a one-element spacer +kernel after every copy, as the ensure kernel precedes every decode copy. On the RTX 5080, in +graph replay, a copy node that follows another kernel's node runs ~8 us longer than one that +follows a copy node, for every variant (CUPTI puts it in the copy's own duration; eager +launches do not show it; cause not identified). Cold GB/s subtracts the spacer's alone time. + +Run: CUDA_VISIBLE_DEVICES=0 PYTHONPATH=python python benchmarks/bench_expert_copy_corun.py +""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import statistics +import tempfile +from dataclasses import dataclass +from typing import Callable + +import torch +import torch.nn.functional as F + +from freetoken.gpu_select import assign_gpu, bind_assigned_gpu, single_gpu_arg +from freetoken.kernel import fast_index_copy as fic +from freetoken.kernel.pinned import device_ptr +from freetoken.moe.host_banks import HostBank +from freetoken.moe.offload_cache import _BANK_BYTES_PER_EXPERT + +L2_BYTES = 64 << 20 +SLEEP_CYCLES = 2_000_000 # ~1 ms head start so the host enqueues every graph before the GPU starts +LMAX = 16 # plan buffer length; the kernels read the live count from num_indices + + +def nvfp4_bank_bytes(h: int, i: int) -> list[int]: + # triton NVFP4 layout() order: gate_up, gate_up_scale, gate_up_global, down, down_scale, down_global + return [2 * i * (h // 2), 2 * i * (h // 16), 2 * i * 2, h * (i // 2), h * (i // 16), h * 2] + + +@dataclass(frozen=True) +class CopyVariant: + name: str + launch: Callable[..., None] + + +def multi_variant(threads: int, blocks_per_bank: int) -> CopyVariant: + def launch(dst, src, feat, di, si, num): + fic.fast_index_copy_multi_jit( + dst, src, feat, di, si, num, num_threads=threads, blocks_per_bank=blocks_per_bank + ) + + return CopyVariant(f"multi {threads}x{blocks_per_bank}", launch) + + +def slim_variant(blocks: int, threads: int, unroll: int) -> CopyVariant: + def launch(dst, src, feat, di, si, num): + fic.fast_index_copy_multi_slim_jit( + dst, src, feat, di, si, num, num_blocks=blocks, num_threads=threads, unroll=unroll + ) + + return CopyVariant(f"slim {blocks}x{threads}u{unroll}", launch) + + +class CopyRig: + """Registered host banks, GPU slot caches and the fused-copy descriptor tensors.""" + + def __init__(self, feats: list[int], n_src: int, n_slots: int, device: torch.device, seed: int): + self.feats = feats + self.expert_bytes = sum(feats) + self.n_src = n_src + self.n_slots = n_slots + self.device = device + gen = torch.Generator().manual_seed(seed) + self.host = [] + for feat in feats: + bank = HostBank((n_src, feat), torch.uint8) + bank.tensor.copy_(torch.randint(0, 256, (n_src, feat), dtype=torch.uint8, generator=gen)) + bank.pin() + self.host.append(bank) + self.slots = [torch.empty((n_slots, feat), dtype=torch.uint8, device=device) for feat in feats] + self.dst_ptrs = torch.tensor([c.data_ptr() for c in self.slots], dtype=torch.int64, device=device) + self.src_ptrs = torch.tensor([device_ptr(b.tensor) for b in self.host], dtype=torch.int64, device=device) + self.feat_bytes = torch.tensor(feats, dtype=torch.int64, device=device) + self.spacer_buf = torch.zeros(1, dtype=torch.int32, device=device) + self.gen = gen + self._cursor = 0 + self._perm = torch.randperm(n_src, generator=gen) + + def plans(self, n_launch: int, experts: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """[n_launch, LMAX] dst/src plans: source rows walk a permutation of the host rows + (no reuse until all are used), destination slots are distinct within a plan.""" + src = torch.zeros((n_launch, LMAX), dtype=torch.int32) + dst = torch.zeros((n_launch, LMAX), dtype=torch.int32) + for r in range(n_launch): + for j in range(experts): + if self._cursor == self.n_src: + self._perm = torch.randperm(self.n_src, generator=self.gen) + self._cursor = 0 + src[r, j] = self._perm[self._cursor] + self._cursor += 1 + dst[r, :experts] = torch.randperm(self.n_slots, generator=self.gen)[:experts] + num = torch.full((1,), experts, dtype=torch.int64, device=self.device) + return dst.to(self.device), src.to(self.device), num + + def copy_loop(self, variant: CopyVariant, n_launch: int, experts: int, cold: bool) -> Callable[[], None]: + dst, src, num = self.plans(n_launch, experts) + + def run(): + for r in range(n_launch): + variant.launch(self.dst_ptrs, self.src_ptrs, self.feat_bytes, dst[r], src[r], num) + if cold: + self.spacer_buf.add_(1) + + return run + + def spacer_us(self, stream: torch.cuda.Stream, trials: int) -> float: + n = 200 + + def run(): + for _ in range(n): + self.spacer_buf.add_(1) + + g = capture(run, stream) + return statistics.median(time_alone(g, n, stream) for _ in range(trials)) + + +@dataclass +class ComputeOp: + name: str + make: Callable[[int], Callable[[], None]] # n_launch -> loop body + single: Callable[[], None] + + +def _weight_copies(nbytes: int) -> int: + return max(2, math.ceil(2 * L2_BYTES / nbytes)) + + +def gemv_bf16(h: int, n: int, device) -> ComputeOp: + x = torch.randn(1, h, dtype=torch.bfloat16, device=device) + ws = [torch.randn(n, h, dtype=torch.bfloat16, device=device) for _ in range(_weight_copies(n * h * 2))] + + def make(n_launch): + def run(): + for r in range(n_launch): + F.linear(x, ws[r % len(ws)]) + + return run + + return ComputeOp(f"gemv_bf16 1x{h}x{n}", make, lambda: F.linear(x, ws[0])) + + +def gemv_fp8(h: int, n: int, device) -> ComputeOp: + fp8 = torch.float8_e4m3fn + x = torch.randn(1, h, device=device).to(fp8) + ws = [torch.randn(n, h, device=device).to(fp8) for _ in range(_weight_copies(n * h))] + one = torch.ones((), dtype=torch.float32, device=device) + + def mm(w): + torch._scaled_mm(x, w.t(), scale_a=one, scale_b=one, out_dtype=torch.bfloat16) + + def make(n_launch): + def run(): + for r in range(n_launch): + mm(ws[r % len(ws)]) + + return run + + return ComputeOp(f"gemv_fp8 1x{h}x{n}", make, lambda: mm(ws[0])) + + +def matmul_bf16(m: int, h: int, n: int, device) -> ComputeOp: + a = torch.randn(m, h, dtype=torch.bfloat16, device=device) + ws = [torch.randn(n, h, dtype=torch.bfloat16, device=device) for _ in range(_weight_copies(n * h * 2))] + + def make(n_launch): + def run(): + for r in range(n_launch): + F.linear(a, ws[r % len(ws)]) + + return run + + return ComputeOp(f"mm_bf16 {m}x{h}x{n}", make, lambda: F.linear(a, ws[0])) + + +def nvfp4_routed(h: int, i: int, top_k: int, n_slots: int, device) -> ComputeOp: + """The decode routed-expert GEMV pair (gate_up -> act -> down) over its own slot banks.""" + from freetoken.moe.fused_nvfp4 import fused_experts_decode_nvfp4_marlin + + fp8 = torch.float8_e4m3fn + gu = torch.randint(0, 256, (n_slots, 2 * i, h // 2), dtype=torch.uint8, device=device) + gus = torch.full((n_slots, 2 * i, h // 16), 1.0, device=device).to(fp8) + gug = torch.full((n_slots, 2 * i), 1.0, dtype=torch.float16, device=device) + dn = torch.randint(0, 256, (n_slots, h, i // 2), dtype=torch.uint8, device=device) + dns = torch.full((n_slots, h, i // 16), 1.0, device=device).to(fp8) + dng = torch.full((n_slots, h), 1.0, dtype=torch.float16, device=device) + x = torch.randn(1, h, dtype=torch.bfloat16, device=device) * 0.01 + w = torch.full((1, top_k), 1.0 / top_k, dtype=torch.float32, device=device) + groups = max(1, n_slots // top_k) + ids = [ + (torch.arange(top_k, dtype=torch.int32, device=device) + g * top_k).view(1, top_k) % n_slots + for g in range(groups) + ] + + def call(t): + fused_experts_decode_nvfp4_marlin(x, gu, gus, gug, dn, dns, dng, w, t, "silu", False) + + def make(n_launch): + def run(): + for r in range(n_launch): + call(ids[r % groups]) + + return run + + return ComputeOp(f"nvfp4_routed k={top_k}", make, lambda: call(ids[0])) + + +def capture(fn: Callable[[], None], stream: torch.cuda.Stream) -> torch.cuda.CUDAGraph: + with torch.cuda.stream(stream): + fn() # warm-up (JIT / Triton compile, cuBLAS heuristics) outside capture + stream.synchronize() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g, stream=stream): + fn() + torch.cuda.synchronize() + g.keepalive = fn # the closure owns the plan tensors the graph reads + return g + + +def _event() -> torch.cuda.Event: + return torch.cuda.Event(enable_timing=True) + + +def time_alone(g: torch.cuda.CUDAGraph, n: int, stream: torch.cuda.Stream) -> float: + """us per launch of a graph holding ``n`` launches, replayed alone.""" + start, end = _event(), _event() + with torch.cuda.stream(stream): + torch.cuda._sleep(SLEEP_CYCLES) + start.record(stream) + g.replay() + end.record(stream) + end.synchronize() + return start.elapsed_time(end) * 1e3 / n + + +def time_corun( + g_fg: torch.cuda.CUDAGraph, n_fg: int, s_fg: torch.cuda.Stream, + g_bg: torch.cuda.CUDAGraph, s_bg: torch.cuda.Stream, +) -> tuple[float, bool]: + """us per launch of the foreground graph while the background graph runs on another + stream, and whether the background outlasted the foreground (a valid co-run).""" + ev0, bg_end, fg_start, fg_end = _event(), _event(), _event(), _event() + s0 = torch.cuda.current_stream() + torch.cuda._sleep(SLEEP_CYCLES) + ev0.record(s0) + s_bg.wait_event(ev0) + s_fg.wait_event(ev0) + with torch.cuda.stream(s_bg): + g_bg.replay() + bg_end.record(s_bg) + with torch.cuda.stream(s_fg): + torch.cuda._sleep(SLEEP_CYCLES // 20) # let the background reach steady state + fg_start.record(s_fg) + g_fg.replay() + fg_end.record(s_fg) + torch.cuda.synchronize() + covered = ev0.elapsed_time(bg_end) >= ev0.elapsed_time(fg_end) + return fg_start.elapsed_time(fg_end) * 1e3 / n_fg, covered + + +def kernel_info(fn: Callable[[], None]) -> list[dict]: + """(name, grid, block, regs/thread, smem) of every kernel one call launches.""" + from torch.profiler import ProfilerActivity, profile + + fn() + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CUDA]) as prof: + fn() + torch.cuda.synchronize() + with tempfile.TemporaryDirectory() as d: + path = os.path.join(d, "trace.json") + prof.export_chrome_trace(path) + with open(path) as f: + events = json.load(f)["traceEvents"] + out = [] + for e in events: + if e.get("cat") != "kernel": + continue + a = e.get("args", {}) + out.append({ + "name": e["name"][:80], + "grid": a.get("grid"), + "block": a.get("block"), + "regs": a.get("registers per thread"), + "smem": a.get("shared memory"), + }) + return out + + +def _fmt_info(info: list[dict]) -> str: + return "; ".join( + f"{k['name'][:48]} grid={k['grid']} block={k['block']} regs={k['regs']} smem={k['smem']}" for k in info + ) + + +def _stats(xs: list[float]) -> tuple[float, float]: + return min(xs), statistics.median(xs) + + +def parse_pairs(text: str, n: int) -> list[tuple[int, ...]]: + out = [] + for item in text.split(","): + item = item.strip() + if not item: + continue + parts = tuple(int(p) for p in item.lower().replace("u", "x").split("x")) + if len(parts) != n: + raise argparse.ArgumentTypeError(f"{item!r}: expected {n} 'x'-separated ints") + out.append(parts) + return out + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--gpu", type=single_gpu_arg, default=None, + help="GPU UUID or nvidia-smi index (default: the first visible GPU)") + p.add_argument("--hidden", type=int, default=2560) + p.add_argument("--inter", type=int, default=640) + p.add_argument("--src-experts", type=int, default=256, help="host rows per bank (pinned)") + p.add_argument("--slots", type=int, default=64, help="GPU slot-cache rows per bank") + p.add_argument("--multi", default="1024x8,1024x4,512x8,512x4,256x16,256x8,128x32", + help="fast_index_copy_multi configs, THREADSxBLOCKS_PER_BANK") + p.add_argument("--slim", default="8x256x4,16x256x2,12x256x4,16x256x4,8x256x2,4x256x4", + help="fast_index_copy_multi_slim configs, BLOCKSxTHREADSxUNROLL " + "(skipped if the kernel is absent)") + p.add_argument("--experts", default="1,2,3,4,5,6,7,8,9,10", help="plan sizes for the sweep") + p.add_argument("--corun-experts", default="4", help="plan sizes for the co-run copy loop") + p.add_argument("--link", choices=["hot", "cold", "both"], default="both", help="sweep link mode(s)") + p.add_argument("--corun-link", choices=["hot", "cold"], default="cold", help="co-run copy loop link mode") + p.add_argument("--trials", type=int, default=5) + p.add_argument("--copy-ms", type=float, default=6.0, help="target length of one timed copy loop") + p.add_argument("--compute-ms", type=float, default=3.0, help="target length of one timed compute loop") + p.add_argument("--cover", type=float, default=3.0, + help="background loop length as a multiple of the foreground's alone time") + p.add_argument("--skip-sweep", action="store_true") + p.add_argument("--skip-corun", action="store_true") + p.add_argument("--ops", default="all", + help="comma list of co-run ops by prefix: gemv_bf16,gemv_fp8,mm_bf16,nvfp4_routed") + p.add_argument("--json", default=None, help="write all results to this JSON path") + p.add_argument("--seed", type=int, default=0) + return p.parse_args() + + +def build_variants(args) -> list[CopyVariant]: + variants = [multi_variant(t, b) for t, b in parse_pairs(args.multi, 2)] + if hasattr(fic, "fast_index_copy_multi_slim_jit"): + variants += [slim_variant(g, t, u) for g, t, u in parse_pairs(args.slim, 3)] + else: + print("fast_index_copy_multi_slim_jit not available: slim configs skipped", flush=True) + return variants + + +def run_sweep(rig: CopyRig, variants, experts_list, args, stream, spacer_us, results) -> None: + modes = ["hot", "cold"] if args.link == "both" else [args.link] + for mode in modes: + cold = mode == "cold" + sub = spacer_us if cold else 0.0 + print(f"\n== sweep ({mode} link{f', spacer {sub:.2f} us subtracted' if cold else ''}): copy alone, " + f"expert={rig.expert_bytes} B, trials={args.trials}; cells are GB/s at min-time / " + f"median-time (us/copy median)", flush=True) + print("variant".ljust(18) + "".join(f"{e:>20d}" for e in experts_list), flush=True) + for v in variants: + row = v.name.ljust(18) + for e in experts_list: + n_launch = max(8, math.ceil(args.copy_ms * 1e3 / (e * 55.0))) + g = capture(rig.copy_loop(v, n_launch, e, cold), stream) + ts = [time_alone(g, n_launch, stream) - sub for _ in range(args.trials)] + t_min, t_med = _stats(ts) + nbytes = e * rig.expert_bytes + results["sweep"].append({ + "link": mode, "variant": v.name, "experts": e, "us_min": t_min, "us_median": t_med, + "gbps_best": nbytes / (t_min * 1e3), "gbps_median": nbytes / (t_med * 1e3), "us_all": ts, + }) + row += f"{nbytes / (t_min * 1e3):>7.1f}/{nbytes / (t_med * 1e3):5.1f} ({t_med:5.0f})" + del g + print(row, flush=True) + + +def run_corun(rig: CopyRig, variants, ops, corun_experts, args, s_fg, s_bg, spacer_us, results) -> None: + cold = args.corun_link == "cold" + sub = spacer_us if cold else 0.0 + print(f"\n== co-run ({args.corun_link} link copy loop): compute on one stream, copy loop on another; " + f"trials={args.trials} (interleaved alone/co-run); sigma = t_corun / t_alone - 1", flush=True) + for op in ops: + n_fg_probe = 20 + g_probe = capture(op.make(n_fg_probe), s_fg) + t_probe = statistics.median(time_alone(g_probe, n_fg_probe, s_fg) for _ in range(3)) + del g_probe + n_op = max(20, math.ceil(args.compute_ms * 1e3 / t_probe)) + g_op = capture(op.make(n_op), s_fg) + print(f"\n-- {op.name}: ~{t_probe:.1f} us/launch alone, loop of {n_op}", flush=True) + results["kernels"][op.name] = kernel_info(op.single) + print(f" kernels: {_fmt_info(results['kernels'][op.name])}", flush=True) + print(" " + "copy variant".ljust(18) + "E".rjust(3) + + " op_alone_us(min/med) op_corun_us(min/med) sigma(min/med)" + + " copy_GBps alone(med) -> under op(med) valid (! = alone median >10% over the probe:" + + " other GPU load, trust the min column)", flush=True) + for v in variants: + for e in corun_experts: + copy_launch_us = e * 55.0 + n_copy_long = max(4, math.ceil(args.cover * n_op * t_probe / copy_launch_us)) + g_copy_long = capture(rig.copy_loop(v, n_copy_long, e, cold), s_bg) + n_copy = max(8, math.ceil(args.copy_ms * 1e3 / copy_launch_us)) + g_copy = capture(rig.copy_loop(v, n_copy, e, cold), s_bg) + n_op_long = max(n_op, math.ceil(args.cover * n_copy * copy_launch_us / t_probe)) + g_op_long = capture(op.make(n_op_long), s_fg) if n_op_long > n_op else g_op + op_alone, op_co, cp_alone, cp_co = [], [], [], [] + valid_op = valid_cp = 0 + for _ in range(args.trials): + op_alone.append(time_alone(g_op, n_op, s_fg)) + t, ok = time_corun(g_op, n_op, s_fg, g_copy_long, s_bg) + op_co.append(t) + valid_op += ok + cp_alone.append(time_alone(g_copy, n_copy, s_bg) - sub) + t, ok = time_corun(g_copy, n_copy, s_bg, g_op_long, s_fg) + cp_co.append(t - sub) + valid_cp += ok + a_min, a_med = _stats(op_alone) + c_min, c_med = _stats(op_co) + sig_min, sig_med = c_min / a_min - 1, c_med / a_med - 1 + nbytes = e * rig.expert_bytes + gb_alone = nbytes / (statistics.median(cp_alone) * 1e3) + gb_co = nbytes / (statistics.median(cp_co) * 1e3) + suspect = a_med > 1.1 * t_probe + results["corun"].append({ + "link": args.corun_link, "op": op.name, "variant": v.name, "experts": e, "suspect": suspect, + "op_alone_us": op_alone, "op_corun_us": op_co, + "sigma_min": sig_min, "sigma_median": sig_med, + "copy_alone_us": cp_alone, "copy_corun_us": cp_co, + "copy_gbps_alone_median": gb_alone, "copy_gbps_corun_median": gb_co, + "valid_op_trials": valid_op, "valid_copy_trials": valid_cp, + }) + print(f" {v.name.ljust(18)}{e:>3d} {a_min:8.2f}/{a_med:8.2f} {c_min:8.2f}/{c_med:8.2f}" + f" {sig_min:+6.3f}/{sig_med:+6.3f} {gb_alone:6.1f} -> {gb_co:6.1f}" + f" {valid_op}/{valid_cp}{' !' if suspect else ''}", flush=True) + del g_copy_long, g_copy, g_op_long + del g_op + + +def main() -> None: + args = parse_args() + assert torch.cuda.is_available(), "CUDA is required" + try: + assign_gpu(args.gpu) + device = bind_assigned_gpu() + except (ValueError, RuntimeError) as e: + raise SystemExit(f"error: {e}") from e + + feats = nvfp4_bank_bytes(args.hidden, args.inter) + assert sum(feats) == _BANK_BYTES_PER_EXPERT["nvfp4"](args.hidden, args.inter) + if (args.hidden, args.inter) == (2560, 640): + assert sum(feats) == 2_772_480 + print("gpu", torch.cuda.get_device_name(device), "| banks", feats, "| expert", sum(feats), "B", flush=True) + rig = CopyRig(feats, args.src_experts, args.slots, device, args.seed) + variants = build_variants(args) + results = {"feats": feats, "sweep": [], "corun": [], "kernels": {}} + + s_fg = torch.cuda.Stream(device=device) + s_bg = torch.cuda.Stream(device=device) + spacer_us = rig.spacer_us(s_bg, args.trials) + results["spacer_us"] = spacer_us + print(f"spacer kernel alone: {spacer_us:.2f} us", flush=True) + for v in variants: + dst, src, num = rig.plans(1, 4) + info = kernel_info(lambda: v.launch(rig.dst_ptrs, rig.src_ptrs, rig.feat_bytes, dst[0], src[0], num)) + results["kernels"][v.name] = info + print(f"{v.name.ljust(18)} {_fmt_info(info)}", flush=True) + + if not args.skip_sweep: + experts_list = [int(x) for x in args.experts.split(",") if x] + run_sweep(rig, variants, experts_list, args, s_bg, spacer_us, results) + + if not args.skip_corun: + h, i = args.hidden, args.inter + all_ops = [ + gemv_bf16(h, 5120, device), gemv_bf16(h, 12288, device), + gemv_fp8(h, 5120, device), gemv_fp8(h, 12288, device), + matmul_bf16(256, h, h, device), + nvfp4_routed(h, i, 10, args.slots, device), + ] + wanted = [w.strip() for w in args.ops.split(",") if w.strip()] + ops = all_ops if wanted == ["all"] else [o for o in all_ops if any(o.name.startswith(w) for w in wanted)] + corun_experts = [int(x) for x in args.corun_experts.split(",") if x] + run_corun(rig, variants, ops, corun_experts, args, s_fg, s_bg, spacer_us, results) + + if args.json: + with open(args.json, "w") as f: + json.dump(results, f, indent=1) + print(f"\nwrote {args.json}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_moe_copy_overlap.py b/benchmarks/bench_moe_copy_overlap.py new file mode 100644 index 0000000000..535f7b1bd2 --- /dev/null +++ b/benchmarks/bench_moe_copy_overlap.py @@ -0,0 +1,233 @@ +"""Decode MoE copy/compute overlap bench (FREETOKEN_MOE_COPY_OVERLAP). + +Qwen4ExpMoE layers at the Qwen3.8-Flash-Next expert geometry (hidden 2560, moe and shared +intermediate 640, top-10, NVFP4 triton banks: 2,772,480 B per expert) over one offload cache, +with fewer experts per layer than the model so host and GPU memory stay small. Each arm +(overlap off / on) captures one decode step through every layer as a CUDA graph, then: + 1. bitwise check of the two arms over the same fresh inputs, + 2. replay time per step with CUDA events, arms interleaved ABAB, fresh inputs every round, + 3. a torch.profiler (CUPTI) trace of graph replays: per layer, the copy kernel's start slip + after the ensure, the shared expert's span (side stream) and whether it ends inside the copy, the join + gap, and the copy's duration against the off arm. +No attention or hyper-connection work runs, so the percentages overstate the model's. + +Run: CUDA_VISIBLE_DEVICES=0 PYTHONPATH=python python benchmarks/bench_moe_copy_overlap.py +""" + +from __future__ import annotations + +import argparse +import statistics +from types import SimpleNamespace + +import torch + +H, I, K = 2560, 640, 10 + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--layers", type=int, default=4) + parser.add_argument("--experts", type=int, default=64, help="experts per layer") + parser.add_argument("--slots", type=int, default=160, help="cache slots (160 of 4 x 64 gives ~4.4 misses/layer)") + parser.add_argument("--bs", type=int, default=1) + parser.add_argument("--steps", type=int, default=40, help="decode steps per timing round") + parser.add_argument("--rounds", type=int, default=6) + return parser.parse_args() + + +def nvfp4_quant(): + from freetoken.layers.quantization import QuantBackend, QuantConfig, set_quant_backend + + set_quant_backend(QuantBackend.parse("moe.nvfp4=triton")) + return QuantConfig.from_hf( + {"quantization_config": {"quant_method": "modelopt", "quant_algo": "NVFP4", "ignore": ["lm_head"]}} + ) + + +def build_arm(args, overlap: bool, banks) -> SimpleNamespace: + import freetoken.core as core + from freetoken.core import Context, set_global_ctx + from freetoken.layers.moe import OffloadMoELayer + from freetoken.models.qwen4_exp.moe import Qwen4ExpMoE + from freetoken.moe.offload_cache import OffloadMoeCache + from freetoken.utils.torch_utils import torch_dtype + + E, L = args.experts, args.layers + config = SimpleNamespace( + hidden_size=H, num_experts=E, num_experts_per_tok=K, moe_intermediate_size=I, + shared_expert_intermediate_size=I, norm_topk_prob=True, quant=None, + moe_strategy="offload", decode_target="gpu", + ) + device = torch.device("cuda") + with torch.device(device), torch_dtype(torch.bfloat16): + moes = [Qwen4ExpMoE(config, l, prefix=f"model.layers.{l}.mlp") for l in range(L)] + gen = torch.Generator(device=device).manual_seed(11) + for moe in moes: + for t in moe.state_dict().values(): + if t.is_floating_point(): + t.normal_(0.0, 0.02, generator=gen) + quant = nvfp4_quant() + experts = [ + OffloadMoELayer(l, E, K, H, I, quant_config=quant, prefix=f"model.layers.{l}.mlp.experts") + for l in range(L) + ] + cache = OffloadMoeCache( + num_layers=L, num_experts=E, cache_size=args.slots, device=device, + quant_format=banks.quant_format, layout=banks.layout, + max_slots=experts[0].quant_method.slot_limit(), decode_copy_overlap=overlap, + ) + cache.set_bank_sources(banks.sources) + cache.set_alphas(banks.gate_up_alpha, banks.down_alpha) + cache.collect_stats = True + for moe, ex in zip(moes, experts): + ex.offload_cache = cache + moe.experts = ex + cache.reset() + core._GLOBAL_CTX = None + ctx = Context(page_size=1) + set_global_ctx(ctx) + ctx._batch = SimpleNamespace(is_prefill=False) + + static = torch.randn(args.bs, H, device=device, dtype=torch.bfloat16) + out = torch.empty(L, args.bs, H, device=device, dtype=torch.bfloat16) + + def step(): + for l, moe in enumerate(moes): + out[l].copy_(moe.forward(static.clone())) + + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + step() # the eager warm-up the graph runner also does + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + step() + cache.reset() + cache.reset_stats() + # the graph holds raw pointers only: the modules must outlive it or replays read freed weights + return SimpleNamespace(graph=graph, static=static, out=out, cache=cache, moes=moes) + + +def fresh_inputs(args, n: int, seed: int) -> list[torch.Tensor]: + gen = torch.Generator(device="cuda").manual_seed(seed) + return [torch.randn(args.bs, H, device="cuda", dtype=torch.bfloat16, generator=gen) for _ in range(n)] + + +def replay(arm, inputs, keep: bool = False) -> list[torch.Tensor]: + outs = [] + for x in inputs: + arm.static.copy_(x) + arm.graph.replay() + if keep: + outs.append(arm.out.clone()) + return outs + + +def time_arm(arm, inputs) -> float: + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + replay(arm, inputs) + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / len(inputs) + + +def trace(arm, inputs) -> list[tuple[float, float, str]]: + from torch.profiler import ProfilerActivity, profile + + with profile(activities=[ProfilerActivity.CUDA]) as prof: + replay(arm, inputs) + torch.cuda.synchronize() + return sorted( + (e.time_range.start, e.time_range.end, e.name) for e in prof.events() if e.device_type.name == "CUDA" + ) + + +def layer_calls(ev, overlap: bool) -> list[dict]: + """Per MoE layer call, located by kernel name (graph replay spreads branches over streams).""" + rows = [] + for i, (_, ens_end, name) in enumerate(ev): + if "_lru_ensure_kernel" not in name: + continue + copy = next((k for k in ev[i + 1:] if "fast_index_copy" in k[2]), None) + gemm = next((k for k in ev[i + 1:] if "_decode_nvfp4" in k[2]), None) + if copy is None or gemm is None: + continue + # the shared expert chain starts at the gate kernel before the ensure; with the overlap it runs on + # the side stream beside the routed path, without it serially ahead of that path's topk + j = max(j for j in range(i) if "_gate_sigmoid_kernel" in ev[j][2]) + routed_path = ("router", "topk", "ensure", "fast_index_copy") + chain = [k for k in (ev[j:] if overlap else ev[j:i]) if k[1] <= gemm[0] + and not any(m in k[2] for m in routed_path)] + rows.append({ + "copy_us": copy[1] - copy[0], + "slip_us": copy[0] - ens_end, + "join_gap_us": gemm[0] - copy[1], + "shared_span_us": max(k[1] for k in chain) - chain[0][0], + "shared_ends_in_copy": max(k[1] for k in chain) <= copy[1], + }) + return rows + + +def pct(xs, q: float) -> float: + xs = sorted(xs) + return xs[min(len(xs) - 1, int(q * len(xs)))] + + +def main() -> None: + args = parse_args() + from freetoken.distributed import set_tp_info + from freetoken.layers.moe import OffloadMoELayer + from freetoken.moe.expert_banks import build_expert_banks + + set_tp_info(rank=0, size=1) + torch.manual_seed(0) + method = OffloadMoELayer(0, args.experts, K, H, I, quant_config=nvfp4_quant(), prefix="model.layers.0.mlp.experts").quant_method + banks = build_expert_banks(method, args.layers, None, device=torch.device("cuda"), dummy=True) + expert_bytes = sum(layers[0][0].numel() * layers[0].element_size() for layers in banks.sources.values()) + print(f"expert {expert_bytes} B; {args.layers} layers x {args.experts} experts, {args.slots} slots, bs {args.bs}") + arms = {False: build_arm(args, False, banks), True: build_arm(args, True, banks)} + + seed = 100 + inputs = fresh_inputs(args, args.steps, seed) + outs = {k: replay(arm, inputs, keep=True) for k, arm in arms.items()} + same = all(torch.equal(a, b) for a, b in zip(outs[True], outs[False])) + active, missed, calls = arms[False].cache.lru_stats.sum(0).tolist() + print(f"bitwise identical on/off over {args.steps} steps: {same}; misses per layer call {missed / calls:.2f}") + + times = {False: [], True: []} + for r in range(args.rounds): + seed += 1 + inputs = fresh_inputs(args, args.steps, seed) + for k in (False, True) if r % 2 == 0 else (True, False): + times[k].append(time_arm(arms[k], inputs)) + for k in (False, True): + t = times[k] + print(f"overlap {'on ' if k else 'off'}: {statistics.median(t):.4f} ms/step median (min {min(t):.4f}, max {max(t):.4f}, n={len(t)})") + per_layer = [1000 * (a - b) / args.layers for a, b in zip(times[False], times[True])] + print(f"saving per layer, paired rounds (us): {[round(x, 1) for x in per_layer]}, median {statistics.median(per_layer):.1f}") + + inputs = fresh_inputs(args, 12, seed + 1) + rows = {k: layer_calls(trace(arms[k], inputs), k) for k in (False, True)} + for k in (False, True): + r = [x for x in rows[k] if x["copy_us"] > 5.0] # layer calls with at least one miss + if not r: + continue + slips = [x["slip_us"] for x in r] + print( + f"overlap {'on ' if k else 'off'}: {len(r)} layer calls with misses; copy {statistics.median(x['copy_us'] for x in r):.1f} us; " + f"copy start - ensure end median {statistics.median(slips):.2f} us (p90 {pct(slips, 0.9):.2f}, <=5 us in {100 * sum(s <= 5 for s in slips) / len(slips):.0f}%); " + f"GEMM start - copy end median {statistics.median(x['join_gap_us'] for x in r):.2f} us; " + f"shared expert span median {statistics.median(x['shared_span_us'] for x in r):.1f} us" + + (f", ends inside the copy in {100 * sum(x['shared_ends_in_copy'] for x in r) / len(r):.0f}%" if k else "") + ) + pairs = [(a["copy_us"], b["copy_us"]) for a, b in zip(rows[False], rows[True]) if a["copy_us"] > 5.0] + if pairs: + ratio = [b / a for a, b in pairs] + print(f"copy duration on/off: median {statistics.median(ratio):.3f} (p10 {pct(ratio, 0.1):.3f}, p90 {pct(ratio, 0.9):.3f})") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_scored_ensure.py b/benchmarks/bench_scored_ensure.py new file mode 100644 index 0000000000..5aa852716e --- /dev/null +++ b/benchmarks/bench_scored_ensure.py @@ -0,0 +1,158 @@ +"""Decode ensure cost per call: flashlib lru_ensure vs the scored --moe-cache-policy kernels. + +One ensure call per MoE layer per decode step, on a full cache (S slots over L x E experts). +Each timed call routes K ids of which exactly m miss: the hits are hot, recently used experts +and the victims are cold filler, under every policy. A CUDA graph of 2 decode steps (2 * L +calls; 1 step when that many fresh ids would not fit) is replayed from a restored snapshot, +so the number is the graph-node cost the decode graph pays. The miss count is verified +through the kernels' stats before timing. + +Run: CUDA_VISIBLE_DEVICES=0 PYTHONPATH=python python benchmarks/bench_scored_ensure.py +""" + +from __future__ import annotations + +import argparse +import statistics + +import numpy as np +import torch +from flashlib.kernels.slot_cache import N_STATS, Stat, lru_ensure + +from freetoken.moe import scored_ensure as se + +L, E, TOP_K = 48, 512, 10 +VARIANTS = ("flashlib-lru", "lru", "kd", "kdfb", "rule") + + +def build(variant: str, slots: int, k: int, m: int, dev: torch.device): + """Full-cache state plus 1-2 steps of L queries of k ids, the first k - m hot, the rest fresh misses.""" + policy = 0 if variant.endswith("lru") else se.POLICY_IDS[variant] + hot = k - m + calls = 2 * L if hot * L + 2 * L * m <= slots else L + step0, tok0 = 1_000_000, 1000 + slot_of_id = np.full(L * E, -1, np.int32) + id_of_slot = np.full(slots, -1, np.int32) + usage = np.zeros(slots, np.int64) + last_tok = np.full(L * E, se.LAST_TOK_NEVER, np.int64) + lc = np.full(L * E, se.LC_NEVER, np.int32) + ct = np.zeros(L * E, np.int64) + ids = [layer * E + e for layer in range(L) for e in range(hot)] + assert len(ids) + calls * m <= slots, "fresh ids would evict hot ones" + filler = (layer * E + e for e in range(k + m, E) for layer in range(L)) + ids += [next(filler) for _ in range(slots - len(ids))] + for s, key in enumerate(ids): + id_of_slot[s], slot_of_id[key] = key, s + is_hot = key % E < hot + usage[s] = step0 - 1 if is_hot else 1 + s + last_tok[key] = tok0 if is_hot else tok0 - 500 - s % 7 + if is_hot: + lc[key], ct[key] = 5 << se.Q, tok0 * L + key // E + queries = np.empty((calls, k), np.int32) + for i in range(calls): + visit = i // L + queries[i] = np.r_[np.arange(hot), hot + visit * m + np.arange(m)] + rng = np.random.default_rng(0) + logits = rng.normal(size=(calls, k // TOP_K, E)).astype(np.float32) + rows = queries.reshape(calls, k // TOP_K, -1).astype(np.int64) + np.put_along_axis(logits, rows, np.take_along_axis(logits, rows, -1) + 3.0, -1) + held = id_of_slot.astype(np.int64) + t = lambda a, dtype=None: torch.from_numpy(a).to(dev, dtype) # noqa: E731 + state = dict( + owner=t(id_of_slot.copy()), slot_last=t(last_tok[held]), slot_lc=t(lc[held]), slot_ct=t(ct[held]), + slot_of_id=t(slot_of_id), id_of_slot=t(id_of_slot), usage=t(usage), + step=torch.tensor(step0, dtype=torch.int64, device=dev), + tok=torch.tensor(tok0, dtype=torch.int64, device=dev), + last_tok=t(last_tok), lc=t(lc), ct=t(ct), queries=t(queries), + src=torch.zeros(slots, dtype=torch.int32, device=dev), + dst=torch.zeros(slots, dtype=torch.int32, device=dev), + num=torch.zeros((), dtype=torch.int64, device=dev), + stats=torch.zeros(L, N_STATS, dtype=torch.int64, device=dev), + ) + return policy, state, t(logits, torch.bfloat16), se.softplus2_table(dev) + + +def step_fn(variant, policy, st, logits, g_table, stats: bool): + def run(): + for i in range(st["queries"].shape[0]): + layer = i % L + q = st["queries"][i] + s = st["stats"][layer] if stats else None + if variant == "flashlib-lru": + lru_ensure(q, st["slot_of_id"], st["id_of_slot"], st["usage"], st["step"], q, + st["src"], st["dst"], st["num"], stats=s, id_base=layer * E) + continue + se.scored_ensure( + q, st["slot_of_id"], st["id_of_slot"], st["usage"], st["step"], q, + st["src"], st["dst"], st["num"], policy=policy, num_layers=L, num_experts=E, + tok=st["tok"], last_tok=st["last_tok"], lc=st["lc"], ct=st["ct"], g_table=g_table, + slot_owner=st["owner"], slot_last_tok=st["slot_last"], slot_lc=st["slot_lc"], slot_ct=st["slot_ct"], + router_logits=logits[i] if policy == 3 else None, near_miss_thr=0.25, + stats=s, id_base=layer * E, bump_tok=layer == 0, + ) + return run + + +def bench(variant, slots, k, m, dev, reps, stats): + policy, st, logits, g_table = build(variant, slots, k, m, dev) + snap = {name: v.clone() for name, v in st.items()} + + def restore(): + for name, v in snap.items(): + st[name].copy_(v) + + # verify the miss count (a stats build of the same calls), then time the requested build + restore() + step_fn(variant, policy, st, logits, g_table, True)() + got = int(st["stats"][:, Stat.MISS].sum()) + calls = st["queries"].shape[0] + assert got == calls * m, f"{variant} K={k} m={m}: {got} misses, expected {calls * m}" + restore() + run = step_fn(variant, policy, st, logits, g_table, stats) + run() # compile outside capture + restore() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + times = [] + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + for rep in range(reps + 3): + restore() + torch.cuda.synchronize() + start.record() + graph.replay() + end.record() + end.synchronize() + if rep >= 3: + times.append(start.elapsed_time(end) * 1000.0 / calls) + return statistics.median(times) + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--slots", type=int, default=1650) + ap.add_argument("--k", default="10,20", help="routed ids per call (10 = bs 1, 20 = bs 2)") + ap.add_argument("--misses", default="0-10", help="list or range; capped at K per call") + ap.add_argument("--variants", default=",".join(VARIANTS)) + ap.add_argument("--reps", type=int, default=30) + ap.add_argument("--stats", action="store_true", help="time with --moe-collect-stats accumulation on") + args = ap.parse_args() + dev = torch.device("cuda") + variants = args.variants.split(",") + misses = [] + for part in args.misses.split(","): + lo, _, hi = part.partition("-") + misses += range(int(lo), int(hi or lo) + 1) + print(f"# {torch.cuda.get_device_name(dev)}, S={args.slots}, us per ensure call " + f"(median of {args.reps} graph replays), stats={'on' if args.stats else 'off'}") + print(f"{'K':>3} {'m':>3} " + " ".join(f"{v:>13}" for v in variants) + " max delta vs flashlib-lru") + for k in (int(x) for x in args.k.split(",")): + for m in (x for x in misses if x <= k): + us = {v: bench(v, args.slots, k, m, dev, args.reps, args.stats) for v in variants} + base = us.get("flashlib-lru") + delta = "" if base is None else f" {max(us[v] - base for v in variants):+.2f}" + print(f"{k:>3} {m:>3} " + " ".join(f"{us[v]:>13.2f}" for v in variants) + delta, flush=True) + + +if __name__ == "__main__": + main() diff --git a/python/freetoken/env.py b/python/freetoken/env.py index c8eb1e7513..0e56633f48 100644 --- a/python/freetoken/env.py +++ b/python/freetoken/env.py @@ -77,6 +77,8 @@ class EnvClassSingleton: MAMBA_SSM_DTYPE = EnvStr("float32") # cap the longest prefill warmup forward; 0 = the engine's own chunk cap WARMUP_MAX_LEN = EnvInt(0) + # GPU-decode MoE: run a layer's expert miss copy on a side stream under the shared expert + MOE_COPY_OVERLAP = EnvBool(False) def __new__(cls): # single instance diff --git a/python/freetoken/kernel/aot.py b/python/freetoken/kernel/aot.py index bbc2d2c669..54539d952c 100644 --- a/python/freetoken/kernel/aot.py +++ b/python/freetoken/kernel/aot.py @@ -124,6 +124,23 @@ def build(build_directory: pathlib.Path) -> object: return KernelSpec(name=_make_name("fast_index_copy_multi", *args), build=build) +def _fast_index_copy_multi_slim_spec() -> KernelSpec: + from .fast_index_copy import SLIM_COPY_BLOCKS, SLIM_COPY_THREADS, SLIM_COPY_UNROLL + + args = make_cpp_args(SLIM_COPY_BLOCKS, SLIM_COPY_THREADS, SLIM_COPY_UNROLL) + + def build(build_directory: pathlib.Path) -> object: + return load_jit( + "fast_index_copy_multi_slim", + *args, + cuda_files=["fast_index_copy.cuh"], + cuda_wrappers=[("launch", f"&MultiIndexCopySlimKernel<{args}>::run")], + build_directory=str(build_directory), + ) + + return KernelSpec(name=_make_name("fast_index_copy_multi_slim", *args), build=build) + + def _batch_memcpy_spec() -> KernelSpec: def build(build_directory: pathlib.Path) -> object: return load_jit( @@ -156,6 +173,8 @@ def default_kernel_specs() -> tuple[KernelSpec, ...]: specs.append(_fast_index_copy_multi_spec(num_threads=1024, blocks_per_bank=32)) # prefill hit-D2D gather (HBM-bound: wide grid) + its miss-side batch H2D binding. specs.append(_fast_index_copy_multi_spec(num_threads=1024, blocks_per_bank=64)) + # opt-in decode copy (FREETOKEN_MOE_SLIM_COPY=1) + specs.append(_fast_index_copy_multi_slim_spec()) specs.append(_batch_memcpy_spec()) specs.append(_radix_spec()) diff --git a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh index 43f3fc1ba5..e37e3d48ad 100644 --- a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh +++ b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh @@ -595,3 +595,112 @@ struct MultiIndexCopyKernel { device.unwrap())(kernel, params); } }; + + +// --------------------------------------------------------------------------- +// Slim multi-bank copy: the fused copy from a small grid, so it can share the GPU with compute. The +// grid strides each (row, bank) segment in turn, carrying the phase across, so no bank runs alone at the tail. +template +__always_inline __device__ void slim_copy_segment( + uint4* __restrict__ dst, const uint4* __restrict__ src, int32_t count, int32_t i, int32_t stride +) { + constexpr int32_t kU = static_cast(kUnroll); + // kUnroll independent loads in flight per thread: host-link latency, not issue rate, bounds a small grid + for (; i + (kU - 1) * stride < count; i += kU * stride) { + uint4 v[kUnroll]; +#pragma unroll + for (int32_t k = 0; k < kU; ++k) { + v[k] = details::load_nc(src + i + k * stride); + } +#pragma unroll + for (int32_t k = 0; k < kU; ++k) { + dst[i + k * stride] = v[k]; + } + } + uint4 v[kUnroll]; +#pragma unroll + for (int32_t k = 0; k < kU; ++k) { + if (i + k * stride < count) { + v[k] = details::load_nc(src + i + k * stride); + } + } +#pragma unroll + for (int32_t k = 0; k < kU; ++k) { + if (i + k * stride < count) { + dst[i + k * stride] = v[k]; + } + } +} + +template +__global__ __launch_bounds__(kNumThreads) void fast_index_copy_multi_slim( + const __grid_constant__ MultiIndexCopyParams p +) { + const int64_t n = p.valid_length ? p.valid_length[0] : p.length; + const auto stride = static_cast(gridDim.x * kNumThreads); + const auto tid = static_cast(blockIdx.x * kNumThreads + threadIdx.x); + const auto* di = static_cast(p.dst_indices); + const auto* si = static_cast(p.src_indices); + int32_t phase = 0; // grid thread that takes the next segment's first chunk + for (int64_t row = 0; row < n; ++row) { + const auto ps = static_cast(si[row]); + const auto pd = static_cast(di[row]); + for (int b = 0; b < p.num_banks; ++b) { + const int64_t feat = p.feat_bytes[b]; + const auto units = static_cast(feat >> 4); // feat % 16 == 0; int32 chunk index: rows < 32 GiB + const auto* src = reinterpret_cast(reinterpret_cast(p.src_ptrs[b]) + ps * feat); + auto* dst = reinterpret_cast(reinterpret_cast(p.dst_ptrs[b]) + pd * feat); + const int32_t first = tid >= phase ? tid - phase : tid - phase + stride; + slim_copy_segment(dst, src, units, first, stride); + phase = static_cast((static_cast(phase) + units) % stride); + } + } +} + +template +struct MultiIndexCopySlimKernel { + static void run( + tvm::ffi::TensorView dst_ptrs, + tvm::ffi::TensorView src_ptrs, + tvm::ffi::TensorView feat_bytes, + tvm::ffi::TensorView dst_indices, + tvm::ffi::TensorView src_indices, + tvm::ffi::Optional num_indices + ) { + using namespace host; + auto device = SymbolicDevice{}; + auto B = SymbolicSize{"num_banks"}; + auto L = SymbolicSize{"indices length"}; + auto ptr_dtype = SymbolicDType{}; + auto indices_dtype = SymbolicDType{}; + auto num_indices_dtype = SymbolicDType{}; + + TensorMatcher({B}).with_dtype(ptr_dtype).with_device(device) + .verify(dst_ptrs).verify(src_ptrs).verify(feat_bytes); + TensorMatcher({L}).with_dtype(indices_dtype).with_device(device) + .verify(dst_indices).verify(src_indices); + + const int64_t* valid_length = nullptr; + if (num_indices.has_value()) { + TensorMatcher({1}).with_dtype(num_indices_dtype).with_device(device) + .verify(num_indices.value()); + valid_length = static_cast(num_indices.value().data_ptr()); + } + + const auto params = MultiIndexCopyParams{ + static_cast(dst_ptrs.data_ptr()), + static_cast(src_ptrs.data_ptr()), + static_cast(feat_bytes.data_ptr()), + dst_indices.data_ptr(), + src_indices.data_ptr(), + valid_length, + static_cast(L.unwrap()), + static_cast(B.unwrap()), + }; + const auto use_int32 = indices_dtype.unwrap().bits == 32; + const auto kernel = use_int32 + ? fast_index_copy_multi_slim + : fast_index_copy_multi_slim; + LaunchKernel(kNumBlocks, kNumThreads, device.unwrap())(kernel, params); + } +}; diff --git a/python/freetoken/kernel/fast_index_copy.py b/python/freetoken/kernel/fast_index_copy.py index b2030620ee..3047947c78 100644 --- a/python/freetoken/kernel/fast_index_copy.py +++ b/python/freetoken/kernel/fast_index_copy.py @@ -14,6 +14,11 @@ DEFAULT_NUM_BLOCKS = 4 +# fast_index_copy_multi_slim grid (8K loads in flight, the least that holds ~51 GB/s on an +# RTX 5080 at 4 experts); the AOT spec (kernel/aot.py) builds exactly this config +SLIM_COPY_BLOCKS = 8 +SLIM_COPY_THREADS = 256 +SLIM_COPY_UNROLL = 4 SKIP_FAST_INDEX_COPY_ENV = "FREETOKEN_SKIP_FAST_INDEX_COPY" _TRUE_VALUES = {"1", "true", "yes", "on"} @@ -202,6 +207,43 @@ def fast_index_copy_multi_jit( module.launch(dst_ptrs, src_ptrs, feat_bytes, dst_indices, src_indices, num_indices) +@lru_cache(maxsize=None) +def _jit_fast_index_copy_multi_slim_module(*, num_blocks: int, num_threads: int, unroll: int) -> Module: + args = make_cpp_args(num_blocks, num_threads, unroll) + return load_jit( + "fast_index_copy_multi_slim", + *args, + cuda_files=["fast_index_copy.cuh"], + cuda_wrappers=[("launch", f"&MultiIndexCopySlimKernel<{args}>::run")], + ) + + +def fast_index_copy_multi_slim_jit( + dst_ptrs: torch.Tensor, + src_ptrs: torch.Tensor, + feat_bytes: torch.Tensor, + dst_indices: torch.Tensor, + src_indices: torch.Tensor, + num_indices: torch.Tensor | None = None, + *, + num_blocks: int = SLIM_COPY_BLOCKS, + num_threads: int = SLIM_COPY_THREADS, + unroll: int = SLIM_COPY_UNROLL, +) -> None: + """:func:`fast_index_copy_multi_jit` from a small fixed grid, for copies that run beside compute. + + ``num_blocks`` CTAs split the concatenated (row, bank) 16-byte chunk range evenly and each + thread keeps ``unroll`` loads in flight, so a few SMs still saturate the host link while + the rest stay free for concurrent kernels. Same arguments and alignment contract. + """ + if _skip_fast_index_copy_enabled(): + return + module = _jit_fast_index_copy_multi_slim_module( + num_blocks=num_blocks, num_threads=num_threads, unroll=unroll + ) + module.launch(dst_ptrs, src_ptrs, feat_bytes, dst_indices, src_indices, num_indices) + + def update_copy_flag_jit(sync_flag: torch.Tensor, delta: int) -> None: assert sync_flag.is_cuda assert sync_flag.numel() == 1 diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index c49eab87eb..63ab8bb8f2 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -245,7 +245,19 @@ def decode_forward( topk=self.top_k, renormalize=self.renormalize, ) - return self._decode_routed(hidden_states, topk_weights, topk_ids) + return self._decode_routed(hidden_states, topk_weights, topk_ids, router_logits=router_logits) + + def decode_side_applies(self) -> bool: + """Whether FREETOKEN_MOE_COPY_OVERLAP applies to this layer's next forward: a GPU decode on a + cache with a side stream (not prefill, CPU/hybrid decode or TP > 1); see decode_side.""" + cache = self.offload_cache + return not ( + cache is None + or cache.decode_copy_stream is None + or self.tp_size > 1 + or get_global_ctx().batch.is_prefill + or cache.is_cpu_layer(self.layer_id) + ) def prefill_forward( self, @@ -273,6 +285,7 @@ def _decode_routed( hidden_states: torch.Tensor, topk_weights: torch.Tensor, topk_ids: torch.Tensor, + router_logits: torch.Tensor | None = None, ) -> torch.Tensor: """On-demand load: ``ensure_experts`` rewrites ``topk_ids`` into cache slot ids in place (loading missing experts), then the GEMM reads the full slot @@ -292,7 +305,7 @@ def _decode_routed( return executor.decode(self.layer_id, hidden_states, topk_weights, topk_ids) if cache.decode_target == "hybrid": return self._decode_hybrid(cache, hidden_states, topk_weights, topk_ids) - cache.ensure_experts(self.layer_id, topk_ids) + cache.ensure_experts(self.layer_id, topk_ids, router_logits=router_logits) cache.copy_missing() return self._expert_gemm( cache, @@ -423,9 +436,11 @@ def _ensure_unique(self, cache: OffloadMoeCache, topk_ids: torch.Tensor) -> torc if n > cache.cache_size // 2: # leave room so a chunk never evicts the one before it return None uniq = uniq.to(torch.int32) + # scored eviction ranks an earlier chunk's fresh installs coldest, so pin them explicitly + pin_since = cache.step.clone() if cache.cache_policy_id else None for start in range(0, n, _ENSURE_CHUNK_IDS): part = uniq[start : start + _ENSURE_CHUNK_IDS] # a contiguous view: ensure rewrites it in place - cache.ensure_experts(self.layer_id, part) + cache.ensure_experts(self.layer_id, part, update_state=False, pin_since=pin_since) # not a decode step cache.copy_missing() return uniq[inverse].view_as(topk_ids) diff --git a/python/freetoken/models/deepseek_v4/moe.py b/python/freetoken/models/deepseek_v4/moe.py index aadcdc0481..65de7f85fe 100644 --- a/python/freetoken/models/deepseek_v4/moe.py +++ b/python/freetoken/models/deepseek_v4/moe.py @@ -126,7 +126,7 @@ def _prefill_routed( or cache.is_unpinned_layer(self.layer_id) ): return super()._prefill_routed(hidden_states, topk_weights, topk_ids) - cache.ensure_experts(self.layer_id, topk_ids) # in-place expert-id -> slot + cache.ensure_experts(self.layer_id, topk_ids, update_state=False) # in-place expert-id -> slot; not a decode step cache.copy_missing() if cache.collect_stats: cache.record_decode_stats(self.layer_id) diff --git a/python/freetoken/models/qwen4_exp/moe.py b/python/freetoken/models/qwen4_exp/moe.py index c10407861f..6f729231f1 100644 --- a/python/freetoken/models/qwen4_exp/moe.py +++ b/python/freetoken/models/qwen4_exp/moe.py @@ -35,9 +35,19 @@ def __init__( def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: num_tokens, hidden_dim = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_dim) + se, ex = self.shared_expert, self.experts + if isinstance(ex, OffloadMoELayer) and ex.decode_side_applies(): + # FREETOKEN_MOE_COPY_OVERLAP: the shared expert runs on a side stream while the router, + # ensure and miss copy stay back to back on this one; joined before the combine + cache = ex.offload_cache + with cache.decode_side(ex.layer_id): + gate = shared_gate_sigmoid(hidden_states, self.shared_expert_gate.weight.view(-1)) + shared = se.forward(hidden_states) + routed = ex.forward(hidden_states=hidden_states, router_logits=self.gate.forward(hidden_states)) + cache.join_decode_side(ex.layer_id) + return shared_gate_mul_add(routed, shared, gate).view(num_tokens, hidden_dim) router_logits = self.gate.forward(hidden_states) gate = shared_gate_sigmoid(hidden_states, self.shared_expert_gate.weight.view(-1)) - se, ex = self.shared_expert, self.experts if ( self._tp_size > 1 and isinstance(se.down_proj, LinearRowParallel) diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 5f91e70e0e..96b13f33de 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -2,6 +2,7 @@ import math import os +from contextlib import contextmanager from dataclasses import dataclass from typing import TYPE_CHECKING, Iterator @@ -13,6 +14,9 @@ # (kept for A/B profiling). Falls back to per-bank automatically if a bank's row bytes or # base address are not 16-byte aligned. _FUSED_COPY = os.getenv("FREETOKEN_FUSED_COPY", "1").strip().lower() not in {"0", "false", "no", "off"} +# Opt-in: run the fused copy_missing on the slim kernel (a few small CTAs, leaves most SMs to +# concurrent compute) instead of the 8-CTAs-per-bank one. Off by default. +_SLIM_COPY = os.getenv("FREETOKEN_MOE_SLIM_COPY", "0").strip().lower() in {"1", "true", "yes", "on"} # cudaMemcpyBatchAsync silently degrades to a SYNCHRONOUS copy when a batch mixes # large entries with sub-~256KB entries on registered host memory (H100 + CUDA 13.0, @@ -26,6 +30,7 @@ _SMALL_BANK_FEAT_BYTES = 256 * 1024 from freetoken.kernel.fast_index_copy import default_blocks_per_bank +from freetoken.env import ENV from freetoken.utils import init_logger if TYPE_CHECKING: @@ -159,11 +164,20 @@ class OffloadMoeCache: # Each layer belongs to one pool: a contiguous slot range with its own LRU scope; pool p owns slots [sum(pool_sizes[:p]), +pool_sizes[p]). Slot ids in slot_for_id are local to the pool. None = one pool for all layers. pool_of_layer: list[int] | None = None pool_sizes: list[int] | None = None + # GPU decode only: the caller's work that is independent of the routed experts (qwen4_exp: the + # shared expert) runs on decode_copy_stream beside the miss copy. None = FREETOKEN_MOE_COPY_OVERLAP. + decode_copy_overlap: bool | None = None def __post_init__(self) -> None: - policy_ids = {"lru": 0} - assert self.cache_policy in policy_ids + from freetoken.moe.scored_ensure import POLICY_IDS + + assert self.cache_policy in POLICY_IDS, self.cache_policy assert self.decode_target in ("gpu", "cpu", "hybrid"), self.decode_target + if self.cache_policy != "lru" and self.decode_target != "gpu": + # the hybrid and CPU decode paths run their own kernels and would silently stay LRU + raise ValueError( + f"--moe-cache-policy {self.cache_policy} needs GPU decode, not {self.decode_target!r}" + ) if self.layout is None: assert self.quant_format in _BANK_SCHEMAS, f"unknown quant_format {self.quant_format!r}" if self.pool_sizes is None: @@ -196,7 +210,11 @@ def __post_init__(self) -> None: "cache, so cache_size must be at least 2 * num_experts " "(raise moe_cache_size or disable moe_prefill_overlap)" ) - self.cache_policy_id = policy_ids[self.cache_policy] + self.cache_policy_id = POLICY_IDS[self.cache_policy] + # rule: an unrouted expert within this many logits of the top-k-th is a near miss + self.near_miss_thr = ( + float(os.getenv("FREETOKEN_MOE_NEAR_MISS_THR", "0.25")) if self.cache_policy == "rule" else None + ) self.slot_for_id = torch.full( (self.num_layers, self.num_experts), -1, @@ -230,6 +248,7 @@ def __post_init__(self) -> None: self.expert_recency = torch.full( (self.num_layers, self.num_experts), -1, dtype=torch.int64, device=self.device ) + self._init_evict_state() # Host source banks (one [num_experts, ...] tensor per layer, so layers can # carry independent host attributes -- see layer_residency) and their GPU # slot caches, keyed by the format's bank schema (attached by @@ -323,6 +342,14 @@ def __post_init__(self) -> None: self.prefill_hit_rows = 0 self.prefill_total_rows = 0 self._slice_pools() + # Decode copy overlap: one side stream plus a (forked, done) event pair per layer. + # Neither is cache_size-shaped, so rebuild() keeps them. + if self.decode_copy_overlap is None: + self.decode_copy_overlap = bool(ENV.MOE_COPY_OVERLAP) + self.decode_copy_stream: torch.cuda.Stream | None = None + self._decode_copy_events: list[tuple[torch.cuda.Event, torch.cuda.Event]] = [] + if self.decode_copy_overlap and self.device.type == "cuda" and self.decode_target == "gpu": + self._init_decode_copy_overlap() def _pool_layout(self, pool: int) -> int: return self.layer_layout[self.pool_of_layer.index(pool)] @@ -426,6 +453,78 @@ def pool_bank_views(self, layer_id: int) -> tuple[torch.Tensor, ...]: """Per-bank slot views of the pool that serves ``layer_id``, indexed by its local slot ids.""" assert self.banks, "set_bank_sources must register the banks first" return self._pool_bank_views[self.pool_of_layer[layer_id]] + def _init_decode_copy_overlap(self) -> None: + # a dedicated stream, not one of torch's 32 pooled ones: a pooled stream can alias the engine + # or capture stream, and the side work would then run serially (correct, but no overlap) + from cuda.bindings import runtime as cudart + + with torch.cuda.device(self.device): + err, handle = cudart.cudaStreamCreateWithFlags(cudart.cudaStreamNonBlocking) + if err != cudart.cudaError_t.cudaSuccess: + raise RuntimeError(f"cudaStreamCreateWithFlags failed: {err}") + self.decode_copy_stream = torch.cuda.ExternalStream(int(handle), device=self.device) + self._decode_copy_events = [ + (torch.cuda.Event(), torch.cuda.Event()) for _ in range(self.num_layers) + ] + # torch creates the CUDA event on its first record: do that now, never inside a capture + for pair in self._decode_copy_events: + for event in pair: + event.record(self.decode_copy_stream) + logger.info("MoE decode copy overlap on: expert misses copy on a side stream") + + def _init_evict_state(self) -> None: + """Scored policies' per-(layer, expert) state (moe/scored_ensure.py); None under LRU. + + It is not slot-shaped: it survives eviction and prefill invalidation, and a rebuild + only resets it. ``evict_tok`` is the decode-step clock the first ensure of each step bumps. + """ + from freetoken.moe import scored_ensure as se + + self.evict_tok = self.evict_last_tok = self.evict_lc = self.evict_ct = None + self.evict_g_table = None + n, dev = self.num_layers * self.num_experts, self.device + if self.cache_policy_id >= 1: + self.evict_tok = torch.zeros((), dtype=torch.int64, device=dev) + self.evict_last_tok = torch.full((n,), se.LAST_TOK_NEVER, dtype=torch.int64, device=dev) + if self.cache_policy_id >= 2: + self.evict_lc = torch.full((n,), se.LC_NEVER, dtype=torch.int32, device=dev) + self.evict_ct = torch.zeros((n,), dtype=torch.int64, device=dev) + self.evict_g_table = se.softplus2_table(dev) + self._alloc_evict_mirrors() + + def _alloc_evict_mirrors(self) -> None: + """Slot-shaped copies of each slot owner's state, read coalesced by the victim scan. + + A mirror is valid while ``evict_slot_owner`` equals ``id_of_slot``; the kernel keeps it + so, and anything else that installs an id into a slot must set the owner to -1. + """ + self.evict_slot_owner = self.evict_slot_last_tok = self.evict_slot_lc = self.evict_slot_ct = None + if self.evict_tok is None: + return + size, dev = self.cache_size, self.device + self.evict_slot_owner = torch.full((size,), -1, dtype=torch.int32, device=dev) + self.evict_slot_last_tok = torch.zeros((size,), dtype=torch.int64, device=dev) + if self.evict_lc is not None: + self.evict_slot_lc = torch.zeros((size,), dtype=torch.int32, device=dev) + self.evict_slot_ct = torch.zeros((size,), dtype=torch.int64, device=dev) + + def reset_evict_state(self) -> None: + if self.evict_tok is None: + return + from freetoken.moe import scored_ensure as se + + self.evict_tok.zero_() + self.evict_last_tok.fill_(se.LAST_TOK_NEVER) + self.evict_slot_owner.fill_(-1) + if self.evict_lc is not None: + self.evict_lc.fill_(se.LC_NEVER) + self.evict_ct.zero_() + + def evict_clock_layer(self) -> int: + """The first layer that ensures on the GPU each decode step; its call bumps ``evict_tok``.""" + if not self.cpu_layer_ids: + return 0 + return min(set(range(self.num_layers)) - self.cpu_layer_ids, default=0) def set_bank_sources( self, @@ -685,6 +784,8 @@ def rebuild(self, cache_size: int, pool_sizes: list[int] | None = None) -> None: self.num_indices.zero_() self.num_missing_full.zero_() self.expert_recency.fill_(-1) + self._alloc_evict_mirrors() + self.reset_evict_state() self.stat_missing.zero_() self.stat_active.zero_() self.stat_calls.zero_() @@ -1027,8 +1128,21 @@ def release_prefill_layer(self, layer_id: int) -> None: self._prefill_buffer_has_release_event[buffer_id] = True self._prefill_buffer_released[buffer_id] = True - def ensure_experts(self, layer_id: int, expert_ids: torch.Tensor) -> None: - from freetoken.moe.offload_kernels import ensure_experts + def ensure_experts( + self, + layer_id: int, + expert_ids: torch.Tensor, + *, + update_state: bool = True, + router_logits: torch.Tensor | None = None, + pin_since: torch.Tensor | None = None, + ) -> None: + """``update_state=False`` (small-prefill ensures) keeps a scored policy's state and + decode-step clock still; ``router_logits`` (``[rows, num_experts]``) feed the ``rule`` + policy's near misses; ``pin_since`` (a ``step`` value) keeps every slot touched after it + (a small prefill's earlier chunks). All are ignored under LRU, whose recency order + already keeps those slots.""" + from freetoken.moe.offload_kernels import ensure_experts, ensure_experts_scored if self.collect_decode_freq: # ``expert_ids`` still holds raw expert ids here (the kernel rewrites them to @@ -1037,7 +1151,18 @@ def ensure_experts(self, layer_id: int, expert_ids: torch.Tensor) -> None: self.decode_freq[layer_id].scatter_add_(0, ids, torch.ones_like(ids)) self._pending_src_layer = layer_id self._pending_whole_layer = False - ensure_experts(self, layer_id, expert_ids) + if self.cache_policy_id == 0: + ensure_experts(self, layer_id, expert_ids) + return + ensure_experts_scored( + self, + layer_id, + expert_ids, + bump_tok=update_state and layer_id == self.evict_clock_layer(), + update_state=update_state, + router_logits=router_logits, + pin_since=pin_since, + ) def ensure_experts_hybrid( self, layer_id: int, expert_ids: torch.Tensor, weights: torch.Tensor, clamp: bool = False @@ -1069,6 +1194,8 @@ def materialize_layer(self, layer_id: int) -> None: self._pending_src_layer = layer_id self._pending_whole_layer = True materialize_layer(self, layer_id) + if self.evict_slot_owner is not None: + self.evict_slot_owner[: self.num_experts].fill_(-1) # it installs the layer into these slots def reset(self) -> None: from freetoken.moe.offload_kernels import reset_cache @@ -1077,6 +1204,7 @@ def reset(self) -> None: # Per-expert recency is not cache_size-shaped, so reset_cache leaves it alone; wipe # it here so a new sequence starts with cold hybrid fetch priorities. self.expert_recency.fill_(-1) + self.reset_evict_state() def reset_stats(self) -> None: self.prefill_hit_rows = 0 @@ -1183,9 +1311,12 @@ def decode_routing_stats(self) -> dict: "norm_entropy": norm_ent, } - def copy_missing(self) -> None: + def copy_missing(self, layer_id: int | None = None) -> None: + """Copy the misses the last ensure_experts/materialize_layer staged, from ``layer_id``'s + host banks (default: the layer that staged them).""" assert self.banks, "set_bank_sources must register the banks first" - layer_id = self._pending_src_layer + if layer_id is None: + layer_id = self._pending_src_layer assert layer_id is not None, "no staged misses (ensure_experts/materialize_layer first)" if layer_id in self._unpinned_layers: if not self._pending_whole_layer: @@ -1202,13 +1333,14 @@ def copy_missing(self) -> None: # materialize stages rows of staging buffer 0; ensure stages slots local to the layer's pool whole = self._pending_whole_layer if self._copy_fused_ok: - from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit + from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit, fast_index_copy_multi_slim_jit # One launch copies the missing rows for every bank (instead of one launch per # bank). evict_slots/src_indices/num_indices are shared across banks; # src_indices holds layer-local expert rows, resolved against this layer's # source pointers (layer_id is a static int per captured graph node). - fast_index_copy_multi_jit( + copy = fast_index_copy_multi_slim_jit if _SLIM_COPY else fast_index_copy_multi_jit + copy( self._copy_dst_ptrs if whole else self._copy_dst_ptrs_pool[self.pool_of_layer[layer_id]], self._copy_src_ptrs[layer_id], self._copy_feat_bytes[self.layer_layout[layer_id]], @@ -1231,6 +1363,24 @@ def copy_missing(self) -> None: self.num_indices, ) + @contextmanager + def decode_side(self, layer_id: int) -> Iterator[None]: + """Run the body on ``decode_copy_stream``, ordered after the work already on the current + stream; ``join_decode_side`` must follow before its results are read. + + The copy stays on the current stream right after its ensure, so it never waits on a + cross-stream edge; the side work (slack-rich) takes that latency and any SM contention. + Under capture the event pair becomes a fork and a join edge of the graph.""" + forked, done = self._decode_copy_events[layer_id] + forked.record(torch.cuda.current_stream(self.device)) + self.decode_copy_stream.wait_event(forked) + with torch.cuda.stream(self.decode_copy_stream): + yield + done.record(self.decode_copy_stream) + + def join_decode_side(self, layer_id: int) -> None: + torch.cuda.current_stream(self.device).wait_event(self._decode_copy_events[layer_id][1]) + def iter_offload_moe_layers(model) -> Iterator: from freetoken.layers import OffloadMoELayer, iter_moe_layers diff --git a/python/freetoken/moe/offload_kernels.py b/python/freetoken/moe/offload_kernels.py index 1d863a2f54..72f8bd5905 100644 --- a/python/freetoken/moe/offload_kernels.py +++ b/python/freetoken/moe/offload_kernels.py @@ -50,6 +50,51 @@ class HybridRoutes(NamedTuple): cpu_ids: torch.Tensor # expert ids, -1 for a GPU route +def ensure_experts_scored( + cache, + layer_id: int, + expert_ids: torch.Tensor, + *, + bump_tok: bool, + update_state: bool, + router_logits: torch.Tensor | None = None, + pin_since: torch.Tensor | None = None, +) -> None: + """``ensure_experts`` for the scored ``--moe-cache-policy`` kinds (kd, kdfb, rule).""" + from freetoken.moe.scored_ensure import scored_ensure + + scored_ensure( + expert_ids, + cache.slot_for_id.view(-1), + cache.id_of_slot, + cache.usage, + cache.step, + expert_ids, + cache.src_indices, + cache.evict_slots, + cache.num_indices, + policy=cache.cache_policy_id, + num_layers=cache.num_layers, + num_experts=cache.num_experts, + tok=cache.evict_tok, + last_tok=cache.evict_last_tok, + lc=cache.evict_lc, + ct=cache.evict_ct, + g_table=cache.evict_g_table, + slot_owner=cache.evict_slot_owner, + slot_last_tok=cache.evict_slot_last_tok, + slot_lc=cache.evict_slot_lc, + slot_ct=cache.evict_slot_ct, + router_logits=router_logits if cache.near_miss_thr is not None else None, + near_miss_thr=cache.near_miss_thr or 0.0, + stats=cache.lru_stats[layer_id] if cache.collect_stats else None, + id_base=layer_id * cache.num_experts, + bump_tok=bump_tok, + update_state=update_state, + pin_since=pin_since, + ) + + def ensure_experts_hybrid( cache, layer_id: int, diff --git a/python/freetoken/moe/scored_ensure.py b/python/freetoken/moe/scored_ensure.py new file mode 100644 index 0000000000..123086ed51 --- /dev/null +++ b/python/freetoken/moe/scored_ensure.py @@ -0,0 +1,395 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright 2026 The FlashLib Authors +# Adapted from flashlib 0.3.0 (kernels/slot_cache/triton/lru_ensure.py: _phase1, _stats, +# _lru_ensure_kernel, _num_warps_for), https://github.com/FlashML-org/flashlib +"""Scored-victim slot-cache admission for the MoE expert cache (``--moe-cache-policy``). + +flashlib's sequential-argmin ``lru_ensure`` with the victim key made pluggable. POLICY 0 is +LRU and matches flashlib bit for bit; the others evict the lowest packed score key: + +- 1 ``kd``: ``-(k + beta * d / L)``. ``k`` is the tokens the expert has been passed over + since its last use, ``d`` the layer-steps until its layer comes up again. +- 2 ``kdfb``: kd plus ``w * log2(decayed use count)`` with a ``halflife``-token decay. +- 3 ``rule``: kdfb plus the near-miss refresh: this layer's experts whose router logit is + within a margin of the top-k-th count as used one token ago (``last_tok``), unpinned. + +The per-(layer, expert) state survives eviction; slot-indexed mirrors of it (tagged with the +owning id) let the victim scan read it coalesced instead of gathering by id. Scores and +counts are int32 Q16 fixed point and the count update ``log2(2^x + 1)`` reads a host-built +table, so a CPU reference can reproduce every victim exactly (``tl.exp2``/``tl.log2`` are +approximate). +""" +from __future__ import annotations + +import math + +import torch +import triton +import triton.language as tl + +POLICY_IDS = {"lru": 0, "kd": 1, "kdfb": 2, "rule": 3} + +# Replay-tuned on sep5_s01 at 1650 slots; the plateau is flat for beta 1-8, w 2-4, halflife 32-128. +BETA = 4.0 +W = 3.0 +HALFLIFE = 64 + +Q = 16 # fraction bits of scores and counts +LAST_TOK_NEVER = -1_000_000 +LC_NEVER = -(1 << 30) +LC_FLOOR = -(30 << Q) +K_MAX = 1 << 14 # tokens; with the caps on beta and w below, a score fits int32 +MAX_BETA, MAX_W = 64.0, 16.0 +DECAY_SHIFT = 8 +SCORE_BIAS = 1 << 31 +MAX_SLOT_BITS = 30 +# g(u) = log2(1 + 2^-u) sampled every 2^-G_STEP_BITS for u in [0, G_RANGE]; zero beyond +G_STEP_BITS = 6 +G_RANGE = 24 +G_LAST = G_RANGE << G_STEP_BITS +# jit code may only read globals that are constexpr +_Q = tl.constexpr(Q) +_LC_NEVER = tl.constexpr(LC_NEVER) +_LC_FLOOR = tl.constexpr(LC_FLOOR) +_K_MAX = tl.constexpr(K_MAX) +_DECAY_SHIFT = tl.constexpr(DECAY_SHIFT) +_SCORE_BIAS = tl.constexpr(SCORE_BIAS) +_G_SHIFT = tl.constexpr(Q - G_STEP_BITS) +_G_LAST = tl.constexpr(G_LAST) + + +def softplus2_table(device=None) -> torch.Tensor: + """Q16 ``log2(1 + 2^-u)`` at ``u = i / 64``, two entries past ``G_LAST`` for the interpolation.""" + step = 1 << G_STEP_BITS + vals = [round(math.log2(1.0 + 2.0 ** (-i / step)) * (1 << Q)) for i in range(G_LAST + 2)] + return torch.tensor(vals, dtype=torch.int32, device=device) + + +def score_params(beta: float, w: float, halflife: int, num_layers: int) -> tuple[int, int, int, int]: + """``(BETA_STEP, W_Q4, DECAY_MUL, DT_MAX)`` shared by the kernel and the CPU reference. + + ``beta * d / L`` is ``d * BETA_STEP`` in Q16; the decay over ``dt`` layer-steps is + ``(dt * DECAY_MUL) >> DECAY_SHIFT``, with ``dt`` capped at 64 half-lives (past the -30 floor). + """ + assert 0 <= beta <= MAX_BETA and 0 <= w <= MAX_W and halflife >= 1 + span = num_layers * int(halflife) + decay_mul = round((1 << (Q + DECAY_SHIFT)) / span) + dt_max = 64 * span + assert dt_max * decay_mul < (1 << 31) + return round(beta * (1 << Q) / num_layers), round(w * 16), decay_mul, dt_max + + +@triton.jit +def _phase1(query_ptr, slot_of_id_ptr, lru_usage_ptr, num_copy_ptr, step, K, + BLOCK_K: tl.constexpr, id_base): + """Vendored unchanged: dedup the query, split hit/miss, rank the misses, bump the hits.""" + k = tl.arange(0, BLOCK_K) + kmask = k < K + q = tl.load(query_ptr + k, mask=kmask, other=-1) + id_base + s = tl.load(slot_of_id_ptr + q, mask=kmask, other=-1) + hit = kmask & (s >= 0) + miss = kmask & (s == -1) + same = ( + (q[:, None] == q[None, :]) + & (k[:, None] > k[None, :]) + & kmask[:, None] + & kmask[None, :] + ) + first = kmask & (tl.sum(same.to(tl.int32), axis=1) == 0) + first_miss = miss & first + smaller = (q[None, :] < q[:, None]) & first_miss[None, :] + rank = tl.sum(smaller.to(tl.int32), axis=1) + num_missing = tl.sum(first_miss.to(tl.int32)) + tl.store(num_copy_ptr, num_missing.to(tl.int64)) + # Duplicated hits write the same value to the same slot -- idempotent. + tl.store(lru_usage_ptr + s, step, mask=hit) + return q, kmask, miss, first, first_miss, rank, num_missing, tl.where(hit, s, -1) + + +@triton.jit +def _stats(stats_ptr, first, num_missing): + # One vectorized atomic over 3 lanes, not three scalar ones: a scalar atomic serializes the CTA. + si = tl.arange(0, 4) + v = tl.where(si == 0, tl.sum(first.to(tl.int32)), tl.where(si == 1, num_missing, 1)) + # relaxed: only later kernels read the counters; flashlib's acq_rel waits on this call's stores (~0.3 us) + tl.atomic_add(stats_ptr + si, v.to(tl.int64), mask=si < 3, sem="relaxed") + + +@triton.jit +def _decayed(lc, ct, now, DECAY_MUL: tl.constexpr, DT_MAX: tl.constexpr): + """int32 Q16 log2 count decayed from ``ct`` to ``now`` (int64 layer-steps).""" + dt = tl.minimum(tl.maximum(now - ct, 0), DT_MAX).to(tl.int32) + return lc - ((dt * DECAY_MUL) >> _DECAY_SHIFT) + + +@triton.jit +def _softplus2(x, g_ptr, mask): + """Q16 ``log2(2^x + 1)`` = ``max(x, 0) + g(|x|)``, g linearly interpolated from the table.""" + u = tl.abs(x) + i = tl.minimum(u >> _G_SHIFT, _G_LAST) + frac = u & ((1 << _G_SHIFT) - 1) + g0 = tl.load(g_ptr + i, mask=mask, other=0) + g1 = tl.load(g_ptr + i + 1, mask=mask, other=0) + g = tl.where(i < _G_LAST, g0 + (((g1 - g0) * frac) >> _G_SHIFT), 0) + return tl.maximum(x, 0) + g + + +@triton.jit +def _owner_state(mirror_ptr, state_ptr, c, oid, mirrored, stale, other): + """The slot owner's state: its coalesced mirror, or a gather by id where the mirror is stale.""" + return tl.where( + mirrored, + tl.load(mirror_ptr + c, mask=mirrored, other=other), + tl.load(state_ptr + oid, mask=stale, other=other), + ) + + +@triton.jit(do_not_specialize=["K", "num_cached", "id_base", "nm_topk", "nm_stride"]) +def _scored_ensure_kernel( + query_ptr, slot_of_id_ptr, id_of_slot_ptr, lru_usage_ptr, lru_step_ptr, + out_ptr, src_ptr, dst_ptr, num_copy_ptr, stats_ptr, + tok_ptr, last_tok_ptr, lc_ptr, ct_ptr, g_ptr, logits_ptr, + owner_ptr, slot_last_ptr, slot_lc_ptr, slot_ct_ptr, pin_ptr, + K, num_cached, id_base, nm_topk, nm_stride, nm_thr, + BLOCK_K: tl.constexpr, BLOCK_C: tl.constexpr, BLOCK_E: tl.constexpr, BLOCK_TOPK: tl.constexpr, + NM_ROWS: tl.constexpr, + USAGE_MAX: tl.constexpr, COLLECT_STATS: tl.constexpr, + POLICY: tl.constexpr, BUMP_TOK: tl.constexpr, UPDATE_STATE: tl.constexpr, PIN_SINCE: tl.constexpr, + NUM_LAYERS: tl.constexpr, NUM_EXPERTS: tl.constexpr, + BETA_STEP: tl.constexpr, W_Q4: tl.constexpr, DECAY_MUL: tl.constexpr, DT_MAX: tl.constexpr, + SLOT_BITS: tl.constexpr, +): + step = tl.load(lru_step_ptr) + 1 + tl.store(lru_step_ptr, step) + if POLICY != 0: + tok = tl.load(tok_ptr) + if BUMP_TOK: + tok = tok + 1 + tl.store(tok_ptr, tok) + layer = id_base // NUM_EXPERTS + now = tok * NUM_LAYERS + layer + if POLICY == 3: + if UPDATE_STATE: + # Routed ids qualify too; they are pinned now and set to tok below, so that is harmless. + ex = tl.arange(0, BLOCK_E) + j = tl.arange(0, BLOCK_TOPK) + prev = tl.load(last_tok_ptr + id_base + ex, mask=ex < NUM_EXPERTS, other=0) + held_at = tl.load(slot_of_id_ptr + id_base + ex, mask=ex < NUM_EXPERTS, other=-1) + near = ex < 0 + for r in tl.static_range(NM_ROWS): + routed = tl.load(query_ptr + r * nm_topk + j, mask=j < nm_topk, other=-1) + row_ptr = logits_ptr + r * nm_stride + kth = tl.min(tl.load(row_ptr + routed, mask=routed >= 0, other=float("inf")).to(tl.float32), axis=0) + row = tl.load(row_ptr + ex, mask=ex < NUM_EXPERTS, other=float("-inf")).to(tl.float32) + near = near | (row >= kth - nm_thr) + refreshed = tl.maximum(prev, tok - 1) + tl.store(last_tok_ptr + id_base + ex, refreshed, mask=near) + tl.store(slot_last_ptr + held_at, refreshed, mask=near & (held_at >= 0)) + # the victim scan and the routed update below read and rewrite what this wrote + tl.debug_barrier() + q, kmask, miss, first, first_miss, rank, num_missing, out = _phase1( + query_ptr, slot_of_id_ptr, lru_usage_ptr, num_copy_ptr, step, K, BLOCK_K, + id_base) + + if num_missing > 0: + # REQUIRED (flashlib): the reload below must see _phase1's hit bump, or a hit is evicted. + tl.debug_barrier() + c = tl.arange(0, BLOCK_C) + cmask = c < num_cached + u = tl.load(lru_usage_ptr + c, mask=cmask, other=USAGE_MAX) + if POLICY == 0: + umax = tl.full([BLOCK_C], USAGE_MAX, u.dtype) + u = tl.where((u == step) | (~cmask), umax, u) + else: + oid = tl.load(id_of_slot_ptr + c, mask=cmask, other=-1) + held = cmask & (oid >= 0) + owner = tl.load(owner_ptr + c, mask=cmask, other=-1) + mirrored = held & (owner == oid) + stale = held & (owner != oid) + lk = oid // NUM_EXPERTS + ahead = lk > layer + d = tl.where(ahead, lk - layer, NUM_LAYERS - layer + lk) + last = _owner_state(slot_last_ptr, last_tok_ptr, c, oid, mirrored, stale, 0) + k = tl.minimum(tl.maximum(tl.where(ahead, tok - 1, tok) - last, 0), _K_MAX).to(tl.int32) + # int32 throughout: |score| < 2^31 by the caps on k, beta and w + score = -((k << _Q) + d * BETA_STEP) + if POLICY >= 2: + lc = _owner_state(slot_lc_ptr, lc_ptr, c, oid, mirrored, stale, _LC_NEVER) + ct = _owner_state(slot_ct_ptr, ct_ptr, c, oid, mirrored, stale, 0) + lcv = _decayed(lc, ct, now, DECAY_MUL, DT_MAX) + lcv = tl.where(lc == _LC_NEVER, _LC_FLOOR, tl.maximum(lcv, _LC_FLOOR)) + score += (W_Q4 * lcv) >> 4 + # Every key is distinct and ties go to the lowest slot; an empty slot keys below any held one. + key = tl.where(held, ((score.to(tl.int64) + _SCORE_BIAS) << SLOT_BITS) | c, c.to(tl.int64)) + evictable = cmask & (u != step) + if PIN_SINCE: + # the score ignores recency, so slots an earlier call touched since *pin_ptr stay pinned too + evictable = evictable & (u <= tl.load(pin_ptr)) + # the pinned key's slot bits are 0, so an all-pinned scan falls back to slot 0 like flashlib's argmin + pinned_key = 0x7FFFFFFFFFFFFFFF ^ ((1 << SLOT_BITS) - 1) + key = tl.where(evictable, key, pinned_key) + for i in tl.range(num_missing): + if POLICY == 0: + victim = tl.argmin(u, axis=0).to(tl.int32) + else: + victim = (tl.min(key, axis=0) & ((1 << SLOT_BITS) - 1)).to(tl.int32) + # Scalar load: victims are distinct, so no earlier iteration wrote this slot. + old = tl.load(id_of_slot_ptr + victim) + if old >= 0: + tl.store(slot_of_id_ptr + old, -1) + e = tl.sum(tl.where((rank == i) & first_miss, q, 0)) + tl.store(id_of_slot_ptr + victim, e) + tl.store(slot_of_id_ptr + e, victim) + tl.store(lru_usage_ptr + victim, step) + tl.store(dst_ptr + i, victim) + tl.store(src_ptr + i, e - id_base) # back to the caller's id space + out = tl.where((rank == i) & miss, victim, out) + if POLICY == 0: + u = tl.where(c == victim, umax, u) # claim in-register + else: + key = tl.where(c == victim, pinned_key, key) + + # Written from registers, never re-read from slot_of_id, so out_ptr may alias query_ptr. + tl.store(out_ptr + tl.arange(0, BLOCK_K), out, mask=kmask) + if POLICY != 0: + if UPDATE_STATE: + # distinct ids only, so a duplicate at bs > 1 counts once per call; out is each id's final slot + tl.store(last_tok_ptr + q, tok, mask=first) + tl.store(owner_ptr + out, q, mask=first) + tl.store(slot_last_ptr + out, tok, mask=first) + if POLICY >= 2: + lc = tl.load(lc_ptr + q, mask=first, other=_LC_NEVER) + x = _decayed(lc, tl.load(ct_ptr + q, mask=first, other=0), now, DECAY_MUL, DT_MAX) + lc_new = tl.where(lc == _LC_NEVER, 0, _softplus2(x, g_ptr, first)) + tl.store(lc_ptr + q, lc_new, mask=first) + tl.store(ct_ptr + q, now, mask=first) + tl.store(slot_lc_ptr + out, lc_new, mask=first) + tl.store(slot_ct_ptr + out, now, mask=first) + else: + # a slot filled without a state update reads the per-id state until a decode hit re-mirrors it + tl.store(owner_ptr + out, -1, mask=first_miss) + if COLLECT_STATS: + _stats(stats_ptr, first, num_missing) + + +def _num_warps_for(block_c: int) -> int: + # flashlib's rule: 8 warps from 2048 slots, more only past ~32 elements per thread. + warps = 8 if block_c >= 2048 else 4 + return max(warps, min(triton.next_power_of_2(max(block_c // 1024, 1)), 32)) + + +def scored_ensure( + query: torch.Tensor, + slot_of_id: torch.Tensor, + id_of_slot: torch.Tensor, + lru_usage: torch.Tensor, + lru_step: torch.Tensor, + out_indices: torch.Tensor, + src_indices: torch.Tensor, + dst_indices: torch.Tensor, + num_copy: torch.Tensor, + *, + policy: int, + num_layers: int, + num_experts: int, + tok: torch.Tensor | None = None, + last_tok: torch.Tensor | None = None, + lc: torch.Tensor | None = None, + ct: torch.Tensor | None = None, + g_table: torch.Tensor | None = None, + slot_owner: torch.Tensor | None = None, + slot_last_tok: torch.Tensor | None = None, + slot_lc: torch.Tensor | None = None, + slot_ct: torch.Tensor | None = None, + router_logits: torch.Tensor | None = None, + near_miss_thr: float = 0.0, + stats: torch.Tensor | None = None, + id_base: int = 0, + bump_tok: bool = False, + update_state: bool = True, + pin_since: torch.Tensor | None = None, + beta: float = BETA, + w: float = W, + halflife: int = HALFLIFE, +) -> None: + """``flashlib.lru_ensure`` (sequential strategy) with the victim picked by ``policy``. + + Same contract as ``lru_ensure``; ``policy`` > 0 also reads and updates the per-(layer, + expert) state ``last_tok``/``lc``/``ct`` (flat ``layer * num_experts + expert``) against + the device token counter ``tok``, and keeps its ``[num_cached]`` mirrors ``slot_last_tok``/ + ``slot_lc``/``slot_ct`` valid where ``slot_owner`` equals ``id_of_slot`` (-1 marks a stale + mirror). ``bump_tok`` advances the counter first (once per decode step); + ``update_state=False`` leaves the state untouched (small-prefill ensures). + ``pin_since`` (a ``lru_step`` value) also pins every slot touched after it, so a small + prefill's chunked calls cannot evict each other's experts; policy > 0 only. + Policy 3 refreshes the experts whose ``router_logits`` (``[rows, num_experts]``, rows of + ``query``) are within ``near_miss_thr`` of their row's lowest routed logit. + """ + k = query.numel() + num_cached = id_of_slot.numel() + assert query.dtype == torch.int32 and query.is_contiguous() + assert slot_of_id.dtype == torch.int32 and id_of_slot.dtype == torch.int32 + assert out_indices.dtype == torch.int32 and out_indices.numel() == k + assert lru_usage.dtype == lru_step.dtype and lru_usage.numel() == num_cached + plan = min(k, num_cached) + assert src_indices.numel() >= plan and dst_indices.numel() >= plan + block_c = triton.next_power_of_2(num_cached) + slot_bits = max(block_c.bit_length() - 1, 1) + beta_step, w_q4, decay_mul, dt_max = score_params(beta, w, halflife, num_layers) + if policy != 0: + assert slot_bits <= MAX_SLOT_BITS, f"{num_cached} slots overflow the packed score key" + assert tok is not None and last_tok is not None + assert last_tok.dtype == torch.int64 and last_tok.numel() == num_layers * num_experts + assert slot_owner is not None and slot_last_tok is not None + assert slot_owner.dtype == torch.int32 and slot_owner.numel() == num_cached + assert slot_last_tok.dtype == torch.int64 and slot_last_tok.numel() == num_cached + if pin_since is not None: + assert policy != 0 and pin_since.dtype == lru_step.dtype and pin_since.numel() == 1 + if policy >= 2: + assert lc is not None and ct is not None and g_table is not None + assert lc.dtype == torch.int32 and ct.dtype == torch.int64 + assert slot_lc is not None and slot_lc.dtype == torch.int32 and slot_lc.numel() == num_cached + assert slot_ct is not None and slot_ct.dtype == torch.int64 and slot_ct.numel() == num_cached + nm_rows = nm_topk = nm_stride = 0 + if policy == 3 and router_logits is not None: + assert router_logits.ndim == 2 and router_logits.shape[1] == num_experts + assert router_logits.stride(1) == 1 and k % router_logits.shape[0] == 0 + nm_rows, nm_stride = router_logits.shape[0], router_logits.stride(0) + nm_topk = k // nm_rows + # unused pointers still need a tensor argument + dummy = lru_step + _scored_ensure_kernel[(1,)]( + query, slot_of_id, id_of_slot, lru_usage, lru_step, + out_indices, src_indices, dst_indices, num_copy, stats, + dummy if tok is None else tok, + dummy if last_tok is None else last_tok, + dummy if lc is None else lc, + dummy if ct is None else ct, + dummy if g_table is None else g_table, + dummy if nm_rows == 0 else router_logits, + dummy if slot_owner is None else slot_owner, + dummy if slot_last_tok is None else slot_last_tok, + dummy if slot_lc is None else slot_lc, + dummy if slot_ct is None else slot_ct, + dummy if pin_since is None else pin_since, + k, num_cached, id_base, nm_topk, nm_stride, float(near_miss_thr), + BLOCK_K=triton.next_power_of_2(k), + BLOCK_C=block_c, + BLOCK_E=triton.next_power_of_2(num_experts), + BLOCK_TOPK=triton.next_power_of_2(max(nm_topk, 1)), + NM_ROWS=nm_rows, + USAGE_MAX=torch.iinfo(lru_usage.dtype).max, + COLLECT_STATS=stats is not None, + POLICY=policy, + BUMP_TOK=bool(bump_tok) and policy != 0, + UPDATE_STATE=bool(update_state) and policy != 0, + PIN_SINCE=pin_since is not None, + NUM_LAYERS=num_layers, + NUM_EXPERTS=num_experts, + BETA_STEP=beta_step, + W_Q4=w_q4, + DECAY_MUL=decay_mul, + DT_MAX=dt_max, + SLOT_BITS=slot_bits, + num_warps=_num_warps_for(block_c), + ) diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index e42dc59cd4..a87882a6dc 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -845,8 +845,13 @@ def _infer_reasoning_parser(model_path: str) -> str | None: parser.add_argument( "--moe-cache-policy", default=ServerArgs.moe_cache_policy, - choices=["lru"], - help="The unified MoE cache eviction policy.", + choices=["lru", "kd", "kdfb", "rule"], + help=( + "The unified MoE cache eviction policy. lru (default); kd: evict the expert passed " + "over for the most tokens, then the one whose layer comes up last; kdfb: kd plus a " + "decayed per-expert use count; rule: kdfb plus router near misses counted as recent " + "(margin FREETOKEN_MOE_NEAR_MISS_THR logits). GPU decode only." + ), ) parser.add_argument( diff --git a/tests/README.md b/tests/README.md index adc87035d9..d892c36dfa 100644 --- a/tests/README.md +++ b/tests/README.md @@ -65,6 +65,7 @@ checkpoint is set: | `FREETOKEN_TEST_MODELS_ROOT` | `models/test_quant_config.py` — directory whose checkpoints the stored-tensor scheme scan checks (default `/mnt/nvme/models`) | | `FREETOKEN_DSV4_VL_MODEL` | `tokenizer/test_mm_tokenize.py` - local DeepSeek-V4 vision checkpoint | | `FREETOKEN_DSV4_VISION_CKPT` | `models/test_dsv4_vision.py` - the same checkpoint, for the parity tests against its `inference/` reference code | +| `FREETOKEN_MOE_ROUTING_TRACE` | `moe/test_scored_ensure.py` — a compact decode routing trace directory (`experts.npy`, `logit_idx.npy`, `logit_val.npy`); its rows are replayed through the scored eviction kernel and the CPU reference | `test_aime.py` takes its sampling protocol from the checkpoint's own `generation_config.json` (pass@N at the recommended temperature, or a single greedy diff --git a/tests/models/qwen4_exp/test_skeleton.py b/tests/models/qwen4_exp/test_skeleton.py index 8fb9112df2..a2412f8346 100644 --- a/tests/models/qwen4_exp/test_skeleton.py +++ b/tests/models/qwen4_exp/test_skeleton.py @@ -432,6 +432,130 @@ def test_shared_gate_kernels_match_torch(num_tokens, hidden, dtype): assert (fused.float() - ref).abs().max() <= (eager.float() - ref).abs().max() + 1e-6 +def _offload_experts(kind: str, layer_id: int, num_experts: int, top_k: int, hidden: int, inter: int): + from freetoken.layers.moe import OffloadMoELayer + from freetoken.layers.quantization import NoQuantConfig, QuantBackend, QuantConfig, set_quant_backend + + if kind == "bf16": + quant = NoQuantConfig() + else: + set_quant_backend(QuantBackend.parse("moe.nvfp4=triton")) # the sm_120 decode kernel + quant = QuantConfig.from_hf({"quantization_config": {"quant_method": "modelopt", "quant_algo": "NVFP4", "ignore": ["lm_head"]}}) + return OffloadMoELayer( + layer_id, num_experts, top_k, hidden, inter, quant_config=quant, + prefix=f"model.layers.{layer_id}.mlp.experts", + ) + + +def _copy_overlap_run(kind, banks, overlap, bs, graph, inputs, monkeypatch): + """Two Qwen4ExpMoE layers sharing one offload cache as small as a layer, so every step evicts + and copies: warm up eagerly (and capture), reset the cache, then decode ``inputs``.""" + import freetoken.kernel.fast_index_copy as fic + from flashlib.kernels.slot_cache import Stat + from freetoken.models.qwen4_exp.moe import Qwen4ExpMoE + from freetoken.moe.offload_cache import OffloadMoeCache + from freetoken.utils.torch_utils import torch_dtype + + num_experts, top_k, hidden, inter = 8, 2, 256, 128 + config = parse_config(toy_hf_config( + hidden_size=hidden, num_experts=num_experts, num_experts_per_tok=top_k, + moe_intermediate_size=inter, shared_expert_intermediate_size=inter, + )) + device = torch.device("cuda") + with torch.device(device), torch_dtype(torch.bfloat16): + moes = [Qwen4ExpMoE(config, layer_id) for layer_id in range(2)] + gen = torch.Generator(device=device).manual_seed(41) + for moe in moes: + _fill(moe, gen, scale=0.2) + experts = [_offload_experts(kind, i, num_experts, top_k, hidden, inter) for i in range(2)] + cache = OffloadMoeCache( + num_layers=2, num_experts=num_experts, cache_size=num_experts, device=device, + quant_format=banks.quant_format, layout=banks.layout, + max_slots=experts[0].quant_method.slot_limit(), decode_copy_overlap=overlap, + ) + cache.set_bank_sources(banks.sources) + cache.set_alphas(banks.gate_up_alpha, banks.down_alpha) + cache.collect_stats = True + for moe, ex in zip(moes, experts): + ex.offload_cache = cache + moe.experts = ex + cache.reset() + + import freetoken.models.qwen4_exp.moe as qmoe + + copy_streams, gate_streams = [], [] + real_gate = qmoe.shared_gate_sigmoid + for name in ("fast_index_copy_multi_jit", "fast_index_copy_multi_slim_jit"): # FREETOKEN_MOE_SLIM_COPY + real_copy = getattr(fic, name) + monkeypatch.setattr( + fic, name, + lambda *a, _real=real_copy, **k: (copy_streams.append(torch.cuda.current_stream().cuda_stream), _real(*a, **k))[1], + ) + monkeypatch.setattr( + qmoe, "shared_gate_sigmoid", + lambda *a, **k: (gate_streams.append(torch.cuda.current_stream().cuda_stream), real_gate(*a, **k))[1], + ) + + def step(x0, x1): + return moes[0].forward(x0.clone()), moes[1].forward(x1.clone()) + + _fresh_ctx(_batch=SimpleNamespace(is_prefill=False)) + static = [torch.randn(bs, hidden, device=device, dtype=torch.bfloat16) for _ in range(2)] + step(*static) # the eager warm-up the graph runner also does + if graph: + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + outs = step(*static) + cache.reset() + cache.reset_stats() + results = [] + for x0, x1 in inputs: + if graph: + static[0].copy_(x0) + static[1].copy_(x1) + g.replay() + else: + outs = step(x0, x1) + results.append(tuple(o.clone() for o in outs)) + torch.cuda.synchronize() + state = [t.clone() for t in (cache.slot_for_id, cache.id_of_slot, cache.usage, *cache.bank_views())] + misses = int(cache.lru_stats[:, Stat.MISS].sum()) + side = getattr(cache.decode_copy_stream, "cuda_stream", None) + return results, state, misses, [s == side for s in copy_streams], [s == side for s in gate_streams] + + +@requires_cuda +@pytest.mark.parametrize("kind", ["bf16", "nvfp4"]) +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +def test_moe_copy_overlap_is_bitwise_identical(kind, graph, bs, monkeypatch): + """FREETOKEN_MOE_COPY_OVERLAP moves the shared expert onto a side stream beside the decode miss + copy; the MoE outputs, the slot map and the slot contents must not change by a single bit.""" + from freetoken.moe.expert_banks import build_expert_banks + + torch.manual_seed(7) + method = _offload_experts(kind, 0, 8, 2, 256, 128).quant_method + banks = build_expert_banks(method, 2, None, device=torch.device("cuda"), dummy=True) + gen = torch.Generator(device="cuda").manual_seed(bs) + inputs = [ + tuple(torch.randn(bs, 256, device="cuda", dtype=torch.bfloat16, generator=gen) * 0.5 for _ in range(2)) + for _ in range(6) + ] + off = _copy_overlap_run(kind, banks, False, bs, graph, inputs, monkeypatch) + on = _copy_overlap_run(kind, banks, True, bs, graph, inputs, monkeypatch) + + assert off[2] > 0 and on[2] == off[2], "the sequence must miss, and copy, in both runs" + assert off[3] and not any(off[3]) and not any(off[4]), "flag off: everything stays on the compute stream" + assert on[3] and not any(on[3]), "flag on: the miss copy stays right behind its ensure" + assert on[4] and all(on[4]), "flag on: the shared expert runs on the cache's side stream" + for got, want in zip(on[0], off[0]): + for a, b in zip(got, want): + assert torch.isfinite(b.float()).all() + assert torch.equal(a, b) + for a, b in zip(on[1], off[1]): + assert torch.equal(a, b) + + @requires_cuda def test_decoder_stack_prefill_and_decode(monkeypatch): """Ragged bs=3 prefill then a bs=3 decode step through the whole model with dummy weights.""" diff --git a/tests/moe/ref_scored_cache.py b/tests/moe/ref_scored_cache.py new file mode 100644 index 0000000000..d7d8b0f2fb --- /dev/null +++ b/tests/moe/ref_scored_cache.py @@ -0,0 +1,159 @@ +"""CPU reference for the scored slot-cache ensure (freetoken.moe.scored_ensure). + +A port of the adversarial-synthesis replay (global slot cache with lru_ensure semantics: the +call's ids pinned, empty slots first, ties to the lowest slot, missing ids ascending paired with +ascending victims) onto the kernel's Q16 integer arithmetic, so victims match the GPU exactly. +""" +from __future__ import annotations + +import numpy as np + +from freetoken.moe import scored_ensure as se + +_KEY_MAX = np.iinfo(np.int64).max + + +class RefScoredCache: + def __init__(self, num_layers, num_experts, cache_size, policy, + beta=se.BETA, w=se.W, halflife=se.HALFLIFE): + self.L, self.E, self.S, self.policy = num_layers, num_experts, cache_size, policy + self.beta_step, self.w_q4, self.decay_mul, self.dt_max = se.score_params(beta, w, halflife, num_layers) + block_c = 1 << max(cache_size - 1, 1).bit_length() + self.slot_bits = max(block_c.bit_length() - 1, 1) + self.g = se.softplus2_table().numpy().astype(np.int64) + self.reset() + + def reset(self): + n = self.L * self.E + self.slot_of_id = np.full(n, -1, np.int64) + self.id_of_slot = np.full(self.S, -1, np.int64) + self.usage = np.zeros(self.S, np.int64) + self.step = 0 + self.tok = 0 + self.last_tok = np.full(n, se.LAST_TOK_NEVER, np.int64) + self.lc = np.full(n, se.LC_NEVER, np.int64) + self.ct = np.zeros(n, np.int64) + + def _decayed(self, ids, now): + dt = np.minimum(np.maximum(now - self.ct[ids], 0), self.dt_max) + return self.lc[ids] - ((dt * self.decay_mul) >> se.DECAY_SHIFT) + + def _softplus2(self, x): + shift = se.Q - se.G_STEP_BITS + u = np.abs(x) + i = np.minimum(u >> shift, se.G_LAST) + frac = u & ((1 << shift) - 1) + g0, g1 = self.g[i], self.g[i + 1] + g = np.where(i < se.G_LAST, g0 + (((g1 - g0) * frac) >> shift), 0) + return np.maximum(x, 0) + g + + def _keys(self, layer, tok, now, pin_since=None): + c = np.arange(self.S, dtype=np.int64) + oid = self.id_of_slot + held = oid >= 0 + if self.policy == 0: + key = (self.usage << self.slot_bits) | c + else: + ids = np.where(held, oid, 0) + lk = ids // self.E + ahead = lk > layer + d = np.where(ahead, lk - layer, self.L - layer + lk) + k = np.where(ahead, tok - 1, tok) - self.last_tok[ids] + k = np.minimum(np.maximum(k, 0), se.K_MAX) + score = -((k << se.Q) + d * self.beta_step) + if self.policy >= 2: + lc = self.lc[ids] + lcv = self._decayed(ids, now) + lcv = np.where(lc == se.LC_NEVER, se.LC_FLOOR, np.maximum(lcv, se.LC_FLOOR)) + score = score + ((self.w_q4 * lcv) >> 4) + assert np.abs(score).max() < 2**31 # the kernel computes it in int32 + key = np.where(held, ((score + se.SCORE_BIAS) << self.slot_bits) | c, c) + evictable = self.usage != self.step + if pin_since is not None: + evictable &= self.usage <= pin_since + return np.where(evictable, key, _KEY_MAX) + + def ensure(self, layer, ids, *, bump_tok=False, update_state=True, logits=None, thr=0.0, near_miss=None, + pin_since=None): + """One ensure call; returns ``(out_slots, src_ids, dst_slots)`` like the kernel's plan. + + ``pin_since`` (a ``step`` value) also pins every slot touched after it. + + Policy 3 refreshes the ids whose ``logits`` (``[rows, E]``) are within ``thr`` of their + row's lowest routed logit, in fp32 as the kernel does, or the ids of a given ``near_miss`` mask. + """ + self.step += 1 + base = layer * self.E + tok = now = 0 + if self.policy != 0: + if bump_tok: + self.tok += 1 + tok = self.tok + now = tok * self.L + layer + if self.policy == 3 and update_state and logits is not None: + lg = np.asarray(logits, dtype=np.float32).reshape(-1, self.E) + routed = np.asarray(ids, dtype=np.int64).reshape(lg.shape[0], -1) + kth = lg[np.arange(lg.shape[0])[:, None], routed].min(axis=1) + near_miss = lg >= (kth - np.float32(thr))[:, None] + if self.policy == 3 and update_state and near_miss is not None: + near = base + np.flatnonzero(np.asarray(near_miss).reshape(-1, self.E).any(0)) + self.last_tok[near] = np.maximum(self.last_tok[near], tok - 1) + q = np.asarray(ids, dtype=np.int64).reshape(-1) + base + slots = self.slot_of_id[q] + self.usage[slots[slots >= 0]] = self.step + uniq = np.unique(q) + missing = uniq[self.slot_of_id[uniq] < 0] # ascending, as the kernel ranks them + src = missing - base + dst = np.empty(missing.size, np.int64) + if missing.size: + keys = self._keys(layer, tok, now, pin_since) + part = np.argpartition(keys, missing.size - 1)[: missing.size] + victims = part[np.argsort(keys[part])] + for i, (e, v) in enumerate(zip(missing, victims)): + old = self.id_of_slot[v] + if old >= 0: + self.slot_of_id[old] = -1 + self.id_of_slot[v] = e + self.slot_of_id[e] = v + self.usage[v] = self.step + dst[i] = v + out = self.slot_of_id[q] + if self.policy != 0 and update_state: + self.last_tok[uniq] = tok + if self.policy >= 2: + lc = self.lc[uniq] + x = self._decayed(uniq, now) + self.lc[uniq] = np.where(lc == se.LC_NEVER, 0, self._softplus2(x)) + self.ct[uniq] = now + return out, src, dst + + + def materialize(self, layer): + """The prefill materialize kernel: the whole layer into slots [0, E), position == expert id.""" + base = layer * self.E + same = (self.id_of_slot >= base) & (self.id_of_slot < base + self.E) + self.id_of_slot[same] = -1 + self.usage[same] = 0 + old = self.id_of_slot[: self.E].copy() + self.slot_of_id[old[old >= 0]] = -1 + self.step += 1 + self.id_of_slot[: self.E] = base + np.arange(self.E) + self.slot_of_id[base : base + self.E] = np.arange(self.E) + self.usage[: self.E] = self.step + + +def replay(cache: RefScoredCache, rows, logits=None, thr=0.0, near_miss=None): + """Decode ``rows`` ``[N, L, K]`` (one step per row, all layers); per-row miss counts. + + ``rows`` may be ``[N, L, B, K]`` for a batch of B rows per step, one K*B call per layer. + Policy 3 takes the matching ``[N, L, B, E]`` router ``logits`` (or a ``near_miss`` mask). + """ + rows = np.asarray(rows) + misses = np.zeros(rows.shape[0], np.int64) + for r in range(rows.shape[0]): + for layer in range(rows.shape[1]): + lg = None if logits is None else logits[r, layer] + nm = None if near_miss is None else near_miss[r, layer] + _, src, _ = cache.ensure(layer, rows[r, layer], bump_tok=layer == 0, logits=lg, thr=thr, near_miss=nm) + misses[r] += src.size + return misses diff --git a/tests/moe/test_fused_copy.py b/tests/moe/test_fused_copy.py index 7a0c086444..c5336ea117 100644 --- a/tests/moe/test_fused_copy.py +++ b/tests/moe/test_fused_copy.py @@ -1,6 +1,8 @@ """The fused multi-bank ``copy_missing`` path must move exactly the same bytes as the legacy per-bank ``fast_index_copy_jit`` loop, for every miss count (including the -zero-copy case), across banks of differing per-row sizes. +zero-copy case), across banks of differing per-row sizes. The opt-in slim kernel +(``FREETOKEN_MOE_SLIM_COPY``) must move exactly the same bytes as the fused one, also when +replayed from a CUDA graph with a device-side plan. """ from __future__ import annotations @@ -8,6 +10,7 @@ import pytest import torch +from freetoken.moe import offload_cache from freetoken.moe.offload_cache import _BANK_SCHEMAS, OffloadMoeCache CUDA = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") @@ -17,22 +20,30 @@ # kernel/aot_models.py must be listed in its TEST_FEATURE_SIZES so the per-bank kernels # stay prebuilt under FREETOKEN_DISABLE_JIT=1. FEATS = [8192, 512, 256, 4096, 512, 256] +# qwen4_exp (Qwen3.8-Flash-Next) triton-NVFP4 rows in _BANK_SCHEMAS["nvfp4"] order: gate_up +# packed/scale/global, down packed/scale/global -- 2,772,480 B per expert. +NVFP4_FEATS = [1638400, 204800, 2560, 819200, 102400, 5120] +# 16 * odd, several rows shorter than one grid pass, so segments end mid-round and the +# stride phase wraps across banks and rows. +ODD_FEATS = [16 * 1031, 16 * 7, 16 * 250, 16, 16 * 4099, 16 * 33] -def _build_cache(num_layers, num_experts, cache_size): +def _build_cache(num_layers, num_experts, cache_size, quant_format="mxfp4_triton", feats=FEATS, pinned=False): dev = torch.device("cuda") cache = OffloadMoeCache( num_layers=num_layers, num_experts=num_experts, cache_size=cache_size, - device=dev, cache_policy="lru", prefill_overlap=False, quant_format="mxfp4_triton", + device=dev, cache_policy="lru", prefill_overlap=False, quant_format=quant_format, ) - schema = _BANK_SCHEMAS["mxfp4_triton"] + schema = _BANK_SCHEMAS[quant_format] # Views into one flat tensor: only per-layer addressing matters here, not # independent allocations. - sources = { - name: list(torch.randint(0, 256, (num_layers * num_experts, feat), dtype=torch.uint8, device=dev) - .split(num_experts)) - for name, feat in zip(schema, FEATS) - } + sources = {} + for name, feat in zip(schema, feats): + if pinned: + flat = torch.randint(0, 256, (num_layers * num_experts, feat), dtype=torch.uint8).pin_memory() + else: + flat = torch.randint(0, 256, (num_layers * num_experts, feat), dtype=torch.uint8, device=dev) + sources[name] = list(flat.split(num_experts)) cache.set_bank_sources(sources) # also builds the fused-copy descriptor return cache @@ -74,3 +85,147 @@ def test_fused_copy_matches_per_bank(num_indices): for b, (r, (_, c)) in enumerate(zip(ref, cache.banks)): assert torch.equal(r, c), f"bank {b} (feat={FEATS[b]}) fused != per-bank at num_indices={num_indices}" + + +def _host_banks(feats, n_src, seed): + from freetoken.kernel.pinned import device_ptr + + gen = torch.Generator().manual_seed(seed) + src = [torch.randint(0, 256, (n_src, f), dtype=torch.uint8, generator=gen).pin_memory() for f in feats] + dev = torch.device("cuda") + src_ptrs = torch.tensor([device_ptr(t) for t in src], dtype=torch.int64, device=dev) + feat_bytes = torch.tensor(feats, dtype=torch.int64, device=dev) + return src, src_ptrs, feat_bytes + + +def _slot_banks(feats, n_slots): + dst = [torch.full((n_slots, f), 0xAB, dtype=torch.uint8, device="cuda") for f in feats] + return dst, torch.tensor([t.data_ptr() for t in dst], dtype=torch.int64, device="cuda") + + +def _expected(src, feats, n_slots, di, si, num): + # untouched slots keep the 0xAB fill; plan rows land at their dst slots + out = [] + for b, f in enumerate(feats): + e = torch.full((n_slots, f), 0xAB, dtype=torch.uint8) + e[di[:num].long().cpu()] = src[b][si[:num].long().cpu()] + out.append(e.cuda()) + return out + + +def _random_plan(gen, n_src, n_slots, plan, dtype): + di = torch.randperm(n_slots, generator=gen)[:plan].to(dtype).cuda() # distinct: no write races + si = torch.randint(0, n_src, (plan,), generator=gen).to(dtype).cuda() + return di, si + + +def _check_slim_matches_multi(feats, idx_dtype, **grid): + from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit, fast_index_copy_multi_slim_jit + + n_src, n_slots, plan = 24, 16, 16 + src, src_ptrs, feat_bytes = _host_banks(feats, n_src, seed=0) + gen = torch.Generator().manual_seed(1) + for num in [0, 1, 3, 10, plan, None]: # None: no num_indices, the whole index length + di, si = _random_plan(gen, n_src, n_slots, plan, idx_dtype) + count = None if num is None else torch.tensor([num], dtype=torch.int64, device="cuda") + ref, ref_ptrs = _slot_banks(feats, n_slots) + fast_index_copy_multi_jit(ref_ptrs, src_ptrs, feat_bytes, di, si, count) + out, out_ptrs = _slot_banks(feats, n_slots) + fast_index_copy_multi_slim_jit(out_ptrs, src_ptrs, feat_bytes, di, si, count, **grid) + torch.cuda.synchronize() + exp = _expected(src, feats, n_slots, di, si, plan if num is None else num) + for b, f in enumerate(feats): + assert torch.equal(ref[b], exp[b]), f"multi reference wrong: bank {b} (feat={f}) num={num}" + assert torch.equal(out[b], ref[b]), f"slim != multi: bank {b} (feat={f}) num={num} grid={grid}" + + +@CUDA +@pytest.mark.parametrize("idx_dtype", [torch.int32, torch.int64]) +@pytest.mark.parametrize("feats", [NVFP4_FEATS, ODD_FEATS], ids=["nvfp4", "odd"]) +def test_slim_copy_matches_multi(feats, idx_dtype): + _check_slim_matches_multi(feats, idx_dtype) + + +@CUDA +@pytest.mark.slow +@pytest.mark.parametrize("grid", [(3, 128, 2), (16, 256, 2), (5, 64, 8)]) +@pytest.mark.parametrize("feats", [NVFP4_FEATS, ODD_FEATS], ids=["nvfp4", "odd"]) +def test_slim_copy_matches_multi_other_grids(feats, grid): + blocks, threads, unroll = grid + _check_slim_matches_multi(feats, torch.int32, num_blocks=blocks, num_threads=threads, unroll=unroll) + + +@CUDA +def test_slim_copy_graph_replay_follows_device_plan(): + from freetoken.kernel.fast_index_copy import fast_index_copy_multi_slim_jit + + feats, n_src, n_slots, plan = NVFP4_FEATS, 24, 16, 16 + src, src_ptrs, feat_bytes = _host_banks(feats, n_src, seed=2) + dst, dst_ptrs = _slot_banks(feats, n_slots) + di = torch.zeros(plan, dtype=torch.int32, device="cuda") + si = torch.zeros(plan, dtype=torch.int32, device="cuda") + count = torch.zeros(1, dtype=torch.int64, device="cuda") + fast_index_copy_multi_slim_jit(dst_ptrs, src_ptrs, feat_bytes, di, si, count) # JIT outside capture + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + fast_index_copy_multi_slim_jit(dst_ptrs, src_ptrs, feat_bytes, di, si, count) + + gen = torch.Generator().manual_seed(3) + for num in [4, 0, plan, 1, 10]: + new_di, new_si = _random_plan(gen, n_src, n_slots, plan, torch.int32) + di.copy_(new_di) + si.copy_(new_si) + count.fill_(num) + for d in dst: + d.fill_(0xAB) + graph.replay() + torch.cuda.synchronize() + exp = _expected(src, feats, n_slots, di, si, num) + for b, f in enumerate(feats): + assert torch.equal(dst[b], exp[b]), f"graph replay: bank {b} (feat={f}) num={num}" + + +def test_slim_copy_flag_defaults_off(monkeypatch): + import importlib.util + + monkeypatch.delenv("FREETOKEN_MOE_SLIM_COPY", raising=False) + spec = importlib.util.find_spec("freetoken.moe.offload_cache") + fresh = importlib.util.module_from_spec(spec) + spec.loader.exec_module(fresh) # a private copy: the imported module and its classes stay untouched + assert fresh._SLIM_COPY is False + monkeypatch.setenv("FREETOKEN_MOE_SLIM_COPY", "1") + spec.loader.exec_module(fresh) + assert fresh._SLIM_COPY is True + + +@CUDA +@pytest.mark.parametrize("slim", [False, True], ids=["multi", "slim"]) +def test_copy_missing_slim_switch(monkeypatch, slim): + from freetoken.kernel import fast_index_copy as fic + + calls = [] + for name in ("fast_index_copy_multi_jit", "fast_index_copy_multi_slim_jit"): + real = getattr(fic, name) + monkeypatch.setattr(fic, name, lambda *a, _n=name, _r=real, **k: (calls.append(_n), _r(*a, **k))) + monkeypatch.setattr(offload_cache, "_SLIM_COPY", slim) + + num_layers, num_experts, cache_size, layer_id, num = 2, 8, 16, 1, 5 + cache = _build_cache(num_layers, num_experts, cache_size, "nvfp4", NVFP4_FEATS, pinned=True) + assert cache._copy_fused_ok + cache._pending_src_layer = layer_id + dev = torch.device("cuda") + cache.num_indices.fill_(num) + cache.evict_slots[:num] = torch.tensor([9, 2, 15, 0, 7], dtype=torch.int32, device=dev) + cache.src_indices[:num] = torch.tensor([3, 3, 0, 7, 5], dtype=torch.int32, device=dev) + for _, c in cache.banks: + c.fill_(0xAB) + cache.copy_missing() + torch.cuda.synchronize() + + assert calls == ["fast_index_copy_multi_slim_jit" if slim else "fast_index_copy_multi_jit"] + slots, rows = cache.evict_slots[:num].long().cpu(), cache.src_indices[:num].long().cpu() + for b, (per_layer, c) in enumerate(cache.banks): + exp = torch.full(c.shape, 0xAB, dtype=torch.uint8) + exp[slots] = per_layer[layer_id][rows] + assert torch.equal(c.cpu(), exp), f"bank {b} (feat={NVFP4_FEATS[b]})" diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index ca8a2aa4d0..5dc7dff602 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -387,7 +387,7 @@ def test_offload_moe_layer_decode_forward_uses_remapped_slot_ids(monkeypatch): lambda *, hidden_states, gating_output, topk, renormalize: (topk_weights, topk_ids), ) - def fake_ensure(layer_id, expert_ids): + def fake_ensure(layer_id, expert_ids, **kw): calls["ensure_layer_id"] = layer_id calls["ensure_expert_ids"] = expert_ids.clone() expert_ids.copy_(torch.tensor([[5, 0]], dtype=torch.int32)) @@ -427,6 +427,121 @@ def fake_fused_decode( assert calls["topk_ids"].tolist() == [[5, 0]] +def test_decode_copy_overlap_follows_the_flag_and_the_decode_target(monkeypatch): + from freetoken.env import ENV + from freetoken.moe.offload_cache import OffloadMoeCache + + def cache(device, **kw): + return OffloadMoeCache(num_layers=3, num_experts=4, cache_size=8, device=torch.device(device), **kw) + + monkeypatch.setattr(ENV.MOE_COPY_OVERLAP, "value", True) + assert cache("cpu").decode_copy_overlap is True # unset -> FREETOKEN_MOE_COPY_OVERLAP + assert cache("cpu").decode_copy_stream is None + monkeypatch.setattr(ENV.MOE_COPY_OVERLAP, "value", False) + assert cache("cpu").decode_copy_overlap is False + if not torch.cuda.is_available(): + return + assert cache("cuda").decode_copy_stream is None + for target in ("cpu", "hybrid"): + assert cache("cuda", decode_copy_overlap=True, decode_target=target).decode_copy_stream is None + on = cache("cuda", decode_copy_overlap=True) + assert on.decode_copy_stream is not None and len(on._decode_copy_events) == 3 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_copy_missing_takes_an_explicit_layer_id(): + # a second copy path (the overlap fork, later prefetch) names its source layer instead of trusting whichever call staged last + from freetoken.moe.offload_cache import OffloadMoeCache + + num_layers, num_experts = 2, 4 + fingerprint = torch.arange(num_layers * num_experts, dtype=torch.float32).view(num_layers, num_experts, 1, 1) + sources = { + "gate_up": [fingerprint[l].expand(num_experts, 32, 8).contiguous().pin_memory() for l in range(num_layers)], + "down": [fingerprint[l].expand(num_experts, 8, 16).contiguous().pin_memory() for l in range(num_layers)], + } + cache = OffloadMoeCache( + num_layers=num_layers, num_experts=num_experts, cache_size=num_experts, device=torch.device("cuda") + ) + cache.set_bank_sources(sources) + cache.reset() + ids = torch.tensor([2, 3], dtype=torch.int32, device="cuda") + cache.ensure_experts(1, ids) + cache._pending_src_layer = 0 + cache.copy_missing(1) + torch.cuda.synchronize() + for bank in cache.bank_views(): + assert [int(bank[s].flatten()[0]) for s in ids.tolist()] == [num_experts + 2, num_experts + 3] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_decode_side_applies_only_to_gpu_decode(monkeypatch): + # prefill, CPU-routed layers, TP > 1 and the flag off keep the shared expert on the compute stream + from types import SimpleNamespace + + import freetoken.core as core + from freetoken.core import Context + from freetoken.moe.offload_cache import OffloadMoeCache + + _init_tp() + layer = _bf16_offload_layer(0, 4, 2, 8, 16) + ctx = Context(page_size=1) + monkeypatch.setattr(core, "_GLOBAL_CTX", ctx) + + def applies(is_prefill=False, cpu_layers=(), **cache_kw): + layer.offload_cache = OffloadMoeCache( + num_layers=1, num_experts=4, cache_size=4, device=torch.device("cuda"), **cache_kw + ) + layer.offload_cache.cpu_layer_ids = frozenset(cpu_layers) + ctx._batch = SimpleNamespace(is_prefill=is_prefill) + return layer.decode_side_applies() + + assert not applies(decode_copy_overlap=False) + assert not applies(decode_copy_overlap=True, decode_target="hybrid") + assert not applies(is_prefill=True, decode_copy_overlap=True) + assert not applies(cpu_layers=(0,), decode_copy_overlap=True) + monkeypatch.setattr(layer, "tp_size", 2) + assert not applies(decode_copy_overlap=True) + monkeypatch.setattr(layer, "tp_size", 1) + assert applies(decode_copy_overlap=True) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_decode_side_stream_never_aliases_a_pooled_stream(): + # torch.cuda.Stream() hands out 32 pooled streams per priority round robin; a pooled side stream + # can come back as the engine or capture stream, and then the side work silently runs serially + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache( + num_layers=1, num_experts=4, cache_size=4, device=torch.device("cuda"), decode_copy_overlap=True + ) + side = cache.decode_copy_stream.cuda_stream + pooled = {torch.cuda.Stream().cuda_stream for _ in range(64)} + assert side not in pooled and side != torch.cuda.current_stream().cuda_stream + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_decode_side_survives_a_rebuild(): + # the side stream and events are not cache_size-shaped, so a runtime rebuild keeps them working + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache( + num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"), decode_copy_overlap=True + ) + cache.set_bank_sources({ + "gate_up": [torch.zeros(4, 32, 8, dtype=torch.bfloat16).pin_memory() for _ in range(2)], + "down": [torch.zeros(4, 8, 16, dtype=torch.bfloat16).pin_memory() for _ in range(2)], + }) + stream = cache.decode_copy_stream + cache.rebuild(6) + assert cache.decode_copy_stream is stream is not None + x = torch.arange(1024, dtype=torch.float32, device="cuda") + for layer_id in range(2): + with cache.decode_side(layer_id): + assert torch.cuda.current_stream() == stream + y = x * 2 + cache.join_decode_side(layer_id) + assert torch.equal(y + 0, x * 2) + def test_lru_gpu_cache_assigns_unique_slots_for_large_miss_batch(): import pytest diff --git a/tests/moe/test_scored_ensure.py b/tests/moe/test_scored_ensure.py new file mode 100644 index 0000000000..db734180b6 --- /dev/null +++ b/tests/moe/test_scored_ensure.py @@ -0,0 +1,622 @@ +"""--moe-cache-policy kd / kdfb / rule: the vendored scored-victim ensure kernel. + +The kernel must reproduce flashlib's LRU bit for bit under POLICY=lru, and match the CPU +reference (ref_scored_cache.py, a port of the eviction replay) victim for victim under the +scored policies. Set FREETOKEN_MOE_ROUTING_TRACE to a compact routing trace directory +(experts.npy, logit_idx.npy, logit_val.npy) to also replay real decode rows. +""" +import os +from pathlib import Path + +import numpy as np +import pytest +import torch + +L, E, S, TOP_K = 48, 512, 1650, 10 +NEAR_MISS_THR = 0.25 +_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +_TRACE = os.environ.get("FREETOKEN_MOE_ROUTING_TRACE", "") + + +def _trace(): + if not _TRACE or not (Path(_TRACE) / "experts.npy").exists(): + pytest.skip("FREETOKEN_MOE_ROUTING_TRACE is not set to a routing trace directory") + load = lambda name: np.load(Path(_TRACE) / name, mmap_mode="r") # noqa: E731 + return load("experts.npy"), load("logit_idx.npy"), load("logit_val.npy") + + +def _trace_logits(li, lv): + """Router logits rebuilt from the trace's recorded top-32 (the rest far below any margin).""" + li, lv = np.asarray(li).astype(np.int64), np.asarray(lv).astype(np.float32) + logits = np.full(li.shape[:-1] + (E,), -1e4, np.float32) + np.put_along_axis(logits, li, lv, axis=-1) + return logits + + +def _zipf_rows(n, rows_per_step=1, seed=0, num_layers=L, num_experts=E, top_k=TOP_K, skew=0.9): + """Skewed decode routing ``[n, L, rows_per_step, top_k]``: distinct ids within a row, duplicates across rows.""" + rng = np.random.default_rng(seed) + p = 1.0 / np.arange(1, num_experts + 1) ** skew + perm = [rng.permutation(num_experts) for _ in range(num_layers)] + out = np.empty((n, num_layers, rows_per_step, top_k), np.int64) + for r in range(n): + for layer in range(num_layers): + for b in range(rows_per_step): + pick = rng.choice(num_experts, top_k, replace=False, p=p / p.sum()) + out[r, layer, b] = perm[layer][pick] + return out + + +def _random_logits(rows, seed=1, num_experts=E): + """bf16-exact router logits shaped like ``rows`` with the routed ids on top (~10% near misses at 0.25).""" + rng = np.random.default_rng(seed) + logits = rng.normal(size=rows.shape[:-1] + (num_experts,)).astype(np.float32) + np.put_along_axis(logits, rows, np.take_along_axis(logits, rows, axis=-1) + 3.0, axis=-1) + return torch.from_numpy(logits).bfloat16().float().numpy() + + +def _cache(policy, cache_size=S, num_layers=L, num_experts=E, device="cuda", **kw): + from freetoken.moe.offload_cache import OffloadMoeCache + + return OffloadMoeCache( + num_layers=num_layers, num_experts=num_experts, cache_size=cache_size, + device=torch.device(device), cache_policy=policy, **kw, + ) + + +def _ref(policy, cache_size=S, num_layers=L, num_experts=E): + from freetoken.moe.scored_ensure import POLICY_IDS + + from .ref_scored_cache import RefScoredCache + + return RefScoredCache(num_layers, num_experts, cache_size, POLICY_IDS[policy]) + + +def _run_gpu(cache, rows, logits=None, dtype=torch.bfloat16): + """Decode ``rows`` ``[N, L, B, K]`` through ``cache.ensure_experts``; per-row miss counts.""" + dev = cache.device + ids = torch.from_numpy(np.ascontiguousarray(rows, dtype=np.int32)).to(dev) + lg = None if logits is None else torch.from_numpy(np.ascontiguousarray(logits)).to(dev, dtype) + cache.collect_stats = True + cache.lru_stats.zero_() + marks = torch.zeros(rows.shape[0], dtype=torch.int64, device=dev) + for r in range(rows.shape[0]): + for layer in range(rows.shape[1]): + cache.ensure_experts(layer, ids[r, layer], router_logits=None if lg is None else lg[r, layer]) + marks[r] = cache.lru_stats[:, 1].sum() + return np.diff(marks.cpu().numpy(), prepend=0) + + +def _run_ref(ref, rows, logits=None): + from .ref_scored_cache import replay + + return replay(ref, rows, logits, NEAR_MISS_THR) + + +def _assert_same_tables(cache, ref): + got = lambda t: t.detach().cpu().numpy().astype(np.int64) # noqa: E731 + np.testing.assert_array_equal(got(cache.slot_for_id.view(-1)), ref.slot_of_id) + np.testing.assert_array_equal(got(cache.id_of_slot), ref.id_of_slot) + np.testing.assert_array_equal(got(cache.usage), ref.usage) + assert int(cache.step) == ref.step + if ref.policy: + assert int(cache.evict_tok) == ref.tok + np.testing.assert_array_equal(got(cache.evict_last_tok), ref.last_tok) + if ref.policy >= 2: + np.testing.assert_array_equal(got(cache.evict_lc), ref.lc) + np.testing.assert_array_equal(got(cache.evict_ct), ref.ct) + if ref.policy: + _assert_mirrors_hold(cache) + + +def _assert_mirrors_hold(cache): + """Every slot whose owner tag matches its id mirrors that id's state.""" + got = lambda t: t.detach().cpu().numpy().astype(np.int64) # noqa: E731 + ids, owner = got(cache.id_of_slot), got(cache.evict_slot_owner) + tagged = (ids >= 0) & (owner == ids) + np.testing.assert_array_equal(got(cache.evict_slot_last_tok)[tagged], got(cache.evict_last_tok)[ids[tagged]]) + if cache.evict_lc is not None: + np.testing.assert_array_equal(got(cache.evict_slot_lc)[tagged], got(cache.evict_lc)[ids[tagged]]) + np.testing.assert_array_equal(got(cache.evict_slot_ct)[tagged], got(cache.evict_ct)[ids[tagged]]) + return tagged + + +def _check_against_ref(policy, rows, logits=None, dtype=torch.bfloat16): + cache, ref = _cache(policy), _ref(policy) + assert policy != "rule" or cache.near_miss_thr == NEAR_MISS_THR + got = _run_gpu(cache, rows, logits, dtype) + want = _run_ref(ref, rows, logits) + np.testing.assert_array_equal(got, want) + _assert_same_tables(cache, ref) + return got + + +# (a) POLICY=lru is flashlib's kernel ------------------------------------------------------- + + +@_cuda +@pytest.mark.parametrize("rows_per_step", [1, 2]) +def test_vendored_lru_is_bit_identical_to_flashlib(rows_per_step): + from flashlib.kernels.slot_cache import N_STATS, lru_ensure + + from freetoken.moe.scored_ensure import scored_ensure + + dev = torch.device("cuda") + k = rows_per_step * TOP_K + # skewed enough that some calls hit every id + stream = _zipf_rows(63, rows_per_step, seed=rows_per_step, skew=1.6).reshape(-1, k)[:3000] + layers = np.arange(stream.shape[0]) % L + + def state(): + return dict( + slot_of_id=torch.full((L * E,), -1, dtype=torch.int32, device=dev), + id_of_slot=torch.full((S,), -1, dtype=torch.int32, device=dev), + usage=torch.zeros(S, dtype=torch.int64, device=dev), + step=torch.zeros((), dtype=torch.int64, device=dev), + src=torch.zeros(S, dtype=torch.int32, device=dev), + dst=torch.zeros(S, dtype=torch.int32, device=dev), + num=torch.zeros((), dtype=torch.int64, device=dev), + stats=torch.zeros(N_STATS, dtype=torch.int64, device=dev), + ) + + a, b = state(), state() + queries = torch.from_numpy(stream.astype(np.int32)).to(dev) + qa, qb = queries.clone(), queries.clone() + plans = ([], []) + for i in range(stream.shape[0]): + base = int(layers[i]) * E + lru_ensure(qa[i], a["slot_of_id"], a["id_of_slot"], a["usage"], a["step"], qa[i], + a["src"], a["dst"], a["num"], stats=a["stats"], id_base=base) + scored_ensure(qb[i], b["slot_of_id"], b["id_of_slot"], b["usage"], b["step"], qb[i], + b["src"], b["dst"], b["num"], policy=0, num_layers=L, num_experts=E, + stats=b["stats"], id_base=base) + for st, plan in zip((a, b), plans): + plan.append(torch.cat([st["num"].view(1), st["src"][:k].long(), st["dst"][:k].long()])) + assert torch.equal(qa, qb), "out slots differ" + pa, pb = (torch.stack(plan) for plan in plans) + n = pa[:, 0] + assert int(n.sum()) > 1000 and int((n == 0).sum()) > 0 # the stream exercises misses and all-hit calls + for i in range(pa.shape[0]): + m = int(n[i]) + assert torch.equal(pa[i, : 1 + m], pb[i, : 1 + m]) and torch.equal(pa[i, 1 + k : 1 + k + m], pb[i, 1 + k : 1 + k + m]) + for key in ("slot_of_id", "id_of_slot", "usage", "step", "stats"): + assert torch.equal(a[key], b[key]), key + + +def test_lru_policy_stays_on_flashlib(monkeypatch): + import freetoken.moe.offload_kernels as ok + + cache = _cache("lru", device="cpu") + calls = [] + monkeypatch.setattr(ok, "lru_ensure", lambda *a, **kw: calls.append("flashlib")) + monkeypatch.setattr(ok, "ensure_experts_scored", lambda *a, **kw: calls.append("scored")) + cache.ensure_experts(3, torch.zeros(TOP_K, dtype=torch.int32)) + cache.ensure_experts(3, torch.zeros(TOP_K, dtype=torch.int32), update_state=False) + assert calls == ["flashlib", "flashlib"] + assert cache.evict_tok is None and cache.evict_last_tok is None and cache.near_miss_thr is None + + +# (b) scored policies == CPU reference ------------------------------------------------------ + + +@_cuda +@pytest.mark.parametrize("policy", ["lru", "kd", "kdfb", "rule"]) +def test_scored_kernel_matches_cpu_reference_on_synthetic_routing(policy): + rows = _zipf_rows(40, seed=3) + got = _check_against_ref(policy, rows, _random_logits(rows) if policy == "rule" else None) + assert got.sum() > 0 + + +@_cuda +@pytest.mark.slow +@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"]) +def test_scored_kernel_matches_cpu_reference_on_trace_rows(policy): + experts, li, lv = _trace() + rows = np.asarray(experts[:2000]).astype(np.int64)[:, :, None, :] + logits = _trace_logits(li[:2000], lv[:2000])[:, :, None, :] if policy == "rule" else None + got = _check_against_ref(policy, rows, logits, torch.float32) + lru = _run_ref(_ref("lru"), rows) + assert got.sum() < lru.sum() # the point of the policy + + +# (c) victims: never the call's own ids, empty slots first ---------------------------------- + + +@_cuda +@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"]) +@pytest.mark.parametrize("rows_per_step", [1, 2]) +def test_victims_skip_the_call_and_take_empty_slots_first(policy, rows_per_step): + from freetoken.moe import scored_ensure as se + + num_layers, num_experts, cache_size = 4, 64, 100 + cache = _cache(policy, cache_size, num_layers, num_experts) + ref = _ref(policy, cache_size, num_layers, num_experts) + warm = _zipf_rows(6, rows_per_step, seed=5, num_layers=num_layers, num_experts=num_experts) + _run_gpu(cache, warm) + _run_ref(ref, warm) + assert int((cache.id_of_slot >= 0).sum()) == cache_size + + layer = 2 + held = cache.id_of_slot.cpu().numpy() + resident = sorted(int(i) - layer * num_experts for i in held if layer * num_experts <= i < (layer + 1) * num_experts) + absent = sorted(set(range(num_experts)) - set(resident)) + k = rows_per_step * TOP_K + hits = resident[: k // 2] + ids = np.array((hits + absent)[:k], np.int64) + if rows_per_step == 2: + ids[-1] = ids[0] # a duplicate hit + # make this call's hits the stalest ids in the cache, so only the pin keeps them + hit_keys = [layer * num_experts + e for e in hits] + for key in hit_keys: + cache.evict_last_tok[key] = se.LAST_TOK_NEVER + ref.last_tok[key] = se.LAST_TOK_NEVER + if cache.evict_lc is not None: + cache.evict_lc[key] = se.LC_NEVER + ref.lc[key] = se.LC_NEVER + hit_slots = {int(cache.slot_for_id[layer, e]) for e in hits} + # invalidate three slots (as the prefill overlap does): they must be the first victims + empty = sorted([s for s in (7, 40, 91, 13, 55, 77) if s not in hit_slots][:3]) + for slot in empty: + for t, v in ((cache.id_of_slot, held[slot]), (ref.id_of_slot, held[slot])): + t[slot] = -1 + cache.slot_for_id.view(-1)[int(held[slot])] = -1 + ref.slot_of_id[int(held[slot])] = -1 + cache.usage[slot] = 0 + ref.usage[slot] = 0 + + q = torch.tensor(ids, dtype=torch.int32, device="cuda") + lg = _random_logits(ids.reshape(rows_per_step, TOP_K), seed=9, num_experts=num_experts) + cache.ensure_experts(layer, q, router_logits=torch.from_numpy(lg).cuda()) + out, src, dst = ref.ensure(layer, ids, logits=lg, thr=NEAR_MISS_THR) + + n = int(cache.num_indices) + got_dst = cache.evict_slots[:n].tolist() + assert n == len(set(ids.tolist()) - set(hits)) == src.size + assert got_dst[: len(empty)] == empty + assert not hit_slots & set(got_dst) + assert q.tolist() == out.tolist() + assert all(int(cache.slot_for_id[layer, e]) >= 0 for e in ids.tolist()) + assert got_dst == dst.tolist() + _assert_same_tables(cache, ref) + + +# (d) the decode-step clock ----------------------------------------------------------------- + + +@_cuda +@pytest.mark.parametrize("rows_per_step", [1, 2]) +def test_token_clock_moves_once_per_decode_step(rows_per_step): + cache = _cache("kdfb") + rows = _zipf_rows(5, rows_per_step, seed=7) + ids = torch.from_numpy(rows.astype(np.int32)).cuda() + for r in range(rows.shape[0]): + for layer in range(L): + cache.ensure_experts(layer, ids[r, layer]) + assert int(cache.evict_tok) == r + 1 + # a CPU-decoded first layer moves the clock to the first GPU layer + cache.reset() + cache.cpu_layer_ids = frozenset({0, 1}) + ids = torch.from_numpy(rows.astype(np.int32)).cuda() + for layer in range(2, L): + cache.ensure_experts(layer, ids[0, layer]) + assert int(cache.evict_tok) == 1 + + +@_cuda +def test_small_prefill_ensures_leave_the_state_alone(monkeypatch): + import freetoken.layers.moe as moe_mod + from freetoken.distributed import set_tp_info, try_get_tp_info + from freetoken.layers.moe import OffloadMoELayer + from freetoken.layers.quantization import NoQuantConfig + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + dev = torch.device("cuda") + experts, top_k, hidden, inter = 64, 4, 16, 32 + layer = OffloadMoELayer(0, experts, top_k, hidden, inter, quant_config=NoQuantConfig(), + prefix="model.layers.0.mlp.experts") + cache = _cache("kdfb", 96, 1, experts) + g = torch.Generator().manual_seed(0) + cache.set_bank_sources({ + "gate_up": [torch.randn(experts, 2 * inter, hidden, generator=g, dtype=torch.bfloat16) * 0.1], + "down": [torch.randn(experts, hidden, inter, generator=g, dtype=torch.bfloat16) * 0.1], + }) + layer.offload_cache = cache + decode = torch.tensor([[1, 2, 3, 4]], dtype=torch.int32, device=dev) + cache.ensure_experts(0, decode) + before = [t.clone() for t in (cache.evict_tok, cache.evict_last_tok, cache.evict_lc, cache.evict_ct)] + + monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", 64) + tokens = 40 + ids = torch.randint(0, experts // 2, (tokens, top_k), dtype=torch.int32, device=dev) + w = torch.full((tokens, top_k), 1.0 / top_k, device=dev) + layer._prefill_routed(torch.randn(tokens, hidden, dtype=torch.bfloat16, device=dev), w, ids.clone()) + assert bool((cache.slot_for_id[0, ids.long().unique()] >= 0).all()) # it did make them resident + assert int((cache.slot_for_id[0] >= 0).sum()) < experts # through the on-demand path, not the whole layer + after = (cache.evict_tok, cache.evict_last_tok, cache.evict_lc, cache.evict_ct) + for b, a in zip(before, after): + assert torch.equal(b, a) + + +@_cuda +def test_dsv4_on_demand_prefill_leaves_the_state_alone(): + from types import SimpleNamespace + + from freetoken.models.deepseek_v4.moe import DSV4OffloadMoELayer + + experts, top_k = 64, 4 + cache = _cache("kdfb", 96, 1, experts) + cache.set_bank_sources({"gate_up": [torch.zeros(experts, 2, 4)], "down": [torch.zeros(experts, 4, 2)]}) + cache.ensure_experts(0, torch.tensor([1, 2, 3, 4], dtype=torch.int32, device="cuda")) + before = [t.clone() for t in (cache.evict_tok, cache.evict_last_tok, cache.evict_lc, cache.evict_ct)] + # few enough routes (3 x 4 < 64) for the slot path rather than whole-layer streaming + layer = SimpleNamespace(offload_cache=cache, num_experts=experts, top_k=top_k, layer_id=0, + _expert_gemm=lambda *a, **kw: None) + ids = torch.tensor([[5, 6, 7, 8], [9, 10, 11, 12], [5, 13, 14, 15]], dtype=torch.int32, device="cuda") + DSV4OffloadMoELayer._prefill_routed(layer, torch.zeros(3, 8, device="cuda"), torch.full((3, top_k), 0.25), ids) + assert bool((cache.slot_for_id[0, 5:16] >= 0).all()) + for b, a in zip(before, (cache.evict_tok, cache.evict_last_tok, cache.evict_lc, cache.evict_ct)): + assert torch.equal(b, a) + + +# (e) reset() and rebuild() ----------------------------------------------------------------- + + +@_cuda +@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"]) +def test_reset_and_rebuild_clear_the_state(policy): + from freetoken.moe import scored_ensure as se + + cache = _cache(policy, 1480) + cache.set_bank_sources({"gate_up": [torch.zeros(E, 2, 4)] * L, "down": [torch.zeros(E, 4, 2)] * L}) + rows = _zipf_rows(3, seed=11) + fresh = _ref(policy) + + def assert_cold(): + assert int(cache.evict_tok) == 0 + assert bool((cache.evict_last_tok == se.LAST_TOK_NEVER).all()) + assert bool((cache.evict_slot_owner == -1).all()) and cache.evict_slot_owner.numel() == cache.cache_size + if cache.evict_lc is not None: + assert bool((cache.evict_lc == se.LC_NEVER).all()) and bool((cache.evict_ct == 0).all()) + + for clear in (cache.reset, lambda: cache.rebuild(S)): + _run_gpu(cache, rows) + assert int(cache.evict_tok) == rows.shape[0] + clear() + assert_cold() + assert cache.cache_size == S + # a rebuilt cache replays exactly like a fresh one + np.testing.assert_array_equal(_run_gpu(cache, rows), _run_ref(fresh, rows)) + _assert_same_tables(cache, fresh) + + +# (f) CUDA graph replay == eager ------------------------------------------------------------ + + +@_cuda +@pytest.mark.parametrize("policy", ["kdfb", "rule"]) +@pytest.mark.parametrize("rows_per_step", [1, 2]) +def test_graph_replay_matches_eager(policy, rows_per_step): + steps = 50 + rows = _zipf_rows(steps + 1, rows_per_step, seed=13) + ids = torch.from_numpy(rows.astype(np.int32)).cuda() + logits = torch.from_numpy(_random_logits(rows, seed=14)).cuda().bfloat16() + + eager = _cache(policy) + for r in range(1, steps + 1): + for layer in range(L): + eager.ensure_experts(layer, ids[r, layer].clone(), router_logits=logits[r, layer]) + + graphed = _cache(policy) + static_ids = ids[0].clone() + static_logits = logits[0].clone() + + def step(): + for layer in range(L): + graphed.ensure_experts(layer, static_ids[layer], router_logits=static_logits[layer]) + + step() # warm-up compiles the kernel variants outside capture + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + step() + graphed.reset() # as graph.py does after capture + for r in range(1, steps + 1): + static_ids.copy_(ids[r]) + static_logits.copy_(logits[r]) + graph.replay() + torch.cuda.synchronize() + + assert int(graphed.evict_tok) == steps + for name in ("slot_for_id", "id_of_slot", "usage", "step", "evict_tok", "evict_last_tok", "evict_lc", "evict_ct", + "evict_slot_owner", "evict_slot_last_tok", "evict_slot_lc", "evict_slot_ct"): + assert torch.equal(getattr(eager, name), getattr(graphed, name)), name + + +# (g) bs=2: two rows per call --------------------------------------------------------------- + + +@_cuda +@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"]) +def test_two_rows_per_call_match_cpu_reference(policy): + rows = _zipf_rows(30, 2, seed=17) + assert any(len(set(rows[r, layer].ravel())) < 2 * TOP_K for r in range(30) for layer in range(L)) + _check_against_ref(policy, rows, _random_logits(rows, seed=18) if policy == "rule" else None) + + +@_cuda +@pytest.mark.slow +@pytest.mark.parametrize("policy", ["kdfb", "rule"]) +def test_two_trace_streams_per_call_match_cpu_reference(policy): + experts, li, lv = _trace() + a, b = slice(0, 600), slice(45000, 45600) # two different requests decoding together + rows = np.stack([np.asarray(experts[a]), np.asarray(experts[b])], axis=2).astype(np.int64) + logits = None + if policy == "rule": + logits = np.stack([_trace_logits(li[a], lv[a]), _trace_logits(li[b], lv[b])], axis=2) + _check_against_ref(policy, rows, logits, torch.float32) + + +# the slot mirrors: foreign installs and state-free ensures fall back to the per-id state ----- + + +@_cuda +@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"]) +def test_materialized_prefill_between_decodes_matches_cpu_reference(policy): + cache, ref = _cache(policy), _ref(policy) + rows = _zipf_rows(24, seed=21) + logits = _random_logits(rows, seed=22) if policy == "rule" else None + # layer 5 comes back to slots it lost to layer 9: its tags there would otherwise look current + for part, layer in ((slice(0, 6), 5), (slice(6, 12), 9), (slice(12, 18), 5), (slice(18, 24), None)): + lg = None if logits is None else logits[part] + np.testing.assert_array_equal(_run_gpu(cache, rows[part], lg), _run_ref(ref, rows[part], lg)) + _assert_same_tables(cache, ref) + if layer is not None: + # the whole-layer prefill path re-seats this layer's experts at slots [0, E) + cache.materialize_layer(layer) + ref.materialize(layer) + assert not bool((cache.evict_slot_owner[:E] >= 0).any()) + _assert_same_tables(cache, ref) + + +@_cuda +@pytest.mark.parametrize("policy", ["kd", "kdfb"]) +def test_materialize_untags_the_slots_it_fills(policy): + # expert 3 of layer 0 leaves slot 3, is used elsewhere, then materialize puts it back at slot 3 + cache = _cache(policy, 24, 2, 8) + dev = torch.device("cuda") + ids = lambda *e: torch.tensor(e, dtype=torch.int32, device=dev) # noqa: E731 + cache.materialize_layer(0) + cache.ensure_experts(0, ids(3, 4)) # a hit at slot 3 tags it with expert 3 + assert int(cache.evict_slot_owner[3]) == 3 + cache.materialize_layer(1) # slot 3 now holds layer 1's expert 3 + cache.ensure_experts(1, ids(0, 1)) + cache.ensure_experts(0, ids(3, 5)) # expert 3 comes back elsewhere, with new state + assert int(cache.slot_for_id[0, 3]) != 3 + cache.materialize_layer(0) + assert int(cache.slot_for_id[0, 3]) == 3 + _assert_mirrors_hold(cache) + + +@_cuda +@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"]) +@pytest.mark.parametrize("chunks", [1, 2]) +def test_state_free_ensures_between_decodes_match_cpu_reference(policy, chunks): + cache, ref = _cache(policy), _ref(policy) + rows = _zipf_rows(12, seed=23) + extra = _zipf_rows(4, 3, seed=24, skew=0.3) + logits = _random_logits(rows, seed=25) if policy == "rule" else None + for r in range(rows.shape[0]): + lg = None if logits is None else logits[r : r + 1] + np.testing.assert_array_equal(_run_gpu(cache, rows[r : r + 1], lg), _run_ref(ref, rows[r : r + 1], lg)) + if r % 3 == 2: + for layer in range(0, L, 7): # a small prefill's chunked ensures + ids = np.unique(extra[r // 3, layer]) + # two chunks pin each other's slots as _ensure_unique does; one chunk keeps the plain pin + pin = cache.step.clone() if chunks > 1 else None + ref_pin = ref.step if chunks > 1 else None + for part in np.array_split(ids, chunks): + q = torch.from_numpy(part.astype(np.int32)).cuda() + cache.ensure_experts(layer, q, update_state=False, pin_since=pin) + out, _, dst = ref.ensure(layer, part, update_state=False, pin_since=ref_pin) + assert q.tolist() == out.tolist() + assert bool((cache.evict_slot_owner[torch.from_numpy(dst).cuda()] == -1).all()) + assert bool((cache.slot_for_id[layer, torch.from_numpy(ids).cuda()] >= 0).all()) + _assert_same_tables(cache, ref) + assert not _assert_mirrors_hold(cache).all() # some slots still read the per-id state + + +# rule: the near misses come from the router logits ---------------------------------------- + + +@_cuda +@pytest.mark.parametrize("rows_per_step", [1, 2]) +def test_rule_refreshes_the_router_near_misses(rows_per_step): + from freetoken.moe import scored_ensure as se + + cache = _cache("rule") + layer = 5 + g = torch.Generator(device="cuda").manual_seed(rows_per_step) + logits = (torch.randn(rows_per_step, E, device="cuda", generator=g) * 1.5).bfloat16() + ids = torch.topk(logits.float(), TOP_K, dim=-1).indices.to(torch.int32) + cache.ensure_experts(0, torch.arange(TOP_K, dtype=torch.int32, device="cuda")) # step 1 starts + cache.ensure_experts(layer, ids.clone(), router_logits=logits) + kth = logits.float().gather(1, ids.long()).min(dim=1, keepdim=True).values + near = (logits.float() >= kth - NEAR_MISS_THR).any(0) + routed = torch.zeros(E, dtype=torch.bool, device="cuda") + routed[ids.long().view(-1)] = True + last = cache.evict_last_tok[layer * E : (layer + 1) * E] + assert bool((last[routed] == 1).all()) + assert bool((last[near & ~routed] == 0).all()) + assert bool((last[~near & ~routed] == se.LAST_TOK_NEVER).all()) + assert 0 < int((near & ~routed).sum()) < E - TOP_K * rows_per_step + + +@pytest.mark.parametrize("policy", ["lru", "kdfb", "rule"]) +def test_decode_forward_hands_the_router_logits_to_ensure(monkeypatch, policy): + from freetoken.distributed import set_tp_info, try_get_tp_info + from freetoken.layers.moe import OffloadMoELayer + from freetoken.layers.quantization import NoQuantConfig + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + layer = OffloadMoELayer(0, 8, 2, 8, 16, quant_config=NoQuantConfig(), prefix="model.layers.0.mlp.experts") + layer.offload_cache = cache = _cache(policy, 12, 1, 8, device="cpu") + seen = {} + monkeypatch.setattr( + "freetoken.layers.moe.fused_topk", + lambda **kw: (torch.full((2, 2), 0.5), torch.zeros((2, 2), dtype=torch.int32)), + ) + monkeypatch.setattr(cache, "ensure_experts", lambda lid, ids, **kw: seen.update(kw)) + monkeypatch.setattr(cache, "copy_missing", lambda: None) + monkeypatch.setattr(cache, "bank_views", lambda: ()) + monkeypatch.setattr(layer, "_expert_gemm", lambda *a, **kw: None) + router_logits = torch.randn(2, 8) + layer.decode_forward(torch.randn(2, 8), router_logits) + assert seen["router_logits"] is router_logits + assert cache.near_miss_thr == (NEAR_MISS_THR if policy == "rule" else None) + + +def test_policy_flag_and_gpu_only(): + import contextlib + import io + + from freetoken.server.args import parse_args + + buf = io.StringIO() + with contextlib.redirect_stdout(buf), contextlib.suppress(SystemExit): + parse_args(["--help"]) + assert "{lru,kd,kdfb,rule}" in buf.getvalue() + for policy in ("kd", "kdfb", "rule"): + _cache(policy, 12, 1, 8, device="cpu") + with pytest.raises(ValueError, match="GPU decode"): + _cache(policy, 12, 1, 8, device="cpu", decode_target="hybrid") + + +@_cuda +@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"]) +def test_an_all_pinned_call_never_stores_past_the_slot_arrays(policy): + # 5 slots pad to an 8-lane block; with every slot pinned the sentinel key must not decode to lane 7. + # The cache enforces cache_size >= num_experts, so the kernel gets 5-slot views of a 16-slot cache. + from freetoken.moe.offload_kernels import ensure_experts_scored + + s, pad = 5, 3 + cache = _cache(policy, cache_size=16, num_layers=1, num_experts=16) + guarded = {} + for name in ("id_of_slot", "usage", "evict_slot_owner", "evict_slot_last_tok", "evict_slot_lc", "evict_slot_ct"): + t = getattr(cache, name) + if t is None: # kd keeps no lc/ct mirrors + continue + buf = torch.full((s + pad,), 7, dtype=t.dtype, device=t.device) + buf[:s] = t[:s] + setattr(cache, name, buf[:s]) + guarded[name] = buf + dev = cache.id_of_slot.device + ensure_experts_scored(cache, 0, torch.arange(s, dtype=torch.int32, device=dev), bump_tok=True, update_state=True) + # 6 distinct ids against 5 slots breaks the caller contract; the kernel must still stay in bounds + q = torch.arange(s + 1, dtype=torch.int32, device=dev) + ensure_experts_scored(cache, 0, q, bump_tok=True, update_state=True) + torch.cuda.synchronize() + for name, buf in guarded.items(): + assert torch.all(buf[s:] == 7), f"{name} written past slot {s - 1}" + assert int(cache.num_indices) == 1 and 0 <= int(cache.evict_slots[0]) < s diff --git a/tests/moe/test_small_prefill_on_demand.py b/tests/moe/test_small_prefill_on_demand.py index eea5842d23..70cec00019 100644 --- a/tests/moe/test_small_prefill_on_demand.py +++ b/tests/moe/test_small_prefill_on_demand.py @@ -8,7 +8,7 @@ from freetoken.distributed import set_tp_info, try_get_tp_info -def _layer_and_cache(device, num_experts=8, top_k=2, hidden=16, inter=32, cache_size=None): +def _layer_and_cache(device, num_experts=8, top_k=2, hidden=16, inter=32, cache_size=None, policy="lru", num_layers=1): from freetoken.layers.moe import OffloadMoELayer from freetoken.layers.quantization import NoQuantConfig from freetoken.moe.offload_cache import OffloadMoeCache @@ -17,11 +17,14 @@ def _layer_and_cache(device, num_experts=8, top_k=2, hidden=16, inter=32, cache_ set_tp_info(rank=0, size=1) layer = OffloadMoELayer(0, num_experts, top_k, hidden, inter, quant_config=NoQuantConfig(), prefix="model.layers.0.mlp.experts") - cache = OffloadMoeCache(num_layers=1, num_experts=num_experts, cache_size=cache_size or num_experts, device=device) + cache = OffloadMoeCache(num_layers=num_layers, num_experts=num_experts, cache_size=cache_size or num_experts, + device=device, cache_policy=policy) g = torch.Generator().manual_seed(0) cache.set_bank_sources({ - "gate_up": [torch.randn(num_experts, 2 * inter, hidden, generator=g, dtype=torch.bfloat16) * 0.1], - "down": [torch.randn(num_experts, hidden, inter, generator=g, dtype=torch.bfloat16) * 0.1], + "gate_up": [torch.randn(num_experts, 2 * inter, hidden, generator=g, dtype=torch.bfloat16) * 0.1 + for _ in range(num_layers)], + "down": [torch.randn(num_experts, hidden, inter, generator=g, dtype=torch.bfloat16) * 0.1 + for _ in range(num_layers)], }) layer.offload_cache = cache return layer, cache @@ -34,7 +37,7 @@ def test_small_prefill_dispatch(monkeypatch, tokens, threshold, on_demand): layer, cache = _layer_and_cache(torch.device("cpu")) monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", threshold) calls = [] - monkeypatch.setattr(cache, "ensure_experts", lambda layer_id, ids: calls.append("ensure")) + monkeypatch.setattr(cache, "ensure_experts", lambda layer_id, ids, **kw: calls.append("ensure")) monkeypatch.setattr(cache, "materialize_layer", lambda layer_id: calls.append("materialize")) monkeypatch.setattr(cache, "copy_missing", lambda: calls.append("copy")) monkeypatch.setattr(cache, "begin_prefill", lambda: calls.append("begin_prefill"), raising=False) @@ -96,7 +99,35 @@ def test_a_real_sized_small_prefill_fits_the_lru_kernel(monkeypatch): monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", threshold) layer, cache = _layer_and_cache(dev, num_experts=experts, top_k=top_k, cache_size=1100) widths, ensure = [], cache.ensure_experts - monkeypatch.setattr(cache, "ensure_experts", lambda lid, q: (widths.append(q.numel()), ensure(lid, q))[1]) + monkeypatch.setattr(cache, "ensure_experts", lambda lid, q, **kw: (widths.append(q.numel()), ensure(lid, q, **kw))[1]) outs[threshold] = _run_and_check_path(layer, hs, w, ids, on_demand=threshold > 0) assert widths and max(widths) <= 32, f"LRU ensure widths {sorted(set(widths))}: wider than decode's" torch.testing.assert_close(outs[1024], outs[0], rtol=2e-2, atol=2e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +@pytest.mark.parametrize("policy", ["lru", "kd", "kdfb", "rule"]) +def test_a_later_chunk_keeps_the_experts_an_earlier_chunk_loaded(monkeypatch, policy): + # Decode fills every slot, then 40 never-decoded experts take two ensure chunks. A scored + # policy ranks the first chunk's fresh installs as its coldest slots, so they need a pin. + dev = torch.device("cuda") + experts, top_k, slots, tokens = 64, 4, 80, 10 + g = torch.Generator(device=dev).manual_seed(2) + hs = torch.randn(tokens, 16, dtype=torch.bfloat16, device=dev, generator=g) * 4 # a wrong expert is ~1.0 off + ids = torch.arange(16, 16 + tokens * top_k, dtype=torch.int32, device=dev).view(tokens, top_k) + w = torch.softmax(torch.randn(tokens, top_k, device=dev, generator=g), dim=-1) + outs = {} + for threshold in (0, 64): + monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", threshold) + layer, cache = _layer_and_cache(dev, num_experts=experts, top_k=top_k, cache_size=slots, policy=policy, + num_layers=2) + for layer_id, first, last in ((1, 0, experts), (0, 0, 16)): + for start in range(first, last, 8): + cache.ensure_experts(layer_id, torch.arange(start, start + 8, dtype=torch.int32, device=dev)) + cache.copy_missing() + assert bool((cache.id_of_slot >= 0).all()) + outs[threshold] = _run_and_check_path(layer, hs, w, ids, on_demand=threshold > 0) + if threshold: + slot_of = cache.slot_for_id[0, ids.long()] + assert bool((slot_of >= 0).all()), f"{int((slot_of < 0).sum())} routed experts were evicted" + torch.testing.assert_close(outs[64], outs[0], rtol=2e-2, atol=0.1) From 3b4893cb2182e81687148fe66ff842ac4cd3b205 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Sun, 27 Sep 2026 03:33:34 -0400 Subject: [PATCH 29/94] merge fork perf: expert prefetch into low-priority slots with single-writer slot maps (FREETOKEN_MOE_PREFETCH), plus decode verify and slot audit debug modes --- benchmarks/README.md | 23 + benchmarks/bench_moe_prefetch.py | 580 ++++++ python/freetoken/engine/engine.py | 72 +- python/freetoken/env.py | 14 + python/freetoken/layers/moe.py | 48 +- python/freetoken/models/qwen4_exp/model.py | 3 +- python/freetoken/models/qwen4_exp/moe.py | 17 +- python/freetoken/moe/offload_cache.py | 239 ++- python/freetoken/moe/offload_kernels.py | 44 +- python/freetoken/moe/prefetch.py | 438 +++++ python/freetoken/moe/scored_ensure.py | 168 +- python/freetoken/moe/slot_audit.py | 688 +++++++ python/freetoken/moe/verify.py | 516 ++++++ tests/models/qwen4_exp/conftest.py | 13 + tests/models/qwen4_exp/test_build_cpu.py | 29 + tests/models/qwen4_exp/test_skeleton.py | 1911 ++++++++++++++++++++ tests/moe/ref_prefetch.py | 22 + tests/moe/ref_scored_cache.py | 44 +- tests/moe/test_prefetch.py | 424 +++++ tests/moe/test_scored_ensure.py | 319 ++++ 20 files changed, 5544 insertions(+), 68 deletions(-) create mode 100644 benchmarks/bench_moe_prefetch.py create mode 100644 python/freetoken/moe/prefetch.py create mode 100644 python/freetoken/moe/slot_audit.py create mode 100644 python/freetoken/moe/verify.py create mode 100644 tests/moe/ref_prefetch.py create mode 100644 tests/moe/test_prefetch.py diff --git a/benchmarks/README.md b/benchmarks/README.md index 3f0df72f81..817f99da18 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -45,6 +45,29 @@ start after the ensure and how much of the side-stream shared expert runs inside python benchmarks/bench_moe_copy_overlap.py --bs 1 --slots 160 ``` +**`bench_moe_prefetch.py`** — synthetic (no checkpoint): `FREETOKEN_MOE_PREFETCH=measure` +cost. The router-lookahead predictor alone per layer (GEMV `[bs, 2560] x [2560, 512]`, top-K +select, and the count kernel), then prefetch off vs measure on a real-geometry Qwen4ExpMoE decode +stack (512 experts, rule eviction, copy overlap on) with forced routing that misses exactly m +experts per layer, as CUDA graphs. `--on` times whole 48-layer decode steps with prefetch off vs +on instead: a GEMV per layer stands in for attention (`--filler-us`), each layer prefetches 3 or +4 candidates of which `--useful` are its misses, and `--delay-us` holds every prefetch copy back +to price a late one. It reports ms/step, demand misses/step and the late fraction. `--steady +--trace DIR` runs steady-state decode on a recorded routing trace instead (experts.npy, +logit_idx.npy, logit_val.npy, req_index.npy): no cache restore, so wrong prefetches and evictions +carry over, and a noisy copy of each layer's traced logits (`--sigma`) stands in for the +lookahead. Arms such as `off,measure,on,on:4/4` (budgets before GDN / QSA layers) replay step by +step in turn, each on its own copy of the cache state; it adds pollution (non-resident demands +per layer against off's misses) and the late fraction per layer kind. Slots and sigma were +calibrated to production's measure line (misses/layer ~4.3, precision ~0.56) at 24 layers. + +```bash +python benchmarks/bench_moe_prefetch.py --bs 1,2 --misses 0,4 +python benchmarks/bench_moe_prefetch.py --on --bs 1,2 --misses 4 --delay-us 0,50,150,300 +python benchmarks/bench_moe_prefetch.py --steady --trace DIR --bs 1,2 --steady-slots 740 --sigma 0.55 \ + --arms off,measure,on,on:empty,on:3/3,on:4/4,on:4/5 +``` + **`bench_expert_copy_corun.py`** — synthetic (no checkpoint): the fused expert copy kernel alone and beside compute, on registered host banks with the qwen4_exp NVFP4 row sizes. Sweeps `fast_index_copy_multi` (threads x blocks/bank) and `fast_index_copy_multi_slim` diff --git a/benchmarks/bench_moe_prefetch.py b/benchmarks/bench_moe_prefetch.py new file mode 100644 index 0000000000..be5d38ebd7 --- /dev/null +++ b/benchmarks/bench_moe_prefetch.py @@ -0,0 +1,580 @@ +"""FREETOKEN_MOE_PREFETCH cost and gain: the router-lookahead predictor alone, its effect on the +decode compute stream in measure mode, and (--on) whole decode steps with prefetch off vs on. + +1. Predictor alone, per layer: the lookahead GEMV ([bs, 2560] x [2560, 512] bf16, a distinct router + weight per layer as in the model), the top-K + select kernel, and both; one stream, CUDA graph. +2. The count kernel alone (it runs on the compute stream in measure mode). +3. A real-geometry decode stack (Qwen3.8-Flash-Next: hidden 2560, 512 experts, top-10, moe and + shared intermediate 640, NVFP4 triton banks) of N Qwen4ExpMoE blocks: router GEMV, shared + expert (on the FREETOKEN_MOE_COPY_OVERLAP side stream by default), top-k, rule ensure, miss + copy of exactly m experts from pinned host, grouped GEMM. The routing is forced (10 - m hot + experts + m cold ones per row) and the cache restored before every replay, so both arms copy + the same bytes. One CUDA graph per step, prefetch off vs measure. No attention runs, so the + percentages overstate the model's. +4. --on: the same stack at the model's depth (48 layers), with a bf16 GEMV "filler" before every + block standing in for attention and the rest of the layer (--filler-us, ~225 us between one + demand copy and the next in production). Prefetch on installs a fixed candidate list per layer: + 3 before GDN / 4 before attention layers, of which --useful are the layer's cold experts and the + rest never-routed ids (production: ~3.24 issued, ~1.75 useful, ~4.2 misses per layer at bs 1). + Reports ms/step, demand misses/step and the late fraction; --delay-us adds a sleep ahead of every + prefetch copy to price the compute stream's wait on a late copy. +5. --steady --trace DIR: steady-state decode on a recorded routing trace (no restore, so wrong + prefetches and evictions carry over), --steady-layers layers of it with --steady-slots slots, + a GEMV filler of --filler-gdn-us / --filler-qsa-us before GDN / QSA layers, and a lookahead + that is each layer's traced top-32 logits plus N(0, --sigma) noise. Every arm (off, measure, + on[:gdn/qsa budgets | :empty]) keeps its own cache state and all arms replay step by step in + turn, so drift hits them alike. Reports ms/step (mean, p50, p90), demand misses per layer, + issued/useful/precision, non-resident demands per layer (misses + prefetch hits: above off's + misses is pollution) and the late fraction before GDN and QSA layers. + +Run: CUDA_VISIBLE_DEVICES=0 PYTHONPATH=python python benchmarks/bench_moe_prefetch.py + ... benchmarks/bench_moe_prefetch.py --on --bs 1,2 --misses 4 + ... benchmarks/bench_moe_prefetch.py --steady --trace DIR --steady-slots 740 --sigma 0.55 +""" + +from __future__ import annotations + +import argparse +import gc +import statistics + +import torch +import torch.nn.functional as F + +H, E, TOP_K, INTER = 2560, 512, 10, 640 + + +def time_graph(run, reps: int, before=None) -> float: + """Median ms of one replay of ``run`` captured as a CUDA graph (after an eager warm-up).""" + if before: + before() + run() + torch.cuda.synchronize() + if before: + before() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + times = [] + for rep in range(reps + 3): + if before: + before() + torch.cuda.synchronize() + start.record() + graph.replay() + end.record() + end.synchronize() + if rep >= 3: + times.append(start.elapsed_time(end)) + return statistics.median(times) + + +def bench_predictor(bs: int, layers: int, k: int, budget: int, reps: int) -> dict: + from freetoken.moe.prefetch import lookahead_select, prefetch_count + + dev = torch.device("cuda") + gen = torch.Generator(device=dev).manual_seed(bs) + weights = [torch.randn(E, H, device=dev, dtype=torch.bfloat16, generator=gen) * 0.02 for _ in range(layers)] + x = torch.randn(bs, H, device=dev, dtype=torch.bfloat16, generator=gen) + resident = torch.where(torch.rand(E, device=dev, generator=gen) < 0.4, 1, -1).to(torch.int32) + sel = torch.full((layers, 2 * k), -1, dtype=torch.int32, device=dev) + res = sel.clone() + logits = [F.linear(x, w) for w in weights] + + def gemv(): + for w in weights: + F.linear(x, w) + + def select(): + for layer in range(layers): + lookahead_select(logits[layer], resident, sel[layer], res[layer], k=k, budget=budget) + + def both(): + for layer, w in enumerate(weights): + lookahead_select(F.linear(x, w), resident, sel[layer], res[layer], k=k, budget=budget) + + slots = torch.randint(0, 1024, (bs * TOP_K,), device=dev, dtype=torch.int32, generator=gen) + id_of_slot = torch.randint(0, E, (1024,), device=dev, dtype=torch.int32, generator=gen) + misses = torch.ones(1, dtype=torch.int64, device=dev) + counters = torch.zeros((layers, 6), dtype=torch.int64, device=dev) + + def count(): + for layer in range(layers): + prefetch_count(slots, id_of_slot, misses, sel[layer], res[layer], counters[layer], id_base=0, rows=bs) + + per = 1000.0 / layers + return {name: time_graph(fn, reps) * per for name, fn in + (("gemv", gemv), ("select", select), ("predictor", both), ("count", count))} + + +class ForcedGate: + """A layer's real router GEMV whose logits are then replaced, so every replay routes (and + misses) the same experts; the next layer's lookahead runs through it too.""" + + def __init__(self, gate, logits: torch.Tensor): + self.gate, self.logits = gate, logits + + def forward(self, x: torch.Tensor) -> torch.Tensor: + self.gate.forward(x) + return self.logits + + +class DecodeStack: + """N Qwen4ExpMoE blocks (router, shared expert, nvfp4 offload experts) over one shared cache, + wired like qwen4_exp for the lookahead.""" + + def __init__(self, layers: int, slots: int, bs: int, misses: int, mode: str, overlap: bool, filler=None, + routing=None, budgets=(3, 4)): + import freetoken.core as core + from types import SimpleNamespace + + from freetoken.core import Context, set_global_ctx + from freetoken.distributed import set_tp_info, try_get_tp_info + from freetoken.layers.moe import OffloadMoELayer + from freetoken.layers.quantization import QuantBackend, QuantConfig, set_quant_backend + from freetoken.models.qwen4_exp.moe import Qwen4ExpMoE + from freetoken.moe.expert_banks import build_expert_banks + from freetoken.moe.offload_cache import OffloadMoeCache + from freetoken.utils.torch_utils import torch_dtype + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + set_quant_backend(QuantBackend.parse("moe.nvfp4=triton")) + quant = QuantConfig.from_hf({"quantization_config": {"quant_method": "modelopt", "quant_algo": "NVFP4", "ignore": ["lm_head"]}}) + dev = torch.device("cuda") + config = SimpleNamespace( + hidden_size=H, num_experts=E, num_experts_per_tok=TOP_K, moe_intermediate_size=INTER, + shared_expert_intermediate_size=INTER, norm_topk_prob=True, quant=None, + moe_strategy="offload", decode_target="gpu", + ) + with torch.device(dev), torch_dtype(torch.bfloat16): + self.moes = [Qwen4ExpMoE(config, i, prefix=f"model.layers.{i}.mlp") for i in range(layers)] + gen = torch.Generator(device=dev).manual_seed(0) + for moe in self.moes: + for t in moe.state_dict().values(): + if t.is_floating_point(): + t.normal_(0.0, 0.02, generator=gen) + experts = [ + OffloadMoELayer(i, E, TOP_K, H, INTER, quant_config=quant, prefix=f"model.layers.{i}.mlp.experts") + for i in range(layers) + ] + method = experts[0].quant_method + if not hasattr(DecodeStack, "_banks"): + DecodeStack._banks = build_expert_banks(method, 1, None, device=dev, dummy=True) # one layer, shared + banks = DecodeStack._banks + self.cache = cache = OffloadMoeCache( + num_layers=layers, num_experts=E, cache_size=slots, device=dev, cache_policy="rule", + quant_format=banks.quant_format, layout=banks.layout, max_slots=method.slot_limit(), + decode_copy_overlap=overlap, prefetch_mode=mode, + ) + cache.set_bank_sources({role: per_layer * layers for role, per_layer in banks.sources.items()}) + cache.set_alphas(banks.gate_up_alpha, banks.down_alpha) + core._GLOBAL_CTX = None + ctx = Context(page_size=1) + set_global_ctx(ctx) + ctx._batch = SimpleNamespace(is_prefill=False) + self.x = [torch.randn(bs, H, device=dev, dtype=torch.bfloat16, generator=gen) * 0.5 for _ in range(layers)] + self.filler = filler + if routing is not None: + # --steady: every layer routes this step's row of route_buf and layer L's lookahead (from + # layer L-1) returns look_buf[L]; both run the real router GEMV first + route_buf, look_buf = routing + gates = [moe.gate for moe in self.moes] + for i, (moe, ex) in enumerate(zip(self.moes, experts)): + ex.offload_cache = cache + moe.experts = ex + moe.gate = ForcedGate(gates[i], route_buf[i]) + for i in range(1, layers): + experts[i - 1].set_lookahead(ForcedGate(gates[i], look_buf[i]), i, budgets[(i + 1) % 4 == 0]) + self.bs = bs + cache.reset() + torch.cuda.synchronize() + return + # forced routing: row r routes hot experts [r*hot, (r+1)*hot) and m fresh ones past all rows' hot ones + hot = TOP_K - misses + for moe, ex in zip(self.moes, experts): + ex.offload_cache = cache + moe.experts = ex + lg = torch.zeros(bs, E, device=dev, dtype=torch.bfloat16) + for r in range(bs): + lg[r, list(range(r * hot, (r + 1) * hot))] = 10.0 + lg[r, list(range(bs * hot + r * misses, bs * hot + (r + 1) * misses))] = 10.0 + moe.gate = ForcedGate(moe.gate, lg) + for i in range(1, layers): + experts[i - 1].set_lookahead(self.moes[i].gate, i, budgets[(i + 1) % 4 == 0]) + self.bs, self.hot, self.misses = bs, hot, misses + # warm the hot experts in, then drop the fresh ones: every replay then misses exactly those + cache.reset() + self.step() + cache.slot_for_id[:, bs * hot:].fill_(-1) + fresh = (cache.id_of_slot >= 0) & (cache.id_of_slot % E >= bs * hot) + cache.id_of_slot[fresh] = -1 + cache.usage[fresh] = 0 + torch.cuda.synchronize() + self.state = [cache.slot_for_id, cache.id_of_slot, cache.usage, cache.step] + for name in ("evict_tok", "evict_last_tok", "evict_lc", "evict_ct", "evict_slot_owner", + "evict_slot_last_tok", "evict_slot_lc", "evict_slot_ct"): + if getattr(cache, name) is not None: + self.state.append(getattr(cache, name)) + self.snap = [t.clone() for t in self.state] + + def restore(self): + for dst, src in zip(self.state, self.snap): + dst.copy_(src) + + def step(self): + for i, (moe, x) in enumerate(zip(self.moes, self.x)): + filler = self.filler[i] if isinstance(self.filler, list) else self.filler + if filler is not None: + F.linear(x, filler) + moe.forward(x.clone()) + + def cold(self, layer: int) -> list[int]: + """The experts ``layer`` misses every replay, merged rank by rank across rows like the select.""" + per_row = [list(range(self.bs * self.hot + r * self.misses, self.bs * self.hot + (r + 1) * self.misses)) + for r in range(self.bs)] + return [row[i] for i in range(self.misses) for row in per_row] + + +def bench_stack(layers, slots, bs, misses, reps, overlap) -> dict: + from flashlib.kernels.slot_cache import Stat + + out = {} + for mode in ("off", "measure"): + stack = DecodeStack(layers, slots, bs, misses, mode, overlap) + stack.cache.collect_stats = True + stack.restore() + stack.cache.lru_stats.zero_() + stack.step() + got = stack.cache.lru_stats[:, Stat.MISS].tolist() + assert got == [bs * misses] * layers, f"expected {bs * misses} misses per layer, got {got}" + stack.cache.collect_stats = False + out[mode] = time_graph(stack.step, reps, before=stack.restore) * 1000.0 / layers + if stack.cache.prefetch is not None: + stack.cache.prefetch.counters.zero_() + stack.restore() + stack.step() + c = stack.cache.prefetch.counters.sum(0).tolist() + out["issued/layer"] = c[0] / max(c[3], 1) + del stack + torch.cuda.empty_cache() + return out + + +def budget_of(layer: int) -> int: + return 3 if (layer + 1) % 4 else 4 + + +def bench_on(layers, slots, bs, misses, useful, reps, filler_us, delays_us) -> list[dict]: + """ms per decode step, prefetch off vs on (and on with every prefetch copy delayed).""" + from flashlib.kernels.slot_cache import Stat + + from freetoken.moe.prefetch import COPIED, ISSUED, LATE, USEFUL + + dev = torch.device("cuda") + x = torch.randn(1, H, device=dev, dtype=torch.bfloat16) + probe = torch.randn(32768, H, device=dev, dtype=torch.bfloat16) + per_row_us = time_graph(lambda: F.linear(x, probe), reps) * 1000.0 / 32768 + del probe + filler = torch.randn(max(int(filler_us / per_row_us), 1), H, device=dev, dtype=torch.bfloat16) * 0.01 + cycles_per_us = 1_000_000 / (time_graph(lambda: torch.cuda._sleep(1_000_000), 5) * 1000.0) + out = [] + for arm in ["off", "on", *[f"on+{d}us" for d in delays_us if d]]: + stack = DecodeStack(layers, slots, bs, misses, "off" if arm == "off" else "on", True, filler=filler) + cache = stack.cache + want = [bs * misses] * layers + if arm != "off": + pf = cache.prefetch + pf.sel_override = torch.full_like(pf.sel, -1) + for layer in range(1, layers): + u = min(useful[layer % len(useful)], bs * misses) + wrong = list(range(E - 1, E - 1 - (budget_of(layer) - u), -1)) + ids = stack.cold(layer)[:u] + wrong + pf.sel_override[layer, : len(ids)] = torch.tensor(ids, dtype=torch.int32) + want[layer] -= u + if arm.startswith("on+"): + pf.delay_copy_cycles = int(float(arm[3:-2]) * cycles_per_us) + cache.collect_stats = True + stack.restore() + cache.lru_stats.zero_() + stack.step() + got = cache.lru_stats[:, Stat.MISS].tolist() + assert got == want, f"{arm}: expected demand misses {want}, got {got}" + cache.collect_stats = False + stack.restore() + stack.step() + torch.cuda.synchronize() + stack.restore() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + stack.step() + if cache.prefetch is not None: + cache.prefetch.stats.zero_() + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + times = [] + for rep in range(reps + 3): + stack.restore() + torch.cuda.synchronize() + start.record() + graph.replay() + end.record() + end.synchronize() + if rep >= 3: + times.append(start.elapsed_time(end)) + row = {"arm": arm, "ms": statistics.median(times), "misses": sum(want)} + if cache.prefetch is not None: + st = cache.prefetch.stats.sum(0).tolist() + calls = max(reps + 3, 1) * (layers - 1) + row.update(issued=st[ISSUED] / calls, useful=st[USEFUL] / calls, late=st[LATE] / max(st[COPIED], 1)) + out.append(row) + del graph, stack, cache + gc.collect() # the module graph holds cycles: free this arm's slot banks before the next one + torch.cuda.empty_cache() + return out + + +def load_trace(path: str, layers: int, offset: int, bs: int, steps: int, skip: int, sigma: float, seed: int): + """Per-step routing and lookahead logits ``[steps, layers, bs, E]`` from a routing trace. + + ``path`` holds experts.npy ``[N, 48, 10]`` (routed ids), logit_idx.npy / logit_val.npy + ``[N, 48, 32]`` (each layer's own top-32 router logits) and req_index.npy ``[N]``. Row r of the + batch follows the r-th longest request from its ``skip``-th token on. The routing logits pick + exactly the traced top-10; the lookahead of layer L is layer L's traced top-32 logits plus + N(0, sigma) noise (everything else far below), a stand-in for W_gate[L] . x_{L-1}.""" + import os + from collections import Counter + + import numpy as np + + ex = np.load(os.path.join(path, "experts.npy"), mmap_mode="r") + li = np.load(os.path.join(path, "logit_idx.npy"), mmap_mode="r") + lv = np.load(os.path.join(path, "logit_val.npy"), mmap_mode="r") + req = np.load(os.path.join(path, "req_index.npy")) + reqs = [r for r, n in Counter(req.tolist()).most_common(bs)] + rng = np.random.default_rng(seed) + route = np.full((steps, layers, bs, E), -10.0, np.float32) + look = np.full((steps, layers, bs, E), -30.0, np.float32) + lsl = slice(offset, offset + layers) + for r, rid in enumerate(reqs): + rows = np.flatnonzero(req == rid)[skip: skip + steps] + assert len(rows) == steps, f"request {rid} has {len(rows) + skip} tokens, fewer than skip + steps" + routed = np.asarray(ex[rows, lsl]).astype(np.int64) # [steps, layers, 10] + np.put_along_axis(route[:, :, r], routed, 10.0 - 0.1 * np.arange(TOP_K), -1) + idx = np.asarray(li[rows, lsl]).astype(np.int64) + val = np.asarray(lv[rows, lsl]).astype(np.float32) + rng.normal(0.0, sigma, idx.shape).astype(np.float32) + np.put_along_axis(look[:, :, r], idx, val, -1) + dev = torch.device("cuda") + return (torch.from_numpy(route).to(dev, torch.bfloat16), torch.from_numpy(look).to(dev, torch.bfloat16), reqs) + + +def bench_steady(args, bs: int, arms: list[str]) -> list[dict]: + """ms per decode step in steady state on traced routing: no restore, the cache evolves. + + One stack, one graph per arm, each arm with its own copy of the cache's slot maps, eviction + state and counters: every step replays all arms back to back (rotating their order), so + clock or link drift hits every arm alike. Slot bytes are shared (timing does not read them).""" + from flashlib.kernels.slot_cache import Stat + + from freetoken.moe.prefetch import COPIED, ISSUED, LATE, MISSES, USEFUL + + dev = torch.device("cuda") + steps = args.warm + args.timed + route_all, look_all, reqs = load_trace(args.trace, args.steady_layers, args.layer_offset, bs, steps, + args.skip, args.sigma, args.seed) + layers = args.steady_layers + x = torch.randn(1, H, device=dev, dtype=torch.bfloat16) + probe = torch.randn(32768, H, device=dev, dtype=torch.bfloat16) + per_row_us = time_graph(lambda: F.linear(x, probe), args.reps) * 1000.0 / 32768 + del probe + fill = {us: torch.randn(max(int(us / per_row_us), 1), H, device=dev, dtype=torch.bfloat16) * 0.01 + for us in (args.filler_gdn_us, args.filler_qsa_us)} + # the filler of layer i stands in for its attention: QSA layers are 3, 7, ... (i % 4 == 3) + fillers = [fill[args.filler_qsa_us if i % 4 == 3 else args.filler_gdn_us] for i in range(layers)] + qsa = torch.tensor([i % 4 == 3 for i in range(layers)]) + route_buf = torch.empty_like(route_all[0]) + look_buf = torch.empty_like(look_all[0]) + stack = DecodeStack(layers, args.steady_slots, bs, 0, "on", True, filler=fillers, routing=(route_buf, look_buf)) + cache, pf = stack.cache, stack.cache.prefetch + cache.collect_stats = True + experts = [moe.experts for moe in stack.moes] + state = [cache.slot_for_id, cache.id_of_slot, cache.usage, cache.step, cache.lru_stats, pf.stats] + for name in ("evict_tok", "evict_last_tok", "evict_lc", "evict_ct", "evict_slot_owner", + "evict_slot_last_tok", "evict_slot_lc", "evict_slot_ct"): + if getattr(cache, name) is not None: + state.append(getattr(cache, name)) + init = [t.clone() for t in state] + + def load(snap): + for dst, src in zip(state, snap): + dst.copy_(src) + + def save(snap): + for dst, src in zip(snap, state): + dst.copy_(src) + + route_buf.copy_(route_all[0]) + look_buf.copy_(look_all[0]) + graphs, snaps, keep = {}, {}, [] + for arm in dict.fromkeys(arms): + mode, _, spec = arm.partition(":") + budgets = tuple(int(b) for b in spec.split("/")) if spec and spec != "empty" else (3, 4) + for i in range(1, layers): + experts[i - 1].set_lookahead(experts[i - 1]._lookahead[0], i, budgets[(i + 1) % 4 == 0]) + # every mode switch is read at enqueue time, so each arm's graph bakes its own path in + pf.mode = "on" if mode == "off" else mode + pf.sel_override = torch.full_like(pf.sel, -1) if spec == "empty" else None + keep.append(pf.sel_override) + cache.prefetch = None if mode == "off" else pf + load(init) + stack.step() # compile outside capture + torch.cuda.synchronize() + load(init) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + stack.step() + graphs[arm], snaps[arm] = graph, [t.clone() for t in init] + cache.prefetch, pf.mode, pf.sel_override = pf, "on", None + names = list(graphs) + events = {arm: [] for arm in names} + for t in range(steps): + if t == args.warm: + for arm in names: + snaps[arm][4].zero_() + snaps[arm][5].zero_() + route_buf.copy_(route_all[t]) + look_buf.copy_(look_all[t]) + for j in range(len(names)): + arm = names[(t + j) % len(names)] + load(snaps[arm]) + ev = (torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)) + ev[0].record() + graphs[arm].replay() + ev[1].record() + save(snaps[arm]) + if t >= args.warm: + events[arm].append(ev) + if t % 50 == 49: + torch.cuda.synchronize() # bound the event backlog + torch.cuda.synchronize() + out = [] + n = args.timed + for arm in names: + mode = arm.partition(":")[0] + times = sorted(a.elapsed_time(b) for a, b in events[arm]) + miss = snaps[arm][4][:, Stat.MISS].double().cpu() + row = {"arm": arm, "mean": statistics.fmean(times), "p50": times[len(times) // 2], + "p90": times[int(len(times) * 0.9)], "miss": miss.sum().item() / (n * layers), + "miss1": miss[1:].sum().item() / (n * (layers - 1))} + if mode != "off": + st = snaps[arm][5].double().cpu() + calls = n * (layers - 1) + row.update(issued=st[:, ISSUED].sum().item() / calls, useful=st[:, USEFUL].sum().item() / calls) + row["precision"] = row["useful"] / row["issued"] if row["issued"] else 0.0 + if mode == "measure": + row["coverage"] = st[:, USEFUL].sum().item() / max(st[:, MISSES].sum().item(), 1) + if mode == "on": + # every routed id that was not a normal resident: residual misses plus prefetch hits + row["nonres1"] = row["miss1"] + row["useful"] + for name, mask in (("late_gdn", ~qsa), ("late_qsa", qsa)): + m = mask.clone() + m[0] = False + row[name] = st[m, LATE].sum().item() / max(st[m, COPIED].sum().item(), 1) + row["late"] = st[:, LATE].sum().item() / max(st[:, COPIED].sum().item(), 1) + out.append(row) + print(f"# peak allocated {torch.cuda.max_memory_allocated(dev) / 2**30:.2f} GiB, requests {reqs}", flush=True) + del graphs, stack, cache, pf, events + gc.collect() + torch.cuda.empty_cache() + return out + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--bs", default="1,2") + ap.add_argument("--k", type=int, default=16) + ap.add_argument("--budget", type=int, default=4) + ap.add_argument("--layers", type=int, default=8, help="MoE layers in the decode stack") + ap.add_argument("--slots", type=int, default=512, help="expert slots (2.77 MB each)") + ap.add_argument("--misses", default="0,4", help="demand misses per layer per row") + ap.add_argument("--reps", type=int, default=30) + ap.add_argument("--overlap", type=int, default=1, help="FREETOKEN_MOE_COPY_OVERLAP in the stack (production: 1)") + ap.add_argument("--on", action="store_true", help="time whole decode steps with prefetch off vs on instead") + ap.add_argument("--on-layers", type=int, default=48) + ap.add_argument("--on-slots", type=int, default=0, help="0 = every layer's routed experts + 64 (the warm-up must not evict)") + ap.add_argument("--filler-us", type=float, default=160.0, help="per-layer compute standing in for attention") + ap.add_argument("--useful", default="2,1,2,2", help="useful prefetches per layer, cycled over layers") + ap.add_argument("--delay-us", default="0,150,300", help="sleep ahead of every prefetch copy (0 = none)") + ap.add_argument("--steady", action="store_true", help="steady-state decode on a routing trace instead") + ap.add_argument("--trace", default=None, help="--steady: dir with experts/logit_idx/logit_val/req_index .npy") + ap.add_argument("--steady-layers", type=int, default=24, help="--steady: layers (VRAM: slots * 2.77 MB)") + ap.add_argument("--layer-offset", type=int, default=0, help="--steady: first traced layer") + ap.add_argument("--steady-slots", type=int, default=825, help="--steady: slots (1650 * 24 / 48 by default)") + ap.add_argument("--warm", type=int, default=300, help="--steady: untimed steps before timing") + ap.add_argument("--timed", type=int, default=500) + ap.add_argument("--skip", type=int, default=200, help="--steady: tokens skipped at each request's start") + ap.add_argument("--sigma", type=float, default=0.3, help="--steady: lookahead noise on the traced logits") + ap.add_argument("--seed", type=int, default=0) + ap.add_argument("--filler-gdn-us", type=float, default=171.0, help="--steady: compute before a GDN layer's MoE") + ap.add_argument("--filler-qsa-us", type=float, default=212.0, help="--steady: compute before a QSA layer's MoE") + ap.add_argument("--arms", default="off,measure,on,on:empty,on:2/3,on:3/3,on:4/4", + help="--steady: mode[:gdn/qsa budgets | :empty (on, nothing selected)]") + args = ap.parse_args() + from freetoken.env import ENV + + ENV.MOE_PREFETCH_K.value, ENV.MOE_PREFETCH_BUDGET.value = args.k, 0 if args.steady else args.budget + dev = torch.device("cuda") + if args.steady: + print(f"# {torch.cuda.get_device_name(dev)}, steady state on {args.trace}: {args.steady_layers} layers from " + f"{args.layer_offset}, {args.steady_slots} slots, rule, copy overlap on, filler {args.filler_gdn_us:.0f}/" + f"{args.filler_qsa_us:.0f} us (GDN/QSA), sigma {args.sigma}, {args.warm} warm + {args.timed} timed steps") + print(f"{'bs':>3} {'arm':>10} {'mean ms':>8} {'vs off':>7} {'p50':>7} {'p90':>7} {'miss/L':>7} {'miss/L1+':>8} " + f"{'iss/L':>6} {'use/L':>6} {'prec':>6} {'nonres/L1+':>10} {'late':>6} {'lateGDN':>7} {'lateQSA':>7}") + for bs in (int(b) for b in args.bs.split(",")): + rows = bench_steady(args, bs, args.arms.split(",")) + base = rows[0]["mean"] + for r in rows: + g = r.get + print(f"{bs:>3} {r['arm']:>10} {r['mean']:>8.3f} {r['mean'] / base - 1:>+7.1%} {r['p50']:>7.3f} " + f"{r['p90']:>7.3f} {r['miss']:>7.3f} {r['miss1']:>8.3f} {g('issued', 0):>6.2f} {g('useful', 0):>6.2f} " + f"{g('precision', 0):>6.3f} {g('nonres1', 0):>10.3f} {g('late', 0):>6.3f} {g('late_gdn', 0):>7.3f} " + f"{g('late_qsa', 0):>7.3f}", flush=True) + return + if args.on: + useful = [int(u) for u in args.useful.split(",")] + delays = [float(d) for d in args.delay_us.split(",")] + print(f"# {torch.cuda.get_device_name(dev)}, {args.on_layers} Qwen4ExpMoE nvfp4 layers, " + f"{args.on_slots or 'layers * bs * top-k + 64'} slots, " + f"rule, copy overlap on, filler {args.filler_us:.0f} us/layer, useful {useful}, median of {args.reps} replays") + print(f"{'bs':>3} {'m':>3} {'arm':>10} {'ms/step':>9} {'vs off':>8} {'misses/step':>12} " + f"{'issued/layer':>13} {'useful/layer':>13} {'late':>6}") + for bs in (int(b) for b in args.bs.split(",")): + for m in (int(x) for x in args.misses.split(",")): + slots = args.on_slots or args.on_layers * bs * TOP_K + 64 + rows = bench_on(args.on_layers, slots, bs, m, useful, args.reps, args.filler_us, delays) + base = rows[0]["ms"] + for r in rows: + extra = (f"{r['issued']:>13.2f} {r['useful']:>13.2f} {r['late']:>6.3f}" if "late" in r + else f"{'':>13} {'':>13} {'':>6}") + print(f"{bs:>3} {m:>3} {r['arm']:>10} {r['ms']:>9.3f} {r['ms'] / base - 1:>+8.1%} " + f"{r['misses']:>12d} {extra}", flush=True) + print(f"# peak allocated {torch.cuda.max_memory_allocated(dev) / 2**30:.2f} GiB") + return + print(f"# {torch.cuda.get_device_name(dev)}, K={args.k}, budget={args.budget}, median of {args.reps} graph replays") + print("# predictor alone, us per layer (47 layers, one stream)") + print(f"{'bs':>3} {'gemv':>8} {'select':>8} {'predictor':>10} {'count':>8}") + for bs in (int(b) for b in args.bs.split(",")): + r = bench_predictor(bs, 47, args.k, args.budget, args.reps) + print(f"{bs:>3} {r['gemv']:>8.2f} {r['select']:>8.2f} {r['predictor']:>10.2f} {r['count']:>8.2f}", flush=True) + print(f"# decode stack ({args.layers} Qwen4ExpMoE nvfp4 layers, {args.slots} slots, rule, overlap={args.overlap}), " + "us per layer, prefetch off vs measure") + print(f"{'bs':>3} {'m':>3} {'off':>9} {'measure':>9} {'delta':>8} {'issued/layer':>13}") + for bs in (int(b) for b in args.bs.split(",")): + for m in (int(x) for x in args.misses.split(",")): + r = bench_stack(args.layers, args.slots, bs, m, args.reps, bool(args.overlap)) + print(f"{bs:>3} {m:>3} {r['off']:>9.2f} {r['measure']:>9.2f} {r['measure'] - r['off']:>+8.2f} " + f"{r.get('issued/layer', 0.0):>13.2f}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 11e13229c9..c46156ab7e 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -1280,6 +1280,9 @@ def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput: use_graph = self.graph_runner.can_use_cuda_graph(batch) with self.ctx.forward_batch(batch), self.model.forward_host_ctx(batch, use_graph): logits = self.graph_runner.replay(batch) if use_graph else self.model.forward() + audit = self.moe_offload_cache.audit if self.moe_offload_cache is not None else None + if audit is not None and batch.is_decode: + audit.end_decode_step(self.moe_offload_cache) # every side stream has joined the compute stream if self.cpu_moe_executor is not None: # One pinned read: surfaces a fired flag-handshake watchdog (dead coordinator # -> stale expert outputs) as a loud error instead of silent corruption. @@ -1294,11 +1297,21 @@ def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput: logprobs_out = self.sampler.compute_logprobs(batch_logits, next_tokens_gpu, args) copy_done_event = torch.cuda.Event() copy_done_event.record(self.stream) - if self.moe_offload_cache is not None and self.moe_offload_cache.collect_stats and batch.is_decode: + cache = self.moe_offload_cache + if cache is not None and batch.is_decode and ( + cache.collect_stats or cache.prefetch is not None or cache.verify is not None or cache.audit is not None + ): self._moe_stats_step += 1 if self._moe_stats_step >= MOE_STATS_INTERVAL: self._moe_stats_step = 0 - self._emit_moe_stats() + if cache.prefetch is not None: + self._emit_prefetch_stats() + if cache.verify is not None: + self._emit_verify_stats() + if cache.audit is not None: + self._emit_audit_stats() + if cache.collect_stats: + self._emit_moe_stats() if logprobs_out is None: return ForwardOutput(next_tokens_gpu, next_tokens_cpu, copy_done_event) chosen_logprobs, top_ids, top_logprobs = logprobs_out @@ -1341,6 +1354,40 @@ def _warmup_prefill_lens(self) -> list[int]: if mc.has_linear_attention and cap >= 8193: lens.add(8193) # chunk_delta_h NT bucket 2 (NT > 128) return sorted(n for n in lens if 2 <= n <= cap) + + def _emit_prefetch_stats(self) -> None: + """Report one window of the FREETOKEN_MOE_PREFETCH counters (one host sync), then restart them.""" + from freetoken.moe.prefetch import format_per_layer, format_summary, summarize + + prefetch = self.moe_offload_cache.prefetch + window = prefetch.take_window() + stats = summarize(window, prefetch.mode) + if not stats["layer_calls"]: + return + logger.info_rank0(f"MoE prefetch {prefetch.mode} ({MOE_STATS_INTERVAL} decode steps): {format_summary(stats)}") + if ENV.MOE_PREFETCH_DEBUG: + logger.info_rank0(f"MoE prefetch per layer (useful/issued/misses per call): {format_per_layer(window)}") + + def _emit_verify_stats(self) -> None: + """Report one window of the FREETOKEN_MOE_PREFETCH_VERIFY counts and the new bad records (host syncs).""" + _log_lines(self.moe_offload_cache.verify.report_window(MOE_STATS_INTERVAL)) + + def _emit_audit_stats(self) -> None: + """Report one window of FREETOKEN_MOE_SLOT_AUDIT totals and the new bad slots (host syncs).""" + _log_lines(self.moe_offload_cache.audit.report_window(MOE_STATS_INTERVAL)) + + def _log_prefetch_totals(self) -> None: + from freetoken.moe.prefetch import format_per_layer, format_summary, summarize + + prefetch = self.moe_offload_cache.prefetch if self.moe_offload_cache is not None else None + if prefetch is None: + return + totals = prefetch.totals + prefetch.stats.cpu() + stats = summarize(totals, prefetch.mode) + if stats["layer_calls"]: + logger.info_rank0(f"MoE prefetch {prefetch.mode} (session, {stats['tokens']} tokens): {format_summary(stats)}") + logger.info_rank0(f"MoE prefetch per layer (useful/issued/misses per call): {format_per_layer(totals)}") + def _emit_moe_stats(self) -> None: """Report one window of expert-cache behaviour, then reset the miss counters. @@ -1463,11 +1510,32 @@ def _warmup_prefill(self) -> None: ) def shutdown(self) -> None: + try: + self._log_prefetch_totals() + except Exception as exc: # noqa: BLE001 -- a diagnostic must never block shutdown + logger.warning(f"MoE prefetch totals unavailable at shutdown: {exc}") + try: + verify = self.moe_offload_cache.verify if self.moe_offload_cache is not None else None + if verify is not None: + _log_lines(verify.report_session()) + except Exception as exc: # noqa: BLE001 + logger.warning(f"MoE verify totals unavailable at shutdown: {exc}") + try: + audit = self.moe_offload_cache.audit if self.moe_offload_cache is not None else None + if audit is not None: + _log_lines(audit.report_session()) + except Exception as exc: # noqa: BLE001 + logger.warning(f"MoE slot audit totals unavailable at shutdown: {exc}") self.graph_runner.destroy_cuda_graphs() torch.distributed.destroy_process_group() destroy_distributed() +def _log_lines(lines: list[tuple[str, str]]) -> None: + for level, line in lines: + (logger.warning_rank0 if level == "warning" else logger.info_rank0)(line) + + def _profile_gpu(index: "int | None" = None) -> Tuple[str | None, str | None]: """(name, uuid) of visible device ``index`` (default: the current, i.e. bound, device); (None, None) without CUDA.""" if not torch.cuda.is_available(): diff --git a/python/freetoken/env.py b/python/freetoken/env.py index 0e56633f48..9a06b1c029 100644 --- a/python/freetoken/env.py +++ b/python/freetoken/env.py @@ -79,6 +79,20 @@ class EnvClassSingleton: WARMUP_MAX_LEN = EnvInt(0) # GPU-decode MoE: run a layer's expert miss copy on a side stream under the shared expert MOE_COPY_OVERLAP = EnvBool(False) + # GPU-decode MoE cross-layer prefetch: off | measure (predict and count, no copies) | on (predict and copy) + MOE_PREFETCH = EnvStr("off") + # router-lookahead candidates per layer, and the non-resident ones kept per target layer + # (0 = 3 before GDN layers, 4 before full-attention layers) + MOE_PREFETCH_K = EnvInt(16) + MOE_PREFETCH_BUDGET = EnvInt(0) + # measure mode: also log the per-layer breakdown with every summary + MOE_PREFETCH_DEBUG = EnvBool(False) + # debug: check every GPU decode layer call's slot maps (meta) and slot bytes before/after the GEMM (1) + MOE_PREFETCH_VERIFY = EnvStr("0") + # routed experts a call byte-checks (scratch rows, ~2.8 MB each on Qwen3.8-Flash) + MOE_PREFETCH_VERIFY_ROWS = EnvInt(20) + # debug: record every expert slot writer and audit every held slot each N decode steps (0 = off) + MOE_SLOT_AUDIT = EnvInt(0) def __new__(cls): # single instance diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index 63ab8bb8f2..e1b4b8517a 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -194,6 +194,13 @@ def __init__( prefix=prefix, ) self.offload_cache: OffloadMoeCache | None = None + # FREETOKEN_MOE_PREFETCH: (next MoE layer's router op, its layer id, its budget); the + # underscore keeps the borrowed router out of this layer's state dict + self._lookahead: tuple[BaseOP, int, int] | None = None + + def set_lookahead(self, gate: BaseOP, target_layer: int, budget: int) -> None: + """Let this layer's GPU decode predict ``target_layer``'s experts with that layer's router.""" + self._lookahead = (gate, target_layer, budget) @property def skips_inactive_slots(self) -> bool: @@ -305,9 +312,25 @@ def _decode_routed( return executor.decode(self.layer_id, hidden_states, topk_weights, topk_ids) if cache.decode_target == "hybrid": return self._decode_hybrid(cache, hidden_states, topk_weights, topk_ids) + prefetch, verify = cache.prefetch, cache.verify + if verify is not None: + verify.begin(cache, self.layer_id, topk_ids, prefetch and prefetch.plan_of(self.layer_id)) cache.ensure_experts(self.layer_id, topk_ids, router_logits=router_logits) + target = None + if prefetch is not None: + # on: install the next layer's prediction (forked at this block's start) on this stream, the slot + # maps' only writer; before this layer's copy join, so a late prefetch copy never delays it + target = prefetch.install(self._lookahead and self._lookahead[1]) + prefetch.join_copy(self.layer_id) cache.copy_missing() - return self._expert_gemm( + if prefetch is not None: + prefetch.join_and_count(self.layer_id, topk_ids, cache.id_of_slot, cache.num_indices) + if target is not None: + prefetch.issue_copy(target) + if verify is not None: + # after issue_copy: the next layer's prefetch copy starts when it would without the checks + verify.before_gemm(cache, topk_ids, None if target is None else prefetch.plan_of(target)) + out = self._expert_gemm( cache, hidden_states, topk_weights, @@ -317,6 +340,29 @@ def _decode_routed( alphas=cache.alphas_for_slots(self.layer_id), is_prefill=False, ) + if verify is not None: + verify.after_gemm(cache) + return out + + def fork_lookahead(self, hidden_states: torch.Tensor) -> None: + """FREETOKEN_MOE_PREFETCH: at the start of this layer's MoE block, predict the next layer's + experts from this block's router input on the predictor stream, beside this layer's router + and ensure. Only on the GPU decode path that ``_decode_routed`` takes and not under TP, and + never toward a CPU-decoded layer, whose ensure would not use it.""" + cache = self.offload_cache + if ( + self._lookahead is None + or cache is None + or cache.prefetch is None + or self.tp_size > 1 + or get_global_ctx().batch.is_prefill + or cache.is_cpu_layer(self.layer_id) + or cache.decode_target == "hybrid" + ): + return + gate, target, budget = self._lookahead + if not cache.is_cpu_layer(target): + cache.prefetch.fork(target, hidden_states, gate, cache.slot_for_id[target], budget) def _decode_hybrid( self, diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py index 5d4c6f9d2f..7f4f826807 100644 --- a/python/freetoken/models/qwen4_exp/model.py +++ b/python/freetoken/models/qwen4_exp/model.py @@ -26,7 +26,7 @@ from .attention import Qwen4ExpAttention from .hc import GatedResidual -from .moe import Qwen4ExpMoE +from .moe import Qwen4ExpMoE, wire_router_lookahead from .ple import PLELayer from freetoken.models.blocks import embed_input_ids from freetoken.models.qwen3_vl.vision import Qwen3VLVisionModel, QwenVLVisionMixin @@ -101,6 +101,7 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model") -> None: for layer_id in range(config.num_layers) ] ) + wire_router_lookahead([layer.mlp for layer in self.layers.op_list], config) self.hyper_connection_mixer = GatedResidual(config, use_combine=False, prefix=f"{prefix}.hyper_connection_mixer") # plain tuple (not an OP child), so it never shows up in the state dict self._ple = tuple(layer.ple for layer in self.layers.op_list if layer.ple is not None) diff --git a/python/freetoken/models/qwen4_exp/moe.py b/python/freetoken/models/qwen4_exp/moe.py index 6f729231f1..4a1907d2ce 100644 --- a/python/freetoken/models/qwen4_exp/moe.py +++ b/python/freetoken/models/qwen4_exp/moe.py @@ -36,6 +36,10 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: num_tokens, hidden_dim = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_dim) se, ex = self.shared_expert, self.experts + if isinstance(ex, OffloadMoELayer): + # FREETOKEN_MOE_PREFETCH: the next layer's prediction needs only this block's input, so it + # starts here and runs beside this layer's router and ensure + ex.fork_lookahead(hidden_states) if isinstance(ex, OffloadMoELayer) and ex.decode_side_applies(): # FREETOKEN_MOE_COPY_OVERLAP: the shared expert runs on a side stream while the router, # ensure and miss copy stay back to back on this one; joined before the combine @@ -65,4 +69,15 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return shared_gate_mul_add(routed, shared, gate).view(num_tokens, hidden_dim) -__all__ = ["Qwen4ExpMoE"] +def wire_router_lookahead(moes: list[Qwen4ExpMoE], config: ModelConfig) -> None: + """Give each offloaded MoE block the next layer's router, for FREETOKEN_MOE_PREFETCH's + lookahead (inert when it is off). ``moes[i]`` is decoder layer ``i``'s block.""" + from freetoken.moe.prefetch import default_budget + + for target in range(1, len(moes)): + experts = moes[target - 1].experts + if isinstance(experts, OffloadMoELayer): + experts.set_lookahead(moes[target].gate, target, default_budget(config.is_linear_layer(target))) + + +__all__ = ["Qwen4ExpMoE", "wire_router_lookahead"] diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 96b13f33de..e897fff3d4 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -2,6 +2,7 @@ import math import os +import weakref from contextlib import contextmanager from dataclasses import dataclass from typing import TYPE_CHECKING, Iterator @@ -167,6 +168,13 @@ class OffloadMoeCache: # GPU decode only: the caller's work that is independent of the routed experts (qwen4_exp: the # shared expert) runs on decode_copy_stream beside the miss copy. None = FREETOKEN_MOE_COPY_OVERLAP. decode_copy_overlap: bool | None = None + # GPU decode only: cross-layer expert prefetch (moe/prefetch.py). None = FREETOKEN_MOE_PREFETCH. + prefetch_mode: str | None = None + # GPU decode only: debug slot checks (moe/verify.py). None = FREETOKEN_MOE_PREFETCH_VERIFY. + verify_mode: str | None = None + # GPU decode only: debug slot writer history + audit every N decode steps (moe/slot_audit.py). + # None = FREETOKEN_MOE_SLOT_AUDIT, 0 = off. + slot_audit: int | None = None def __post_init__(self) -> None: from freetoken.moe.scored_ensure import POLICY_IDS @@ -322,6 +330,9 @@ def __post_init__(self) -> None: self._stage_bytes: list[int] = [] self._stage_views: list[list[tuple[torch.Tensor, ...] | None]] = [] self._inval: list[slice | torch.Tensor] = [] + # the slot audit's label for the staged copy, and the most entries its plan can have + self._pending_kind = 0 + self._pending_width: int | None = None self.prefill_copy_stream: torch.cuda.Stream | None = None self.prefill_begin_event: torch.cuda.Event | None = None self.prefill_ready_events: list[torch.cuda.Event] = [] @@ -350,6 +361,70 @@ def __post_init__(self) -> None: self._decode_copy_events: list[tuple[torch.cuda.Event, torch.cuda.Event]] = [] if self.decode_copy_overlap and self.device.type == "cuda" and self.decode_target == "gpu": self._init_decode_copy_overlap() + self.prefetch = self._init_prefetch() + self.verify = self._init_verify() + self.audit = self._init_audit() + + def _init_audit(self): + from freetoken.kernel.fast_index_copy import _skip_fast_index_copy_enabled + from freetoken.moe.slot_audit import SlotAuditor, resolve_interval + + from_env = self.slot_audit is None + self.slot_audit = resolve_interval(self.slot_audit) + if not self.slot_audit or self.device.type != "cuda": + return None + if self.decode_target != "gpu": + if not from_env: + raise ValueError(f"FREETOKEN_MOE_SLOT_AUDIT={self.slot_audit} needs GPU decode, not {self.decode_target!r}") + logger.warning(f"FREETOKEN_MOE_SLOT_AUDIT={self.slot_audit} ignored: decode target is {self.decode_target!r}") + self.slot_audit = 0 + return None + if _skip_fast_index_copy_enabled(): + logger.warning("FREETOKEN_MOE_SLOT_AUDIT ignored: FREETOKEN_SKIP_FAST_INDEX_COPY makes every copy a no-op") + self.slot_audit = 0 + return None + return SlotAuditor(self.num_layers, self.num_experts, self.device, interval=self.slot_audit) + + def _init_verify(self): + from freetoken.moe.verify import ExpertVerifier, resolve_mode + + from_env = self.verify_mode is None + self.verify_mode = resolve_mode(self.verify_mode) + if self.verify_mode == "off" or self.device.type != "cuda": + return None + if self.decode_target != "gpu": + if not from_env: + raise ValueError(f"FREETOKEN_MOE_PREFETCH_VERIFY={self.verify_mode} needs GPU decode, not {self.decode_target!r}") + logger.warning(f"FREETOKEN_MOE_PREFETCH_VERIFY={self.verify_mode} ignored: decode target is {self.decode_target!r}") + self.verify_mode = "off" + return None + return ExpertVerifier(self.num_layers, self.num_experts, self.device, mode=self.verify_mode) + + def _init_prefetch(self): + from freetoken.moe.prefetch import ExpertPrefetcher, resolve_mode + + from_env = self.prefetch_mode is None + self.prefetch_mode = resolve_mode(self.prefetch_mode) + if self.prefetch_mode == "off": + return None + if self.decode_target != "gpu": + if not from_env: + raise ValueError(f"FREETOKEN_MOE_PREFETCH={self.prefetch_mode} needs GPU decode, not {self.decode_target!r}") + # the process-wide flag skips hybrid/CPU-decoded caches, like FREETOKEN_MOE_COPY_OVERLAP + logger.warning(f"FREETOKEN_MOE_PREFETCH={self.prefetch_mode} ignored: decode target is {self.decode_target!r}") + self.prefetch_mode = "off" + return None + if self.device.type != "cuda": + return None + prefetch = ExpertPrefetcher(self.num_layers, self.num_experts, self.device, mode=self.prefetch_mode) + # a proxy, not a cycle: dropping the cache must free its slot banks without waiting for gc + prefetch.cache = weakref.proxy(self) + return prefetch + + @property + def prefetch_on(self) -> bool: + """FREETOKEN_MOE_PREFETCH=on: every ensure keys prefetched slots low priority.""" + return self.prefetch is not None and self.prefetch.mode == "on" def _pool_layout(self, pool: int) -> int: return self.layer_layout[self.pool_of_layer.index(pool)] @@ -611,6 +686,10 @@ def set_bank_sources( self.prefill_overlap = False self._slice_pools() self._build_copy_plan() + if self.verify is not None: + self.verify.bind(self) + if self.audit is not None: + self.audit.bind(self) if self.prefill_overlap: self._init_prefill_overlap_buffers() @@ -776,6 +855,8 @@ def rebuild(self, cache_size: int, pool_sizes: list[int] | None = None) -> None: self.usage = torch.zeros((cache_size,), dtype=torch.int64, device=self.device) self._slice_pools() self._build_copy_plan() # slot caches were reallocated -> refresh fused-copy addrs + if self.verify is not None: + self.verify.bind(self) plan_slots = max(self.num_experts, cache_size) self.evict_slots = torch.empty((plan_slots,), dtype=torch.int32, device=self.device) self.src_indices = torch.empty((plan_slots,), dtype=torch.int32, device=self.device) @@ -797,6 +878,10 @@ def rebuild(self, cache_size: int, pool_sizes: list[int] | None = None) -> None: self.stat_fetched_layer.zero_() self.stat_steps_layer.zero_() self.decode_freq.zero_() + if self.prefetch is not None: + self.prefetch.reset() + if self.audit is not None: + self.audit.bind(self) # new slot arrays: a new history self.prefill_hit_rows = 0 self.prefill_total_rows = 0 self._hit_d2d_fallback_logged = False # geometry changed; re-log if still unusable @@ -897,9 +982,13 @@ def _init_prefill_overlap_buffers(self) -> None: ) self._prefill_hit_num = torch.zeros((1,), dtype=torch.int64, device=self.device) - def _invalidate_prefill_buffer(self, buffer_id: int) -> None: + def _invalidate_prefill_buffer(self, buffer_id: int, layer_id: int = -1) -> None: slots = self._inval[buffer_id] if isinstance(slots, slice): + if self.audit is not None: + from freetoken.moe.slot_audit import INVALIDATE + + self.audit.record_clear(self, INVALIDATE, layer_id, slots.start, slots.stop - slots.start) # One fixed-shape launch. The boolean-mask index # (slot_for_id[old_ids[old_ids >= 0]] = -1) has a data-dependent shape, so every # call hid a device-to-host sync -- and with two buffer reuses per chunk x 48 @@ -960,9 +1049,13 @@ def prefetch_prefill_layer(self, layer_id: int) -> None: ) def copy() -> None: - self._invalidate_prefill_buffer(buffer_id) + self._invalidate_prefill_buffer(buffer_id, layer_id) for (per_layer, _), view in zip(self.banks, self.stage_views(layer_id, buffer_id)): view.copy_(per_layer[layer_id], non_blocking=True) + if self.audit is not None: + from freetoken.moe.slot_audit import PREFILL_BUFFER + + self.audit.record_range(self, PREFILL_BUFFER, layer_id, buffer_id * self.num_experts, self.num_experts, meta=False) if self._prefill_hit_d2d_active: self._prefetch_split(layer_id, buffer_id) @@ -1064,11 +1157,16 @@ def _prefetch_split(self, layer_id: int, buffer_id: int) -> None: self._prefill_hit_num, blocks_per_bank=64, ) + if self.audit is not None: + from freetoken.moe.slot_audit import PREFILL_HIT_D2D + + self.audit.record_plan(self, PREFILL_HIT_D2D, layer_id, self._prefill_hit_dst, None, self._prefill_hit_num, + meta=False, aux=self._prefill_hit_src, row_base=buffer_id * E) miss = np.nonzero(~hit_mask)[0] with torch.cuda.stream(self.prefill_copy_stream): if self._prefill_buffer_has_release_event[buffer_id]: self.prefill_copy_stream.wait_event(self.prefill_release_events[buffer_id]) - self._invalidate_prefill_buffer(buffer_id) + self._invalidate_prefill_buffer(buffer_id, layer_id) if miss.size: run_starts = np.concatenate(([0], np.nonzero(np.diff(miss) != 1)[0] + 1)) starts = miss[run_starts] @@ -1094,6 +1192,11 @@ def _prefetch_split(self, layer_id: int, buffer_id: int) -> None: torch.tensor(nbytes, dtype=torch.int64), torch.cuda.current_stream(self.device).cuda_stream, ) + if self.audit is not None: + from freetoken.moe.slot_audit import PREFILL_SPLIT_H2D + + # every row gets at least its small banks from the host, the misses all of them + self.audit.record_range(self, PREFILL_SPLIT_H2D, layer_id, buffer_id * E, E, meta=False) self.prefill_ready_events[buffer_id].record(self.prefill_copy_stream) def wait_prefill_layer(self, layer_id: int) -> tuple[torch.Tensor, ...]: @@ -1151,18 +1254,49 @@ def ensure_experts( self.decode_freq[layer_id].scatter_add_(0, ids, torch.ones_like(ids)) self._pending_src_layer = layer_id self._pending_whole_layer = False - if self.cache_policy_id == 0: + audit = self.audit + if audit is not None and update_state and layer_id == self.evict_clock_layer(): + audit.begin_step() + lowpri = self.prefetch_on + width = expert_ids.numel() + if self.cache_policy_id == 0 and not lowpri: ensure_experts(self, layer_id, expert_ids) - return - ensure_experts_scored( - self, - layer_id, - expert_ids, - bump_tok=update_state and layer_id == self.evict_clock_layer(), - update_state=update_state, - router_logits=router_logits, - pin_since=pin_since, - ) + else: + # prefetch on: LRU runs the vendored kernel (flashlib's cannot key low-priority slots) + ensure_experts_scored( + self, + layer_id, + expert_ids, + bump_tok=update_state and layer_id == self.evict_clock_layer(), + update_state=update_state, + router_logits=router_logits, + pin_since=pin_since, + lowpri=lowpri, + pf_count=self.prefetch.count_args(layer_id) if lowpri and update_state else None, + ) + if audit is not None: + from freetoken.moe import slot_audit as sa + + kind = sa.DEMAND_INSTALL if update_state else sa.PREFILL_INSTALL + self._pending_kind = sa.DEMAND_COPY if update_state else sa.PREFILL_COPY + self._pending_width = width + audit.record_plan(self, kind, layer_id, self.evict_slots, self.src_indices, self.num_indices, meta=True, width=width) + + def prefetch_ensure( + self, layer_id: int, query: torch.Tensor, dst_slots: torch.Tensor, src_rows: torch.Tensor, + num: torch.Tensor, stats_row: torch.Tensor, ready: torch.Tensor, budget: int | None = None, + ) -> None: + """Install the first ``budget`` of ``query``'s non-resident experts of ``layer_id`` as + low-priority slots and plan their copy into ``dst_slots``/``src_rows``/``num`` (never the + demand plan, which copy(L-1) still reads). Runs on the compute stream like every other slot + map writer.""" + from freetoken.moe.offload_kernels import prefetch_ensure_experts + + prefetch_ensure_experts(self, layer_id, query, dst_slots, src_rows, num, stats_row, ready, budget=budget) + if self.audit is not None: + from freetoken.moe.slot_audit import PREFETCH_INSTALL + + self.audit.record_plan(self, PREFETCH_INSTALL, layer_id, dst_slots, src_rows, num, meta=True) def ensure_experts_hybrid( self, layer_id: int, expert_ids: torch.Tensor, weights: torch.Tensor, clamp: bool = False @@ -1193,18 +1327,36 @@ def materialize_layer(self, layer_id: int) -> None: self._pending_src_layer = layer_id self._pending_whole_layer = True + audit = self.audit + if audit is not None: + from freetoken.moe import slot_audit as sa + + # the kernel also frees this layer's slots above E; record them while they still name it + audit.record_clear(self, sa.MATERIALIZE_CLEAR, layer_id, 0, self.cache_size, layer_above_e=True) materialize_layer(self, layer_id) if self.evict_slot_owner is not None: self.evict_slot_owner[: self.num_experts].fill_(-1) # it installs the layer into these slots + if audit is not None: + self._pending_kind = sa.MATERIALIZE_COPY + self._pending_width = self.num_experts + audit.record_plan(self, sa.MATERIALIZE_INSTALL, layer_id, self.evict_slots, self.src_indices, self.num_indices, + meta=True, width=self.num_experts) def reset(self) -> None: from freetoken.moe.offload_kernels import reset_cache + if self.audit is not None: + from freetoken.moe.slot_audit import RESET + + self.audit.record_clear(self, RESET, -1, 0, self.cache_size) reset_cache(self) # Per-expert recency is not cache_size-shaped, so reset_cache leaves it alone; wipe # it here so a new sequence starts with cold hybrid fetch priorities. self.expert_recency.fill_(-1) self.reset_evict_state() + if self.prefetch is not None: + # the graph runner resets after each capture: drop the capture-time counts too + self.prefetch.reset_counters() def reset_stats(self) -> None: self.prefill_hit_rows = 0 @@ -1329,39 +1481,56 @@ def copy_missing(self, layer_id: int | None = None) -> None: # never CUDA-graph captured: prefill is not captured, and decode never reaches this branch (it routes to the CPU executor) for (per_layer, _), view in zip(self.banks, self.stage_views(layer_id, 0)): view.copy_(per_layer[layer_id]) + if self.audit is not None: + from freetoken.moe.slot_audit import PAGEABLE_COPY + + self.audit.record_range(self, PAGEABLE_COPY, layer_id, 0, self.num_experts, meta=False) return # materialize stages rows of staging buffer 0; ensure stages slots local to the layer's pool - whole = self._pending_whole_layer + self.copy_rows(layer_id, self.evict_slots, self.src_indices, self.num_indices, slim=_SLIM_COPY, + kind=self._pending_kind, width=self._pending_width, whole=self._pending_whole_layer) + + def copy_rows( + self, layer_id: int, dst_slots: torch.Tensor, src_rows: torch.Tensor, num: torch.Tensor, *, slim: bool, + kind: int = 0, width: int | None = None, whole: bool = False, + ) -> None: + """Copy host rows ``src_rows[:num]`` of ``layer_id``'s banks into slots ``dst_slots[:num]`` + on the current stream (the demand plan, or a prefetch plan on the prefetch copy stream); + ``kind`` labels the copy for the slot audit and ``width`` bounds ``num`` for it. ``dst_slots`` + are local to the layer's pool, or rows of staging buffer 0 when ``whole`` (materialize).""" + self._copy_rows(layer_id, dst_slots, src_rows, num, slim=slim, whole=whole) + if self.audit is not None: + from freetoken.moe.slot_audit import OTHER_COPY + + self.audit.record_plan(self, kind or OTHER_COPY, layer_id, dst_slots, src_rows, num, meta=False, width=width) + + def _copy_rows( + self, layer_id: int, dst_slots: torch.Tensor, src_rows: torch.Tensor, num: torch.Tensor, *, slim: bool, + whole: bool = False, + ) -> None: if self._copy_fused_ok: from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit, fast_index_copy_multi_slim_jit # One launch copies the missing rows for every bank (instead of one launch per - # bank). evict_slots/src_indices/num_indices are shared across banks; - # src_indices holds layer-local expert rows, resolved against this layer's - # source pointers (layer_id is a static int per captured graph node). - copy = fast_index_copy_multi_slim_jit if _SLIM_COPY else fast_index_copy_multi_jit - copy( - self._copy_dst_ptrs if whole else self._copy_dst_ptrs_pool[self.pool_of_layer[layer_id]], - self._copy_src_ptrs[layer_id], - self._copy_feat_bytes[self.layer_layout[layer_id]], - self.evict_slots, - self.src_indices, - self.num_indices, - blocks_per_bank=self._copy_blocks_per_bank, - ) + # bank). The slot/row plan is shared across banks; src_rows holds layer-local + # expert rows, resolved against this layer's source pointers (layer_id is a + # static int per captured graph node). + dst_ptrs = self._copy_dst_ptrs if whole else self._copy_dst_ptrs_pool[self.pool_of_layer[layer_id]] + feat_bytes = self._copy_feat_bytes[self.layer_layout[layer_id]] + if slim: + fast_index_copy_multi_slim_jit(dst_ptrs, self._copy_src_ptrs[layer_id], feat_bytes, dst_slots, src_rows, num) + else: + fast_index_copy_multi_jit( + dst_ptrs, self._copy_src_ptrs[layer_id], feat_bytes, dst_slots, src_rows, num, + blocks_per_bank=self._copy_blocks_per_bank, + ) return from freetoken.kernel import fast_index_copy_jit views = self.stage_views(layer_id, 0) if whole else self.pool_bank_views(layer_id) for (per_layer, _), cache in zip(self.banks, views): - fast_index_copy_jit( - cache, - self.evict_slots, - per_layer[layer_id], - self.src_indices, - self.num_indices, - ) + fast_index_copy_jit(cache, dst_slots, per_layer[layer_id], src_rows, num) @contextmanager def decode_side(self, layer_id: int) -> Iterator[None]: diff --git a/python/freetoken/moe/offload_kernels.py b/python/freetoken/moe/offload_kernels.py index 72f8bd5905..e34a591f76 100644 --- a/python/freetoken/moe/offload_kernels.py +++ b/python/freetoken/moe/offload_kernels.py @@ -59,10 +59,15 @@ def ensure_experts_scored( update_state: bool, router_logits: torch.Tensor | None = None, pin_since: torch.Tensor | None = None, + lowpri: bool = False, + pf_count: tuple | None = None, ) -> None: - """``ensure_experts`` for the scored ``--moe-cache-policy`` kinds (kd, kdfb, rule).""" + """``ensure_experts`` for the scored ``--moe-cache-policy`` kinds (kd, kdfb, rule), and for + every policy under FREETOKEN_MOE_PREFETCH=on (``lowpri``; ``pf_count`` is the prefetcher's + ``(stats_row, ready, num, rows)`` when this layer was prefetched this step).""" from freetoken.moe.scored_ensure import scored_ensure + stats_row, ready, num, rows = pf_count if pf_count is not None else (None, None, None, 0) scored_ensure( expert_ids, cache.slot_for_id.view(-1), @@ -73,6 +78,23 @@ def ensure_experts_scored( cache.src_indices, cache.evict_slots, cache.num_indices, + router_logits=router_logits if cache.near_miss_thr is not None else None, + near_miss_thr=cache.near_miss_thr or 0.0, + stats=cache.lru_stats[layer_id] if cache.collect_stats else None, + bump_tok=bump_tok, + update_state=update_state, + pin_since=pin_since, + lowpri=lowpri, + pf_stats=stats_row, + pf_ready=ready, + pf_num=num, + pf_rows=rows, + **_evict_state(cache, layer_id), + ) + + +def _evict_state(cache, layer_id: int) -> dict: + return dict( policy=cache.cache_policy_id, num_layers=cache.num_layers, num_experts=cache.num_experts, @@ -85,13 +107,21 @@ def ensure_experts_scored( slot_last_tok=cache.evict_slot_last_tok, slot_lc=cache.evict_slot_lc, slot_ct=cache.evict_slot_ct, - router_logits=router_logits if cache.near_miss_thr is not None else None, - near_miss_thr=cache.near_miss_thr or 0.0, - stats=cache.lru_stats[layer_id] if cache.collect_stats else None, id_base=layer_id * cache.num_experts, - bump_tok=bump_tok, - update_state=update_state, - pin_since=pin_since, + ) + + +def prefetch_ensure_experts( + cache, layer_id: int, query: torch.Tensor, dst_slots: torch.Tensor, src_rows: torch.Tensor, + num: torch.Tensor, stats_row: torch.Tensor, ready: torch.Tensor, budget: int | None = None, +) -> None: + """FREETOKEN_MOE_PREFETCH=on: install the first ``budget`` of ``query``'s non-resident experts + of ``layer_id`` as low-priority slots and plan their copy (scored_ensure.prefetch_ensure).""" + from freetoken.moe.scored_ensure import prefetch_ensure + + prefetch_ensure( + query, cache.slot_for_id.view(-1), cache.id_of_slot, cache.usage, cache.step, + dst_slots, src_rows, num, stats_row, ready, budget=budget, **_evict_state(cache, layer_id), ) diff --git a/python/freetoken/moe/prefetch.py b/python/freetoken/moe/prefetch.py new file mode 100644 index 0000000000..01129d3997 --- /dev/null +++ b/python/freetoken/moe/prefetch.py @@ -0,0 +1,438 @@ +"""Cross-layer expert prefetch for GPU decode (``FREETOKEN_MOE_PREFETCH``). + +A training-free router lookahead predicts layer L's experts from the tensor layer L-1's router +consumed (``W_gate[L] . x_{L-1}``, top-K) and keeps the first BUDGET candidates that are not +resident. At bs 2 both rows' top-K lists merge rank by rank (row 0's first, then row 1's +first, ...), duplicates dropped, under one budget. + +``measure``: after layer L's ensure, count how many of them layer L actually routed. Nothing is +copied and the cache is never touched. The predictor reads ``slot_for_id[L]`` without ordering +against the ensures; a late read only skews the counters, never the model. + +The prediction forks at the start of layer L-1's MoE block, where x_{L-1} first exists, so it runs +beside layer L-1's own router and ensure: + + compute: -fork(L)-> router_{L-1} ensure(L-1) copy(L-1) GEMM(L-1) ... ensure(L) copy(L) -join-> count(L) GEMM(L) + predict: \\-> router_L(x_{L-1}) -> top-K -> select(L) -rec sel(L) -/ + +``on``: the predictor only merges the candidates (it reads no cache state); the compute stream +installs the first BUDGET of them that are not resident as low-priority slots (prefetch_ensure: +held, usage 0, keyed between the empty slots and every resident) and a third stream copies them +from layer L's host banks with the slim copy kernel: + + compute: -fork(L)-> ... ensure(L-1) [wait sel(L)] prefetch_ensure(L) [wait pfcopy(L-1)] copy(L-1) -rec dcopy(L)-> GEMM(L-1) ... ensure(L) [wait pfcopy(L)] copy(L) GEMM(L) + predict: \\-> router_L -> top-K -> merge(L) -rec sel(L) + copy: [wait dcopy(L)] slim copy(L) -> ready(L) -rec pfcopy(L) + +Every writer of the slot maps and their mirrors runs on the compute stream, so the maps have one +writer in stream order and no cross-stream edge guards them, and the installs follow layer L's +residency right after ensure(L-1) however late the predictor runs. dcopy(L) follows +prefetch_ensure(L) and the demand copy(L-1) on the compute stream, so the copy stream reads a +finished plan and owns the host link second. prefetch_ensure(L) keeps usage == step, i.e. layer +L-1's routed slots, pinned while GEMM(L-1) may read them; the compute stream joins the prefetch copy +before its own copy(L), so ensure(L) may reuse a prefetched slot as a demand victim and GEMM(L) only +ever reads landed bytes. A demand hit on a prefetched slot sets its usage (a normal resident) and +counts as useful. Layer 0 is never predicted. +""" +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from freetoken.env import ENV +from freetoken.moe.slot_audit import PREFETCH_COPY +from freetoken.utils import init_logger + +logger = init_logger(__name__) + +MODES = ("off", "measure", "on") +# non-resident candidates per target layer: GDN layers leave a shorter copy window than QSA ones +GDN_BUDGET = 3 +ATTN_BUDGET = 4 +# counters[layer, col]; on mode adds LATE and COPIED (calls whose plan had a copy) in a wider stats row +ISSUED, USEFUL, RESIDENT_HITS, CALLS, ROWS, MISSES, LATE, COPIED = range(8) +NUM_COLS = 6 +NUM_STAT_COLS = 8 +# the select kernel's dedup tile is (rows * k)^2; larger (eager-only) decode batches skip the prediction +MAX_ROWS = 8 + + +def resolve_mode(mode: str | None = None) -> str: + mode = (ENV.MOE_PREFETCH.value if mode is None else mode).strip().lower() + if mode not in MODES: + raise ValueError(f"FREETOKEN_MOE_PREFETCH={mode!r}: expected one of {', '.join(MODES)}") + return mode + + +def default_budget(before_linear: bool) -> int: + """Non-resident candidates kept for a target layer, by the target's attention kind.""" + return GDN_BUDGET if before_linear else ATTN_BUDGET + + +def dedicated_stream(device: torch.device, *, highest_priority: bool = False) -> torch.cuda.ExternalStream: + # not one of torch's 32 pooled streams: a pooled one can alias the engine or capture stream + from cuda.bindings import runtime as cudart + + with torch.cuda.device(device): + if highest_priority: + err, _, greatest = cudart.cudaDeviceGetStreamPriorityRange() + if err == cudart.cudaError_t.cudaSuccess: + err, handle = cudart.cudaStreamCreateWithPriority(cudart.cudaStreamNonBlocking, greatest) + else: + err, handle = cudart.cudaStreamCreateWithFlags(cudart.cudaStreamNonBlocking) + if err != cudart.cudaError_t.cudaSuccess: + raise RuntimeError(f"creating a dedicated CUDA stream failed: {err}") + return torch.cuda.ExternalStream(int(handle), device=device) + + +@triton.jit(do_not_specialize=["stride", "budget"]) +def _lookahead_select_kernel( + logits_ptr, stride, resident_ptr, sel_ptr, res_ptr, budget, + ROWS: tl.constexpr, ROWS_P: tl.constexpr, E: tl.constexpr, E_P: tl.constexpr, + K: tl.constexpr, N_P: tl.constexpr, W: tl.constexpr, W_P: tl.constexpr, +): + b = tl.arange(0, ROWS_P) + e = tl.arange(0, E_P) + live = (b < ROWS)[:, None] & (e < E)[None, :] + v = tl.load(logits_ptr + b[:, None] * stride + e[None, :], mask=live, other=float("-inf")).to(tl.float32) + # candidates in merged order: position r * ROWS + b holds row b's rank-r id + j = tl.arange(0, N_P) + cand = tl.full([N_P], -1, tl.int32) + for r in tl.static_range(K): + top = tl.argmax(v, axis=1).to(tl.int32) # ties go to the lowest id + at = (j[:, None] == r * ROWS + b[None, :]) & (b < ROWS)[None, :] + cand = tl.where(tl.sum(at.to(tl.int32), axis=1) > 0, tl.sum(tl.where(at, top[None, :], 0), axis=1), cand) + v = tl.where(e[None, :] == top[:, None], float("-inf"), v) + valid = j < ROWS * K + earlier = (cand[:, None] == cand[None, :]) & (j[None, :] < j[:, None]) & valid[None, :] + first = valid & (tl.sum(earlier.to(tl.int32), axis=1) == 0) + held = tl.load(resident_ptr + cand, mask=first, other=-1) >= 0 + pick = first & (~held) + rank = tl.cumsum(pick.to(tl.int32), axis=0) - 1 + kept = first & held & (tl.cumsum(first.to(tl.int32), axis=0) <= budget) + krank = tl.cumsum(kept.to(tl.int32), axis=0) - 1 + p = tl.arange(0, W_P) + m_sel = pick[None, :] & (rank[None, :] == p[:, None]) & (p < budget)[:, None] + m_res = kept[None, :] & (krank[None, :] == p[:, None]) + # at most one candidate matches each output position + sel = tl.sum(tl.where(m_sel, cand[None, :] + 1, 0), axis=1) - 1 + res = tl.sum(tl.where(m_res, cand[None, :] + 1, 0), axis=1) - 1 + tl.store(sel_ptr + p, sel, mask=p < W) + tl.store(res_ptr + p, res, mask=p < W) + + +@triton.jit(do_not_specialize=["id_base"]) +def _prefetch_count_kernel( + slots_ptr, id_of_slot_ptr, misses_ptr, sel_ptr, res_ptr, counters_ptr, id_base, + N: tl.constexpr, N_P: tl.constexpr, W: tl.constexpr, W_P: tl.constexpr, ROWS: tl.constexpr, +): + i = tl.arange(0, N_P) + slot = tl.load(slots_ptr + i, mask=i < N, other=-1) + # ensure rewrote the routed ids to slots, all pinned to this layer's ids until the next ensure + routed = tl.where(slot >= 0, tl.load(id_of_slot_ptr + slot, mask=slot >= 0, other=0) - id_base, -1) + p = tl.arange(0, W_P) + sel = tl.load(sel_ptr + p, mask=p < W, other=-1) + res = tl.load(res_ptr + p, mask=p < W, other=-1) + sel_hit = tl.sum(((sel[:, None] == routed[None, :]) & (i < N)[None, :]).to(tl.int32), axis=1) > 0 + res_hit = tl.sum(((res[:, None] == routed[None, :]) & (i < N)[None, :]).to(tl.int32), axis=1) > 0 + issued = tl.sum((sel >= 0).to(tl.int32)) + useful = tl.sum((sel_hit & (sel >= 0)).to(tl.int32)) + res_hits = tl.sum((res_hit & (res >= 0)).to(tl.int32)) + misses = tl.load(misses_ptr) + c = tl.arange(0, 8) + add = tl.where(c == 0, issued, tl.where(c == 1, useful, tl.where(c == 2, res_hits, tl.where(c == 3, 1, ROWS)))) + add = tl.where(c == 5, misses, add.to(tl.int64)) + # one writer per layer row, stream-ordered: no atomics needed + old = tl.load(counters_ptr + c, mask=c < 6, other=0) + tl.store(counters_ptr + c, old + add, mask=c < 6) + + +def lookahead_select( + logits: torch.Tensor, resident: torch.Tensor, sel: torch.Tensor, res: torch.Tensor, *, k: int, budget: int +) -> None: + """Top-``k`` of each ``logits`` row, merged rank by rank; ``sel`` gets the first ``budget`` + distinct ids with ``resident[id] < 0`` and ``res`` the resident ids among the first ``budget`` + distinct ones, both in merge order and padded with -1. One CTA, no host sync.""" + rows, num_experts = logits.shape + width = sel.numel() + assert logits.stride(1) == 1 and resident.numel() == num_experts and resident.dtype == torch.int32 + assert sel.dtype == res.dtype == torch.int32 and res.numel() == width and 1 <= k <= num_experts + _lookahead_select_kernel[(1,)]( + logits, logits.stride(0), resident, sel, res, min(budget, width), + ROWS=rows, ROWS_P=triton.next_power_of_2(rows), + E=num_experts, E_P=triton.next_power_of_2(num_experts), + K=k, N_P=triton.next_power_of_2(rows * k), + W=width, W_P=triton.next_power_of_2(width), + num_warps=4, + ) + + +def prefetch_count( + slots: torch.Tensor, id_of_slot: torch.Tensor, misses: torch.Tensor, + sel: torch.Tensor, res: torch.Tensor, counters: torch.Tensor, *, id_base: int, rows: int, +) -> None: + """Add one call's issued / useful / resident-hit / call / row / miss counts to ``counters``.""" + assert slots.dtype == torch.int32 and slots.is_contiguous() and counters.dtype == torch.int64 + n, width = slots.numel(), sel.numel() + _prefetch_count_kernel[(1,)]( + slots, id_of_slot, misses, sel, res, counters, id_base, + N=n, N_P=triton.next_power_of_2(n), W=width, W_P=triton.next_power_of_2(width), ROWS=rows, + num_warps=1, + ) + + +@triton.jit +def _mark_ready_kernel(ready_ptr): + tl.store(ready_ptr, 1) + + +class ExpertPrefetcher: + """Predictor stream, per-layer buffers and device counters for the router lookahead. + + ``fork`` runs layer L's prediction on the dedicated stream at the start of layer L-1's MoE + block. ``measure``: ``join_and_count`` waits for it after layer L's copy is enqueued and counts + on the compute stream. ``on``: ``install`` (right after layer L-1's ensure) waits for it and + runs prefetch_ensure(L) on the compute stream; ``issue_copy`` puts its copy on a second + dedicated stream after layer L-1's demand copy; the compute stream joins the copy before its + copy(L) (``join_copy``) and ensure(L) counts the call (``count_args``). The events are created + up front, the shapes are fixed and nothing syncs the host, so every side captures into the + decode graph. Nothing here is cache_size-shaped: a rebuild keeps it and only resets the counters. + """ + + def __init__( + self, num_layers: int, num_experts: int, device: torch.device, *, + mode: str = "measure", k: int | None = None, budget: int | None = None, + ) -> None: + self.mode = mode + self.num_layers = num_layers + self.num_experts = num_experts + self.device = device + k = ENV.MOE_PREFETCH_K.value if k is None else int(k) + if k < 1: + raise ValueError(f"FREETOKEN_MOE_PREFETCH_K={k} must be >= 1") + self.k = min(k, num_experts) + # 0 = the per-layer default the model wires in + self.budget_override = ENV.MOE_PREFETCH_BUDGET.value if budget is None else int(budget) + if self.budget_override < 0: + raise ValueError(f"FREETOKEN_MOE_PREFETCH_BUDGET={self.budget_override} must be >= 0") + self.stream = dedicated_stream(device) + # (forked, selected): the predictor starts after forked and records selected after the select + self._events = [(torch.cuda.Event(), torch.cuda.Event()) for _ in range(num_layers)] + # torch creates the CUDA event on its first record: do that now, never inside a capture + for pair in self._events: + for event in pair: + event.record(self.stream) + # on mode's select sees nothing resident: prefetch_ensure applies residency and the budget + self._none_resident = torch.full((num_experts,), -1, dtype=torch.int32, device=device) + self._budget = [0] * num_layers + width = 2 * self.k # room for the whole bs-2 union + self.sel = torch.full((num_layers, width), -1, dtype=torch.int32, device=device) + self.res = torch.full((num_layers, width), -1, dtype=torch.int32, device=device) + self.stats = torch.zeros((num_layers, NUM_STAT_COLS), dtype=torch.int64, device=device) + # measure mode's six columns, a view: its count kernel and readers keep the stage-A layout + self.counters = self.stats[:, :NUM_COLS] + self.totals = torch.zeros((num_layers, NUM_STAT_COLS), dtype=torch.int64) + # the router input of each forked prediction, alive until the compute stream has waited for + # the select so the allocator cannot hand its block to later work the predictor still reads + self._inflight: list[torch.Tensor | None] = [None] * num_layers + # on mode: whether install ran prefetch_ensure for the forked prediction (its copy is then due) + self._installed = [False] * num_layers + # on mode: the cache whose slot maps prefetch_ensure writes and whose banks the copy reads + self.cache = None + self.copy_stream = None + if mode == "on": + # the slim copy is an SM kernel: at default priority a late copy queues for SMs behind the next + # compute grid while the compute stream waits on it (bench --on: a 50 us slip cost 14% per step) + self.copy_stream = dedicated_stream(device, highest_priority=True) + # (dcopy: after prefetch_ensure(L) and layer L-1's demand copy on the compute stream, + # pfcopy: after layer L's prefetch copy) + self._copy_events = [(torch.cuda.Event(), torch.cuda.Event()) for _ in range(num_layers)] + for pair in self._copy_events: + for event in pair: + event.record(self.copy_stream) + # one plan per target layer: prefetch_ensure(L+1) must not overwrite the plan copy(L) reads + self.pf_slots = torch.full((num_layers, width), -1, dtype=torch.int32, device=device) + self.pf_src = torch.full((num_layers, width), -1, dtype=torch.int32, device=device) + self.pf_num = torch.zeros((num_layers, 1), dtype=torch.int64, device=device) + self.pf_ready = torch.zeros((num_layers, 1), dtype=torch.int32, device=device) + # test hooks, read at enqueue (capture) time: a [num_layers, width] device buffer replacing + # sel (on mode still installs at most the budget), and sleeps put on the predictor / copy + # stream ahead of their work + self.sel_override: torch.Tensor | None = None + self.delay_predict_cycles = 0 + self.delay_copy_cycles = 0 + logger.info( + f"MoE prefetch {mode}: router lookahead top-{self.k}, budget " + + (str(self.budget_override) if self.budget_override else f"{GDN_BUDGET} before GDN / {ATTN_BUDGET} before attention layers") + + (", slim copies on a dedicated high-priority stream" if mode == "on" else "") + ) + + def budget_for(self, wired: int) -> int: + return min(self.budget_override or wired, self.sel.shape[1]) + + def fork(self, target: int, x: torch.Tensor, gate, resident: torch.Tensor, budget: int) -> bool: + """Predict ``target``'s experts from ``x`` (the previous layer's router input) on the + predictor stream, ordered after the work already on the current stream. Measure mode skips + the ids ``resident`` holds (a racy read of the cache); on mode reads no cache state and + leaves residency and the budget to ``install``. Returns whether it forked.""" + if x.shape[0] > MAX_ROWS: + return False + budget = self.budget_for(budget) + if self.mode == "on": + self._budget[target] = budget + resident, budget = self._none_resident, self.sel.shape[1] + forked, selected = self._events[target] + forked.record(torch.cuda.current_stream(self.device)) + self.stream.wait_event(forked) + with torch.cuda.stream(self.stream): + if self.delay_predict_cycles: + torch.cuda._sleep(self.delay_predict_cycles) + logits = gate.forward(x) + lookahead_select(logits, resident, self.sel[target], self.res[target], k=self.k, budget=budget) + if self.sel_override is not None: + self.sel[target].copy_(self.sel_override[target]) + selected.record(self.stream) + self._inflight[target] = x + self._installed[target] = False + return True + + def install(self, target: int | None) -> int | None: + """On mode, right after the previous layer's ensure: wait for ``target``'s prediction and + install its first budget non-resident candidates with prefetch_ensure on the current + (compute) stream. Returns ``target`` when its copy is due (``issue_copy``), else None.""" + if target is None or self.mode != "on" or self._inflight[target] is None: + return None + torch.cuda.current_stream(self.device).wait_event(self._events[target][1]) + self.cache.prefetch_ensure( + target, self.sel[target], self.pf_slots[target], self.pf_src[target], + self.pf_num[target], self.stats[target], self.pf_ready[target], budget=self._budget[target], + ) + self._installed[target] = True + return target + + def _pending(self, layer_id: int) -> bool: + """On mode: ``layer_id``'s prediction was installed and its copy is not joined yet.""" + return self.mode == "on" and self._installed[layer_id] and self._inflight[layer_id] is not None + + def count_args(self, layer_id: int) -> tuple | None: + """On mode: what ``layer_id``'s demand ensure counts into, or None when it was not prefetched.""" + if not self._pending(layer_id): + return None + return self.stats[layer_id], self.pf_ready[layer_id], self.pf_num[layer_id], self._inflight[layer_id].shape[0] + + def plan_of(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor] | None: + """On mode: ``layer_id``'s prefetch copy plan (slots, count) while its prefetch is in flight.""" + if not self._pending(layer_id): + return None + return self.pf_slots[layer_id], self.pf_num[layer_id] + + def join_copy(self, layer_id: int) -> None: + """On mode: the compute stream waits for ``layer_id``'s prefetch copy before its own copy, so + a prefetched slot ensure(L) picked as a demand victim is never written by two copies.""" + if self.mode != "on" or self._inflight[layer_id] is None: + return + # a prediction that was never installed still has to rejoin (a capture needs every fork joined) + done = self._copy_events[layer_id][1] if self._pending(layer_id) else self._events[layer_id][1] + torch.cuda.current_stream(self.device).wait_event(done) + self._inflight[layer_id] = None + self._installed[layer_id] = False + + def issue_copy(self, target: int) -> None: + """On mode: copy ``target``'s prefetch plan on the copy stream after the work already on the + current stream: its prefetch_ensure and the demand copy (which owns the host link first).""" + if not self._pending(target): + return + dcopy, pfcopy = self._copy_events[target] + dcopy.record(torch.cuda.current_stream(self.device)) + self.copy_stream.wait_event(dcopy) + with torch.cuda.stream(self.copy_stream): + if self.delay_copy_cycles: + torch.cuda._sleep(self.delay_copy_cycles) + self.cache.copy_rows(target, self.pf_slots[target], self.pf_src[target], self.pf_num[target], slim=True, + kind=PREFETCH_COPY) + _mark_ready_kernel[(1,)](self.pf_ready[target]) + pfcopy.record(self.copy_stream) + + def join_and_count(self, layer_id: int, slots: torch.Tensor, id_of_slot: torch.Tensor, misses: torch.Tensor) -> None: + """Measure mode: join ``layer_id``'s prediction (if one was forked this step) and count it + against the routed slots ``slots`` of this layer's ensure.""" + x = self._inflight[layer_id] if self.mode == "measure" else None + if x is None: + return + torch.cuda.current_stream(self.device).wait_event(self._events[layer_id][1]) + prefetch_count( + slots.view(-1), id_of_slot, misses, self.sel[layer_id], self.res[layer_id], self.counters[layer_id], + id_base=layer_id * self.num_experts, rows=x.shape[0], + ) + self._inflight[layer_id] = None + + def reset_counters(self) -> None: + self.stats.zero_() + + def reset(self) -> None: + """Counters and in-flight state; the streams, events and buffers stay.""" + self.stats.zero_() + self.totals.zero_() + self._inflight = [None] * self.num_layers + self._installed = [False] * self.num_layers + + def take_window(self) -> torch.Tensor: + """This window's ``[num_layers, NUM_STAT_COLS]`` stats on the host (one sync), added to + ``totals``; the device ones restart.""" + window = self.stats.cpu() + self.stats.zero_() + self.totals += window + return window + + +def summarize(counters: torch.Tensor, mode: str = "measure") -> dict: + """Aggregate ``[num_layers, NUM_COLS or NUM_STAT_COLS]`` counters into per-layer-call and + per-token rates. In on mode the misses are what was left after the prefetch, so coverage is + useful / (useful + misses), and late is the share of copied calls whose copy had not landed + when the compute stream reached it.""" + total = counters.sum(0).tolist() + calls, issued, useful = total[CALLS], total[ISSUED], total[USEFUL] + tokens = int(counters[:, ROWS].max()) if counters.numel() else 0 + wanted = useful + total[MISSES] if mode == "on" else total[MISSES] + stats = { + "layer_calls": calls, + "tokens": tokens, + "issued_per_layer": issued / calls if calls else 0.0, + "useful_per_layer": useful / calls if calls else 0.0, + "resident_hits_per_layer": total[RESIDENT_HITS] / calls if calls else 0.0, + "misses_per_layer": total[MISSES] / calls if calls else 0.0, + "precision": useful / issued if issued else 0.0, + "coverage": useful / wanted if wanted else 0.0, + "useful_per_token": useful / tokens if tokens else 0.0, + } + if mode == "on": + copied = total[COPIED] if len(total) > COPIED else 0 + stats["copied_calls"] = copied + stats["late_fraction"] = total[LATE] / copied if copied else 0.0 + return stats + + +def format_summary(stats: dict) -> str: + if "late_fraction" in stats: + return ( + f"issued/layer={stats['issued_per_layer']:.2f}, useful/layer={stats['useful_per_layer']:.2f}, " + f"precision={stats['precision']:.3f}, useful/token={stats['useful_per_token']:.1f}, " + f"misses/layer={stats['misses_per_layer']:.2f}, coverage={stats['coverage']:.3f}, " + f"late={stats['late_fraction']:.3f} of {stats['copied_calls']} copies" + ) + return ( + f"issued/layer={stats['issued_per_layer']:.2f}, useful/layer={stats['useful_per_layer']:.2f}, " + f"precision={stats['precision']:.3f}, useful/token={stats['useful_per_token']:.1f}, " + f"misses/layer={stats['misses_per_layer']:.2f}, coverage={stats['coverage']:.3f}, " + f"resident_hits/layer={stats['resident_hits_per_layer']:.2f}" + ) + + +def format_per_layer(counters: torch.Tensor) -> str: + parts = [] + for layer, row in enumerate(counters.tolist()): + if row[CALLS]: + c = row[CALLS] + parts.append(f"L{layer}={row[USEFUL] / c:.2f}/{row[ISSUED] / c:.2f}/{row[MISSES] / c:.2f}") + return ", ".join(parts) diff --git a/python/freetoken/moe/scored_ensure.py b/python/freetoken/moe/scored_ensure.py index 123086ed51..ba570b5169 100644 --- a/python/freetoken/moe/scored_ensure.py +++ b/python/freetoken/moe/scored_ensure.py @@ -106,6 +106,22 @@ def _phase1(query_ptr, slot_of_id_ptr, lru_usage_ptr, num_copy_ptr, step, K, return q, kmask, miss, first, first_miss, rank, num_missing, tl.where(hit, s, -1) +@triton.jit +def _prefetch_phase1(query_ptr, slot_of_id_ptr, K, BLOCK_K: tl.constexpr, id_base): + """_phase1 for a prefetch: -1 entries are padding, hits are left alone (a prefetch is not an + access), and the misses rank in query order so the select's best candidates install first.""" + k = tl.arange(0, BLOCK_K) + raw = tl.load(query_ptr + k, mask=k < K, other=-1) + kmask = raw >= 0 + q = tl.where(kmask, raw + id_base, 0) + s = tl.load(slot_of_id_ptr + q, mask=kmask, other=0) + miss = kmask & (s == -1) + same = (q[:, None] == q[None, :]) & (k[:, None] > k[None, :]) & kmask[:, None] & kmask[None, :] + first_miss = miss & (tl.sum(same.to(tl.int32), axis=1) == 0) + rank = tl.cumsum(first_miss.to(tl.int32), axis=0) - 1 + return q, first_miss, rank, tl.sum(first_miss.to(tl.int32)) + + @triton.jit def _stats(stats_ptr, first, num_missing): # One vectorized atomic over 3 lanes, not three scalar ones: a scalar atomic serializes the CTA. @@ -144,13 +160,14 @@ def _owner_state(mirror_ptr, state_ptr, c, oid, mirrored, stale, other): ) -@triton.jit(do_not_specialize=["K", "num_cached", "id_base", "nm_topk", "nm_stride"]) +@triton.jit(do_not_specialize=["K", "num_cached", "id_base", "nm_topk", "nm_stride", "pf_rows", "pf_budget"]) def _scored_ensure_kernel( query_ptr, slot_of_id_ptr, id_of_slot_ptr, lru_usage_ptr, lru_step_ptr, out_ptr, src_ptr, dst_ptr, num_copy_ptr, stats_ptr, tok_ptr, last_tok_ptr, lc_ptr, ct_ptr, g_ptr, logits_ptr, owner_ptr, slot_last_ptr, slot_lc_ptr, slot_ct_ptr, pin_ptr, - K, num_cached, id_base, nm_topk, nm_stride, nm_thr, + pf_stats_ptr, pf_ready_ptr, pf_num_ptr, + K, num_cached, id_base, nm_topk, nm_stride, nm_thr, pf_rows, pf_budget, BLOCK_K: tl.constexpr, BLOCK_C: tl.constexpr, BLOCK_E: tl.constexpr, BLOCK_TOPK: tl.constexpr, NM_ROWS: tl.constexpr, USAGE_MAX: tl.constexpr, COLLECT_STATS: tl.constexpr, @@ -158,15 +175,27 @@ def _scored_ensure_kernel( NUM_LAYERS: tl.constexpr, NUM_EXPERTS: tl.constexpr, BETA_STEP: tl.constexpr, W_Q4: tl.constexpr, DECAY_MUL: tl.constexpr, DT_MAX: tl.constexpr, SLOT_BITS: tl.constexpr, + LOWPRI: tl.constexpr, PF_COUNT: tl.constexpr, PREFETCH: tl.constexpr, ): - step = tl.load(lru_step_ptr) + 1 - tl.store(lru_step_ptr, step) + if PREFETCH: + # not an access: the clock stays, so usage == step still pins the last demand call's slots + # (at step 0, right after a reset, every slot reads pinned: nothing installs before a demand call) + step = tl.load(lru_step_ptr) + tl.store(num_copy_ptr, tl.zeros([], tl.int64)) + tl.store(pf_ready_ptr, 0) + else: + step = tl.load(lru_step_ptr) + 1 + tl.store(lru_step_ptr, step) if POLICY != 0: tok = tl.load(tok_ptr) if BUMP_TOK: tok = tok + 1 tl.store(tok_ptr, tok) layer = id_base // NUM_EXPERTS + if PREFETCH: + # score from where the pass is, right after layer - 1's demand call: the target layer's + # residents are next up, not a whole pass away (else it evicts what the layer is about to route) + layer = layer - 1 now = tok * NUM_LAYERS + layer if POLICY == 3: if UPDATE_STATE: @@ -187,9 +216,18 @@ def _scored_ensure_kernel( tl.store(slot_last_ptr + held_at, refreshed, mask=near & (held_at >= 0)) # the victim scan and the routed update below read and rewrite what this wrote tl.debug_barrier() - q, kmask, miss, first, first_miss, rank, num_missing, out = _phase1( - query_ptr, slot_of_id_ptr, lru_usage_ptr, num_copy_ptr, step, K, BLOCK_K, - id_base) + if PREFETCH: + q, first_miss, rank, num_missing = _prefetch_phase1(query_ptr, slot_of_id_ptr, K, BLOCK_K, id_base) + else: + if PF_COUNT: + kk = tl.arange(0, BLOCK_K) + sk = tl.load(slot_of_id_ptr + tl.load(query_ptr + kk, mask=kk < K, other=0) + id_base, mask=kk < K, other=-1) + # a hit on a held slot with usage 0 consumes a prefetch; read it before _phase1 bumps it + was_pf = (sk >= 0) & (tl.load(lru_usage_ptr + sk, mask=sk >= 0, other=1) == 0) + tl.debug_barrier() + q, kmask, miss, first, first_miss, rank, num_missing, out = _phase1( + query_ptr, slot_of_id_ptr, lru_usage_ptr, num_copy_ptr, step, K, BLOCK_K, + id_base) if num_missing > 0: # REQUIRED (flashlib): the reload below must see _phase1's hit bump, or a hit is evicted. @@ -199,7 +237,13 @@ def _scored_ensure_kernel( u = tl.load(lru_usage_ptr + c, mask=cmask, other=USAGE_MAX) if POLICY == 0: umax = tl.full([BLOCK_C], USAGE_MAX, u.dtype) - u = tl.where((u == step) | (~cmask), umax, u) + if LOWPRI: + held0 = cmask & (tl.load(id_of_slot_ptr + c, mask=cmask, other=-1) >= 0) + # empty slots, then unconsumed prefetches (held, usage 0), then residents by recency + u = tl.where((u == step) | (~cmask), umax, tl.where(held0, u + 1, 0)) + else: + u = tl.where((u == step) | (~cmask), umax, u) + n_free = tl.sum((u != umax).to(tl.int32)) else: oid = tl.load(id_of_slot_ptr + c, mask=cmask, other=-1) held = cmask & (oid >= 0) @@ -221,6 +265,10 @@ def _scored_ensure_kernel( score += (W_Q4 * lcv) >> 4 # Every key is distinct and ties go to the lowest slot; an empty slot keys below any held one. key = tl.where(held, ((score.to(tl.int64) + _SCORE_BIAS) << SLOT_BITS) | c, c.to(tl.int64)) + if LOWPRI: + # an unconsumed prefetch (held, usage 0) keys above every empty slot and below every + # resident: the caps keep score + 2^31 > 2^30 - 2^25 - 2^22, far above this band's 1 + key = tl.where(held & (u == 0), (1 << SLOT_BITS) | c.to(tl.int64), key) evictable = cmask & (u != step) if PIN_SINCE: # the score ignores recency, so slots an earlier call touched since *pin_ptr stay pinned too @@ -228,7 +276,16 @@ def _scored_ensure_kernel( # the pinned key's slot bits are 0, so an all-pinned scan falls back to slot 0 like flashlib's argmin pinned_key = 0x7FFFFFFFFFFFFFFF ^ ((1 << SLOT_BITS) - 1) key = tl.where(evictable, key, pinned_key) - for i in tl.range(num_missing): + n_free = tl.sum(evictable.to(tl.int32)) + if PREFETCH: + # never the all-pinned fallback: that slot may be under the running GEMM; misses rank in + # query order, so the budget keeps the first non-resident candidates + n_iter = tl.minimum(tl.minimum(num_missing, n_free), pf_budget) + tl.store(num_copy_ptr, n_iter.to(tl.int64)) + tl.store(pf_stats_ptr, tl.load(pf_stats_ptr) + n_iter) + else: + n_iter = num_missing + for i in tl.range(n_iter): if POLICY == 0: victim = tl.argmin(u, axis=0).to(tl.int32) else: @@ -240,17 +297,38 @@ def _scored_ensure_kernel( e = tl.sum(tl.where((rank == i) & first_miss, q, 0)) tl.store(id_of_slot_ptr + victim, e) tl.store(slot_of_id_ptr + e, victim) - tl.store(lru_usage_ptr + victim, step) + if PREFETCH: + # usage 0 on a held slot marks it low priority until a demand call touches it + tl.store(lru_usage_ptr + victim, 0) + if POLICY != 0: + tl.store(owner_ptr + victim, -1) + else: + tl.store(lru_usage_ptr + victim, step) tl.store(dst_ptr + i, victim) tl.store(src_ptr + i, e - id_base) # back to the caller's id space - out = tl.where((rank == i) & miss, victim, out) + if not PREFETCH: + out = tl.where((rank == i) & miss, victim, out) if POLICY == 0: u = tl.where(c == victim, umax, u) # claim in-register else: key = tl.where(c == victim, pinned_key, key) - # Written from registers, never re-read from slot_of_id, so out_ptr may alias query_ptr. - tl.store(out_ptr + tl.arange(0, BLOCK_K), out, mask=kmask) + if not PREFETCH: + # Written from registers, never re-read from slot_of_id, so out_ptr may alias query_ptr. + tl.store(out_ptr + tl.arange(0, BLOCK_K), out, mask=kmask) + if PF_COUNT: + # one writer per column: pf_ensure adds ISSUED (column 0), earlier on this stream + pn = tl.load(pf_num_ptr) + late = (pn > 0) & (tl.load(pf_ready_ptr, volatile=True) == 0) + col = tl.arange(0, 8) + add = tl.where(col == 1, tl.sum((first & was_pf).to(tl.int32)).to(tl.int64), 0) + add = tl.where(col == 3, 1, add) + add = tl.where(col == 4, pf_rows, add) + add = tl.where(col == 5, num_missing.to(tl.int64), add) + add = tl.where(col == 6, late.to(tl.int64), add) + add = tl.where(col == 7, (pn > 0).to(tl.int64), add) + cols = (col == 1) | (col >= 3) + tl.store(pf_stats_ptr + col, tl.load(pf_stats_ptr + col, mask=cols, other=0) + add, mask=cols) if POLICY != 0: if UPDATE_STATE: # distinct ids only, so a duplicate at bs > 1 counts once per call; out is each id's final slot @@ -265,7 +343,7 @@ def _scored_ensure_kernel( tl.store(ct_ptr + q, now, mask=first) tl.store(slot_lc_ptr + out, lc_new, mask=first) tl.store(slot_ct_ptr + out, now, mask=first) - else: + elif not PREFETCH: # a slot filled without a state update reads the per-id state until a decode hit re-mirrors it tl.store(owner_ptr + out, -1, mask=first_miss) if COLLECT_STATS: @@ -311,6 +389,13 @@ def scored_ensure( beta: float = BETA, w: float = W, halflife: int = HALFLIFE, + lowpri: bool = False, + pf_stats: torch.Tensor | None = None, + pf_ready: torch.Tensor | None = None, + pf_num: torch.Tensor | None = None, + pf_rows: int = 0, + pf_budget: int | None = None, + prefetch: bool = False, ) -> None: """``flashlib.lru_ensure`` (sequential strategy) with the victim picked by ``policy``. @@ -324,7 +409,21 @@ def scored_ensure( prefill's chunked calls cannot evict each other's experts; policy > 0 only. Policy 3 refreshes the experts whose ``router_logits`` (``[rows, num_experts]``, rows of ``query``) are within ``near_miss_thr`` of their row's lowest routed logit. + + FREETOKEN_MOE_PREFETCH=on: ``lowpri`` keys held slots with usage 0 (unconsumed prefetches) + between the empty slots and every resident. ``pf_stats`` (a ``[8]`` int64 prefetch stats row) + also counts this demand call: hits on such slots, the call, ``pf_rows``, the misses, and whether + the prefetch plan ``pf_num`` had rows whose copy had not set ``pf_ready`` yet. + ``prefetch=True`` is :func:`prefetch_ensure`'s install mode instead, installing at most + ``pf_budget`` ids (default: all). """ + if prefetch: + assert lowpri and pf_stats is not None and pf_ready is not None and stats is None + assert not update_state and not bump_tok and pin_since is None and router_logits is None + if pf_stats is not None: + assert lowpri and pf_stats.dtype == torch.int64 and pf_stats.numel() == 8 and pf_stats.is_contiguous() + assert pf_ready is not None and pf_ready.dtype == torch.int32 and pf_ready.numel() == 1 + assert prefetch or (pf_num is not None and pf_num.dtype == torch.int64 and pf_num.numel() == 1) k = query.numel() num_cached = id_of_slot.numel() assert query.dtype == torch.int32 and query.is_contiguous() @@ -372,7 +471,11 @@ def scored_ensure( dummy if slot_lc is None else slot_lc, dummy if slot_ct is None else slot_ct, dummy if pin_since is None else pin_since, - k, num_cached, id_base, nm_topk, nm_stride, float(near_miss_thr), + dummy if pf_stats is None else pf_stats, + dummy if pf_ready is None else pf_ready, + dummy if pf_num is None else pf_num, + k, num_cached, id_base, nm_topk, nm_stride, float(near_miss_thr), int(pf_rows), + k if pf_budget is None else int(pf_budget), BLOCK_K=triton.next_power_of_2(k), BLOCK_C=block_c, BLOCK_E=triton.next_power_of_2(num_experts), @@ -391,5 +494,40 @@ def scored_ensure( DECAY_MUL=decay_mul, DT_MAX=dt_max, SLOT_BITS=slot_bits, + LOWPRI=bool(lowpri), + PF_COUNT=pf_stats is not None and not prefetch, + PREFETCH=bool(prefetch), num_warps=_num_warps_for(block_c), ) + + +def prefetch_ensure( + query: torch.Tensor, + slot_of_id: torch.Tensor, + id_of_slot: torch.Tensor, + lru_usage: torch.Tensor, + lru_step: torch.Tensor, + dst_slots: torch.Tensor, + src_rows: torch.Tensor, + num_copy: torch.Tensor, + pf_stats: torch.Tensor, + pf_ready: torch.Tensor, + budget: int | None = None, + **state, +) -> None: + """Install the ids of ``query`` (-1 = padding) that are not resident as low-priority slots + (held, usage 0), at most the first ``budget`` of them in query order (default: all), and + write their copy plan to ``dst_slots``/``src_rows``/``num_copy``. + + Unlike a demand call it bumps no clock, touches no hit and no per-id state, and so keeps + every slot of the last demand call (usage == step) pinned; victims follow the demand key + (empty, then earlier unconsumed prefetches, then the coldest residents), scored as of the + previous layer's call, the position the prefetch runs at. With fewer + evictable slots than misses it installs only the first ones, in query order. Adds the + installs to ``pf_stats[0]`` and clears ``pf_ready``. + """ + scored_ensure( + query, slot_of_id, id_of_slot, lru_usage, lru_step, query, src_rows, dst_slots, num_copy, + update_state=False, lowpri=True, pf_stats=pf_stats, pf_ready=pf_ready, pf_budget=budget, prefetch=True, + **state, + ) diff --git a/python/freetoken/moe/slot_audit.py b/python/freetoken/moe/slot_audit.py new file mode 100644 index 0000000000..981239758f --- /dev/null +++ b/python/freetoken/moe/slot_audit.py @@ -0,0 +1,688 @@ +"""Expert slot audit (``FREETOKEN_MOE_SLOT_AUDIT=N``): a debug detector for PERSISTENT slot corruption. + +Two parts, both on the device and off by default: + +- writer history: every writer of the slot maps (demand and prefetch ensures, prefill ensures, the + whole-layer materialize, prefill buffer invalidation, reset) and every writer of the slot banks + (demand, prefetch and prefill copies, the prefill double buffers) is followed on its writer's + stream by a tiny kernel keyed on the same plan (or slot range) that appends ``(seq, kind, decode + step, lru step, layer, expert row, ...)`` to a per-slot ring of the last ``DEPTH`` metadata events + and one of the last ``DEPTH`` byte events. The map writers all run on the compute stream except the + prefill buffer invalidation (on the prefill copy stream, fenced behind it); the byte writers also + run on the prefetch and prefill copy streams. ``seq`` is one global atomic counter, so the two rings + merge into one cross-stream order. Clears are recorded just before their writer, so they carry the + old owner. + Fixed shapes and no host sync: the recorders capture into the decode graph. +- audit: every N decode steps, after the forward on the compute stream, every held slot is checked: + its bytes in every bank against the host row of the (layer, expert) its ``id_of_slot`` names + (host rows gathered with the fused copy into a scratch of fewer than 64 MiB, one layer chunk at a + time), ``slot_for_id[id_of_slot[s]] == s``, and, over every id, ``id_of_slot[slot_for_id[id]] == id`` + (a stale entry routes that id to another expert's bytes). The first ``RING`` bad (slot, owner) + pairs are kept with both history rings; the engine logs them and the totals every + MOE_STATS_INTERVAL decode steps and at shutdown, with a diagnosis read from the history. +""" +from __future__ import annotations + +import math + +import torch +import triton +import triton.language as tl + +from freetoken.env import ENV +from freetoken.utils import init_logger + +logger = init_logger(__name__) + +# events kept per slot and ring (metadata / bytes) +DEPTH = 4 +# event fields +(E_SEQ, E_KIND, E_DSTEP, E_LSTEP, E_LAYER, E_ROW, E_AUX, E_OWNER, E_USAGE, E_POS) = range(10) +EVENT_FIELDS = 10 +# metadata kinds (installs name the new owner; clears the old one, recorded before the writer) +DEMAND_INSTALL, PREFETCH_INSTALL, PREFILL_INSTALL, MATERIALIZE_INSTALL = 1, 2, 3, 4 +MATERIALIZE_CLEAR, INVALIDATE, RESET = 5, 6, 7 +# byte kinds +DEMAND_COPY, PREFETCH_COPY, PREFILL_COPY, MATERIALIZE_COPY, PAGEABLE_COPY = 11, 12, 13, 14, 15 +PREFILL_BUFFER, PREFILL_SPLIT_H2D, PREFILL_HIT_D2D, OTHER_COPY = 16, 17, 18, 19 +KIND_NAMES = { + DEMAND_INSTALL: "demand install", PREFETCH_INSTALL: "prefetch install", PREFILL_INSTALL: "prefill ensure install", + MATERIALIZE_INSTALL: "materialize install", MATERIALIZE_CLEAR: "materialize clear", + INVALIDATE: "prefill buffer invalidate", RESET: "reset", + DEMAND_COPY: "demand copy", PREFETCH_COPY: "prefetch copy", PREFILL_COPY: "prefill ensure copy", + MATERIALIZE_COPY: "materialize copy", PAGEABLE_COPY: "pageable materialize copy", + PREFILL_BUFFER: "prefill buffer copy", PREFILL_SPLIT_H2D: "prefill buffer H2D (misses and small banks)", + PREFILL_HIT_D2D: "prefill buffer D2D", OTHER_COPY: "copy", +} +INSTALLS = (DEMAND_INSTALL, PREFETCH_INSTALL, PREFILL_INSTALL, MATERIALIZE_INSTALL) +# the copy each install expects next +COPY_FOR = { + DEMAND_INSTALL: DEMAND_COPY, PREFETCH_INSTALL: PREFETCH_COPY, PREFILL_INSTALL: PREFILL_COPY, + MATERIALIZE_INSTALL: MATERIALIZE_COPY, +} +# recorder row sources: the plan's src rows, slot - row_base, the range index, the old owner +ROW_PLAN, ROW_FROM_SLOT, ROW_INDEX, ROW_OLD_OWNER = range(4) +# recorder predicates: every slot, held slots, slots >= E held by the layer +PRED_ALL, PRED_HELD, PRED_LAYER_ABOVE_E = range(3) +# totals: audits, held slots seen, slots byte-checked, held slots not byte-checked, bad slots found +# (summed over audits), of them with bad bytes / a forward map miss / named by a stale slot_for_id / +# an owner out of range, slot_for_id entries past the cache, new records, repeat findings, events +# recorded, plan entries naming no slot +(T_AUDITS, T_HELD, T_BYTES_CHECKED, T_UNCHECKED, T_BAD, T_BYTES, T_FWD, T_STALE, T_OWNER_RANGE, T_MAP_RANGE, + T_NEW, T_STILL, T_EVENTS, T_BAD_PLAN) = range(14) +NUM_TOTALS = 16 +# per-slot audit flags +F_BYTES, F_FWD, F_STALE, F_OWNER_RANGE = 1, 2, 4, 8 +# record header, then the metadata ring and the byte ring as stored ([DEPTH, EVENT_FIELDS] each) +(H_AUDIT, H_DSTEP, H_LSTEP, H_SLOT, H_OWNER, H_MAPPED, H_USAGE, H_FLAGS, H_BANKS, H_FIRST_BANK, H_FIRST_BYTE, + H_BAD_WORDS, H_STALE_ID, H_STALE_N, H_META_N, H_BYTES_N) = range(16) +HEADER = 16 +RING_EVENTS = DEPTH * EVENT_FIELDS +RECORD_FIELDS = HEADER + 2 * RING_EVENTS +RING = 32 +SCRATCH_BYTES = 64 << 20 +MAX_CHUNK_ROWS = 64 +COMPARE_WORDS = 4096 +_NO_BAD = tl.constexpr(0x7FFFFFFF) +_D = tl.constexpr(DEPTH) +_F = tl.constexpr(EVENT_FIELDS) + + +def resolve_interval(interval: int | None = None) -> int: + value = ENV.MOE_SLOT_AUDIT.value if interval is None else interval + value = int(value or 0) + if value < 0: + raise ValueError(f"FREETOKEN_MOE_SLOT_AUDIT={value}: expected 0 (off) or the decode steps between audits") + return value + + +@triton.jit +def _put(vals, f, idx: tl.constexpr, v): + return tl.where(f[None, :] == idx, v.to(tl.int64)[:, None], vals) + + +@triton.jit +def _bump_step_kernel(dstep_ptr): + tl.store(dstep_ptr, tl.load(dstep_ptr) + 1) + + +@triton.jit(do_not_specialize=["kind", "layer", "start", "row_base", "width", "num_slots"]) +def _record_kernel( + hist_ptr, cur_ptr, seq_ptr, dstep_ptr, lstep_ptr, id_of_slot_ptr, usage_ptr, totals_ptr, + dst_ptr, src_ptr, aux_ptr, num_ptr, + kind, layer, start, row_base, width, num_slots, + PLAN: tl.constexpr, ROW: tl.constexpr, HAS_AUX: tl.constexpr, HAS_NUM: tl.constexpr, PRED: tl.constexpr, + E: tl.constexpr, BLOCK: tl.constexpr, +): + """Append one event per written slot to that slot's ring: plan entries ``i < min(*num, width)`` + (``PLAN``) or the slots ``start + i`` for ``i < width``.""" + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + n = width + if HAS_NUM: + n = tl.minimum(tl.load(num_ptr).to(tl.int32), width) + valid = i < n + if PLAN: + slot = tl.load(dst_ptr + i, mask=valid, other=-1) + else: + slot = start + i + ok = valid & (slot >= 0) & (slot < num_slots) + bad_plan = valid & ~ok + owner = tl.load(id_of_slot_ptr + slot, mask=ok, other=-3) + usage = tl.load(usage_ptr + slot, mask=ok, other=-3) + if PRED == 1: + ok = ok & (owner >= 0) + if PRED == 2: + ok = ok & (owner >= layer * E) & (owner < layer * E + E) & (slot >= E) + lay = tl.zeros([BLOCK], tl.int32) + layer + aux = tl.zeros([BLOCK], tl.int32) - 1 + if ROW == 0: + row = tl.load(src_ptr + i, mask=ok, other=-1) + elif ROW == 1: + row = slot - row_base + elif ROW == 2: + row = i + else: + row = tl.where(owner >= 0, owner % E, -1) + lay = tl.where(owner >= 0, owner // E, -1) + aux = aux + 1 + layer + if HAS_AUX: + aux = tl.load(aux_ptr + i, mask=ok, other=-1) + cnt = tl.sum(ok.to(tl.int64), axis=0) + seq = tl.atomic_add(seq_ptr, cnt) + tl.cumsum(ok.to(tl.int64), axis=0) - 1 + pos = tl.atomic_add(cur_ptr + slot, 1, mask=ok) + f = tl.arange(0, 16) + z = tl.zeros([BLOCK], tl.int64) + vals = tl.zeros([BLOCK, 16], tl.int64) - 1 + vals = _put(vals, f, 0, seq) + vals = _put(vals, f, 1, z + kind) + vals = _put(vals, f, 2, z + tl.load(dstep_ptr)) + vals = _put(vals, f, 3, z + tl.load(lstep_ptr)) + vals = _put(vals, f, 4, lay) + vals = _put(vals, f, 5, row) + vals = _put(vals, f, 6, aux) + vals = _put(vals, f, 7, owner) + vals = _put(vals, f, 8, usage) + vals = _put(vals, f, 9, i) + at = (slot.to(tl.int64) * _D + (pos % _D)) * _F + tl.store(hist_ptr + at[:, None] + f[None, :], vals, mask=ok[:, None] & (f[None, :] < _F)) + t = tl.arange(0, 2) + add = tl.where(t == 0, cnt, tl.sum(bad_plan.to(tl.int64), axis=0)) + tl.atomic_add(totals_ptr + 12 + t, add, sem="relaxed") + + +@triton.jit(do_not_specialize=["num_slots", "num_ids"]) +def _audit_meta_kernel( + id_of_slot_ptr, slot_for_id_ptr, layer_ok_ptr, count_ptr, rows_ptr, slots_ptr, + flags_ptr, stale_id_ptr, stale_n_ptr, totals_ptr, num_slots, num_ids, + E: tl.constexpr, CAP: tl.constexpr, BLOCK: tl.constexpr, +): + """Both map checks, and each byte-checkable held slot appended to its layer's list.""" + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + ms = i < num_slots + o = tl.load(id_of_slot_ptr + i, mask=ms, other=-1) + held = ms & (o >= 0) + in_range = held & (o < num_ids) + mapped = tl.load(slot_for_id_ptr + o, mask=in_range, other=-2) + fwd = in_range & (mapped != i) + out_of_range = held & ~in_range + fl = tl.where(fwd, 2, 0) | tl.where(out_of_range, 8, 0) + tl.atomic_or(flags_ptr + i, fl, mask=ms & (fl != 0)) + layer = tl.where(in_range, o // E, 0) + ok = in_range & (tl.load(layer_ok_ptr + layer, mask=in_range, other=0) != 0) + p = tl.atomic_add(count_ptr + layer, 1, mask=ok) + take = ok & (p < CAP) + tl.store(rows_ptr + layer * CAP + p, o - layer * E, mask=take) + tl.store(slots_ptr + layer * CAP + p, i, mask=take) + mi = i < num_ids + s = tl.load(slot_for_id_ptr + i, mask=mi, other=-1) + live = mi & (s >= 0) + s_ok = live & (s < num_slots) + owner = tl.load(id_of_slot_ptr + s, mask=s_ok, other=-1) + stale = s_ok & (owner != i) + tl.atomic_or(flags_ptr + s, 4, mask=stale) + tl.store(stale_id_ptr + s, i, mask=stale) + tl.atomic_add(stale_n_ptr + s, 1, mask=stale) + t = tl.arange(0, 16) + add = tl.where(t == 1, tl.sum(held.to(tl.int64), axis=0), 0) + add = tl.where(t == 2, tl.sum(take.to(tl.int64), axis=0), add) + add = tl.where(t == 3, tl.sum((held & ~take).to(tl.int64), axis=0), add) + add = tl.where(t == 6, tl.sum(fwd.to(tl.int64), axis=0), add) + add = tl.where(t == 7, tl.sum(stale.to(tl.int64), axis=0), add) + add = tl.where(t == 8, tl.sum(out_of_range.to(tl.int64), axis=0), add) + add = tl.where(t == 9, tl.sum((live & ~s_ok).to(tl.int64), axis=0), add) + add = tl.where((t == 0) & (tl.program_id(0) == 0), 1, add) + tl.atomic_add(totals_ptr + t, add, mask=add != 0, sem="relaxed") + + +@triton.jit +def _audit_nums_kernel(count_ptr, nums_ptr, L: tl.constexpr, C: tl.constexpr, R: tl.constexpr, BLOCK: tl.constexpr): + j = tl.arange(0, BLOCK) + m = j < L * C + cnt = tl.load(count_ptr + j // C, mask=m, other=0) + tl.store(nums_ptr + j, tl.minimum(tl.maximum(cnt - (j % C) * R, 0), R).to(tl.int64), mask=m) + + +@triton.jit +def _audit_compare_kernel(cache_ptrs, ref_ptrs, words_ptr, slots_ptr, num_ptr, bad_cnt_ptr, bad_first_ptr, + NB: tl.constexpr, BLOCK: tl.constexpr): + """Program (row, bank, chunk): one chunk of a listed slot's bank row against its gathered host row.""" + r = tl.program_id(0) + if r < tl.load(num_ptr): + b = tl.program_id(1) + words = tl.load(words_ptr + b) + off = tl.program_id(2).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + m = off < words + slot = tl.load(slots_ptr + r).to(tl.int64) + got = tl.load(tl.load(cache_ptrs + b).to(tl.pointer_type(tl.int32)) + slot * words + off, mask=m, other=0) + want = tl.load(tl.load(ref_ptrs + b).to(tl.pointer_type(tl.int32)) + r * words + off, mask=m, other=0) + bad = m & (got != want) + nbad = tl.sum(bad.to(tl.int32), axis=0) + if nbad > 0: + tl.atomic_add(bad_cnt_ptr + slot * NB + b, nbad) + tl.atomic_min(bad_first_ptr + slot * NB + b, tl.min(tl.where(bad, off, _NO_BAD), axis=0).to(tl.int32)) + + +@triton.jit(do_not_specialize=["num_slots", "num_ids"]) +def _audit_report_kernel( + id_of_slot_ptr, slot_for_id_ptr, usage_ptr, lstep_ptr, dstep_ptr, + flags_ptr, stale_id_ptr, stale_n_ptr, bad_cnt_ptr, bad_first_ptr, reported_ptr, + meta_hist_ptr, meta_cur_ptr, bytes_hist_ptr, bytes_cur_ptr, + ring_ptr, cursor_ptr, totals_ptr, count_ptr, + num_slots, num_ids, + NB: tl.constexpr, NB_P: tl.constexpr, L: tl.constexpr, L_P: tl.constexpr, RING_N: tl.constexpr, + BLOCK: tl.constexpr, +): + """Keep each new bad (slot, owner) with its history, count the findings and clear the work arrays.""" + s = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + ms = s < num_slots + b = tl.arange(0, NB_P) + cell = ms[:, None] & (b < NB)[None, :] + cnt = tl.load(bad_cnt_ptr + s[:, None] * NB + b[None, :], mask=cell, other=0) + first = tl.load(bad_first_ptr + s[:, None] * NB + b[None, :], mask=cell, other=_NO_BAD) + bank_bad = cnt > 0 + bytes_bad = tl.sum(bank_bad.to(tl.int32), axis=1) > 0 + bank_mask = tl.sum(tl.where(bank_bad, 1 << b[None, :], 0), axis=1) + first_bank = tl.min(tl.where(bank_bad, b[None, :], NB_P), axis=1) + first_word = tl.sum(tl.where(b[None, :] == first_bank[:, None], first, 0), axis=1) + fl = tl.load(flags_ptr + s, mask=ms, other=0) | tl.where(bytes_bad, 1, 0) + bad = ms & (fl != 0) + owner = tl.load(id_of_slot_ptr + s, mask=ms, other=-1) + sid = tl.load(stale_id_ptr + s, mask=ms, other=-1) + key = (owner.to(tl.int64) + 1) * (num_ids + 2) + sid + 1 + new = bad & (tl.load(reported_ptr + s, mask=ms, other=-1) != key) + tl.store(reported_ptr + s, tl.where(bad, key, -1), mask=ms) + pos = tl.atomic_add(cursor_ptr + s * 0, 1, mask=new) + keep = new & (pos < RING_N) + in_range = ms & (owner >= 0) & (owner < num_ids) + f = tl.arange(0, 16) + z = tl.zeros([BLOCK], tl.int64) + vals = tl.zeros([BLOCK, 16], tl.int64) - 1 + vals = _put(vals, f, 0, z + tl.load(totals_ptr) - 1) + vals = _put(vals, f, 1, z + tl.load(dstep_ptr)) + vals = _put(vals, f, 2, z + tl.load(lstep_ptr)) + vals = _put(vals, f, 3, s) + vals = _put(vals, f, 4, owner) + vals = _put(vals, f, 5, tl.load(slot_for_id_ptr + owner, mask=in_range, other=-2)) + vals = _put(vals, f, 6, tl.load(usage_ptr + s, mask=ms, other=-1)) + vals = _put(vals, f, 7, fl) + vals = _put(vals, f, 8, bank_mask) + vals = _put(vals, f, 9, tl.where(bytes_bad, first_bank, -1)) + vals = _put(vals, f, 10, tl.where(bytes_bad, first_word.to(tl.int64) * 4, -1)) + vals = _put(vals, f, 11, tl.sum(cnt, axis=1)) + vals = _put(vals, f, 12, sid) + vals = _put(vals, f, 13, tl.load(stale_n_ptr + s, mask=ms, other=0)) + vals = _put(vals, f, 14, tl.load(meta_cur_ptr + s, mask=ms, other=0)) + vals = _put(vals, f, 15, tl.load(bytes_cur_ptr + s, mask=ms, other=0)) + rec = ring_ptr + pos.to(tl.int64) * (16 + 2 * _D * _F) + tl.store(rec[:, None] + f[None, :], vals, mask=keep[:, None]) + j = tl.arange(0, 64) + jm = keep[:, None] & (j < _D * _F)[None, :] + src = s.to(tl.int64)[:, None] * (_D * _F) + j[None, :] + tl.store(rec[:, None] + 16 + j[None, :], tl.load(meta_hist_ptr + src, mask=jm, other=-1), mask=jm) + tl.store(rec[:, None] + 16 + _D * _F + j[None, :], tl.load(bytes_hist_ptr + src, mask=jm, other=-1), mask=jm) + t = tl.arange(0, 16) + add = tl.where(t == 4, tl.sum(bad.to(tl.int64), axis=0), 0) + add = tl.where(t == 5, tl.sum((ms & bytes_bad).to(tl.int64), axis=0), add) + add = tl.where(t == 10, tl.sum(new.to(tl.int64), axis=0), add) + add = tl.where(t == 11, tl.sum((bad & ~new).to(tl.int64), axis=0), add) + tl.atomic_add(totals_ptr + t, add, mask=add != 0, sem="relaxed") + tl.store(flags_ptr + s, tl.zeros([BLOCK], tl.int32), mask=ms) + tl.store(stale_id_ptr + s, tl.zeros([BLOCK], tl.int32) - 1, mask=ms) + tl.store(stale_n_ptr + s, tl.zeros([BLOCK], tl.int32), mask=ms) + tl.store(bad_cnt_ptr + s[:, None] * NB + b[None, :], tl.zeros([BLOCK, NB_P], tl.int32), mask=cell) + tl.store(bad_first_ptr + s[:, None] * NB + b[None, :], tl.zeros([BLOCK, NB_P], tl.int32) + _NO_BAD, mask=cell) + if tl.program_id(0) == 0: + ll = tl.arange(0, L_P) + tl.store(count_ptr + ll, tl.zeros([L_P], tl.int32), mask=ll < L) + + +class SlotAuditor: + """Per-slot writer history (recorded by ``record_*`` on the writer's stream) and the periodic + audit (``end_decode_step`` -> ``scan`` on the compute stream). ``bind`` sizes everything for a + cache's slots and banks; a rebuild rebinds and starts a new history, the totals and records stay. + Holds no reference to its cache: every call takes it.""" + + def __init__(self, num_layers: int, num_experts: int, device: torch.device, *, interval: int, + ring: int | None = None, scratch_bytes: int = SCRATCH_BYTES) -> None: + self.num_layers = num_layers + self.num_experts = num_experts + self.device = device + self.interval = interval + self.ring_n = RING if ring is None else ring + self.scratch_bytes = scratch_bytes + ring = self.ring_n + self.dstep = torch.zeros((1,), dtype=torch.int64, device=device) + self.seq = torch.zeros((1,), dtype=torch.int64, device=device) + self.totals = torch.zeros((NUM_TOTALS,), dtype=torch.int64, device=device) + self.ring = torch.full((ring, RECORD_FIELDS), -1, dtype=torch.int64, device=device) + self.cursor = torch.zeros((1,), dtype=torch.int64, device=device) + self.count = torch.zeros((num_layers,), dtype=torch.int32, device=device) + self.num_slots = 0 + self.bank_names: list[str] = [] + self._feats: list[int] = [] + self._steps = 0 + self._printed = 0 + self._last = torch.zeros((NUM_TOTALS,), dtype=torch.int64) + self._dummy = torch.zeros((1,), dtype=torch.int32, device=device) + logger.info(f"MoE slot audit: recording every slot writer, auditing every held slot each {interval} decode steps") + + # ------------------------------------------------------------------ setup + + def bind(self, cache) -> None: + """Size the history and audit arrays for ``cache``'s slots and banks (a rebuild starts a new + history) and point the byte checks at its banks.""" + caches = [bank for _, bank in cache.banks] + feats = [math.prod(bank.shape[1:]) * bank.element_size() for bank in caches] + assert all(f % 4 == 0 for f in feats), feats + if self.bank_names: + assert list(cache.bank_schema) == self.bank_names and feats == self._feats, "the bank layout changed" + else: + self.bank_names = list(cache.bank_schema) + self._feats = feats + self._alloc_scratch(cache, feats) + self.cache_ptrs = torch.tensor([bank.data_ptr() for bank in caches], dtype=torch.int64, device=self.device) + S, dev = cache.cache_size, self.device + self.num_slots = S + self.meta_hist = torch.full((S, DEPTH, EVENT_FIELDS), -1, dtype=torch.int64, device=dev) + self.bytes_hist = torch.full((S, DEPTH, EVENT_FIELDS), -1, dtype=torch.int64, device=dev) + self.meta_cur = torch.zeros((S,), dtype=torch.int32, device=dev) + self.bytes_cur = torch.zeros((S,), dtype=torch.int32, device=dev) + self.flags = torch.zeros((S,), dtype=torch.int32, device=dev) + self.stale_id = torch.full((S,), -1, dtype=torch.int32, device=dev) + self.stale_n = torch.zeros((S,), dtype=torch.int32, device=dev) + self.reported = torch.full((S,), -1, dtype=torch.int64, device=dev) + nb = len(feats) + self.bad_cnt = torch.zeros((S * nb,), dtype=torch.int32, device=dev) + self.bad_first = torch.full((S * nb,), 0x7FFFFFFF, dtype=torch.int32, device=dev) + # a correct cache holds at most num_experts slots of a layer; more are counted unchecked + per_layer = min(self.num_experts, S) + self.chunks = triton.cdiv(per_layer, self.rows) + self.cap = self.chunks * self.rows + self.list_rows = torch.zeros((self.num_layers, self.cap), dtype=torch.int32, device=dev) + self.list_slots = torch.zeros((self.num_layers, self.cap), dtype=torch.int32, device=dev) + self.nums = torch.zeros((self.num_layers * self.chunks,), dtype=torch.int64, device=dev) + # layers without a device alias (LOCKED/PAGEABLE host banks) cannot be gathered + self.layer_ok_host = [layer not in cache._unpinned_layers for layer in range(self.num_layers)] + self.layer_ok = torch.tensor(self.layer_ok_host, dtype=torch.int32, device=dev) + self.count.zero_() + + def _alloc_scratch(self, cache, feats: list[int]) -> None: + row = sum(feats) + self.rows = max(1, min(MAX_CHUNK_ROWS, (self.scratch_bytes - 1) // row)) + rows = self.rows + self.scratch = torch.empty((rows * row,), dtype=torch.uint8, device=self.device) + offsets = [rows * sum(feats[:b]) for b in range(len(feats))] + self.scratch_views = [ + self.scratch[o : o + rows * f].view(bank.dtype).view(rows, *bank.shape[1:]) + for o, f, (_, bank) in zip(offsets, feats, cache.banks) + ] + self.scratch_ptrs = torch.tensor([self.scratch.data_ptr() + o for o in offsets], dtype=torch.int64, device=self.device) + self.words = torch.tensor([f // 4 for f in feats], dtype=torch.int64, device=self.device) + self.iota = torch.arange(rows, dtype=torch.int32, device=self.device) + self._word_chunks = triton.cdiv(max(feats) // 4, COMPARE_WORDS) + logger.info(f"MoE slot audit scratch: {self.scratch.numel() / 2**20:.1f} MiB, {rows} experts a chunk") + + # ------------------------------------------------------------------ recording (graph-capturable) + + def begin_step(self) -> None: + """Advance the decode-step clock (the first GPU layer's decode ensure calls this).""" + _bump_step_kernel[(1,)](self.dstep) + + def _launch(self, cache, kind: int, layer: int, *, meta: bool, width: int, dst=None, src=None, aux=None, + num=None, start: int = 0, row: int = ROW_PLAN, row_base: int = 0, pred: int = PRED_ALL) -> None: + if width <= 0 or not self.num_slots: + return + block = min(triton.next_power_of_2(width), 128) + d = self._dummy + _record_kernel[(triton.cdiv(width, block),)]( + self.meta_hist if meta else self.bytes_hist, self.meta_cur if meta else self.bytes_cur, + self.seq, self.dstep, cache.step, cache.id_of_slot, cache.usage, self.totals, + d if dst is None else dst, d if src is None else src, d if aux is None else aux, d if num is None else num, + kind, layer, start, row_base, width, self.num_slots, + PLAN=dst is not None, ROW=row, HAS_AUX=aux is not None, HAS_NUM=num is not None, PRED=pred, + E=self.num_experts, BLOCK=block, num_warps=4, + ) + + def record_plan(self, cache, kind: int, layer: int, dst: torch.Tensor, src: torch.Tensor | None, num: torch.Tensor, + *, meta: bool, width: int | None = None, aux: torch.Tensor | None = None, row_base: int = 0) -> None: + """The writer just wrote slots ``dst[:num]`` (rows ``src[:num]``, else ``dst - row_base``).""" + width = dst.numel() if width is None else min(width, dst.numel()) + self._launch(cache, kind, layer, meta=meta, width=width, dst=dst, src=src, aux=aux, num=num, + row=ROW_PLAN if src is not None else ROW_FROM_SLOT, row_base=row_base) + + def record_range(self, cache, kind: int, layer: int, start: int, n: int, *, meta: bool) -> None: + """The writer just wrote expert ``i`` of ``layer`` into slot ``start + i`` for ``i < n``.""" + self._launch(cache, kind, layer, meta=meta, width=n, start=start, row=ROW_INDEX) + + def record_clear(self, cache, kind: int, context_layer: int, start: int, n: int, *, layer_above_e: bool = False) -> None: + """Before a writer clears held slots in ``[start, start + n)`` (with ``layer_above_e``: only + those >= E held by ``context_layer``); the event names the old owner.""" + self._launch(cache, kind, context_layer, meta=True, width=n, start=start, row=ROW_OLD_OWNER, + pred=PRED_LAYER_ABOVE_E if layer_above_e else PRED_HELD) + + # ------------------------------------------------------------------ audit + + def end_decode_step(self, cache) -> bool: + """After a decode forward, on the compute stream: audit every ``interval`` steps.""" + self._steps += 1 + if self._steps % self.interval: + return False + self.scan(cache) + return True + + def scan(self, cache) -> None: + """Audit every held slot now (enqueued on the current stream; no host sync).""" + S, num_ids = self.num_slots, self.num_layers * self.num_experts + assert S == cache.cache_size and cache.id_of_slot.numel() == S, "bind after a rebuild" + block = 1024 + _audit_meta_kernel[(triton.cdiv(max(S, num_ids), block),)]( + cache.id_of_slot, cache.slot_for_id.view(-1), self.layer_ok, self.count, self.list_rows, self.list_slots, + self.flags, self.stale_id, self.stale_n, self.totals, S, num_ids, + E=self.num_experts, CAP=self.cap, BLOCK=block, num_warps=4, + ) + n_lists = self.num_layers * self.chunks + _audit_nums_kernel[(1,)](self.count, self.nums, L=self.num_layers, C=self.chunks, R=self.rows, + BLOCK=triton.next_power_of_2(n_lists)) + nb = len(self.bank_names) + R = self.rows + for layer in range(self.num_layers): + if not self.layer_ok_host[layer]: + continue + for c in range(self.chunks): + rows = self.list_rows[layer, c * R : (c + 1) * R] + slots = self.list_slots[layer, c * R : (c + 1) * R] + num = self.nums[layer * self.chunks + c : layer * self.chunks + c + 1] + self._gather(cache, layer, rows, num) + _audit_compare_kernel[(R, nb, self._word_chunks)]( + self.cache_ptrs, self.scratch_ptrs, self.words, slots, num, self.bad_cnt, self.bad_first, + NB=nb, BLOCK=COMPARE_WORDS, num_warps=8, + ) + rblock = 256 + _audit_report_kernel[(triton.cdiv(S, rblock),)]( + cache.id_of_slot, cache.slot_for_id.view(-1), cache.usage, cache.step, self.dstep, + self.flags, self.stale_id, self.stale_n, self.bad_cnt, self.bad_first, self.reported, + self.meta_hist, self.meta_cur, self.bytes_hist, self.bytes_cur, + self.ring, self.cursor, self.totals, self.count, S, num_ids, + NB=nb, NB_P=triton.next_power_of_2(nb), L=self.num_layers, L_P=triton.next_power_of_2(self.num_layers), + RING_N=self.ring_n, BLOCK=rblock, num_warps=4, + ) + + def _gather(self, cache, layer: int, rows: torch.Tensor, num: torch.Tensor) -> None: + # the plain fused copy, not the slim one prefetch copies use: a slim-kernel bug cannot hide itself + if cache._copy_fused_ok: + from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit + + fast_index_copy_multi_jit( + self.scratch_ptrs, cache._copy_src_ptrs[layer], cache._copy_feat_bytes, self.iota, rows, num, + ) + return + from freetoken.kernel import fast_index_copy_jit + + for (per_layer, _), view in zip(cache.banks, self.scratch_views): + fast_index_copy_jit(view, self.iota, per_layer[layer], rows, num) + + # ------------------------------------------------------------------ reporting (host syncs) + + def take_window(self) -> tuple[list[int], list[int], list[list[int]], int]: + """(totals now, totals of this window, records kept since the last call, records in all).""" + totals = self.totals.cpu() + window = (totals - self._last).tolist() + self._last = totals + cursor = int(self.cursor.item()) + upto = min(cursor, self.ring_n) + records = self.ring[self._printed : upto].cpu().tolist() if upto > self._printed else [] + self._printed = max(self._printed, upto) + return totals.tolist(), window, records, cursor + + def report_window(self, steps: int) -> list[tuple[str, str]]: + """(level, line) pairs for one MOE_STATS_INTERVAL window; nothing when no audit ran.""" + totals, window, records, cursor = self.take_window() + if not window[T_AUDITS] and not records: + return [] + bad = bool(window[T_BAD] or window[T_MAP_RANGE] or window[T_BAD_PLAN] or records) + lines = [("warning" if bad else "info", + f"MoE slot audit ({steps} decode steps, every {self.interval}): {self.format_counts(window)}; " + f"session: {totals[T_AUDITS]} audits, {totals[T_NEW]} bad (slot, owner) pairs")] + lines += [("warning", "MoE slot audit record: " + self.format_record(r)) for r in records] + if records and cursor > self.ring_n: + lines.append(("warning", f"MoE slot audit: {cursor} bad (slot, owner) pairs in all, the first {self.ring_n} are kept")) + return lines + + def report_session(self) -> list[tuple[str, str]]: + """Session totals and every kept record (shutdown).""" + totals = self.totals.cpu().tolist() + cursor = int(self.cursor.item()) + if not totals[T_AUDITS]: + return [] + bad = bool(totals[T_BAD] or totals[T_MAP_RANGE] or totals[T_BAD_PLAN]) + lines = [("warning" if bad else "info", f"MoE slot audit (session, every {self.interval} decode steps): " + f"{self.format_counts(totals)}")] + records = self.ring[: min(cursor, self.ring_n)].cpu().tolist() + lines += [("warning", "MoE slot audit record: " + self.format_record(r)) for r in records] + if cursor > self.ring_n: + lines.append(("warning", f"MoE slot audit: {cursor} bad (slot, owner) pairs in all, the first {self.ring_n} are kept")) + return lines + + @staticmethod + def format_counts(t: list[int]) -> str: + audits = t[T_AUDITS] + per = (lambda v: f"{v / audits:.0f}") if audits else str + text = ( + f"{audits} audits, {per(t[T_HELD])} held slots an audit ({per(t[T_BYTES_CHECKED])} byte-checked), " + f"bad slots={t[T_BAD]} (bytes={t[T_BYTES]}, map={t[T_FWD]}, stale slot_for_id={t[T_STALE]}, " + f"owner out of range={t[T_OWNER_RANGE]}), new={t[T_NEW]}, repeat={t[T_STILL]}, events={t[T_EVENTS]}" + ) + if t[T_UNCHECKED]: + text += f", held slots not byte-checked={t[T_UNCHECKED]}" + if t[T_MAP_RANGE]: + text += f", slot_for_id entries past the cache={t[T_MAP_RANGE]}" + if t[T_BAD_PLAN]: + text += f", plan entries naming no slot={t[T_BAD_PLAN]}" + return text + + # ------------------------------------------------------------------ record decoding + + def _name(self, flat: int) -> str: + e = self.num_experts + return f"L{flat // e}/e{flat % e}" if flat >= 0 else str(flat) + + @staticmethod + def events(rec: list[int]) -> list[dict]: + """Both rings of a record as event dicts in ``seq`` order (``ring`` 'meta' or 'bytes').""" + out = [] + for ring, base in (("meta", HEADER), ("bytes", HEADER + RING_EVENTS)): + for d in range(DEPTH): + ev = rec[base + d * EVENT_FIELDS : base + (d + 1) * EVENT_FIELDS] + if ev[E_SEQ] >= 0: + out.append({ + "ring": ring, "seq": ev[E_SEQ], "kind": ev[E_KIND], "dstep": ev[E_DSTEP], "lstep": ev[E_LSTEP], + "layer": ev[E_LAYER], "row": ev[E_ROW], "aux": ev[E_AUX], "owner": ev[E_OWNER], + "usage": ev[E_USAGE], "pos": ev[E_POS], + }) + return sorted(out, key=lambda ev: ev["seq"]) + + def _flat(self, ev: dict) -> int: + return ev["layer"] * self.num_experts + ev["row"] if ev["layer"] >= 0 and ev["row"] >= 0 else -1 + + def format_event(self, ev: dict) -> str: + e = self.num_experts + text = f"#{ev['seq']} {KIND_NAMES.get(ev['kind'], ev['kind'])} L{ev['layer']}/e{ev['row']}" + if ev["kind"] == PREFILL_HIT_D2D: + text += f" from slot {ev['aux']}" + elif ev["kind"] in (MATERIALIZE_CLEAR, INVALIDATE) and ev["aux"] >= 0: + text += f" (for layer {ev['aux']})" + seen = "slot held" if ev["ring"] == "bytes" else "then held" + if ev["kind"] in (MATERIALIZE_CLEAR, INVALIDATE, RESET): + seen = "held" + owner = f"L{ev['owner'] // e}/e{ev['owner'] % e}" if ev["owner"] >= 0 else "nothing" + return f"{text} @step {ev['dstep']}/lru {ev['lstep']} ({seen} {owner}, usage {ev['usage']})" + + def _kind(self, ev: dict) -> str: + return KIND_NAMES.get(ev["kind"], str(ev["kind"])) + + def _pending_install(self, copy: dict, meta: list[dict], byte: list[dict], before: dict) -> dict | None: + """The earlier install (before ``before``) that ``copy`` completes: an install of the id it + copied with no copy of that id between them.""" + flat = self._flat(copy) + for inst in reversed([ev for ev in meta if ev["kind"] in INSTALLS and ev["seq"] < before["seq"]]): + if self._flat(inst) == flat: + done = any(self._flat(ev) == flat and inst["seq"] < ev["seq"] < copy["seq"] for ev in byte) + return None if done else inst + return None + + def _bytes_story(self, owner: int, meta: list[dict], byte: list[dict]) -> str: + insts = [ev for ev in meta if ev["kind"] in INSTALLS and self._flat(ev) == owner] + if not insts: + return (f"no recorded install of {self._name(owner)} in the kept history " + f"(an unrecorded map writer, or it is older than the last {DEPTH} events)") + inst = insts[-1] + name = self._kind(inst) + after = [ev for ev in byte if ev["seq"] > inst["seq"]] + own = [ev for ev in after if self._flat(ev) == owner] + if own: + last = after[-1] + if self._flat(last) == owner: + return (f"{name} and {self._kind(last)} of {self._name(owner)} (#{last['seq']}), then the bytes changed " + f"with no recorded write: overwritten after the last recorded write, or that copy did not land") + text = (f"{name} and {self._kind(own[-1])} of {self._name(owner)} (#{own[-1]['seq']}), then " + f"{self._kind(last)} of {self._name(self._flat(last))} (#{last['seq']}) overwrote it") + prior = self._pending_install(last, meta, byte, inst) + if prior is not None: + text += f": the late copy of the earlier {self._kind(prior)} of {self._name(self._flat(prior))} (#{prior['seq']})" + return text + text = f"{name}, no {KIND_NAMES[COPY_FOR[inst['kind']]]}" + if not after: + return text + last = after[-1] + prior = self._pending_install(last, meta, byte, inst) + if prior is not None: + return (f"{text}; the last write, {self._kind(last)} of {self._name(self._flat(last))} (#{last['seq']}), " + f"is the copy of the earlier {self._kind(prior)} (#{prior['seq']}) and landed after this install") + where = (f"from layer {last['layer']}'s source" if last["row"] == owner % self.num_experts + else f"of {self._name(self._flat(last))}") + return f"{name} of {self._name(owner)}, then {self._kind(last)} {where} ({self._name(self._flat(last))}, #{last['seq']})" + + def diagnose(self, rec: list[int]) -> str: + """What the history says about a bad slot, in words.""" + owner, fl = rec[H_OWNER], rec[H_FLAGS] + evs = self.events(rec) + meta = [ev for ev in evs if ev["ring"] == "meta"] + byte = [ev for ev in evs if ev["ring"] == "bytes"] + parts = [] + if fl & F_OWNER_RANGE: + parts.append(f"id_of_slot={owner} names no expert") + if fl & F_BYTES and owner >= 0: + parts.append(self._bytes_story(owner, meta, byte)) + if fl & F_FWD: + parts.append(f"slot_for_id[{self._name(owner)}]={rec[H_MAPPED]}, not this slot") + if fl & F_STALE: + sid = rec[H_STALE_ID] + text = (f"stale slot_for_id: {self._name(sid)} -> this slot, which holds {self._name(owner)}" + + (f" ({rec[H_STALE_N]} such ids)" if rec[H_STALE_N] > 1 else "")) + gone = [ev for ev in meta if ev["kind"] in INSTALLS and self._flat(ev) == sid] + if gone: + later = [ev for ev in meta if ev["seq"] > gone[-1]["seq"]] + text += (f"; {self._name(sid)} was installed here by {self._kind(gone[-1])} #{gone[-1]['seq']}" + + (f" and replaced by {self._kind(later[0])} #{later[0]['seq']} without clearing " + f"its map entry" if later else "")) + parts.append(text) + # an install's recorder reads the slot right after the install on the same stream + for ev in meta: + if ev["kind"] in INSTALLS and ev["owner"] != self._flat(ev): + parts.append(f"right after the {self._kind(ev)} of {self._name(self._flat(ev))} (#{ev['seq']}) the slot " + f"named {self._name(ev['owner'])}: another map writer ran concurrently") + return "; ".join(parts) or "clean" + + def format_record(self, rec: list[int]) -> str: + fl = rec[H_FLAGS] + text = (f"audit={rec[H_AUDIT]} step={rec[H_DSTEP]} lru={rec[H_LSTEP]} slot={rec[H_SLOT]} " + f"holds {self._name(rec[H_OWNER])} (slot_for_id={rec[H_MAPPED]}, usage={rec[H_USAGE]})") + if fl & F_BYTES: + banks = [self.bank_names[b] if b < len(self.bank_names) else str(b) + for b in range(max(len(self.bank_names), 1)) if rec[H_BANKS] >> b & 1] + first = self.bank_names[rec[H_FIRST_BANK]] if 0 <= rec[H_FIRST_BANK] < len(self.bank_names) else rec[H_FIRST_BANK] + text += (f" | bytes differ in {', '.join(banks)} (first at {first}+{rec[H_FIRST_BYTE]}, " + f"{rec[H_BAD_WORDS]} bad words)") + text += f" | diagnosis: {self.diagnose(rec)}" + history = "; ".join(self.format_event(ev) for ev in self.events(rec)) + return text + f" | history ({rec[H_META_N]} map / {rec[H_BYTES_N]} byte events in all): {history or 'none'}" diff --git a/python/freetoken/moe/verify.py b/python/freetoken/moe/verify.py new file mode 100644 index 0000000000..1f096c9c3a --- /dev/null +++ b/python/freetoken/moe/verify.py @@ -0,0 +1,516 @@ +"""Decode expert-cache verifier (``FREETOKEN_MOE_PREFETCH_VERIFY``): a debug detector for slot races. + +For every GPU decode layer call it checks, on the compute stream and in its order, what the +grouped GEMM is about to read and has just read: + +- meta (before and after the GEMM): each routed entry's slot maps to the expert the router chose, + in both slot maps (``id_of_slot[s] == layer * E + id`` and ``slot_for_id[layer, id] == s``); the + router's ids are copied before ensure rewrites them in place; +- bytes before the GEMM (``full``): every routed slot's row in every bank equals the host bank row + of that expert for that layer, gathered into a scratch buffer with the fused copy; +- bytes after the GEMM (``full``): the same slots against the same scratch, so a slot overwritten + while the GEMM ran shows up even when a later copy restores it. + +Nothing here records or waits on an event: the prefetch streams keep exactly the edges they have +without it. The checks do delay the compute stream (``full`` gathers every routed expert over +PCIe), which moves where the other streams land relative to it. ``meta`` skips the byte checks and +costs a few one-CTA launches per layer. Counts accumulate per layer on the device and the first +``RING`` offending records are kept with the slot state seen at check time; the engine reports +them every MOE_STATS_INTERVAL decode steps and at shutdown. +""" +from __future__ import annotations + +import math + +import torch +import triton +import triton.language as tl + +from freetoken.env import ENV +from freetoken.utils import init_logger + +logger = init_logger(__name__) + +MODES = ("off", "meta", "full") +# counters[layer, col]: checked calls, offending records per check, routed experts left unchecked +CHECKS, META_BAD, PRE_BAD, POST_BAD, POST_META_BAD, UNCHECKED = range(6) +NUM_COLS = 6 +COL_NAMES = ("checks", "meta_bad", "pre_bad", "post_bad", "post_meta_bad", "unchecked") +META_PRE, BYTES_PRE, META_POST, BYTES_POST = range(4) +KINDS = ("meta_pre", "bytes_pre", "meta_post", "bytes_post") +# record fields; owner / mapped / usage / step are read when the check runs +(R_CALL, R_LAYER, R_ROW, R_EXPERT, R_SLOT, R_BANK, R_KIND, R_OWNER, R_MAPPED, R_USAGE, R_STEP, + R_SOURCE, R_FIRST_BYTE, R_BAD_WORDS) = range(14) +NUM_FIELDS = 14 +# R_SOURCE bits: the slot is in this call's demand copy plan, this layer's prefetch copy plan, or +# the next layer's prefetch plan (installed earlier on the compute stream) +SRC_DEMAND, SRC_PREFETCH, SRC_NEXT_PREFETCH = 1, 2, 4 +RING = 64 +# routed entries a call may have (bs * top_k); wider eager batches are counted unchecked +MAX_ENTRIES = 64 +CHUNK_WORDS = 4096 +_NF = tl.constexpr(NUM_FIELDS) +_NC = tl.constexpr(NUM_COLS) +_NO_BAD = tl.constexpr(0x7FFFFFFF) + + +def resolve_mode(mode: str | None = None) -> str: + raw = (ENV.MOE_PREFETCH_VERIFY.value if mode is None else mode).strip().lower() + if raw in ("", "0", "off", "false", "no"): + return "off" + if raw in ("1", "on", "true", "yes", "full"): + return "full" + if raw == "meta": + return "meta" + raise ValueError(f"FREETOKEN_MOE_PREFETCH_VERIFY={raw!r}: expected 0, 1 (full) or meta") + + +@triton.jit +def _put(vals, f, idx: tl.constexpr, v): + return tl.where(f[None, :] == idx, v.to(tl.int64)[:, None], vals) + + +@triton.jit +def _in_plan(s, plan_ptr, num_ptr, W: tl.constexpr, W_P: tl.constexpr): + j = tl.arange(0, W_P) + plan = tl.load(plan_ptr + j, mask=j < tl.minimum(tl.load(num_ptr), W), other=-3) + return tl.sum((s[:, None] == plan[None, :]).to(tl.int32), axis=1) > 0 + + +@triton.jit +def _source(s, dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, + DEM_W: tl.constexpr, DEM_P: tl.constexpr, PF_W: tl.constexpr, PF_P: tl.constexpr, + HAS_PF: tl.constexpr, HAS_NX: tl.constexpr): + src = tl.where(_in_plan(s, dem_ptr, dem_num_ptr, DEM_W, DEM_P), 1, 0) + if HAS_PF: + src += tl.where(_in_plan(s, pf_ptr, pf_num_ptr, PF_W, PF_P), 2, 0) + if HAS_NX: + src += tl.where(_in_plan(s, nx_ptr, nx_num_ptr, PF_W, PF_P), 4, 0) + return src + + +@triton.jit +def _write(ring_ptr, cursor, bad, vals, RING_N: tl.constexpr): + """Append the ``bad`` rows of ``vals`` after ``cursor`` while the ring has room; the new cursor.""" + f = tl.arange(0, 16) + pos = cursor + tl.cumsum(bad.to(tl.int64), axis=0) - 1 + keep = bad & (pos < RING_N) + tl.store(ring_ptr + pos[:, None] * _NF + f[None, :], vals, mask=keep[:, None] & (f[None, :] < _NF)) + return cursor + tl.sum(bad.to(tl.int64), axis=0) + + +@triton.jit +def _meta(ids_ptr, slots_ptr, id_of_slot_ptr, slot_for_id_ptr, usage_ptr, step_ptr, + dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, + ring_ptr, cursor, counters_ptr, col: tl.constexpr, kind: tl.constexpr, call, layer, num_slots, + N: tl.constexpr, N_P: tl.constexpr, E: tl.constexpr, + DEM_W: tl.constexpr, DEM_P: tl.constexpr, PF_W: tl.constexpr, PF_P: tl.constexpr, + HAS_PF: tl.constexpr, HAS_NX: tl.constexpr, RING_N: tl.constexpr): + """Check every routed entry's slot against both maps; returns the new ring cursor.""" + i = tl.arange(0, N_P) + m = i < N + ids = tl.load(ids_ptr + i, mask=m, other=0) + s = tl.load(slots_ptr + i, mask=m, other=-1) + ok_s = m & (s >= 0) & (s < num_slots) + ok_id = m & (ids >= 0) & (ids < E) + fid = layer * E + ids + owner = tl.load(id_of_slot_ptr + s, mask=ok_s, other=-2) + mapped = tl.load(slot_for_id_ptr + fid, mask=ok_id, other=-2) + bad = m & ~(ok_s & ok_id & (owner == fid) & (mapped == s)) + usage = tl.load(usage_ptr + s, mask=ok_s, other=-1) + src = _source(s, dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, DEM_W, DEM_P, PF_W, PF_P, HAS_PF, HAS_NX) + f = tl.arange(0, 16) + z = tl.zeros([N_P], tl.int64) + vals = tl.zeros([N_P, 16], tl.int64) - 1 + vals = _put(vals, f, 0, z + call) + vals = _put(vals, f, 1, z + layer) + vals = _put(vals, f, 2, i) + vals = _put(vals, f, 3, ids) + vals = _put(vals, f, 4, s) + vals = _put(vals, f, 6, z + kind) + vals = _put(vals, f, 7, owner) + vals = _put(vals, f, 8, mapped) + vals = _put(vals, f, 9, usage) + vals = _put(vals, f, 10, z + tl.load(step_ptr)) + vals = _put(vals, f, 11, src) + nbad = tl.sum(bad.to(tl.int64), axis=0) + tl.store(counters_ptr + col, tl.load(counters_ptr + col) + nbad) + return _write(ring_ptr, cursor, bad, vals, RING_N) + + +@triton.jit(do_not_specialize=["layer", "num_slots"]) +def _plan_kernel( + ids_ptr, slots_ptr, id_of_slot_ptr, slot_for_id_ptr, usage_ptr, step_ptr, + dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, + ring_ptr, cursor_ptr, counters_ptr, calls_ptr, + u_ids_ptr, u_slots_ptr, u_rows_ptr, u_num_ptr, slots_copy_ptr, + layer, num_slots, + N: tl.constexpr, N_P: tl.constexpr, E: tl.constexpr, CAP: tl.constexpr, + DEM_W: tl.constexpr, DEM_P: tl.constexpr, PF_W: tl.constexpr, PF_P: tl.constexpr, + HAS_PF: tl.constexpr, HAS_NX: tl.constexpr, BYTES: tl.constexpr, RING_N: tl.constexpr, +): + """Before the GEMM: count the call, check the maps, keep the slots for the after-GEMM checks and + list the distinct (expert, slot) pairs to byte-check (first CAP of them).""" + call = tl.load(calls_ptr + layer) + tl.store(calls_ptr + layer, call + 1) + row = counters_ptr + layer * _NC + tl.store(row + 0, tl.load(row + 0) + 1) + cursor = tl.load(cursor_ptr) + cursor = _meta(ids_ptr, slots_ptr, id_of_slot_ptr, slot_for_id_ptr, usage_ptr, step_ptr, + dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, + ring_ptr, cursor, row, 1, 0, call, layer, num_slots, + N, N_P, E, DEM_W, DEM_P, PF_W, PF_P, HAS_PF, HAS_NX, RING_N) + tl.store(cursor_ptr, cursor) + i = tl.arange(0, N_P) + m = i < N + ids = tl.load(ids_ptr + i, mask=m, other=-1) + s = tl.load(slots_ptr + i, mask=m, other=-1) + tl.store(slots_copy_ptr + i, s, mask=m) + if BYTES: + # an unreadable slot or id is already a meta record; a duplicate pair shares its first row + ok = m & (s >= 0) & (s < num_slots) & (ids >= 0) & (ids < E) + same = (ids[:, None] == ids[None, :]) & (s[:, None] == s[None, :]) & (i[None, :] < i[:, None]) & ok[None, :] + first = ok & (tl.sum(same.to(tl.int32), axis=1) == 0) + rank = tl.cumsum(first.to(tl.int32), axis=0) - 1 + take = first & (rank < CAP) + tl.store(u_ids_ptr + rank, ids, mask=take) + tl.store(u_slots_ptr + rank, s, mask=take) + tl.store(u_rows_ptr + rank, i, mask=take) + total = tl.sum(first.to(tl.int64), axis=0) + kept = tl.minimum(total, CAP) + tl.store(u_num_ptr, kept) + tl.store(row + 5, tl.load(row + 5) + total - kept) + + +@triton.jit +def _compare_kernel(cache_ptrs, ref_ptrs, words_ptr, u_slots_ptr, u_num_ptr, bad_cnt_ptr, bad_first_ptr, + NB: tl.constexpr, BLOCK: tl.constexpr): + """Program (row, bank, chunk): compare one chunk of a routed slot's row with its host copy.""" + r = tl.program_id(0) + b = tl.program_id(1) + words = tl.load(words_ptr + b) + off = tl.program_id(2).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + m = (r < tl.load(u_num_ptr)) & (off < words) + slot = tl.load(u_slots_ptr + r).to(tl.int64) + got = tl.load(tl.load(cache_ptrs + b).to(tl.pointer_type(tl.int32)) + slot * words + off, mask=m, other=0) + want = tl.load(tl.load(ref_ptrs + b).to(tl.pointer_type(tl.int32)) + r * words + off, mask=m, other=0) + bad = m & (got != want) + nbad = tl.sum(bad.to(tl.int32), axis=0) + if nbad > 0: + tl.atomic_add(bad_cnt_ptr + r * NB + b, nbad) + tl.atomic_min(bad_first_ptr + r * NB + b, tl.min(tl.where(bad, off, _NO_BAD), axis=0).to(tl.int32)) + + +@triton.jit(do_not_specialize=["layer", "num_slots"]) +def _tally_kernel( + ids_ptr, slots_copy_ptr, id_of_slot_ptr, slot_for_id_ptr, usage_ptr, step_ptr, + dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, + ring_ptr, cursor_ptr, counters_ptr, calls_ptr, + u_ids_ptr, u_slots_ptr, u_rows_ptr, u_num_ptr, bad_cnt_ptr, bad_first_ptr, + layer, num_slots, + N: tl.constexpr, N_P: tl.constexpr, E: tl.constexpr, CAP: tl.constexpr, CAP_P: tl.constexpr, + NB: tl.constexpr, NB_P: tl.constexpr, + DEM_W: tl.constexpr, DEM_P: tl.constexpr, PF_W: tl.constexpr, PF_P: tl.constexpr, + HAS_PF: tl.constexpr, HAS_NX: tl.constexpr, BYTES: tl.constexpr, POST: tl.constexpr, + BYTES_KIND: tl.constexpr, BYTES_COL: tl.constexpr, RING_N: tl.constexpr, +): + """Turn the compare flags into counts and records (and clear them); after the GEMM also recheck + the maps of the slots it read.""" + call = tl.load(calls_ptr + layer) - 1 + row = counters_ptr + layer * _NC + cursor = tl.load(cursor_ptr) + if POST: + cursor = _meta(ids_ptr, slots_copy_ptr, id_of_slot_ptr, slot_for_id_ptr, usage_ptr, step_ptr, + dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, + ring_ptr, cursor, row, 4, 2, call, layer, num_slots, + N, N_P, E, DEM_W, DEM_P, PF_W, PF_P, HAS_PF, HAS_NX, RING_N) + if BYTES: + j = tl.arange(0, CAP_P * NB_P) + r = j // NB_P + b = j % NB_P + cell = (r < CAP) & (b < NB) + live = cell & (r < tl.load(u_num_ptr)) + cnt = tl.load(bad_cnt_ptr + r * NB + b, mask=live, other=0) + first = tl.load(bad_first_ptr + r * NB + b, mask=live, other=_NO_BAD) + tl.store(bad_cnt_ptr + r * NB + b, tl.zeros([CAP_P * NB_P], tl.int32), mask=cell) + tl.store(bad_first_ptr + r * NB + b, tl.zeros([CAP_P * NB_P], tl.int32) + _NO_BAD, mask=cell) + bad = live & (cnt > 0) + ids = tl.load(u_ids_ptr + r, mask=live, other=0) + s = tl.load(u_slots_ptr + r, mask=live, other=-1) + ok_s = live & (s >= 0) & (s < num_slots) + src = _source(s, dem_ptr, dem_num_ptr, pf_ptr, pf_num_ptr, nx_ptr, nx_num_ptr, DEM_W, DEM_P, PF_W, PF_P, HAS_PF, HAS_NX) + f = tl.arange(0, 16) + z = tl.zeros([CAP_P * NB_P], tl.int64) + vals = tl.zeros([CAP_P * NB_P, 16], tl.int64) - 1 + vals = _put(vals, f, 0, z + call) + vals = _put(vals, f, 1, z + layer) + vals = _put(vals, f, 2, tl.load(u_rows_ptr + r, mask=live, other=-1)) + vals = _put(vals, f, 3, ids) + vals = _put(vals, f, 4, s) + vals = _put(vals, f, 5, b) + vals = _put(vals, f, 6, z + BYTES_KIND) + vals = _put(vals, f, 7, tl.load(id_of_slot_ptr + s, mask=ok_s, other=-2)) + vals = _put(vals, f, 8, tl.load(slot_for_id_ptr + layer * E + ids, mask=live, other=-2)) + vals = _put(vals, f, 9, tl.load(usage_ptr + s, mask=ok_s, other=-1)) + vals = _put(vals, f, 10, z + tl.load(step_ptr)) + vals = _put(vals, f, 11, src) + vals = _put(vals, f, 12, first.to(tl.int64) * 4) + vals = _put(vals, f, 13, cnt) + tl.store(row + BYTES_COL, tl.load(row + BYTES_COL) + tl.sum(bad.to(tl.int64), axis=0)) + cursor = _write(ring_ptr, cursor, bad, vals, RING_N) + tl.store(cursor_ptr, cursor) + + +class ExpertVerifier: + """Device counters, the record ring and (``full``) the scratch the byte checks compare against. + + ``begin`` copies the router's ids before ensure rewrites them, ``before_gemm`` runs after the + joins and the demand copy, ``after_gemm`` right after the GEMM; all three enqueue on the current + (compute) stream with fixed shapes and no host sync, so they capture into the decode graph. + ``bind`` sizes the scratch from the banks once and refreshes the slot bank addresses in place + (a rebuild reallocates the banks); counters survive rebuilds and resets. + """ + + def __init__(self, num_layers: int, num_experts: int, device: torch.device, *, mode: str = "full", + rows: int | None = None) -> None: + from freetoken.kernel.fast_index_copy import _skip_fast_index_copy_enabled + + if mode == "full" and _skip_fast_index_copy_enabled(): + logger.warning( + "FREETOKEN_MOE_PREFETCH_VERIFY=1 with FREETOKEN_SKIP_FAST_INDEX_COPY: the host rows cannot be " + "gathered, so only the slot maps are checked (meta)" + ) + mode = "meta" + self.mode = mode + self.num_layers = num_layers + self.num_experts = num_experts + self.device = device + self.rows = ENV.MOE_PREFETCH_VERIFY_ROWS.value if rows is None else int(rows) + if self.rows < 1: + raise ValueError(f"FREETOKEN_MOE_PREFETCH_VERIFY_ROWS={self.rows} must be >= 1") + self.counters = torch.zeros((num_layers, NUM_COLS), dtype=torch.int64, device=device) + self.totals = torch.zeros((num_layers, NUM_COLS), dtype=torch.int64) + self.calls = torch.zeros((num_layers,), dtype=torch.int64, device=device) + self.ring = torch.full((RING, NUM_FIELDS), -1, dtype=torch.int64, device=device) + self.cursor = torch.zeros((1,), dtype=torch.int64, device=device) + self.ids = torch.full((MAX_ENTRIES,), -1, dtype=torch.int32, device=device) + self.slots = torch.full((MAX_ENTRIES,), -1, dtype=torch.int32, device=device) + self.u_ids = torch.zeros((self.rows,), dtype=torch.int32, device=device) + self.u_slots = torch.zeros((self.rows,), dtype=torch.int32, device=device) + self.u_rows = torch.zeros((self.rows,), dtype=torch.int32, device=device) + self.u_num = torch.zeros((1,), dtype=torch.int64, device=device) + self.iota = torch.arange(self.rows, dtype=torch.int32, device=device) + self.bank_names: list[str] = [] + self.scratch: torch.Tensor | None = None + self._printed = 0 + self._top_k = 1 + self._call = None # (layer, entries, this layer's prefetch plan or None) between begin and after_gemm + self._next = None + logger.info( + f"MoE verify {self.mode}: checking every GPU decode layer call's slot maps" + + (f" and slot bytes before and after the GEMM (up to {self.rows} experts a call)" if self.mode == "full" else "") + ) + + def bind(self, cache) -> None: + """Size the scratch for ``cache``'s banks (first call) or point the checks at its new banks.""" + caches = [bank for _, bank in cache.banks] + feats = [math.prod(bank.shape[1:]) * bank.element_size() for bank in caches] + if self.bank_names: + assert list(cache.bank_schema) == self.bank_names and feats == self._feats, "the bank layout changed" + else: + assert all(f % 4 == 0 for f in feats), feats + self.bank_names = list(cache.bank_schema) + self._feats = feats + self.cache_ptrs = torch.zeros((len(feats),), dtype=torch.int64, device=self.device) + if self.mode == "full": + self._alloc_scratch(cache, feats) + self.cache_ptrs.copy_(torch.tensor([bank.data_ptr() for bank in caches], dtype=torch.int64)) + + def _alloc_scratch(self, cache, feats: list[int]) -> None: + rows = self.rows + self.scratch = torch.empty((rows * sum(feats),), dtype=torch.uint8, device=self.device) + offsets = [rows * sum(feats[:b]) for b in range(len(feats))] + self.scratch_views = [ + self.scratch[o : o + rows * f].view(bank.dtype).view(rows, *bank.shape[1:]) + for o, f, (_, bank) in zip(offsets, feats, cache.banks) + ] + self.scratch_ptrs = torch.tensor([self.scratch.data_ptr() + o for o in offsets], dtype=torch.int64, device=self.device) + self.words = torch.tensor([f // 4 for f in feats], dtype=torch.int64, device=self.device) + self.bad_cnt = torch.zeros((rows * len(feats),), dtype=torch.int32, device=self.device) + self.bad_first = torch.full((rows * len(feats),), 0x7FFFFFFF, dtype=torch.int32, device=self.device) + self._chunks = triton.cdiv(max(feats) // 4, CHUNK_WORDS) + logger.info(f"MoE verify scratch: {self.scratch.numel() / 2**20:.1f} MiB for {rows} experts") + + def begin(self, cache, layer_id: int, topk_ids: torch.Tensor, plan: tuple | None) -> None: + """Copy the router's ids of ``layer_id`` (``plan``: this layer's prefetch copy plan, if forked).""" + assert topk_ids.dtype == torch.int32 and topk_ids.is_contiguous() + n = topk_ids.numel() + self._top_k = topk_ids.shape[-1] + if n > MAX_ENTRIES: + self.counters[layer_id, UNCHECKED] += n + self._call = None + return + self.ids[:n].copy_(topk_ids.view(-1)) + self._call = (layer_id, n, plan) + + def _plans(self, cache) -> tuple: + _, n, plan = self._call + pf, nx = plan, self._next + dummy = (self.u_slots, self.u_num) + width = (pf or nx or dummy)[0].numel() + return ( + cache.evict_slots, cache.num_indices, *(pf or dummy), *(nx or dummy), + dict(DEM_W=n, DEM_P=triton.next_power_of_2(n), PF_W=width, PF_P=triton.next_power_of_2(width), + HAS_PF=pf is not None, HAS_NX=nx is not None), + ) + + def before_gemm(self, cache, slots: torch.Tensor, next_plan: tuple | None) -> None: + """Check the maps and (``full``) the bytes of the slots ``slots`` the GEMM is about to read + (``next_plan``: the next layer's prefetch plan, if one was issued).""" + if self._call is None: + return + layer, n, _ = self._call + self._next = next_plan + *plans, plan_kw = self._plans(cache) + full = self.mode == "full" + _plan_kernel[(1,)]( + self.ids, slots, cache.id_of_slot, cache.slot_for_id.view(-1), cache.usage, cache.step, *plans, + self.ring, self.cursor, self.counters, self.calls, + self.u_ids, self.u_slots, self.u_rows, self.u_num, self.slots, + layer, cache.cache_size, + N=n, N_P=triton.next_power_of_2(n), E=self.num_experts, CAP=self.rows, + BYTES=full, RING_N=RING, num_warps=4, **plan_kw, + ) + if not full: + return + self._gather(cache, layer) + self._compare_and_tally(cache, post=False) + + def after_gemm(self, cache) -> None: + """Recheck the maps and (``full``) the bytes of the slots the GEMM just read.""" + if self._call is None: + return + if self.mode == "full": + self._compare_and_tally(cache, post=True) + else: + self._tally(cache, post=True, bytes_=False) + self._call = self._next = None + + def _gather(self, cache, layer: int) -> None: + if cache._copy_fused_ok: + from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit + + fast_index_copy_multi_jit( + self.scratch_ptrs, cache._copy_src_ptrs[layer], cache._copy_feat_bytes, self.iota, self.u_ids, self.u_num, + ) + return + from freetoken.kernel import fast_index_copy_jit + + for (per_layer, _), view in zip(cache.banks, self.scratch_views): + fast_index_copy_jit(view, self.iota, per_layer[layer], self.u_ids, self.u_num) + + def _compare_and_tally(self, cache, *, post: bool) -> None: + nb = len(self.bank_names) + _compare_kernel[(self.rows, nb, self._chunks)]( + self.cache_ptrs, self.scratch_ptrs, self.words, self.u_slots, self.u_num, self.bad_cnt, self.bad_first, + NB=nb, BLOCK=CHUNK_WORDS, num_warps=8, + ) + self._tally(cache, post=post, bytes_=True) + + def _tally(self, cache, *, post: bool, bytes_: bool) -> None: + layer, n, _ = self._call + *plans, plan_kw = self._plans(cache) + nb = max(len(self.bank_names), 1) + _tally_kernel[(1,)]( + self.ids, self.slots, cache.id_of_slot, cache.slot_for_id.view(-1), cache.usage, cache.step, *plans, + self.ring, self.cursor, self.counters, self.calls, + self.u_ids, self.u_slots, self.u_rows, self.u_num, + self.bad_cnt if bytes_ else self.u_rows, self.bad_first if bytes_ else self.u_rows, + layer, cache.cache_size, + N=n, N_P=triton.next_power_of_2(n), E=self.num_experts, CAP=self.rows, + CAP_P=triton.next_power_of_2(self.rows), NB=nb, NB_P=triton.next_power_of_2(nb), + BYTES=bytes_, POST=post, BYTES_KIND=BYTES_POST if post else BYTES_PRE, BYTES_COL=POST_BAD if post else PRE_BAD, + RING_N=RING, num_warps=4, **plan_kw, + ) + + # ------------------------------------------------------------------ reporting (host syncs) + + def take_window(self) -> tuple[torch.Tensor, list[list[int]], int]: + """This window's ``[num_layers, NUM_COLS]`` counts (added to ``totals``; the device ones + restart), the records kept since the last call, and how many records there were in all.""" + window = self.counters.cpu() + self.counters.zero_() + self.totals += window + cursor = int(self.cursor.item()) + upto = min(cursor, RING) + records = self.ring[self._printed : upto].cpu().tolist() if upto > self._printed else [] + self._printed = max(self._printed, upto) + return window, records, cursor + + def report_window(self, steps: int) -> list[tuple[str, str]]: + """(level, line) pairs for one MOE_STATS_INTERVAL window; nothing when no call was checked.""" + window, records, cursor = self.take_window() + if not int(window[:, CHECKS].sum()) and not records: + return [] + bad = self._bad(window) or bool(records) + session = self._bad(self.totals) + lines = [("warning" if bad else "info", + f"MoE verify {self.mode} ({steps} decode steps): {self.format_counts(window)}; session bad={session}")] + lines += [("warning", "MoE verify record: " + self.format_record(r)) for r in records] + if bad and cursor > RING: + lines.append(("warning", f"MoE verify: {cursor} bad records in all, the first {RING} are kept")) + return lines + + def report_session(self) -> list[tuple[str, str]]: + """Session counts and every kept record (shutdown).""" + totals = self.totals + self.counters.cpu() + cursor = int(self.cursor.item()) + if not int(totals[:, CHECKS].sum()): + return [] + bad = self._bad(totals) + lines = [("warning" if bad else "info", f"MoE verify {self.mode} (session): {self.format_counts(totals)}")] + records = self.ring[: min(cursor, RING)].cpu().tolist() + lines += [("warning", "MoE verify record: " + self.format_record(r)) for r in records] + if cursor > RING: + lines.append(("warning", f"MoE verify: {cursor} bad records in all, the first {RING} are kept")) + return lines + + @staticmethod + def _bad(counts: torch.Tensor) -> int: + return int(counts[:, [META_BAD, PRE_BAD, POST_BAD, POST_META_BAD]].sum()) + + def format_counts(self, counts: torch.Tensor) -> str: + total = counts.sum(0).tolist() + text = f"{total[CHECKS]} layer calls checked, " + ", ".join( + f"{COL_NAMES[c]}={total[c]}" for c in (META_BAD, PRE_BAD, POST_BAD, POST_META_BAD) + ) + if total[UNCHECKED]: + text += f", unchecked experts={total[UNCHECKED]}" + per_layer = [ + f"L{layer}=" + "/".join(str(row[c]) for c in (META_BAD, PRE_BAD, POST_BAD, POST_META_BAD)) + for layer, row in enumerate(counts.tolist()) + if any(row[c] for c in (META_BAD, PRE_BAD, POST_BAD, POST_META_BAD)) + ] + if per_layer: + text += " (meta/pre/post/post_meta by layer: " + ", ".join(per_layer) + ")" + return text + + def format_record(self, rec: list[int]) -> str: + e = self.num_experts + kind = KINDS[rec[R_KIND]] if 0 <= rec[R_KIND] < len(KINDS) else str(rec[R_KIND]) + owner = rec[R_OWNER] + owner_text = f"L{owner // e}/e{owner % e}" if owner >= 0 else str(owner) + src = [name for bit, name in ((SRC_DEMAND, "demand copy"), (SRC_PREFETCH, "prefetch copy"), + (SRC_NEXT_PREFETCH, "next layer's prefetch plan")) if rec[R_SOURCE] & bit] + row = rec[R_ROW] + text = ( + f"{kind} call={rec[R_CALL]} layer={rec[R_LAYER]} row={row} (token {row // self._top_k}, rank {row % self._top_k}) " + f"expert={rec[R_EXPERT]} slot={rec[R_SLOT]}" + ) + if rec[R_BANK] >= 0: + bank = self.bank_names[rec[R_BANK]] if rec[R_BANK] < len(self.bank_names) else rec[R_BANK] + text += f" bank={bank} first_bad_byte={rec[R_FIRST_BYTE]} bad_words={rec[R_BAD_WORDS]}" + return text + ( + f" | slot now holds {owner_text}, slot_for_id={rec[R_MAPPED]}, usage={rec[R_USAGE]}, lru_step={rec[R_STEP]}, " + f"slot in: {', '.join(src) or 'no copy plan (a resident hit)'}" + ) diff --git a/tests/models/qwen4_exp/conftest.py b/tests/models/qwen4_exp/conftest.py index 2b9e80214f..4e2bb67bb8 100644 --- a/tests/models/qwen4_exp/conftest.py +++ b/tests/models/qwen4_exp/conftest.py @@ -13,3 +13,16 @@ def _runtime(): core._GLOBAL_CTX = None yield core._GLOBAL_CTX = None + import torch + + if torch.cuda.is_available() and torch.cuda.is_initialized(): + import gc + + # torch keeps a cuBLAS workspace per stream for good: each offload cache's dedicated + # streams would otherwise pin ~16 MiB for the rest of the session + gc.collect() + torch.cuda.synchronize() + clear = getattr(torch._C, "_cuda_clearCublasWorkspaces", None) # private torch API + if clear is not None: + clear() + torch.cuda.empty_cache() diff --git a/tests/models/qwen4_exp/test_build_cpu.py b/tests/models/qwen4_exp/test_build_cpu.py index 58ad0c4dc3..85366638a5 100644 --- a/tests/models/qwen4_exp/test_build_cpu.py +++ b/tests/models/qwen4_exp/test_build_cpu.py @@ -47,3 +47,32 @@ def test_every_layer_gets_its_own_prefixed_weights(): mlp_keys = {k for k in sd if ".mlp." in k} layers = {k.split(".layers.")[1].split(".")[0] for k in mlp_keys if ".layers." in k} assert len(layers) > 1, f"MoE weights landed on a single layer: {sorted(layers)}" + + +def test_offloaded_layers_borrow_the_next_router_for_the_lookahead(): + # FREETOKEN_MOE_PREFETCH: layer L-1's experts hold layer L's router, outside the state dict + # and without the engine's offload-layer walk finding extra layers through it + import dataclasses + + from freetoken.moe.offload_cache import iter_offload_moe_layers + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + previous = rotary._ROPE_DEVICE + set_rope_device(torch.device("cpu")) + try: + config = dataclasses.replace(parse_config(toy_hf_config()), moe_strategy="offload") + with torch.device("meta"): + model = Qwen4ExpForCausalLM(config) + finally: + rotary.get_rope.cache_clear() + set_rope_device(previous) + mlps = [layer.mlp for layer in model.model.layers.op_list] + assert len(list(iter_offload_moe_layers(model))) == len(mlps) == 4 + wired = [mlp.experts._lookahead for mlp in mlps] + assert wired[-1] is None + assert [(gate is mlps[t].gate, t, budget) for gate, t, budget in wired[:-1]] == [ + (True, 1, 3), (True, 2, 3), (True, 3, 4), # layer 3 is the full-attention one + ] + assert not any("lookahead" in k for k in model.state_dict()) + assert sum(k.endswith(".mlp.gate.weight") for k in model.state_dict()) == 4 diff --git a/tests/models/qwen4_exp/test_skeleton.py b/tests/models/qwen4_exp/test_skeleton.py index a2412f8346..2b47748f53 100644 --- a/tests/models/qwen4_exp/test_skeleton.py +++ b/tests/models/qwen4_exp/test_skeleton.py @@ -556,6 +556,1917 @@ def test_moe_copy_overlap_is_bitwise_identical(kind, graph, bs, monkeypatch): assert torch.equal(a, b) +# FREETOKEN_MOE_PREFETCH stack: 4 layers (full attention on layer 3), 16 experts, top-2, and one +# shared cache as small as a layer so every step evicts +_PF_LAYERS, _PF_EXPERTS, _PF_TOPK, _PF_HIDDEN, _PF_INTER = 4, 16, 2, 256, 128 +# ~50-100 us between layers: the predictor must finish before the next ensure even at toy size +_PF_SLEEP_CYCLES = 200_000 + + +def _prefetch_banks(kind): + from freetoken.moe.expert_banks import build_expert_banks + + torch.manual_seed(7) + method = _offload_experts(kind, 0, _PF_EXPERTS, _PF_TOPK, _PF_HIDDEN, _PF_INTER).quant_method + return build_expert_banks(method, _PF_LAYERS, None, device=torch.device("cuda"), dummy=True) + + +def _prefetch_stack(kind, banks, mode, monkeypatch, *, k=6, budget=0, wire=True, overlap=False, policy="rule", + cache_size=_PF_EXPERTS, **cache_kw): + from freetoken.env import ENV + from freetoken.models.qwen4_exp.moe import Qwen4ExpMoE, wire_router_lookahead + from freetoken.moe.offload_cache import OffloadMoeCache + from freetoken.utils.torch_utils import torch_dtype + + monkeypatch.setattr(ENV.MOE_PREFETCH_K, "value", k) + monkeypatch.setattr(ENV.MOE_PREFETCH_BUDGET, "value", budget) + config = parse_config(toy_hf_config( + _PF_LAYERS, hidden_size=_PF_HIDDEN, num_experts=_PF_EXPERTS, num_experts_per_tok=_PF_TOPK, + moe_intermediate_size=_PF_INTER, shared_expert_intermediate_size=_PF_INTER, + )) + device = torch.device("cuda") + with torch.device(device), torch_dtype(torch.bfloat16): + moes = [Qwen4ExpMoE(config, layer_id) for layer_id in range(_PF_LAYERS)] + gen = torch.Generator(device=device).manual_seed(43) + for moe in moes: + _fill(moe, gen, scale=0.2) + experts = [_offload_experts(kind, i, _PF_EXPERTS, _PF_TOPK, _PF_HIDDEN, _PF_INTER) for i in range(_PF_LAYERS)] + cache = OffloadMoeCache( + num_layers=_PF_LAYERS, num_experts=_PF_EXPERTS, cache_size=cache_size, device=device, + cache_policy=policy, quant_format=banks.quant_format, layout=banks.layout, + max_slots=experts[0].quant_method.slot_limit(), prefetch_mode=mode, decode_copy_overlap=overlap, + **cache_kw, + ) + cache.set_bank_sources(banks.sources) + cache.set_alphas(banks.gate_up_alpha, banks.down_alpha) + cache.collect_stats = True + for moe, ex in zip(moes, experts): + ex.offload_cache = cache + moe.experts = ex + if wire: + wire_router_lookahead(moes, config) + cache.reset() + return moes, cache + + +def _prefetch_run(moes, cache, bs, graph, inputs, before_decode=None, sleep=_PF_SLEEP_CYCLES, before_step=None): + """Warm up eagerly (and capture), reset the cache, then decode ``inputs`` (per step, one + ``[bs, hidden]`` input per layer); ``sleep`` cycles between layers, ``before_step(i)`` before step i.""" + from flashlib.kernels.slot_cache import Stat + + _fresh_ctx(_batch=SimpleNamespace(is_prefill=False)) + + def step(xs): + outs = [] + for layer, (moe, x) in enumerate(zip(moes, xs)): + if layer and sleep: + torch.cuda._sleep(sleep) + outs.append(moe.forward(x.clone())) + return outs + + static = [torch.randn(bs, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16) for _ in moes] + step(static) # the eager warm-up the graph runner also does + if graph: + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + outs = step(static) + cache.reset() + cache.reset_stats() + if before_decode is not None: + before_decode() + results = [] + for i, xs in enumerate(inputs): + if before_step is not None: + before_step(i) + if graph: + for s, x in zip(static, xs): + s.copy_(x) + g.replay() + else: + outs = step(xs) + if cache.audit is not None: + cache.audit.end_decode_step(cache) # where the engine calls it: after the forward + results.append([o.clone() for o in outs]) + torch.cuda.synchronize() + state = [t.clone() for t in (cache.slot_for_id, cache.id_of_slot, cache.usage, *cache.bank_views())] + counters = None if cache.prefetch is None else cache.prefetch.counters.cpu() + return results, state, counters, cache.lru_stats[:, Stat.MISS].cpu() + + +def _prefetch_inputs(bs, steps=8, chain=False, seed=0): + gen = torch.Generator(device="cuda").manual_seed(seed * 10 + bs) + + def x(): + return torch.randn(bs, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16, generator=gen) * 0.5 + + if chain: # every layer sees the same input, so the lookahead equals the next router + return [[x()] * _PF_LAYERS for _ in range(steps)] + return [[x() for _ in range(_PF_LAYERS)] for _ in range(steps)] + + +@requires_cuda +@pytest.mark.parametrize("kind", ["bf16", "nvfp4"]) +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("overlap,policy", [(False, "rule"), (True, "rule"), (True, "lru")], ids=["serial", "overlap", "overlap-lru"]) +def test_moe_prefetch_measure_is_bitwise_identical(kind, graph, bs, overlap, policy, monkeypatch): + """FREETOKEN_MOE_PREFETCH=measure predicts and counts beside the decode (with and without + FREETOKEN_MOE_COPY_OVERLAP's side stream); the MoE outputs, the slot map, the usage clock and + the slot contents must not change by a single bit.""" + banks = _prefetch_banks(kind) + inputs = _prefetch_inputs(bs) + off = _prefetch_run(*_prefetch_stack(kind, banks, "off", monkeypatch, overlap=overlap, policy=policy), bs, graph, inputs) + on = _prefetch_run(*_prefetch_stack(kind, banks, "measure", monkeypatch, overlap=overlap, policy=policy), bs, graph, inputs) + + assert off[2] is None and int(off[3].sum()) > 0 and torch.equal(on[3], off[3]) + for got, want in zip(on[0], off[0]): + for a, b in zip(got, want): + assert torch.isfinite(b.float()).all() + assert torch.equal(a, b) + for a, b in zip(on[1], off[1]): + assert torch.equal(a, b) + from freetoken.moe.prefetch import CALLS, ISSUED, ROWS + + counters = on[2] + assert counters[0].tolist() == [0] * counters.shape[1], "layer 0 has no predecessor" + assert counters[1:, CALLS].tolist() == [len(inputs)] * (_PF_LAYERS - 1) + assert counters[1:, ROWS].tolist() == [len(inputs) * bs] * (_PF_LAYERS - 1) + assert int(counters[1:, ISSUED].sum()) > 0 + + +def _reference_counters(selects, ensures, budget_override, k): + """CPU counters from the spied predictor logits and the residency its select read, and each + target layer's routing and residency at its ensure.""" + from freetoken.moe.prefetch import NUM_COLS, default_budget + from tests.moe.ref_prefetch import ref_count, ref_select + + config = parse_config(toy_hf_config(_PF_LAYERS)) + counters = torch.zeros((_PF_LAYERS, NUM_COLS), dtype=torch.int64) + for layer in range(1, _PF_LAYERS): + assert len(selects[layer]) == len(ensures[layer]) > 0 + want_budget = budget_override or default_budget(config.is_linear_layer(layer)) + for (logits, budget, seen), (raw, resident) in zip(selects[layer], ensures[layer]): + assert budget == want_budget + sel, res = ref_select(logits.tolist(), seen.tolist(), k, budget) + routed = raw.view(-1).tolist() + misses = len({e for e in routed if resident[e] < 0}) + counters[layer] += torch.tensor(ref_count(sel, res, routed, misses, logits.shape[0])) + return counters + + +@requires_cuda +@pytest.mark.parametrize("kind", ["bf16", "nvfp4"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("budget", [0, 5]) +def test_moe_prefetch_counters_match_reference(kind, bs, budget, monkeypatch): + """Eager counters equal a CPU reference built from the predictor's own logits and the residency + its select read, and each layer's routing and pre-ensure residency; the captured graph counts + the same. The predictor runs on its dedicated stream, the count on the compute stream.""" + import freetoken.moe.prefetch as pf_mod + + banks = _prefetch_banks(kind) + inputs = _prefetch_inputs(bs, seed=1) + k = 6 + moes, cache = _prefetch_stack(kind, banks, "measure", monkeypatch, k=k, budget=budget) + selects = {layer: [] for layer in range(_PF_LAYERS)} + ensures = {layer: [] for layer in range(_PF_LAYERS)} + streams = {"select": set(), "count": set()} + real_select, real_count, real_ensure = pf_mod.lookahead_select, pf_mod.prefetch_count, cache.ensure_experts + + def spy_select(logits, resident, sel, res, *, k, budget): + # the select reads residency beside layer L-1's ensure: syncing on both sides pins what it saw + torch.cuda.synchronize() + layer = (sel.data_ptr() - cache.prefetch.sel.data_ptr()) // cache.prefetch.sel[0].nbytes + selects[layer].append((logits.float().cpu(), budget, resident.cpu())) + streams["select"].add(torch.cuda.current_stream().cuda_stream) + real_select(logits, resident, sel, res, k=k, budget=budget) + torch.cuda.synchronize() + + def spy_count(*args, **kwargs): + streams["count"].add(torch.cuda.current_stream().cuda_stream) + return real_count(*args, **kwargs) + + def spy_ensure(layer_id, expert_ids, **kwargs): + torch.cuda.synchronize() # the prediction forked after the previous ensure has landed + ensures[layer_id].append((expert_ids.cpu(), cache.slot_for_id[layer_id].cpu())) + return real_ensure(layer_id, expert_ids, **kwargs) + + def install(): + monkeypatch.setattr(pf_mod, "lookahead_select", spy_select) + monkeypatch.setattr(pf_mod, "prefetch_count", spy_count) + monkeypatch.setattr(cache, "ensure_experts", spy_ensure) + + eager = _prefetch_run(moes, cache, bs, False, inputs, before_decode=install) + want = _reference_counters(selects, ensures, budget, k) + assert torch.equal(eager[2], want), (eager[2], want) + assert int(want[:, 1].sum()) > 0, "the sequence must have useful predictions" + assert streams["select"] == {cache.prefetch.stream.cuda_stream} + assert streams["count"] == {torch.cuda.current_stream().cuda_stream} + + monkeypatch.undo() + moes, cache = _prefetch_stack(kind, banks, "measure", monkeypatch, k=k, budget=budget) + graph = _prefetch_run(moes, cache, bs, True, inputs) + assert torch.equal(graph[2], eager[2]) + for got, want_out in zip(graph[0], eager[0]): + for a, b in zip(got, want_out): + assert torch.equal(a, b) + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +def test_moe_prefetch_perfect_lookahead_counts_every_miss(graph, bs, monkeypatch): + """Known answer: when layer L-1's router input equals layer L's, the lookahead is layer L's own + router, so with the budget covering every candidate each demand miss is a useful prediction.""" + from freetoken.moe.prefetch import ISSUED, MISSES, RESIDENT_HITS, USEFUL + + banks = _prefetch_banks("nvfp4") + moes, cache = _prefetch_stack("nvfp4", banks, "measure", monkeypatch, k=4, budget=8) + _, _, counters, lru_misses = _prefetch_run(moes, cache, bs, graph, _prefetch_inputs(bs, chain=True, seed=2)) + assert int(counters[1:, MISSES].sum()) > 0 + assert counters[1:, USEFUL].tolist() == counters[1:, MISSES].tolist() == lru_misses[1:].tolist() + assert (counters[1:, ISSUED] >= counters[1:, USEFUL]).all() + assert int(counters[1:, RESIDENT_HITS].sum()) > 0, "routed hits among the top candidates" + + +@requires_cuda +def test_moe_prefetch_off_launches_nothing_new(monkeypatch): + """With the flag off the wired lookahead is inert: no prefetcher, and the decode launches + exactly the kernels of an unwired stack (measure adds some, so the check can see them).""" + from torch.profiler import ProfilerActivity, profile + + banks = _prefetch_banks("bf16") + inputs = _prefetch_inputs(1, steps=2) + + def kernels(mode, wire): + moes, cache = _prefetch_stack("bf16", banks, mode, monkeypatch, wire=wire) + _prefetch_run(moes, cache, 1, False, inputs[:1]) # warm up and compile outside the trace + with profile(activities=[ProfilerActivity.CUDA]) as prof: + _prefetch_run(moes, cache, 1, False, inputs) + events = [e for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA] + return cache, [e.name for e in sorted(events, key=lambda e: e.time_range.start)] + + cache, wired_off = kernels("off", True) + assert cache.prefetch is None + _, unwired_off = kernels("off", False) + _, measured = kernels("measure", True) + assert wired_off == unwired_off + assert len(measured) > len(wired_off) + + +@requires_cuda +def test_moe_prefetch_survives_a_rebuild(monkeypatch): + """A runtime rebuild resets the counters and keeps the predictor stream; a recaptured graph + counts like an eager run on the rebuilt cache.""" + banks = _prefetch_banks("bf16") + inputs = _prefetch_inputs(1, seed=3) + moes, cache = _prefetch_stack("bf16", banks, "measure", monkeypatch) + before = _prefetch_run(moes, cache, 1, True, inputs)[2] + stream = cache.prefetch.stream + assert int(before.sum()) > 0 + cache.rebuild(_PF_EXPERTS + 4) + assert int(cache.prefetch.counters.abs().sum()) == 0 and cache.prefetch.stream is stream + graph = _prefetch_run(moes, cache, 1, True, inputs) + eager = _prefetch_run(moes, cache, 1, False, inputs) + assert int(graph[2].sum()) > 0 and torch.equal(graph[2], eager[2]) + for a, b in zip(graph[1], eager[1]): + assert torch.equal(a, b) + + +@requires_cuda +@pytest.mark.parametrize("kind", ["bf16", "nvfp4"]) +@pytest.mark.parametrize("bs", [1, 2]) +def test_moe_prefetch_predicts_layer_l_from_layer_l_minus_1_input(kind, bs, monkeypatch): + """With a distinct input per layer, the prediction for layer L reads exactly layer L-1's router + input through layer L's router, and the select sees those logits and layer L's residency.""" + import freetoken.moe.prefetch as pf_mod + + banks = _prefetch_banks(kind) + inputs = _prefetch_inputs(bs, steps=3, seed=4) + moes, cache = _prefetch_stack(kind, banks, "measure", monkeypatch) + forks, selects = [], [] + real_fork, real_select = cache.prefetch.fork, pf_mod.lookahead_select + + def spy_fork(target, x, gate, resident, budget): + forks.append((target, x.clone(), gate, resident.data_ptr())) + return real_fork(target, x, gate, resident, budget) + + def spy_select(logits, resident, sel, res, *, k, budget): + selects.append(logits.clone()) + return real_select(logits, resident, sel, res, k=k, budget=budget) + + def install(): + monkeypatch.setattr(cache.prefetch, "fork", spy_fork) + monkeypatch.setattr(pf_mod, "lookahead_select", spy_select) + + _prefetch_run(moes, cache, bs, False, inputs, before_decode=install) + assert [f[0] for f in forks] == list(range(1, _PF_LAYERS)) * len(inputs) and len(selects) == len(forks) + for call, ((target, x, gate, resident), logits) in enumerate(zip(forks, selects)): + assert torch.equal(x, inputs[call // (_PF_LAYERS - 1)][target - 1]) + assert gate is moes[target].gate and resident == cache.slot_for_id[target].data_ptr() + assert torch.equal(logits, F.linear(x, moes[target].gate.weight)) + + +# FREETOKEN_MOE_PREFETCH=on: the same stack, with copies. No sleep between layers, so the prefetch +# streams race the compute stream at toy size and only the joins keep the GEMMs on landed bytes. +_PF_ON_STEPS = 200 + + +def _assert_slots_hold_their_experts(cache): + """Lossless: every held slot's bytes equal the host rows of the expert its id names.""" + held = (cache.id_of_slot >= 0).nonzero().view(-1).cpu() + ids = cache.id_of_slot[held.cuda()].long().cpu() + assert held.numel() > 0 + for (per_layer, slots), view in zip(cache.banks, cache.bank_views()): + host = torch.stack(per_layer).flatten(0, 1) + want = host[ids].contiguous().view(torch.uint8) + got = view[held.cuda()].cpu().contiguous().view(torch.uint8) + assert torch.equal(got, want) + + +def _routed_ids(moes, inputs): + """[step][layer] -> the experts each layer routes (the router alone decides; no cache input).""" + from freetoken.moe.fused import fused_topk + + out = [] + for xs in inputs: + row = [] + for moe, x in zip(moes, xs): + _, ids = fused_topk(hidden_states=x, gating_output=moe.gate.forward(x), topk=_PF_TOPK, renormalize=True) + row.append(ids.view(-1).tolist()) + out.append(row) + return out + + +def _garbage(routed, width, kind, seed): + """Per step a [layers, width] candidate override: 'random' ids (and -1), 'adversarial' ones + that mix layer L's routed ids with ids it does not route, duplicates and padding, or 'flood': + every expert of layer L, so a prefetch wants more slots than the unpinned ones.""" + gen = torch.Generator().manual_seed(seed) + steps = [] + for row in routed: + sel = torch.full((_PF_LAYERS, width), -1, dtype=torch.int32) + for layer in range(1, _PF_LAYERS): + if kind == "random": + ids = torch.randint(-1, _PF_EXPERTS, (width,), generator=gen) + elif kind == "flood": + ids = torch.randperm(_PF_EXPERTS, generator=gen)[:width] + else: + hit = list(dict.fromkeys(row[layer])) + miss = [e for e in torch.randperm(_PF_EXPERTS, generator=gen).tolist() if e not in hit][:3] + ids = torch.tensor((hit[:3] + miss + hit[:1] + [-1] + miss[:1]) * 2)[:width] + ids = ids[torch.randperm(ids.numel(), generator=gen)] + sel[layer, : ids.numel()] = ids.to(torch.int32) + steps.append(sel) + return steps + + +def _on_stack(kind, banks, monkeypatch, *, policy, overlap, k=6, budget=0): + return _prefetch_stack(kind, banks, "on", monkeypatch, k=k, budget=budget, overlap=overlap, policy=policy) + + +def _with_override(cache, steps, **delays): + """before_step hook that loads step i's candidate override into the prefetcher's static buffer.""" + pf = cache.prefetch + pf.sel_override = torch.full_like(pf.sel, -1) + for name, cycles in delays.items(): + setattr(pf, name, cycles) + + def before_step(i): + pf.sel_override.copy_(steps[i][:, : pf.sel.shape[1]]) + + return before_step + + +def _assert_same_outputs(got, want): + for a_step, b_step in zip(got, want): + for a, b in zip(a_step, b_step): + assert torch.isfinite(b.float()).all() + assert torch.equal(a, b) + + +@requires_cuda +@pytest.mark.parametrize("kind", ["bf16", "nvfp4"]) +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("policy", ["rule", "lru"]) +@pytest.mark.parametrize("overlap", [False, True], ids=["serial", "overlap"]) +def test_moe_prefetch_on_is_bitwise_identical(kind, graph, bs, policy, overlap, monkeypatch): + """FREETOKEN_MOE_PREFETCH=on copies predicted experts into low-priority slots beside the decode: + over 200 steps of a cache that evicts every step the MoE outputs are those of prefetch off to + the bit, every held slot holds its expert's bytes, and the prefetches were used.""" + from freetoken.moe.prefetch import CALLS, COPIED, ISSUED, USEFUL + + banks = _prefetch_banks(kind) + inputs = _prefetch_inputs(bs, steps=_PF_ON_STEPS, seed=5) + moes, cache = _prefetch_stack(kind, banks, "off", monkeypatch, overlap=overlap, policy=policy) + off = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0) + _assert_slots_hold_their_experts(cache) + moes, cache = _on_stack(kind, banks, monkeypatch, policy=policy, overlap=overlap) + on = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0) + _assert_same_outputs(on[0], off[0]) + _assert_slots_hold_their_experts(cache) + stats = cache.prefetch.stats.cpu() + assert stats[0].tolist() == [0] * stats.shape[1], "layer 0 is never prefetched" + assert stats[1:, CALLS].tolist() == [_PF_ON_STEPS] * (_PF_LAYERS - 1) + assert int(stats[:, ISSUED].sum()) > 0 and int(stats[:, USEFUL].sum()) > 0 and int(stats[:, COPIED].sum()) > 0 + assert int(on[3].sum()) < int(off[3].sum()), "useful prefetches remove demand misses" + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("policy", ["rule", "lru"]) +@pytest.mark.parametrize("garbage", ["random", "adversarial"]) +def test_moe_prefetch_on_survives_a_garbage_predictor(graph, bs, policy, garbage, monkeypatch): + """Whatever the candidates (random ids and padding, or the next layer's routed ids mixed with ids + it will not route, duplicates and -1), prefetch on changes no output bit and no slot's bytes.""" + from freetoken.moe.prefetch import ISSUED, USEFUL + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(bs, steps=_PF_ON_STEPS, seed=6) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, policy=policy) + off = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0) + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True, budget=8) + steps = _garbage(_routed_ids(moes, inputs), 12, garbage, seed=bs) + on = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0, before_step=_with_override(cache, steps)) + _assert_same_outputs(on[0], off[0]) + _assert_slots_hold_their_experts(cache) + stats = cache.prefetch.stats.cpu() + assert int(stats[:, ISSUED].sum()) > 0 and int(stats[:, USEFUL].sum()) > 0 + assert int(stats[:, USEFUL].sum()) < int(stats[:, ISSUED].sum()), "some prefetches must be wrong" + + +# ~0.5 ms at the RTX 5080's clock: far longer than a toy layer, so every delayed stage is late +_PF_DELAY_CYCLES = 1_000_000 + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("policy", ["rule", "lru"]) +@pytest.mark.parametrize("delay", ["copy", "predict", "both"]) +def test_moe_prefetch_on_joins_hold_when_the_prefetch_streams_run_late(graph, bs, policy, delay, monkeypatch): + """Race test: a sleep on the prefetch copy stream (every copy lands after its layer's ensure) or on + the predictor stream (prefetch_ensure runs long after the fork) must not change a bit. The + candidates include each layer's routed ids, so GEMMs do read prefetched slots. The joins also + make the slot maps timing-free: the delayed run ends with the undelayed run's maps and counts.""" + from freetoken.moe.prefetch import COPIED, LATE, NUM_STAT_COLS, USEFUL + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(bs, steps=_PF_ON_STEPS, seed=7) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, policy=policy) + off = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0) + runs = [] + for cycles in (0, _PF_DELAY_CYCLES): + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True, budget=8) + steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=bs + 10) + delays = {"delay_copy_cycles": cycles * (delay in ("copy", "both")), + "delay_predict_cycles": cycles * (delay in ("predict", "both"))} + on = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0, before_step=_with_override(cache, steps, **delays)) + _assert_same_outputs(on[0], off[0]) + _assert_slots_hold_their_experts(cache) + runs.append((on[1][:3], cache.prefetch.stats.cpu())) + (maps, stats), (delayed_maps, delayed_stats) = runs + assert int(stats[:, USEFUL].sum()) > 0 + for a, b in zip(maps, delayed_maps): + assert torch.equal(a, b) + cols = [c for c in range(NUM_STAT_COLS) if c != LATE] + assert torch.equal(stats[:, cols], delayed_stats[:, cols]) + if delay == "copy": + # the late counter sees what the join waited for: the copies were still sleeping at ensure(L) + # (all of them under replay; in eager a loaded host can take longer than the sleep per layer) + late, copied = int(delayed_stats[:, LATE].sum()), int(delayed_stats[:, COPIED].sum()) + assert copied > 0 and (late == copied if graph else late > 0) + + +@requires_cuda +def test_moe_prefetch_race_test_catches_a_missing_copy_join(monkeypatch): + """Negative control for the race test: without the compute stream's wait on the prefetch copy, + GEMMs read prefetched slots before their bytes land and the outputs change.""" + from freetoken.moe.prefetch import ExpertPrefetcher + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(1, steps=40, seed=7) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True) + off = _prefetch_run(moes, cache, 1, False, inputs, sleep=0) + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy="rule", overlap=True, budget=8) + + def no_wait(self, layer_id): + self._inflight[layer_id] = None + + monkeypatch.setattr(ExpertPrefetcher, "join_copy", no_wait) + steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=11) + on = _prefetch_run(moes, cache, 1, False, inputs, sleep=0, + before_step=_with_override(cache, steps, delay_copy_cycles=_PF_DELAY_CYCLES)) + assert any(not torch.equal(a, b) for got, want in zip(on[0], off[0]) for a, b in zip(got, want)) + + +def _replay_on_reference(policy, calls, rows): + """Replay the spied prefetch_ensure / ensure sequence on the CPU reference; per-layer stats.""" + from freetoken.moe.prefetch import CALLS, COPIED, ISSUED, MISSES, ROWS, USEFUL + from freetoken.moe.scored_ensure import POLICY_IDS + from tests.moe.ref_scored_cache import RefScoredCache + + ref = RefScoredCache(_PF_LAYERS, _PF_EXPERTS, _PF_EXPERTS, POLICY_IDS[policy]) + stats = torch.zeros((_PF_LAYERS, 8), dtype=torch.int64) + planned = {} + for call in calls: + if call[0] == "prefetch": + _, layer, sel, budget = call + src, _ = ref.prefetch(layer, sel.tolist(), budget) + stats[layer, ISSUED] += src.size + planned[layer] = src.size + else: + _, layer, ids, logits = call + _, src, _ = ref.ensure(layer, ids.view(-1).numpy(), bump_tok=layer == 0, lowpri=True, + logits=None if logits is None else logits.float().numpy(), thr=0.25) + if layer in planned: + stats[layer, [USEFUL, CALLS, ROWS, MISSES, COPIED]] += torch.tensor( + [ref.useful, 1, rows, src.size, int(planned.pop(layer) > 0)]) + return ref, stats + + +@requires_cuda +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("policy", ["rule", "lru"]) +def test_moe_prefetch_on_counters_match_reference(bs, policy, monkeypatch): + """issued / useful / calls / rows / misses / copied equal a CPU replay of the spied candidates + and routings, the slot maps end identical, late is 0 when each ensure waits for the copies + (known answer), and with the candidates fixed a captured graph counts the same as eager. + prefetch_ensure and the demand copy run on the compute stream, the slim prefetch copy on the + copy stream.""" + import freetoken.kernel.fast_index_copy as fic + from freetoken.moe.prefetch import LATE, NUM_STAT_COLS, RESIDENT_HITS + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(bs, steps=60, seed=8) + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True) + calls, streams = [], {"pf_ensure": set(), "slim": set(), "demand": set()} + real_pf, real_ensure = cache.prefetch_ensure, cache.ensure_experts + real_slim, real_multi = fic.fast_index_copy_multi_slim_jit, fic.fast_index_copy_multi_jit + + def spy_pf(layer, query, *args, **kw): + streams["pf_ensure"].add(torch.cuda.current_stream().cuda_stream) + calls.append(("prefetch", layer, query.clone(), kw.get("budget"))) # on the compute stream after its wait on the select + return real_pf(layer, query, *args, **kw) + + def spy_ensure(layer, ids, **kw): + torch.cuda.synchronize() # the copies forked before this ensure have landed: late must be 0 + logits = kw.get("router_logits") + calls.append(("ensure", layer, ids.clone(), None if logits is None or policy != "rule" else logits.clone())) + return real_ensure(layer, ids, **kw) + + def spy_slim(*args): + streams["slim"].add(torch.cuda.current_stream().cuda_stream) + return real_slim(*args) + + def spy_multi(*args): + streams["demand"].add(torch.cuda.current_stream().cuda_stream) + return real_multi(*args) + + def install(): + monkeypatch.setattr(cache, "prefetch_ensure", spy_pf) + monkeypatch.setattr(cache, "ensure_experts", spy_ensure) + monkeypatch.setattr(fic, "fast_index_copy_multi_slim_jit", spy_slim) + monkeypatch.setattr(fic, "fast_index_copy_multi_jit", spy_multi) + + eager = _prefetch_run(moes, cache, bs, False, inputs, before_decode=install, sleep=0) + torch.cuda.synchronize() + calls = [c if c[0] == "prefetch" else (c[0], c[1], c[2].cpu(), None if c[3] is None else c[3].cpu()) for c in calls] + calls = [(c[0], c[1], c[2].cpu(), c[3]) if c[0] == "prefetch" else c for c in calls] + ref, want = _replay_on_reference(policy, calls, bs) + got = cache.prefetch.stats.cpu() + assert got.shape == (_PF_LAYERS, NUM_STAT_COLS) + assert int(got[:, LATE].sum()) == 0 and int(got[:, RESIDENT_HITS].sum()) == 0 + assert torch.equal(got, want), (got, want) + assert int(want[:, 1].sum()) > 0, "the sequence must have useful prefetches" + assert torch.equal(cache.slot_for_id.view(-1).cpu().long(), torch.from_numpy(ref.slot_of_id)) + assert torch.equal(cache.id_of_slot.cpu().long(), torch.from_numpy(ref.id_of_slot)) + assert torch.equal(cache.usage.cpu(), torch.from_numpy(ref.usage)) + assert streams["pf_ensure"] == {torch.cuda.current_stream().cuda_stream} + assert streams["slim"] == {cache.prefetch.copy_stream.cuda_stream} + assert streams["demand"] == {torch.cuda.current_stream().cuda_stream} + + # the select reads residency beside layer L-1's ensure, so two runs may pick different candidates: + # with the candidates fixed, a captured graph must install, count and end like eager + monkeypatch.undo() + runs = [] + for graph in (False, True): + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True) + steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=bs + 8) + run = _prefetch_run(moes, cache, bs, graph, inputs, sleep=_PF_SLEEP_CYCLES, before_step=_with_override(cache, steps)) + _assert_same_outputs(run[0], eager[0]) + runs.append((run[1][:3], cache.prefetch.stats.cpu())) + cols = [c for c in range(NUM_STAT_COLS) if c != LATE] + (maps, stats), (graph_maps, graph_stats) = runs + for a, b in zip(maps, graph_maps): + assert torch.equal(a, b) + assert torch.equal(graph_stats[:, cols], stats[:, cols]) and int(stats[:, 1].sum()) > 0 + + +@requires_cuda +def test_moe_prefetch_on_rebuild_and_reset_clear_the_prefetch_state(monkeypatch): + """A rebuild keeps both prefetch streams, the events and plan buffers, and drops the counters, + the in-flight state and every low-priority slot; a recaptured graph then matches eager and off. + A reset (as after a capture) also drops the counters and low-priority slots.""" + banks = _prefetch_banks("bf16") + inputs = _prefetch_inputs(1, steps=40, seed=9) + moes, cache = _on_stack("bf16", banks, monkeypatch, policy="rule", overlap=True) + _prefetch_run(moes, cache, 1, True, inputs, sleep=0) + pf = cache.prefetch + streams, plan = (pf.stream, pf.copy_stream), (pf.pf_slots, pf.pf_src, pf.pf_num, pf.pf_ready) + assert int(pf.stats.sum()) > 0 and bool(((cache.id_of_slot >= 0) & (cache.usage == 0)).any()) + pf._inflight[2] = torch.zeros(1) + cache.rebuild(_PF_EXPERTS + 4) + assert (pf.stream, pf.copy_stream) == streams and (pf.pf_slots, pf.pf_src, pf.pf_num, pf.pf_ready) == plan + assert int(pf.stats.abs().sum()) == 0 and int(pf.totals.abs().sum()) == 0 and pf._inflight == [None] * _PF_LAYERS + assert not bool(((cache.id_of_slot >= 0) & (cache.usage == 0)).any()) + graph = _prefetch_run(moes, cache, 1, True, inputs, sleep=0) + eager = _prefetch_run(moes, cache, 1, False, inputs, sleep=0) + _assert_same_outputs(graph[0], eager[0]) + _assert_slots_hold_their_experts(cache) + assert int(pf.stats.sum()) > 0 + cache.reset() + assert int(pf.stats.abs().sum()) == 0 and not bool(((cache.id_of_slot >= 0) & (cache.usage == 0)).any()) + moes_off, cache_off = _prefetch_stack("bf16", banks, "off", monkeypatch, overlap=True) + cache_off.rebuild(_PF_EXPERTS + 4) + off = _prefetch_run(moes_off, cache_off, 1, False, inputs, sleep=0) + _assert_same_outputs(eager[0], off[0]) + + +@requires_cuda +def test_moe_prefetch_on_skips_batches_past_the_select_tile(monkeypatch): + """An eager decode wider than MAX_ROWS forks nothing: no join, no count, the same outputs.""" + from freetoken.moe.prefetch import MAX_ROWS + + banks = _prefetch_banks("bf16") + bs = MAX_ROWS + 1 + inputs = _prefetch_inputs(bs, steps=6, seed=10) + off = _prefetch_run(*_prefetch_stack("bf16", banks, "off", monkeypatch), bs, False, inputs, sleep=0) + moes, cache = _on_stack("bf16", banks, monkeypatch, policy="rule", overlap=True) + on = _prefetch_run(moes, cache, bs, False, inputs, sleep=0) + _assert_same_outputs(on[0], off[0]) + assert int(cache.prefetch.stats.abs().sum()) == 0 and cache.prefetch._inflight == [None] * _PF_LAYERS + + +def _stall_hooks(monkeypatch): + """Sleeps (cycles, read at enqueue time) on the compute stream right before each decode GEMM, i.e. + after issue_copy(L+1), and on the stream that runs the shared expert gate (the side stream).""" + import freetoken.models.qwen4_exp.moe as q4moe + from freetoken.layers.moe import OffloadMoELayer + + stall = {"gemm": 0, "shared": 0} + real_gemm, real_gate = OffloadMoELayer._expert_gemm, q4moe.shared_gate_sigmoid + + def gemm(self, cache, *args, is_prefill, **kw): + if stall["gemm"] and not is_prefill: + torch.cuda._sleep(stall["gemm"]) + return real_gemm(self, cache, *args, is_prefill=is_prefill, **kw) + + def gate(*args, **kw): + if stall["shared"]: + torch.cuda._sleep(stall["shared"]) + return real_gate(*args, **kw) + + monkeypatch.setattr(OffloadMoELayer, "_expert_gemm", gemm) + monkeypatch.setattr(q4moe, "shared_gate_sigmoid", gate) + return stall + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("policy", ["rule", "lru"]) +@pytest.mark.parametrize("where", ["gemm", "shared"]) +@pytest.mark.parametrize("cands", ["adversarial", "flood"]) +def test_moe_prefetch_on_holds_when_the_compute_side_stalls(graph, bs, policy, where, cands, monkeypatch): + """Race test from the other side: a sleep on the compute stream between issue_copy(L+1) and + GEMM(L) lets copy(L+1) land before GEMM(L) reads, so a prefetch into any of layer L's routed + slots changes outputs deterministically; a sleep on the shared expert's side stream must not + matter either. 'flood' asks for all 16 experts a layer, more than the unpinned slots, so even + LRU's recency order would reach the pinned ones. Bit-identical to off, and the same final maps + as the unstalled run.""" + k, width = (8, 16) if cands == "flood" else (6, 12) + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(bs, steps=80, seed=21) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, policy=policy) + off = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0) + stall = _stall_hooks(monkeypatch) + runs = [] + for cycles in (0, _PF_DELAY_CYCLES): + stall[where] = cycles + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True, k=k, budget=width) + steps = _garbage(_routed_ids(moes, inputs), width, cands, seed=bs + 20) + on = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0, before_step=_with_override(cache, steps)) + _assert_same_outputs(on[0], off[0]) + _assert_slots_hold_their_experts(cache) + runs.append(on[1][:3]) + for a, b in zip(*runs): + assert torch.equal(a, b) + + +def _assert_held_slots_hold_their_experts(cache): + """_assert_slots_hold_their_experts, but a prefill may leave no slot held.""" + if bool((cache.id_of_slot >= 0).any()): + _assert_slots_hold_their_experts(cache) + + +def _mixed_run(moes, cache, bs, graph, schedule, before_step=None, late=None, settle=True): + """Warm up (and capture) the decode like _prefetch_run, reset, then run ``schedule``: ("decode", xs) + is one decode step (a replay under ``graph``), ("prefill", xs) an eager prefill of xs's tokens. + ``late=(cycles, steps)`` sleeps that long on the predictor stream in those decode steps (a second + graph captured with the sleep, in the same pool, replays them); ``settle`` syncs and checks the + slots after each prefill, else the next decode is enqueued right behind it.""" + ctx = _fresh_ctx(_batch=SimpleNamespace(is_prefill=False)) + + def step(xs): + return [moe.forward(x.clone()) for moe, x in zip(moes, xs)] + + static = [torch.randn(bs, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16) for _ in moes] + step(static) + if graph: + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + graphs = {False: (g, step(static))} + if late is not None: + cache.prefetch.delay_predict_cycles = late[0] + g_late = torch.cuda.CUDAGraph() + with torch.cuda.graph(g_late, pool=g.pool()): + graphs[True] = (g_late, step(static)) + cache.prefetch.delay_predict_cycles = 0 + cache.reset() + results = [] + for i, (phase, xs) in enumerate(schedule): + if before_step is not None: + before_step(i) + ctx._batch.is_prefill = phase == "prefill" + slow = late is not None and i in late[1] + if phase == "decode" and graph: + for s, x in zip(static, xs): + s.copy_(x) + g, outs = graphs[slow] + g.replay() + else: + if late is not None: + cache.prefetch.delay_predict_cycles = late[0] if slow else 0 + outs = step(xs) + if phase == "decode" and cache.audit is not None: + cache.audit.end_decode_step(cache) + results.append([o.clone() for o in outs]) + if phase == "prefill" and settle: + torch.cuda.synchronize() + _assert_held_slots_hold_their_experts(cache) + torch.cuda.synchronize() + _assert_held_slots_hold_their_experts(cache) + return results + + +_PREFILL_PATHS = { + # (cache_size, prefill tokens, cache kwargs, FREETOKEN_MOE_SMALL_PREFILL_TOKENS) + "materialize": (_PF_EXPERTS + 4, 24, {}, 0), + "overlap": (2 * _PF_EXPERTS + 4, 24, {"prefill_overlap": True}, 0), + "overlap_d2d": (2 * _PF_EXPERTS + 4, 24, {"prefill_overlap": True, "prefill_hit_d2d": True}, 0), + "small": (_PF_EXPERTS, 3, {}, 8), +} + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("policy", ["rule", "lru"]) +@pytest.mark.parametrize("path", list(_PREFILL_PATHS)) +def test_moe_prefetch_on_survives_prefills_between_decodes(graph, bs, policy, path, monkeypatch): + """Prefills between decode steps rewrite slots behind the prefetch's back: the whole-layer + materialize installs slots [0, E), the overlap double buffers invalidate and stream [0, 2E) + (with the hit-D2D gather reading resident slots), and small prefills ensure their experts + under the low-priority keys. Decode, prefill, decode, prefill, decode: every output of prefetch + on equals off to the bit, and after each prefill every held slot holds its expert.""" + import freetoken.layers.moe as layers_moe + + cache_size, tokens, cache_kw, small = _PREFILL_PATHS[path] + monkeypatch.setattr(layers_moe, "_SMALL_PREFILL_TOKENS", small) + banks = _prefetch_banks("nvfp4") + decode = _prefetch_inputs(bs, steps=18, seed=31) + gen = torch.Generator(device="cuda").manual_seed(32) + prefills = [[torch.randn(tokens, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16, generator=gen) * 0.5 + for _ in range(_PF_LAYERS)] for _ in range(2)] + schedule = ([("decode", xs) for xs in decode[:6]] + [("prefill", prefills[0])] + + [("decode", xs) for xs in decode[6:12]] + [("prefill", prefills[1])] + + [("decode", xs) for xs in decode[12:]]) + kw = dict(overlap=True, policy=policy, cache_size=cache_size, **cache_kw) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, **kw) + off = _mixed_run(moes, cache, bs, graph, schedule) + moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, budget=8, **kw) + steps = _garbage(_routed_ids(moes, [xs for _, xs in schedule]), 12, "adversarial", seed=bs + 30) + on = _mixed_run(moes, cache, bs, graph, schedule, before_step=_with_override(cache, steps)) + _assert_same_outputs(on, off) + if path == "overlap_d2d": + assert cache.prefill_hit_rows > 0, "the hit-D2D gather must have read resident slots" + assert int(cache.prefetch.stats[:, 0].sum()) > 0 + + +def _record_compute_streams(monkeypatch, moes): + """The streams the MoE blocks are entered on (the default stream, and the capture stream under a + graph): the compute streams of the decode.""" + compute = set() + for moe in moes: + real = moe.forward + + def forward(x, real=real): + compute.add(torch.cuda.current_stream().cuda_stream) + return real(x) + + monkeypatch.setattr(moe, "forward", forward) + return compute + + +def _spy_map_writers(monkeypatch): + """(writer, stream) of every launch that writes the slot maps or their mirrors, at enqueue time.""" + import freetoken.kernel.triton.moe as triton_moe + import freetoken.moe.offload_kernels as ok + import freetoken.moe.scored_ensure as se + + seen = [] + + def spy(name, real): + def launch(*args, **kwargs): + seen.append((name(kwargs) if callable(name) else name, torch.cuda.current_stream().cuda_stream)) + return real(*args, **kwargs) + + return launch + + # scored_ensure is the kernel of both the demand ensures and prefetch_ensure (prefetch=True) + monkeypatch.setattr(se, "scored_ensure", spy(lambda kw: "prefetch_ensure" if kw.get("prefetch") else "ensure", + se.scored_ensure)) + monkeypatch.setattr(ok, "lru_ensure", spy("ensure", ok.lru_ensure)) + monkeypatch.setattr(ok, "materialize_layer", spy("materialize", ok.materialize_layer)) + monkeypatch.setattr(ok, "reset_cache", spy("reset", ok.reset_cache)) + monkeypatch.setattr(triton_moe, "invalidate_prefill_slots", spy("invalidate", triton_moe.invalidate_prefill_slots)) + return seen + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("path", list(_PREFILL_PATHS)) +def test_moe_prefetch_on_writes_the_slot_maps_only_on_the_compute_stream(graph, bs, path, monkeypatch): + """Single writer: every launch that writes the slot maps (the demand and prefetch ensures, prefill + ensures, the whole-layer materialize, reset) runs on the stream the MoE blocks run on (the capture + stream under a graph), never on the predictor, prefetch copy or shared-expert side stream. The only + other writer is the overlap prefill's buffer invalidation, on the prefill copy stream it fences. + The predictor's select reads no cache state.""" + import freetoken.layers.moe as layers_moe + import freetoken.moe.prefetch as pf_mod + from freetoken.moe.prefetch import ISSUED + + cache_size, tokens, cache_kw, small = _PREFILL_PATHS[path] + monkeypatch.setattr(layers_moe, "_SMALL_PREFILL_TOKENS", small) + banks = _prefetch_banks("nvfp4") + decode = _prefetch_inputs(bs, steps=6, seed=95) + gen = torch.Generator(device="cuda").manual_seed(96) + prefill = [torch.randn(tokens, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16, generator=gen) * 0.5 + for _ in range(_PF_LAYERS)] + schedule = [("decode", xs) for xs in decode[:3]] + [("prefill", prefill)] + [("decode", xs) for xs in decode[3:]] + moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, budget=8, overlap=True, cache_size=cache_size, + **cache_kw) + steps = _garbage(_routed_ids(moes, [xs for _, xs in schedule]), 12, "adversarial", seed=bs + 97) + seen = _spy_map_writers(monkeypatch) + compute = _record_compute_streams(monkeypatch, moes) + selects, real_select = [], pf_mod.lookahead_select + + def spy_select(logits, resident, *args, **kwargs): + selects.append((resident.data_ptr(), torch.cuda.current_stream().cuda_stream)) + return real_select(logits, resident, *args, **kwargs) + + monkeypatch.setattr(pf_mod, "lookahead_select", spy_select) + _mixed_run(moes, cache, bs, graph, schedule, before_step=_with_override(cache, steps)) + pf = cache.prefetch + assert not compute & {pf.stream.cuda_stream, pf.copy_stream.cuda_stream, cache.decode_copy_stream.cuda_stream} + assert len(compute) == 1 + graph + writers = {name for name, _ in seen} + want = {"ensure", "prefetch_ensure", "reset"} | {"materialize": {"materialize"}, "overlap": {"invalidate"}, + "overlap_d2d": {"invalidate"}, "small": set()}[path] + assert writers == want, writers + off_compute = [(name, stream) for name, stream in seen if stream not in compute] + fence = {cache.prefill_copy_stream.cuda_stream} if cache.prefill_copy_stream is not None else set() + assert all(name == "invalidate" and stream in fence for name, stream in off_compute), off_compute + assert {stream for name, stream in seen if name == "prefetch_ensure"} == compute # eager, and the capture + assert int(pf.stats[:, ISSUED].sum()) > 0 + # after the writer checks, which are what a predictor-side install fails + assert selects and set(selects) == {(pf._none_resident.data_ptr(), pf.stream.cuda_stream)} + + +# ~4 ms at the RTX 5080's clock: a select far later than its whole layer +_PF_LATE_CYCLES = 8 * _PF_DELAY_CYCLES + + +def _step_starts(times): + """before_step hook recording a timed event on the compute stream at the start of each step.""" + def before_step(i): + times.append(torch.cuda.Event(enable_timing=True)) + times[-1].record() + + return before_step + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("policy", ["rule", "lru"]) +def test_moe_prefetch_on_waits_for_a_late_select(graph, bs, policy, monkeypatch): + """The compute stream waits for the select before it installs, however late the predictor runs: + with every select ~4 ms late (far past its whole layer) a step lasts at least its three late + selects, and the outputs equal off's and the maps and counts (all but LATE) those of the run + without the delay.""" + from freetoken.moe.prefetch import LATE, NUM_STAT_COLS, USEFUL + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(bs, steps=24, seed=100) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, policy=policy) + off = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0) + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + torch.cuda._sleep(_PF_LATE_CYCLES) + end.record() + end.synchronize() + late_ms = start.elapsed_time(end) + runs = [] + for cycles in (0, _PF_LATE_CYCLES): + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True, budget=8) + steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=bs + 101) + load, times = _with_override(cache, steps, delay_predict_cycles=cycles), [] + mark = _step_starts(times) + + def before_step(i, load=load, mark=mark): + mark(i) + load(i) + + on = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0, before_step=before_step) + _assert_same_outputs(on[0], off[0]) + _assert_slots_hold_their_experts(cache) + step_ms = sorted(a.elapsed_time(b) for a, b in zip(times, times[1:])) + runs.append((on[1][:3], cache.prefetch.stats.cpu(), step_ms[len(step_ms) // 2])) + (maps, stats, fast_ms), (late_maps, late_stats, slow_ms) = runs + assert int(stats[:, USEFUL].sum()) > 0 + for a, b in zip(maps, late_maps): + assert torch.equal(a, b) + cols = [c for c in range(NUM_STAT_COLS) if c != LATE] + assert torch.equal(stats[:, cols], late_stats[:, cols]) + # layers 1-3 are predicted; each prediction forks after the previous one was waited for + assert slow_ms >= 0.9 * (_PF_LAYERS - 1) * late_ms > 4 * fast_ms, (slow_ms, late_ms, fast_ms) + + +@requires_cuda +def test_moe_prefetch_on_late_select_without_the_wait_changes_the_maps_but_no_output(monkeypatch): + """Negative control for the test above (eager: a capture needs the predictor joined): without the + compute stream's wait, prefetch_ensure installs stale candidates, so the maps differ from the + undelayed run's, yet every output still equals off's: the slot maps have one writer, so a wrong + prediction only costs a copy.""" + from freetoken.moe.prefetch import ExpertPrefetcher + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(1, steps=24, seed=100) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True) + off = _prefetch_run(moes, cache, 1, False, inputs, sleep=0) + + def no_wait(self, target): + """install without the compute stream's wait on the select.""" + if target is None or self._inflight[target] is None: + return None + self.cache.prefetch_ensure( + target, self.sel[target], self.pf_slots[target], self.pf_src[target], + self.pf_num[target], self.stats[target], self.pf_ready[target], budget=self._budget[target], + ) + self._installed[target] = True + return target + + runs = [] + for skip in (False, True): + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy="rule", overlap=True, budget=8) + if skip: + monkeypatch.setattr(ExpertPrefetcher, "install", no_wait) + steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=102) + on = _prefetch_run(moes, cache, 1, False, inputs, sleep=0, + before_step=_with_override(cache, steps, delay_predict_cycles=_PF_LATE_CYCLES * skip)) + _assert_same_outputs(on[0], off[0]) + _assert_slots_hold_their_experts(cache) + runs.append(on[1][:3]) + assert any(not torch.equal(a, b) for a, b in zip(*runs)) + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("path", list(_PREFILL_PATHS)) +def test_moe_prefetch_on_late_select_right_after_a_prefill(graph, bs, path, monkeypatch): + """The suspected production trigger: a decode step enqueued right behind a prefill (which + rewrote slots behind the prefetch's back, the overlap buffers' [0, 2E) included) with the + predictor ~4 ms late. Every output equals off's to the bit, the maps equal those of the same run + with no late step, every held slot holds its expert and the decode verifier finds nothing.""" + import freetoken.layers.moe as layers_moe + from freetoken.moe.verify import CHECKS + + cache_size, tokens, cache_kw, small = _PREFILL_PATHS[path] + monkeypatch.setattr(layers_moe, "_SMALL_PREFILL_TOKENS", small) + banks = _prefetch_banks("nvfp4") + decode = _prefetch_inputs(bs, steps=15, seed=105) + gen = torch.Generator(device="cuda").manual_seed(106) + prefills = [[torch.randn(tokens, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16, generator=gen) * 0.5 + for _ in range(_PF_LAYERS)] for _ in range(2)] + schedule = ([("decode", xs) for xs in decode[:5]] + [("prefill", prefills[0])] + + [("decode", xs) for xs in decode[5:10]] + [("prefill", prefills[1])] + + [("decode", xs) for xs in decode[10:]]) + after_prefill = {i + 1 for i, (phase, _) in enumerate(schedule) if phase == "prefill"} + kw = dict(overlap=True, cache_size=cache_size, **cache_kw) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, **kw) + off = _mixed_run(moes, cache, bs, graph, schedule) + runs = [] + for late in (None, (_PF_LATE_CYCLES, after_prefill)): + moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, budget=8, verify_mode="meta", **kw) + steps = _garbage(_routed_ids(moes, [xs for _, xs in schedule]), 12, "adversarial", seed=bs + 107) + on = _mixed_run(moes, cache, bs, graph, schedule, before_step=_with_override(cache, steps), late=late, + settle=False) + _assert_same_outputs(on, off) + counts, records = _verify_state(cache) + assert counts[CHECKS] > 0 and counts[CHECKS + 1 :] == [0] * (len(counts) - 1) and records == [], records + runs.append([t.clone() for t in (cache.slot_for_id, cache.id_of_slot, cache.usage)]) + for a, b in zip(*runs): + assert torch.equal(a, b) + if path == "overlap_d2d": + assert cache.prefill_hit_rows > 0, "the hit-D2D gather must have read resident slots" + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("policy", ["rule", "lru"]) +def test_moe_prefetch_on_install_may_evict_a_slot_whose_prefetch_copy_is_in_flight(graph, bs, policy, monkeypatch): + """prefetch_ensure(L) runs on the compute stream before it joins copy(L-1), so it may evict a slot + that plan L-1 installed and whose slim copy is still writing it. The next writer of that slot's + bytes (copy(L) on the same copy stream, or a demand copy after join_copy) is ordered after copy(L-1): + with every copy late and flood candidates the eviction happens (counted on the device, so it holds + under replay too), and still no output bit and no slot's bytes change.""" + from freetoken.moe.prefetch import ExpertPrefetcher + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(bs, steps=40, seed=110) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, policy=policy) + off = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0) + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True, budget=8) + pf = cache.prefetch + evicted_in_flight = torch.zeros((), dtype=torch.int64, device="cuda") + col = torch.arange(pf.pf_slots.shape[1], device="cuda") + real_install = ExpertPrefetcher.install + + def install(self, target): + got = real_install(self, target) + if got is not None and got >= 2: + # on the compute stream right after prefetch_ensure(got), before join_copy(got - 1) + new = torch.where(col < self.pf_num[got], self.pf_slots[got], -2) + old = torch.where(col < self.pf_num[got - 1], self.pf_slots[got - 1], -3) + shared = (new[:, None] == old[None, :]).any() + evicted_in_flight.add_((shared & (self.pf_ready[got - 1][0] == 0)).long()) + return got + + monkeypatch.setattr(ExpertPrefetcher, "install", install) + steps = _garbage(_routed_ids(moes, inputs), 12, "flood", seed=bs + 111) + on = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0, + before_step=_with_override(cache, steps, delay_copy_cycles=_PF_DELAY_CYCLES)) + _assert_same_outputs(on[0], off[0]) + _assert_slots_hold_their_experts(cache) + assert int(evicted_in_flight) > 0, "no install evicted a slot of the previous layer's in-flight copy" + + +@requires_cuda +def test_moe_prefetch_in_flight_eviction_test_catches_an_unordered_next_copy(monkeypatch): + """Negative control for the test above: order each prefetch copy only after its own install (not + after copy(L-1)'s join) and put odd and even layers' copies on two streams, only the even ones + late. A late copy then lands on a slot the next install took over, and the outputs or the held + slots' bytes go wrong.""" + from freetoken.moe.prefetch import ExpertPrefetcher, _mark_ready_kernel, dedicated_stream + from freetoken.moe.slot_audit import PREFETCH_COPY + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(1, steps=40, seed=110) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True) + off = _prefetch_run(moes, cache, 1, False, inputs, sleep=0) + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy="rule", overlap=True, budget=8) + pf = cache.prefetch + streams = [pf.copy_stream, dedicated_stream(torch.device("cuda"), highest_priority=True)] + installed = [torch.cuda.Event() for _ in range(_PF_LAYERS)] + real_install = ExpertPrefetcher.install + + def install(self, target): + got = real_install(self, target) + if got is not None: + installed[got].record(torch.cuda.current_stream()) + return got + + def issue_copy(self, target): + if not self._pending(target): + return + stream = streams[target % 2] + stream.wait_event(installed[target]) + with torch.cuda.stream(stream): + if target % 2 == 0: + torch.cuda._sleep(_PF_DELAY_CYCLES) + self.cache.copy_rows(target, self.pf_slots[target], self.pf_src[target], self.pf_num[target], slim=True, + kind=PREFETCH_COPY) + _mark_ready_kernel[(1,)](self.pf_ready[target]) + self._copy_events[target][1].record(stream) + + monkeypatch.setattr(ExpertPrefetcher, "install", install) + monkeypatch.setattr(ExpertPrefetcher, "issue_copy", issue_copy) + steps = _garbage(_routed_ids(moes, inputs), 12, "flood", seed=112) + on = _prefetch_run(moes, cache, 1, False, inputs, sleep=0, before_step=_with_override(cache, steps)) + try: + _assert_same_outputs(on[0], off[0]) + _assert_slots_hold_their_experts(cache) + except AssertionError: + return + pytest.fail("a prefetch copy landing on a slot the next install took over went unnoticed") + + +@requires_cuda +@pytest.mark.parametrize("kind", ["bf16", "nvfp4"]) +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +def test_moe_prefetch_on_per_bank_copy_fallback_is_bitwise_identical(kind, graph, bs, monkeypatch): + """Without the fused copy plan (FREETOKEN_FUSED_COPY=0 or misaligned banks) the prefetch copy + takes copy_rows' per-bank launches on the copy stream; late copies still change no bit.""" + import freetoken.moe.offload_cache as offload_cache + + monkeypatch.setattr(offload_cache, "_FUSED_COPY", False) + banks = _prefetch_banks(kind) + inputs = _prefetch_inputs(bs, steps=60, seed=5) + moes, cache = _prefetch_stack(kind, banks, "off", monkeypatch, overlap=True) + off = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0) + moes, cache = _on_stack(kind, banks, monkeypatch, policy="rule", overlap=True, budget=8) + assert not cache._copy_fused_ok + steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=bs) + on = _prefetch_run(moes, cache, bs, graph, inputs, sleep=0, + before_step=_with_override(cache, steps, delay_copy_cycles=_PF_DELAY_CYCLES // 2)) + _assert_same_outputs(on[0], off[0]) + _assert_slots_hold_their_experts(cache) + assert int(cache.prefetch.stats[:, 0].sum()) > 0 + + +def _multi_graph_run(moes, cache, schedule, before_step=None): + """The graph runner's sequence: per graph size (largest first) an eager warm-up, a capture into + one shared pool and a reset; then ``schedule`` of (bs, xs, replay) steps, replays and eager + steps interleaved.""" + _fresh_ctx(_batch=SimpleNamespace(is_prefill=False)) + + def step(xs): + return [moe.forward(x.clone()) for moe, x in zip(moes, xs)] + + graphs, pool = {}, None + for bs in (2, 1): + static = [torch.randn(bs, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16) for _ in moes] + step(static) + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g, pool=pool): + outs = step(static) + cache.reset() + pool = g.pool() if pool is None else pool + graphs[bs] = (g, static, outs) + results = [] + for i, (bs, xs, replay) in enumerate(schedule): + if before_step is not None: + before_step(i) + if replay: + g, static, outs = graphs[bs] + for s, x in zip(static, xs): + s.copy_(x) + g.replay() + else: + outs = step(xs) + if cache.audit is not None: + cache.audit.end_decode_step(cache) + results.append([o.clone() for o in outs]) + torch.cuda.synchronize() + return results + + +@requires_cuda +@pytest.mark.parametrize("policy", ["rule", "lru"]) +def test_moe_prefetch_on_interleaves_graph_sizes_and_eager_steps(policy, monkeypatch): + """bs 1 and bs 2 graphs captured into one pool (as the graph runner does), replayed in an + arbitrary order with eager steps between them, including an eager bs 9 step that forks nothing: + every output equals off's to the bit and every held slot holds its expert.""" + from freetoken.moe.prefetch import MAX_ROWS + + banks = _prefetch_banks("nvfp4") + gen = torch.Generator().manual_seed(41) + order = [(int(b), bool(r)) for b, r in zip(torch.randint(0, 4, (60,), generator=gen), torch.randint(0, 2, (60,), generator=gen))] + sizes = {0: 1, 1: 2, 2: 3, 3: MAX_ROWS + 1} + xs_by_bs = {bs: iter(_prefetch_inputs(bs, steps=60, seed=42 + bs)) for bs in sizes.values()} + schedule = [] + for b, r in order: + bs = sizes[b] + schedule.append((bs, next(xs_by_bs[bs]), r and bs <= 2)) + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, policy=policy) + off = _multi_graph_run(moes, cache, schedule) + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True, budget=8) + steps = _garbage(_routed_ids(moes, [xs for _, xs, _ in schedule]), 12, "adversarial", seed=43) + on = _multi_graph_run(moes, cache, schedule, before_step=_with_override(cache, steps)) + _assert_same_outputs(on, off) + _assert_slots_hold_their_experts(cache) + assert cache.prefetch._inflight == [None] * _PF_LAYERS + assert int(cache.prefetch.stats[:, 0].sum()) > 0 and int(cache.prefetch.stats[:, 1].sum()) > 0 + + +@requires_cuda +@pytest.mark.parametrize("seed", range(4)) +@pytest.mark.parametrize("rows", ["bs1", "bs2", "bs2_same"]) +@pytest.mark.parametrize("policy", ["rule", "lru"]) +def test_moe_prefetch_on_fuzzes_per_layer_stream_delays(seed, rows, policy, monkeypatch): + """Eager fuzz: before every layer a random choice of sleeps on the predictor stream, the prefetch + copy stream, the compute stream before the GEMM and the shared expert's side stream, with up to + 12 candidates a layer (most of the 16 slots re-installed every call) and, for bs2_same, two + identical rows. Outputs equal off's, held slots hold their experts, and the final maps and + stats (all but LATE) equal those of the same run without delays.""" + from freetoken.moe.prefetch import LATE, NUM_STAT_COLS + + bs = 1 if rows == "bs1" else 2 + inputs = _prefetch_inputs(bs, steps=60, seed=50 + seed) + if rows == "bs2_same": + inputs = [[x[:1].expand(2, -1).contiguous() for x in xs] for xs in inputs] + banks = _prefetch_banks("nvfp4") + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, policy=policy) + off = _prefetch_run(moes, cache, bs, False, inputs, sleep=0) + stall = _stall_hooks(monkeypatch) + gen = torch.Generator().manual_seed(seed) + choices = torch.randint(0, 16, (len(inputs), _PF_LAYERS), generator=gen).tolist() + runs = [] + for delayed in (False, True): + moes, cache = _on_stack("nvfp4", banks, monkeypatch, policy=policy, overlap=True, budget=12) + steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=60 + seed) + pf = cache.prefetch + load = _with_override(cache, steps) + at = {"step": -1} + + def before_step(i, load=load): + at["step"] = i + load(i) + + for layer, moe in enumerate(moes): + real = moe.forward + + def forward(x, layer=layer, real=real, pf=pf): + bits = choices[at["step"]][layer] if delayed and at["step"] >= 0 else 0 + # layer L's fork runs layer L+1's prediction and copy: these delays hit them + pf.delay_predict_cycles = _PF_DELAY_CYCLES // 2 if bits & 1 else 0 + pf.delay_copy_cycles = _PF_DELAY_CYCLES // 2 if bits & 2 else 0 + stall["gemm"] = _PF_DELAY_CYCLES // 2 if bits & 4 else 0 + stall["shared"] = _PF_DELAY_CYCLES // 2 if bits & 8 else 0 + return real(x) + + monkeypatch.setattr(moe, "forward", forward) + on = _prefetch_run(moes, cache, bs, False, inputs, sleep=0, before_step=before_step) + _assert_same_outputs(on[0], off[0]) + _assert_slots_hold_their_experts(cache) + runs.append((on[1][:3], pf.stats.cpu())) + (maps, stats), (delayed_maps, delayed_stats) = runs + for a, b in zip(maps, delayed_maps): + assert torch.equal(a, b) + cols = [c for c in range(NUM_STAT_COLS) if c != LATE] + assert torch.equal(stats[:, cols], delayed_stats[:, cols]) + + +# FREETOKEN_MOE_PREFETCH_VERIFY: the debug slot checks on the same stacks, quiet on a correct decode +# and loud on each kind of fault they exist for +_VERIFY_STEPS = 60 + + +def _verify_state(cache): + """(counters summed over layers, kept records) of the stack's verifier.""" + from freetoken.moe.verify import RING + + v = cache.verify + torch.cuda.synchronize() + return v.counters.cpu().sum(0).tolist(), v.ring[: min(int(v.cursor), RING)].cpu().tolist() + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("prefetch", ["off", "on"]) +@pytest.mark.parametrize("verify", ["full", "meta"]) +def test_moe_prefetch_verify_is_quiet_on_a_correct_decode(graph, bs, prefetch, verify, monkeypatch): + """FREETOKEN_MOE_PREFETCH_VERIFY checks every decode layer call and finds nothing on a correct + decode: prefetch off, and prefetch on with candidates the GEMMs read, late copies and a late + predictor. The checks only read, so the outputs equal those of verify off to the bit.""" + from freetoken.moe.prefetch import USEFUL + from freetoken.moe.verify import CHECKS + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(bs, steps=_VERIFY_STEPS, seed=70) + runs = [] + for mode in ("off", verify): + moes, cache = _prefetch_stack("nvfp4", banks, prefetch, monkeypatch, overlap=True, budget=8, verify_mode=mode) + before_step = None + if prefetch == "on": + steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=bs + 70) + before_step = _with_override(cache, steps, delay_copy_cycles=_PF_DELAY_CYCLES // 2, + delay_predict_cycles=_PF_DELAY_CYCLES // 4) + runs.append(_prefetch_run(moes, cache, bs, graph, inputs, sleep=0, before_step=before_step)) + _assert_same_outputs(runs[1][0], runs[0][0]) + counts, records = _verify_state(cache) + assert counts[CHECKS] == (_VERIFY_STEPS + 1) * _PF_LAYERS # the eager warm-up is checked too + assert counts[CHECKS + 1 :] == [0] * (len(counts) - 1) and records == [], records + assert (cache.verify.scratch is None) == (verify == "meta") + if prefetch == "on": + assert int(cache.prefetch.stats[:, USEFUL].sum()) > 0, "the GEMMs must have read prefetched slots" + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("prefetch", ["off", "on"]) +@pytest.mark.parametrize("path", list(_PREFILL_PATHS)) +def test_moe_prefetch_verify_is_quiet_across_prefills(graph, prefetch, path, monkeypatch): + """Every request prefills before it decodes: the whole-layer materialize, the overlap double + buffers (with the hit-D2D gather) and small prefills rewrite slots between decode steps, and the + decode checks after each of them still find nothing.""" + import freetoken.layers.moe as layers_moe + from freetoken.moe.verify import CHECKS + + cache_size, tokens, cache_kw, small = _PREFILL_PATHS[path] + monkeypatch.setattr(layers_moe, "_SMALL_PREFILL_TOKENS", small) + banks = _prefetch_banks("nvfp4") + decode = _prefetch_inputs(2, steps=18, seed=75) + gen = torch.Generator(device="cuda").manual_seed(76) + prefills = [[torch.randn(tokens, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16, generator=gen) * 0.5 + for _ in range(_PF_LAYERS)] for _ in range(2)] + schedule = ([("decode", xs) for xs in decode[:6]] + [("prefill", prefills[0])] + + [("decode", xs) for xs in decode[6:12]] + [("prefill", prefills[1])] + + [("decode", xs) for xs in decode[12:]]) + moes, cache = _prefetch_stack("nvfp4", banks, prefetch, monkeypatch, budget=8, overlap=True, cache_size=cache_size, + verify_mode="full", **cache_kw) + before_step = None + if prefetch == "on": + steps = _garbage(_routed_ids(moes, [xs for _, xs in schedule]), 12, "adversarial", seed=77) + before_step = _with_override(cache, steps) + _mixed_run(moes, cache, 2, graph, schedule, before_step=before_step) + counts, records = _verify_state(cache) + if path == "overlap_d2d": + assert cache.prefill_hit_rows > 0, "the hit-D2D gather must have read resident slots" + assert counts[CHECKS] == (len(decode) + 1) * _PF_LAYERS + assert counts[CHECKS + 1 :] == [0] * (len(counts) - 1) and records == [], records + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +def test_moe_prefetch_verify_catches_a_missing_copy_join(graph, monkeypatch): + """Negative control: without the compute stream's wait on the prefetch copy (and with that copy + late), GEMMs read prefetched slots before their bytes land. The byte check before the GEMM flags + them, and the records place the slot in the layer's prefetch copy plan.""" + from freetoken.moe.prefetch import ExpertPrefetcher + from freetoken.moe.verify import BYTES_PRE, META_BAD, PRE_BAD, R_KIND, R_SOURCE, SRC_PREFETCH + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(1, steps=40, seed=7) + moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, overlap=True, budget=8, verify_mode="full") + + def no_wait(self, layer_id): + self._inflight[layer_id] = None + + monkeypatch.setattr(ExpertPrefetcher, "join_copy", no_wait) + if graph: + # a capture needs the copy stream joined back: after the last layer, past every GEMM it races + last, copy_stream = moes[-1].forward, cache.prefetch.copy_stream + + def forward(x): + out = last(x) + torch.cuda.current_stream().wait_stream(copy_stream) + return out + + monkeypatch.setattr(moes[-1], "forward", forward) + steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=11) + # ~2.5 ms a copy: the copy stream falls ever further behind even an eager host's slow enqueue + _prefetch_run(moes, cache, 1, graph, inputs, sleep=0, + before_step=_with_override(cache, steps, delay_copy_cycles=5 * _PF_DELAY_CYCLES)) + counts, records = _verify_state(cache) + assert counts[PRE_BAD] > 0 and counts[META_BAD] == 0 + assert any(r[R_KIND] == BYTES_PRE and r[R_SOURCE] & SRC_PREFETCH for r in records) + + +def _on_call(layer, call, action): + """Wrap ExpertVerifier.before_gemm so ``action(verifier, cache, slots, real)`` replaces it on the + ``call``-th decode call of ``layer`` (call 0 is the eager warm-up).""" + from freetoken.moe.verify import ExpertVerifier + + real, seen = ExpertVerifier.before_gemm, {} + + def before_gemm(self, cache, slots, next_plan): + at = self._call[0] + seen[at] = seen.get(at, -1) + 1 + if (at, seen[at]) == (layer, call): + return action(self, cache, slots, lambda: real(self, cache, slots, next_plan)) + return real(self, cache, slots, next_plan) + + return before_gemm + + +@requires_cuda +@pytest.mark.parametrize("prefetch", ["off", "on"]) +def test_moe_prefetch_verify_catches_a_corrupt_slot(prefetch, monkeypatch): + """Negative control: bytes flipped in one bank of one routed slot right before the checks are + flagged before the GEMM, with the call, layer, slot, expert, bank and byte offset of the flip.""" + from freetoken.moe.verify import (BYTES_PRE, META_BAD, PRE_BAD, R_BANK, R_CALL, R_EXPERT, R_FIRST_BYTE, + R_KIND, R_LAYER, R_SLOT, ExpertVerifier) + + banks = _prefetch_banks("nvfp4") + moes, cache = _prefetch_stack("nvfp4", banks, prefetch, monkeypatch, overlap=True, verify_mode="full") + hit = {} + + def corrupt(verifier, cache, slots, real): + torch.cuda.synchronize() + hit.update(slot=int(slots.view(-1)[0]), expert=int(verifier.ids[0])) + cache.bank_views()[3][hit["slot"]].view(torch.uint8).view(-1)[100:108].bitwise_not_() + real() + + monkeypatch.setattr(ExpertVerifier, "before_gemm", _on_call(2, 10, corrupt)) + _prefetch_run(moes, cache, 1, False, _prefetch_inputs(1, steps=20, seed=71), sleep=0) + counts, records = _verify_state(cache) + assert counts[PRE_BAD] >= 1 and counts[META_BAD] == 0 + first = records[0] + assert [first[f] for f in (R_KIND, R_CALL, R_LAYER, R_SLOT, R_EXPERT, R_BANK, R_FIRST_BYTE)] == [ + BYTES_PRE, 10, 2, hit["slot"], hit["expert"], 3, 100] + + +@requires_cuda +@pytest.mark.parametrize("verify", ["full", "meta"]) +def test_moe_prefetch_verify_catches_a_wrong_slot_owner(verify, monkeypatch): + """Negative control: a routed slot's id_of_slot names another layer's expert while the checks + before the GEMM run (restored right after them): exactly that entry is flagged, with the owner it + saw, and the checks after the GEMM find the maps whole again.""" + from freetoken.moe.verify import (CHECKS, META_BAD, META_PRE, R_CALL, R_EXPERT, R_KIND, R_LAYER, R_MAPPED, + R_OWNER, R_ROW, R_SLOT, ExpertVerifier) + + banks = _prefetch_banks("nvfp4") + moes, cache = _prefetch_stack("nvfp4", banks, "off", monkeypatch, overlap=True, verify_mode=verify) + hit = {} + + def corrupt(verifier, cache, slots, real): + torch.cuda.synchronize() + s, e = int(slots.view(-1)[1]), int(verifier.ids[1]) + hit.update(slot=s, expert=e, saved=int(cache.id_of_slot[s])) + cache.id_of_slot[s] = 3 * _PF_EXPERTS + e # in range: an ensure reading it stays in bounds + real() + cache.id_of_slot[s] = hit["saved"] + + monkeypatch.setattr(ExpertVerifier, "before_gemm", _on_call(1, 7, corrupt)) + _prefetch_run(moes, cache, 1, False, _prefetch_inputs(1, steps=12, seed=72), sleep=0) + counts, records = _verify_state(cache) + assert hit["saved"] == _PF_EXPERTS + hit["expert"] + assert counts[META_BAD] == 1 and sum(counts[CHECKS + 1 :]) == 1 and len(records) == 1 + rec = records[0] + assert [rec[f] for f in (R_KIND, R_CALL, R_LAYER, R_ROW, R_SLOT, R_EXPERT, R_OWNER, R_MAPPED)] == [ + META_PRE, 7, 1, 1, hit["slot"], hit["expert"], 3 * _PF_EXPERTS + hit["expert"], hit["slot"]] + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("prefetch", ["off", "on"]) +def test_moe_prefetch_verify_catches_a_write_during_the_gemm(graph, prefetch, monkeypatch): + """Negative control: another stream rewrites a routed slot while the GEMM runs, after the checks + before it (a sleep on the compute stream holds the GEMM until the write lands). The byte check + after the GEMM flags it; the check before it saw nothing.""" + from freetoken.layers.moe import OffloadMoELayer + from freetoken.moe.verify import (BYTES_POST, BYTES_PRE, POST_BAD, R_BANK, R_CALL, R_FIRST_BYTE, R_KIND, R_LAYER, + R_SLOT, ExpertVerifier) + + banks = _prefetch_banks("nvfp4") + moes, cache = _prefetch_stack("nvfp4", banks, prefetch, monkeypatch, overlap=True, verify_mode="full") + side, real, armed = torch.cuda.Stream(), OffloadMoELayer._expert_gemm, {"on": False} + slot = torch.zeros((1,), dtype=torch.int64, device="cuda") + offset = torch.arange(16, device="cuda") + + def gemm(self, cache, hidden, weights, topk_ids, *, is_prefill, **kw): + if armed["on"] and self.layer_id == 3 and not is_prefill: + slot.copy_(topk_ids.view(-1)[:1]) # device side, so a captured graph rewrites this step's slot + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + torch.cuda._sleep(_PF_DELAY_CYCLES // 20) + row = cache.bank_views()[0].view(cache.cache_size, -1) + row.index_put_((slot, offset), row[slot, offset].bitwise_not()) + torch.cuda._sleep(_PF_DELAY_CYCLES) + out = real(self, cache, hidden, weights, topk_ids, is_prefill=is_prefill, **kw) + torch.cuda.current_stream().wait_stream(side) # a capture must join the side stream back + return out + return real(self, cache, hidden, weights, topk_ids, is_prefill=is_prefill, **kw) + + monkeypatch.setattr(OffloadMoELayer, "_expert_gemm", gemm) + real_run = ExpertVerifier.before_gemm + + def arm_at_call_5(self, cache, slots, next_plan): + # graph: the arm must be set while capturing, so the rewrite replays every step; eager: once at call 5 + armed["on"] = self._call[0] == 3 and (graph or int(self.calls[3]) == 5) + return real_run(self, cache, slots, next_plan) + + monkeypatch.setattr(ExpertVerifier, "before_gemm", arm_at_call_5) + _prefetch_run(moes, cache, 1, graph, _prefetch_inputs(1, steps=8, seed=73), sleep=0) + torch.cuda.synchronize() + counts, records = _verify_state(cache) + assert counts[POST_BAD] >= 1 + first = records[0] + assert [first[f] for f in (R_KIND, R_LAYER, R_BANK, R_FIRST_BYTE)] == [BYTES_POST, 3, 0, 0] + assert not any(r[R_KIND] == BYTES_PRE and (r[R_CALL], r[R_SLOT]) == (first[R_CALL], first[R_SLOT]) for r in records) + + +@requires_cuda +def test_moe_prefetch_verify_off_builds_and_launches_nothing(monkeypatch): + """With the flag off there is no verifier and a decode launches only what it launches without + one; meta and full only add launches (their kernels, the id copy and the host-row gather).""" + from collections import Counter + + from torch.profiler import ProfilerActivity, profile + + from freetoken.env import ENV + + monkeypatch.setattr(ENV.MOE_PREFETCH_VERIFY, "value", "0") + banks = _prefetch_banks("bf16") + inputs = _prefetch_inputs(1, steps=2) + + def kernels(verify): + moes, cache = _prefetch_stack("bf16", banks, "off", monkeypatch, verify_mode=verify) + _prefetch_run(moes, cache, 1, False, inputs[:1]) # warm up and compile outside the trace + with profile(activities=[ProfilerActivity.CUDA]) as prof: + _prefetch_run(moes, cache, 1, False, inputs) + return cache, Counter(e.name for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA) + + cache, off = kernels(None) + assert cache.verify is None and cache.verify_mode == "off" + ours = ("_plan_kernel", "_compare_kernel", "_tally_kernel") + assert not any(k in name for name in off for k in ours) + for mode in ("meta", "full"): + _, on = kernels(mode) + assert not off - on, f"{mode} dropped a launch" + added = on - off + assert all(any(k in name for k in (*ours, "Memcpy", "fast_index_copy_multi")) for name in added), added + assert any("_compare_kernel" in name for name in added) == (mode == "full") + + +@requires_cuda +def test_moe_prefetch_verify_follows_a_rebuild(monkeypatch): + """A rebuild reallocates the slot banks (production rebuilds to its final slot count after start): + the verifier keeps its scratch and counts, points its byte checks at the new banks, and a + recaptured graph checks clean.""" + from freetoken.moe.verify import CHECKS + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(2, steps=30, seed=74) + moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, overlap=True, verify_mode="full") + v = cache.verify + scratch = v.scratch + _prefetch_run(moes, cache, 2, True, inputs, sleep=0) + cache.rebuild(_PF_EXPERTS + 4) + assert cache.verify is v and v.scratch is scratch + assert v.cache_ptrs.tolist() == [bank.data_ptr() for bank in cache.bank_views()] + _prefetch_run(moes, cache, 2, True, inputs, sleep=0) + counts, records = _verify_state(cache) + assert counts[CHECKS] == 2 * (30 + 1) * _PF_LAYERS + assert counts[CHECKS + 1 :] == [0] * (len(counts) - 1) and records == [] + + +# FREETOKEN_MOE_SLOT_AUDIT (moe/slot_audit.py): quiet and output-neutral on a correct decode with +# prefills between the steps, and each injected persistent fault found with the history that explains it + + +def _audit_records(cache): + """(totals, kept records) of the stack's slot audit.""" + a = cache.audit + torch.cuda.synchronize() + return a.totals.cpu().tolist(), a.ring[: min(int(a.cursor), a.ring_n)].cpu().tolist() + + +def _bad_slots(cache): + """Ground truth: the (slot, owner) pairs whose bytes differ from the owner's host rows.""" + torch.cuda.synchronize() + held = (cache.id_of_slot >= 0).nonzero().view(-1).cpu() + ids = cache.id_of_slot[held.cuda()].long().cpu() + bad = torch.zeros(held.numel(), dtype=torch.bool) + for (per_layer, _), view in zip(cache.banks, cache.bank_views()): + want = torch.stack(per_layer).flatten(0, 1)[ids].contiguous().view(torch.uint8).view(held.numel(), -1) + got = view[held.cuda()].cpu().contiguous().view(torch.uint8).view(held.numel(), -1) + bad |= (got != want).any(dim=1) + return {(int(s), int(o)) for s, o, b in zip(held, ids, bad) if b} + + +def _spy_audit_kinds(monkeypatch): + """The event kinds the recorders were asked for, at enqueue time, each with the streams its + recorders ran on (a recorder follows its writer on the writer's stream).""" + from freetoken.moe.slot_audit import SlotAuditor + + kinds, real = {}, SlotAuditor._launch + + def launch(self, cache, kind, *args, **kwargs): + kinds.setdefault(kind, set()).add(torch.cuda.current_stream().cuda_stream) + return real(self, cache, kind, *args, **kwargs) + + monkeypatch.setattr(SlotAuditor, "_launch", launch) + return kinds + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("bs", [1, 2]) +@pytest.mark.parametrize("prefetch", ["off", "on"]) +@pytest.mark.parametrize("path", list(_PREFILL_PATHS)) +def test_moe_slot_audit_is_quiet_and_changes_no_bit(graph, bs, prefetch, path, monkeypatch): + """FREETOKEN_MOE_SLOT_AUDIT=1 records every slot writer (the decode and prefetch ensures and copies, + and between the decode steps the whole-layer materialize, the overlap double buffers with the + hit-D2D gather, or small prefills) and audits every held slot after each decode step: it finds + nothing, and the outputs, slot maps and slot bytes equal those of the audit off to the bit. Every + map writer it records ran on the compute stream but the prefill buffer invalidation (on the + prefill copy stream), and the prefetch copy on the prefetch copy stream.""" + import freetoken.layers.moe as layers_moe + import freetoken.moe.offload_cache as offload_cache + from freetoken.moe import slot_audit as sa + + cache_size, tokens, cache_kw, small = _PREFILL_PATHS[path] + monkeypatch.setattr(layers_moe, "_SMALL_PREFILL_TOKENS", small) + if path == "overlap_d2d": + # every toy bank is under the 256 KiB small-bank floor, which would leave the hit gather idle + monkeypatch.setattr(offload_cache, "_SMALL_BANK_FEAT_BYTES", 8192) + banks = _prefetch_banks("nvfp4") + decode = _prefetch_inputs(bs, steps=18, seed=81) + gen = torch.Generator(device="cuda").manual_seed(82) + prefills = [[torch.randn(tokens, _PF_HIDDEN, device="cuda", dtype=torch.bfloat16, generator=gen) * 0.5 + for _ in range(_PF_LAYERS)] for _ in range(2)] + schedule = ([("decode", xs) for xs in decode[:6]] + [("prefill", prefills[0])] + + [("decode", xs) for xs in decode[6:12]] + [("prefill", prefills[1])] + + [("decode", xs) for xs in decode[12:]]) + kinds = _spy_audit_kinds(monkeypatch) + runs = [] + for interval in (0, 1): + moes, cache = _prefetch_stack("nvfp4", banks, prefetch, monkeypatch, budget=8, overlap=True, + cache_size=cache_size, slot_audit=interval, **cache_kw) + assert (cache.audit is None) == (interval == 0) + before_step = None + if prefetch == "on": + steps = _garbage(_routed_ids(moes, [xs for _, xs in schedule]), 12, "adversarial", seed=bs + 83) + before_step = _with_override(cache, steps) + compute = _record_compute_streams(monkeypatch, moes) + out = _mixed_run(moes, cache, bs, graph, schedule, before_step=before_step) + runs.append((out, [t.clone() for t in (cache.slot_for_id, cache.id_of_slot, cache.usage, *cache.bank_views())])) + _assert_same_outputs(runs[1][0], runs[0][0]) + for a, b in zip(runs[1][1], runs[0][1]): + assert torch.equal(a, b) + totals, records = _audit_records(cache) + assert records == [] and totals[sa.T_BAD] == 0, [cache.audit.format_record(r) for r in records] + assert totals[sa.T_MAP_RANGE] == totals[sa.T_BAD_PLAN] == totals[sa.T_UNCHECKED] == 0 + assert totals[sa.T_AUDITS] == len(decode) + assert totals[sa.T_HELD] > 0 and totals[sa.T_BYTES_CHECKED] == totals[sa.T_HELD] and totals[sa.T_EVENTS] > 0 + expected = {sa.DEMAND_INSTALL, sa.DEMAND_COPY, sa.RESET} + if prefetch == "on": + expected |= {sa.PREFETCH_INSTALL, sa.PREFETCH_COPY} + expected |= { + "materialize": {sa.MATERIALIZE_CLEAR, sa.MATERIALIZE_INSTALL, sa.MATERIALIZE_COPY}, + "overlap": {sa.INVALIDATE, sa.PREFILL_BUFFER}, + "overlap_d2d": {sa.INVALIDATE, sa.PREFILL_SPLIT_H2D, sa.PREFILL_HIT_D2D}, + "small": {sa.PREFILL_INSTALL, sa.PREFILL_COPY}, + }[path] + assert set(kinds) == expected, (set(kinds) - expected, expected - set(kinds)) + fence = set() if cache.prefill_copy_stream is None else {cache.prefill_copy_stream.cuda_stream} + for kind in (sa.DEMAND_INSTALL, sa.PREFETCH_INSTALL, sa.PREFILL_INSTALL, sa.MATERIALIZE_INSTALL, + sa.MATERIALIZE_CLEAR, sa.RESET, sa.INVALIDATE): + assert kinds.get(kind, set()) <= (fence if kind == sa.INVALIDATE else compute), kind + if prefetch == "on": + assert kinds[sa.PREFETCH_COPY] == {cache.prefetch.copy_stream.cuda_stream} + + +@requires_cuda +@pytest.mark.parametrize("policy", ["rule", "lru"]) +def test_moe_slot_audit_is_quiet_across_graph_sizes_and_eager_steps(policy, monkeypatch): + """The graph runner's shape: bs 1 and bs 2 graphs in one pool, replayed in any order with eager + steps between them (a bs 9 one forks no prefetch); audited after every step, with the decode + verifier on beside it, neither finds anything.""" + from freetoken.moe import slot_audit as sa + from freetoken.moe.prefetch import MAX_ROWS + from freetoken.moe.verify import CHECKS + + banks = _prefetch_banks("nvfp4") + gen = torch.Generator().manual_seed(90) + order = [(int(b), bool(r)) for b, r in zip(torch.randint(0, 4, (40,), generator=gen), torch.randint(0, 2, (40,), generator=gen))] + sizes = {0: 1, 1: 2, 2: 3, 3: MAX_ROWS + 1} + xs_by_bs = {bs: iter(_prefetch_inputs(bs, steps=40, seed=91 + bs)) for bs in sizes.values()} + schedule = [(sizes[b], next(xs_by_bs[sizes[b]]), r and sizes[b] <= 2) for b, r in order] + moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, policy=policy, overlap=True, budget=8, slot_audit=1, + verify_mode="meta") + steps = _garbage(_routed_ids(moes, [xs for _, xs, _ in schedule]), 12, "adversarial", seed=92) + _multi_graph_run(moes, cache, schedule, before_step=_with_override(cache, steps)) + totals, records = _audit_records(cache) + assert records == [] and totals[sa.T_BAD] == 0 and totals[sa.T_AUDITS] == len(schedule) + assert totals[sa.T_BYTES_CHECKED] == totals[sa.T_HELD] > 0 + counts, verify_records = _verify_state(cache) + assert counts[CHECKS] > 0 and counts[CHECKS + 1 :] == [0] * (len(counts) - 1) and verify_records == [] + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("policy", ["rule", "lru"]) +def test_moe_slot_audit_explains_a_skipped_prefetch_copy(graph, policy, monkeypatch): + """Fault injection: layer 2's prefetch copies copy nothing. Every slot the audit finds bad holds a + layer-2 expert whose history reads 'prefetch install, no prefetch copy', and every bad slot still + held at the end was found.""" + from freetoken.moe import slot_audit as sa + + monkeypatch.setattr(sa, "RING", 1024) + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(1, steps=40, seed=84) + moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, policy=policy, overlap=True, budget=8, slot_audit=1) + real = cache.copy_rows + + def copy_rows(layer_id, dst, src, num, *, slim, **kw): + if kw.get("kind") == sa.PREFETCH_COPY and layer_id == 2: + return None # the fault: no bytes move (and so nothing is recorded) + return real(layer_id, dst, src, num, slim=slim, **kw) + + monkeypatch.setattr(cache, "copy_rows", copy_rows) + steps = _garbage(_routed_ids(moes, inputs), 12, "adversarial", seed=85) + _prefetch_run(moes, cache, 1, graph, inputs, sleep=0, before_step=_with_override(cache, steps)) + truth = _bad_slots(cache) + totals, records = _audit_records(cache) + assert records and totals[sa.T_BAD] > 0 and totals[sa.T_NEW] == len(records) + a = cache.audit + for rec in records: + assert rec[sa.H_FLAGS] == sa.F_BYTES and rec[sa.H_OWNER] // _PF_EXPERTS == 2, a.format_record(rec) + # a layer-1 prefetch copy still in flight may land after the install and is named as such + assert a.diagnose(rec).startswith("prefetch install, no prefetch copy"), a.format_record(rec) + assert truth <= {(r[sa.H_SLOT], r[sa.H_OWNER]) for r in records} + + +@requires_cuda +@pytest.mark.parametrize("graph", [False, True], ids=["eager", "graph"]) +@pytest.mark.parametrize("prefetch", ["off", "on"]) +def test_moe_slot_audit_explains_a_demand_copy_from_the_wrong_layer(graph, prefetch, monkeypatch): + """Fault injection: layer 2's demand copies read layer 1's host banks (the right rows). The bad + slots are layer-2 demand installs whose last copy came from layer 1's source.""" + from freetoken.moe import slot_audit as sa + from freetoken.moe.offload_cache import OffloadMoeCache + + monkeypatch.setattr(sa, "RING", 1024) + real = OffloadMoeCache.copy_missing + + def copy_missing(self, layer_id=None): + if self._pending_src_layer == 2 and not self._pending_whole_layer: + return self.copy_rows(1, self.evict_slots, self.src_indices, self.num_indices, slim=False, kind=self._pending_kind) + return real(self, layer_id) + + monkeypatch.setattr(OffloadMoeCache, "copy_missing", copy_missing) + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(1, steps=30, seed=86) + moes, cache = _prefetch_stack("nvfp4", banks, prefetch, monkeypatch, overlap=True, budget=8, slot_audit=1) + _prefetch_run(moes, cache, 1, graph, inputs, sleep=0) + truth = _bad_slots(cache) + totals, records = _audit_records(cache) + assert records and truth <= {(r[sa.H_SLOT], r[sa.H_OWNER]) for r in records} + for rec in records: + owner = rec[sa.H_OWNER] + assert owner // _PF_EXPERTS == 2 and rec[sa.H_FLAGS] == sa.F_BYTES + e = owner % _PF_EXPERTS + assert cache.audit.diagnose(rec).startswith( + f"demand install of L2/e{e}, then demand copy from layer 1's source (L1/e{e}, #"), cache.audit.format_record(rec) + + +@requires_cuda +@pytest.mark.parametrize("prefetch", ["off", "on"]) +def test_moe_slot_audit_explains_bytes_overwritten_after_install(prefetch, monkeypatch): + """Fault injection: bytes of a held slot flipped behind every recorder's back. The next audit keeps + one record (bank, byte offset, bad words) whose history ends in the owner's install and copy, a + later audit counts it again without a second record, and a clean slot is forgotten.""" + from freetoken.moe import slot_audit as sa + + banks = _prefetch_banks("nvfp4") + moes, cache = _prefetch_stack("nvfp4", banks, prefetch, monkeypatch, overlap=True, budget=8, slot_audit=1000) + _prefetch_run(moes, cache, 1, False, _prefetch_inputs(1, steps=12, seed=87), sleep=0) + a = cache.audit + a.scan(cache) + assert _audit_records(cache)[1] == [] + slot = int((cache.id_of_slot >= 0).nonzero()[0]) + owner = int(cache.id_of_slot[slot]) + row = cache.bank_views()[3][slot].view(torch.uint8).view(-1) + row[100:108].bitwise_not_() + a.scan(cache) + a.scan(cache) + totals, records = _audit_records(cache) + assert len(records) == 1 and [totals[t] for t in (sa.T_AUDITS, sa.T_BAD, sa.T_NEW, sa.T_STILL)] == [3, 2, 1, 1] + rec = records[0] + assert [rec[f] for f in (sa.H_AUDIT, sa.H_SLOT, sa.H_OWNER, sa.H_FLAGS, sa.H_BANKS, sa.H_FIRST_BANK, + sa.H_FIRST_BYTE, sa.H_BAD_WORDS)] == [1, slot, owner, sa.F_BYTES, 1 << 3, 3, 100, 2] + events = a.events(rec) + install = [ev for ev in events if ev["kind"] in sa.INSTALLS][-1] + copy = [ev for ev in events if ev["ring"] == "bytes"][-1] + assert a._flat(install) == owner == a._flat(copy) and copy["seq"] > install["seq"] + assert "then the bytes changed with no recorded write" in a.diagnose(rec) + row[100:108].bitwise_not_() + a.scan(cache) + row[100:108].bitwise_not_() + a.scan(cache) + totals, records = _audit_records(cache) + assert len(records) == 2 and records[1][sa.H_AUDIT] == 4 and totals[sa.T_BAD] == 3 + + +@requires_cuda +def test_moe_slot_audit_explains_stale_and_missing_map_entries(monkeypatch): + """Fault injection on the maps: slot_for_id still names a slot for its previous owner (an eviction + that forgot to clear it: routing that expert would read another's bytes), and a held slot's owner + has lost its entry. Both are found and the stale one is traced to the install it outlived.""" + from freetoken.moe import slot_audit as sa + + banks = _prefetch_banks("nvfp4") + moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, overlap=True, budget=8, slot_audit=1000) + _prefetch_run(moes, cache, 1, False, _prefetch_inputs(1, steps=12, seed=88), sleep=0) + a = cache.audit + flat = cache.slot_for_id.view(-1) + held = (cache.id_of_slot >= 0).nonzero().view(-1).tolist() + found = None + for s in held: + for ev in sorted(a.meta_hist[s].tolist()): + fid = ev[sa.E_LAYER] * _PF_EXPERTS + ev[sa.E_ROW] + if ev[sa.E_KIND] in sa.INSTALLS and int(flat[fid]) == -1: + found = (s, fid) + if found: + break + assert found, "some slot must have outlived an earlier owner" + stale_slot, stale_id = found + fwd_slot = next(t for t in held if t != stale_slot) + fwd_owner = int(cache.id_of_slot[fwd_slot]) + flat[stale_id] = stale_slot + flat[fwd_owner] = -1 + a.scan(cache) + totals, records = _audit_records(cache) + assert [totals[t] for t in (sa.T_BAD, sa.T_STALE, sa.T_FWD, sa.T_BYTES)] == [2, 1, 1, 0] and len(records) == 2 + by_slot = {r[sa.H_SLOT]: r for r in records} + rec = by_slot[stale_slot] + assert rec[sa.H_FLAGS] == sa.F_STALE and rec[sa.H_STALE_ID] == stale_id + holder = a._name(int(cache.id_of_slot[stale_slot])) + text = a.diagnose(rec) + assert text.startswith(f"stale slot_for_id: {a._name(stale_id)} -> this slot, which holds {holder}; " + f"{a._name(stale_id)} was installed here by "), text + assert text.endswith("without clearing its map entry"), text + rec = by_slot[fwd_slot] + assert rec[sa.H_FLAGS] == sa.F_FWD and rec[sa.H_MAPPED] == -1 + assert a.diagnose(rec) == f"slot_for_id[{a._name(fwd_owner)}]=-1, not this slot" + + +@requires_cuda +def test_moe_slot_audit_off_builds_and_launches_nothing(monkeypatch): + """With the flag off there is no auditor and a decode launches only what it launches without one; + on, a decode only adds the recorders (and the audit its own kernels and host-row gathers).""" + from collections import Counter + + from torch.profiler import ProfilerActivity, profile + + from freetoken.env import ENV + + monkeypatch.setattr(ENV.MOE_SLOT_AUDIT, "value", 0) + banks = _prefetch_banks("bf16") + inputs = _prefetch_inputs(1, steps=2) + + def kernels(interval): + moes, cache = _prefetch_stack("bf16", banks, "on", monkeypatch, slot_audit=interval) + _prefetch_run(moes, cache, 1, False, inputs[:1]) # warm up and compile outside the trace + with profile(activities=[ProfilerActivity.CUDA]) as prof: + _prefetch_run(moes, cache, 1, False, inputs) + return cache, Counter(e.name for e in prof.events() if e.device_type == torch.autograd.DeviceType.CUDA) + + cache, off = kernels(None) + assert cache.audit is None and cache.slot_audit == 0 + ours = ("_record_kernel", "_bump_step_kernel", "_audit_") + assert not any(k in name for name in off for k in ours) + _, on = kernels(1) + assert not off - on, "the audit dropped a launch" + added = on - off + assert all(any(k in name for k in (*ours, "fast_index_copy_multi")) for name in added), added + assert any("_audit_report_kernel" in name for name in added) + + +@requires_cuda +def test_moe_slot_audit_follows_a_rebuild(monkeypatch): + """A rebuild reallocates the slot banks and maps: the auditor keeps its scratch and totals, starts a + new history sized for the new slot count, points its byte checks at the new banks, and a recaptured + graph audits clean.""" + from freetoken.moe import slot_audit as sa + + banks = _prefetch_banks("nvfp4") + inputs = _prefetch_inputs(2, steps=20, seed=89) + moes, cache = _prefetch_stack("nvfp4", banks, "on", monkeypatch, overlap=True, slot_audit=1) + a = cache.audit + scratch = a.scratch + _prefetch_run(moes, cache, 2, True, inputs, sleep=0) + cache.rebuild(_PF_EXPERTS + 4) + assert cache.audit is a and a.scratch is scratch and a.num_slots == _PF_EXPERTS + 4 + assert a.meta_hist.shape[0] == a.bytes_cur.shape[0] == _PF_EXPERTS + 4 and int(a.meta_cur.sum()) == 0 + assert a.cache_ptrs.tolist() == [bank.data_ptr() for bank in cache.bank_views()] + _prefetch_run(moes, cache, 2, True, inputs, sleep=0) + totals, records = _audit_records(cache) + assert records == [] and totals[sa.T_BAD] == 0 and totals[sa.T_AUDITS] == 2 * len(inputs) + + @requires_cuda def test_decoder_stack_prefill_and_decode(monkeypatch): """Ragged bs=3 prefill then a bs=3 decode step through the whole model with dummy weights.""" diff --git a/tests/moe/ref_prefetch.py b/tests/moe/ref_prefetch.py new file mode 100644 index 0000000000..c40d9351a8 --- /dev/null +++ b/tests/moe/ref_prefetch.py @@ -0,0 +1,22 @@ +"""CPU reference for the router-lookahead select and count (freetoken.moe.prefetch).""" +from __future__ import annotations + + +def ref_select(logits, resident, k: int, budget: int): + """Each row's top-k (ties to the lowest id), merged rank by rank, duplicates dropped; the first + ``budget`` non-resident ids, and the resident ids among the first ``budget`` distinct ones.""" + rows = [sorted(range(len(row)), key=lambda e: (-row[e], e))[:k] for row in logits] + merged = [] + for r in range(k): + for row in rows: + if row[r] not in merged: + merged.append(row[r]) + sel = [e for e in merged if resident[e] < 0][:budget] + res = [e for e in merged[:budget] if resident[e] >= 0] + return sel, res + + +def ref_count(sel, res, routed, misses: int, rows: int): + """``[issued, useful, resident_hits, calls, rows, misses]`` for one layer call.""" + routed = set(routed) + return [len(sel), sum(e in routed for e in sel), sum(e in routed for e in res), 1, rows, misses] diff --git a/tests/moe/ref_scored_cache.py b/tests/moe/ref_scored_cache.py index d7d8b0f2fb..15153c86d5 100644 --- a/tests/moe/ref_scored_cache.py +++ b/tests/moe/ref_scored_cache.py @@ -47,12 +47,13 @@ def _softplus2(self, x): g = np.where(i < se.G_LAST, g0 + (((g1 - g0) * frac) >> shift), 0) return np.maximum(x, 0) + g - def _keys(self, layer, tok, now, pin_since=None): + def _keys(self, layer, tok, now, pin_since=None, lowpri=False): c = np.arange(self.S, dtype=np.int64) oid = self.id_of_slot held = oid >= 0 if self.policy == 0: - key = (self.usage << self.slot_bits) | c + # lowpri: empty (0) < unconsumed prefetch (held, usage 0 -> 1) < resident (usage + 1) + key = (np.where(held, self.usage + 1, 0) << self.slot_bits) | c if lowpri else (self.usage << self.slot_bits) | c else: ids = np.where(held, oid, 0) lk = ids // self.E @@ -68,16 +69,20 @@ def _keys(self, layer, tok, now, pin_since=None): score = score + ((self.w_q4 * lcv) >> 4) assert np.abs(score).max() < 2**31 # the kernel computes it in int32 key = np.where(held, ((score + se.SCORE_BIAS) << self.slot_bits) | c, c) + if lowpri: + key = np.where(held & (self.usage == 0), (1 << self.slot_bits) | c, key) evictable = self.usage != self.step if pin_since is not None: evictable &= self.usage <= pin_since return np.where(evictable, key, _KEY_MAX) def ensure(self, layer, ids, *, bump_tok=False, update_state=True, logits=None, thr=0.0, near_miss=None, - pin_since=None): + pin_since=None, lowpri=False): """One ensure call; returns ``(out_slots, src_ids, dst_slots)`` like the kernel's plan. - ``pin_since`` (a ``step`` value) also pins every slot touched after it. + ``pin_since`` (a ``step`` value) also pins every slot touched after it. ``lowpri`` + (FREETOKEN_MOE_PREFETCH=on) keys unconsumed prefetches between the empty slots and the + residents; ``self.useful`` is then the call's distinct hits on such slots. Policy 3 refreshes the ids whose ``logits`` (``[rows, E]``) are within ``thr`` of their row's lowest routed logit, in fp32 as the kernel does, or the ids of a given ``near_miss`` mask. @@ -100,13 +105,15 @@ def ensure(self, layer, ids, *, bump_tok=False, update_state=True, logits=None, self.last_tok[near] = np.maximum(self.last_tok[near], tok - 1) q = np.asarray(ids, dtype=np.int64).reshape(-1) + base slots = self.slot_of_id[q] - self.usage[slots[slots >= 0]] = self.step uniq = np.unique(q) + held_at = self.slot_of_id[uniq] + self.useful = int(((held_at >= 0) & (self.usage[np.maximum(held_at, 0)] == 0)).sum()) + self.usage[slots[slots >= 0]] = self.step missing = uniq[self.slot_of_id[uniq] < 0] # ascending, as the kernel ranks them src = missing - base dst = np.empty(missing.size, np.int64) if missing.size: - keys = self._keys(layer, tok, now, pin_since) + keys = self._keys(layer, tok, now, pin_since, lowpri) part = np.argpartition(keys, missing.size - 1)[: missing.size] victims = part[np.argsort(keys[part])] for i, (e, v) in enumerate(zip(missing, victims)): @@ -127,6 +134,31 @@ def ensure(self, layer, ids, *, bump_tok=False, update_state=True, logits=None, self.ct[uniq] = now return out, src, dst + def prefetch(self, layer, sel, budget=None): + """``prefetch_ensure``: install ``sel``'s non-resident ids (-1 = padding, query order), at most + the first ``budget``, as held slots with usage 0, without bumping a clock or touching a hit; + the last demand call's slots (usage == step) stay pinned. Returns ``(src_ids, dst_slots)``, + the copy plan.""" + base = layer * self.E + tok = self.tok if self.policy else 0 + pos = layer - 1 # the kernel scores from the previous layer's call, where the prefetch runs + now = tok * self.L + pos + order = [] + for e in sel: + if e >= 0 and e + base not in order: + order.append(int(e) + base) + missing = [i for i in order if self.slot_of_id[i] < 0] + keys = self._keys(pos, tok, now, lowpri=True) + n = min(len(missing), int((keys != _KEY_MAX).sum()), len(missing) if budget is None else budget) + victims = np.argsort(keys, kind="stable")[:n] + for e, v in zip(missing[:n], victims): + old = self.id_of_slot[v] + if old >= 0: + self.slot_of_id[old] = -1 + self.id_of_slot[v] = e + self.slot_of_id[e] = v + self.usage[v] = 0 + return np.asarray(missing[:n], np.int64) - base, victims.astype(np.int64) def materialize(self, layer): """The prefill materialize kernel: the whole layer into slots [0, E), position == expert id.""" diff --git a/tests/moe/test_prefetch.py b/tests/moe/test_prefetch.py new file mode 100644 index 0000000000..47c26066a4 --- /dev/null +++ b/tests/moe/test_prefetch.py @@ -0,0 +1,424 @@ +"""FREETOKEN_MOE_PREFETCH stage A: the router-lookahead select / count kernels against a Python +reference, the mode switch, and the predictor stream's lifetime across resets and rebuilds. +The qwen4_exp stack tests (bit-identical outputs, counters over a real decode) live in +tests/models/qwen4_exp/test_skeleton.py.""" + +from __future__ import annotations + +import weakref + +import pytest +import torch + +from .ref_prefetch import ref_count, ref_select + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + + +def _padded(ids, width): + return ids + [-1] * (width - len(ids)) + + +def _select(logits, resident, k, budget, width=16): + from freetoken.moe.prefetch import lookahead_select + + sel = torch.full((width,), -7, dtype=torch.int32, device="cuda") + res = torch.full((width,), -7, dtype=torch.int32, device="cuda") + lookahead_select( + torch.tensor(logits, dtype=torch.bfloat16, device="cuda"), + torch.tensor(resident, dtype=torch.int32, device="cuda"), + sel, res, k=k, budget=budget, + ) + return sel.tolist(), res.tolist() + + +def test_mode_switch(): + from freetoken.moe.prefetch import default_budget, resolve_mode + + assert resolve_mode("off") == "off" and resolve_mode(" Measure ") == "measure" and resolve_mode("ON") == "on" + with pytest.raises(ValueError): + resolve_mode("lookahead") + assert (default_budget(True), default_budget(False)) == (3, 4) + + +@requires_cuda +def test_cache_builds_the_prefetcher_only_when_asked(): + from freetoken.moe.offload_cache import OffloadMoeCache + + def cache(**kw): + return OffloadMoeCache(num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"), **kw) + + assert cache(prefetch_mode="off").prefetch is None + assert cache(prefetch_mode="measure").prefetch.copy_stream is None + on = cache(prefetch_mode="on") + assert on.prefetch.mode == "on" and on.prefetch.cache.prefetch is on.prefetch and on.prefetch_on + # no reference cycle: dropping the cache frees its tensors at once + ref = weakref.ref(on) + del on + assert ref() is None + for mode in ("measure", "on"): + with pytest.raises(ValueError): + cache(prefetch_mode=mode, decode_target="hybrid") + + +@requires_cuda +def test_the_prefetch_env_flag_skips_caches_that_do_not_decode_on_the_gpu(monkeypatch): + # a process-wide flag must not break hybrid or CPU-decoded caches, as FREETOKEN_MOE_COPY_OVERLAP does not + from freetoken.env import ENV + from freetoken.moe.offload_cache import OffloadMoeCache + + monkeypatch.setattr(ENV.MOE_PREFETCH, "value", "on") + for target in ("hybrid", "cpu"): + c = OffloadMoeCache(num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"), decode_target=target) + assert c.prefetch is None and not c.prefetch_on + assert OffloadMoeCache(num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda")).prefetch_on + + +@requires_cuda +def test_select_known_answers(): + # order by logit: 1, 4, 6, 2, 3, 7, 0, 5; experts 4 and 2 are resident + logits = [[0.1, 0.9, 0.5, 0.3, 0.8, 0.0, 0.7, 0.2]] + resident = [-1, -1, 7, -1, 3, -1, -1, -1] + assert _select(logits, resident, 5, 2) == (_padded([1, 6], 16), _padded([4], 16)) + assert _select(logits, resident, 5, 4) == (_padded([1, 6, 3], 16), _padded([4, 2], 16)) + # the top-k cap: 3 and 7 are out of a top-4 list + assert _select(logits, resident, 4, 4) == (_padded([1, 6], 16), _padded([4, 2], 16)) + # all-equal logits rank by id + assert _select([[0.5] * 8], [-1] * 8, 8, 3) == (_padded([0, 1, 2], 16), _padded([], 16)) + # bs 2: row 0 ranks 1, 4, 6 and row 1 ranks 4, 0, 5 -> merged 1, 4, 0, 6, 5; expert 0 resident + rows = [[0.0, 0.9, 0.1, 0.1, 0.8, 0.1, 0.7, 0.1], [0.8, 0.1, 0.1, 0.1, 0.9, 0.7, 0.1, 0.1]] + resident = [2, -1, -1, -1, -1, -1, -1, -1] + assert _select(rows, resident, 3, 3) == (_padded([1, 4, 6], 16), _padded([0], 16)) + assert _select(rows, resident, 3, 16) == (_padded([1, 4, 6, 5], 16), _padded([0], 16)) + + +@requires_cuda +@pytest.mark.parametrize("rows", [1, 2]) +@pytest.mark.parametrize("k,budget", [(3, 3), (16, 4), (16, 3), (8, 16)]) +def test_select_matches_reference(rows, k, budget): + gen = torch.Generator().manual_seed(rows * 100 + k + budget) + for trial in range(20): + # coarse values make bf16 ties common, so the tie rule is exercised + logits = (torch.randint(-24, 24, (rows, 512), generator=gen).float() / 8).tolist() + resident = torch.where(torch.rand(512, generator=gen) < 0.4, 1, -1).tolist() + want = ref_select(logits, resident, k, budget) + width = 2 * k + got = _select(logits, resident, k, budget, width=width) + assert got == (_padded(want[0], width), _padded(want[1], width)), trial + + +@requires_cuda +def test_count_known_answers(): + from freetoken.moe.prefetch import prefetch_count + + num_experts, layer = 8, 2 + base = layer * num_experts + # slots 0, 1, 3 hold this layer's experts 3, 1, 6; slot 4 another layer's expert + id_of_slot = torch.tensor([base + 3, base + 1, -1, base + 6, 5, -1], dtype=torch.int32, device="cuda") + slots = torch.tensor([0, 1, 3], dtype=torch.int32, device="cuda") + sel = torch.tensor(_padded([3, 5], 6), dtype=torch.int32, device="cuda") + res = torch.tensor(_padded([1, 4], 6), dtype=torch.int32, device="cuda") + misses = torch.tensor([2], dtype=torch.int64, device="cuda") + counters = torch.zeros(6, dtype=torch.int64, device="cuda") + for _ in range(2): + prefetch_count(slots, id_of_slot, misses, sel, res, counters, id_base=base, rows=1) + want = ref_count([3, 5], [1, 4], [3, 1, 6], 2, 1) + assert counters.tolist() == [2 * v for v in want] == [4, 2, 2, 2, 2, 4] + # bs 2 with a duplicate routed expert counts it once + counters.zero_() + slots2 = torch.tensor([0, 1, 0, 3], dtype=torch.int32, device="cuda") + prefetch_count(slots2, id_of_slot, misses, sel, res, counters, id_base=base, rows=2) + assert counters.tolist() == ref_count([3, 5], [1, 4], [3, 1, 3, 6], 2, 2) + + +@requires_cuda +def test_prefetch_stream_never_aliases_a_pooled_stream(): + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache( + num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"), + decode_copy_overlap=True, prefetch_mode="measure", + ) + pred = cache.prefetch.stream.cuda_stream + pooled = {torch.cuda.Stream().cuda_stream for _ in range(64)} + assert pred not in pooled and pred != torch.cuda.current_stream().cuda_stream + assert pred != cache.decode_copy_stream.cuda_stream + + +@requires_cuda +def test_prefetch_copy_stream_is_a_third_dedicated_stream(): + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache( + num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"), + decode_copy_overlap=True, prefetch_mode="on", + ) + pf = cache.prefetch + copy = pf.copy_stream.cuda_stream + pooled = {torch.cuda.Stream().cuda_stream for _ in range(64)} + assert copy not in pooled and pf.stream.cuda_stream not in pooled + assert len({copy, pf.stream.cuda_stream, cache.decode_copy_stream.cuda_stream, torch.cuda.current_stream().cuda_stream}) == 4 + # its CTAs go ahead of pending compute ones, so a late copy does not also wait for SMs + from cuda.bindings import runtime as cudart + + _, _, greatest = cudart.cudaDeviceGetStreamPriorityRange() + assert cudart.cudaStreamGetPriority(copy)[1] == greatest < 0 + assert cudart.cudaStreamGetPriority(pf.stream.cuda_stream)[1] == 0 + + +@requires_cuda +def test_rebuild_and_reset_clear_the_counters(): + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache( + num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"), prefetch_mode="measure" + ) + cache.set_bank_sources({ + "gate_up": [torch.zeros(4, 32, 8, dtype=torch.bfloat16).pin_memory() for _ in range(2)], + "down": [torch.zeros(4, 8, 16, dtype=torch.bfloat16).pin_memory() for _ in range(2)], + }) + pf = cache.prefetch + stream, sel = pf.stream, pf.sel + pf.counters.fill_(3) + cache.reset() + assert int(pf.counters.abs().sum()) == 0 + pf.counters.fill_(5) + pf.take_window() + assert int(pf.totals.sum()) > 0 + pf.counters.fill_(2) + pf._inflight[1] = torch.zeros(1) + cache.rebuild(6) + assert cache.prefetch is pf and pf.stream is stream and pf.sel is sel + assert int(pf.counters.abs().sum()) == 0 and int(pf.totals.abs().sum()) == 0 + assert pf._inflight == [None, None] + + +def _emit_prefetch(window, monkeypatch, *, debug=False, mode="measure"): + from types import SimpleNamespace + + from freetoken.engine import engine as engine_mod + from freetoken.env import ENV + + taken, lines = [], [] + prefetch = SimpleNamespace(mode=mode, take_window=lambda: taken.append(1) or window) + engine = SimpleNamespace(moe_offload_cache=SimpleNamespace(prefetch=prefetch)) + monkeypatch.setattr(ENV.MOE_PREFETCH_DEBUG, "value", debug) + # the freetoken loggers do not propagate, so caplog only sees them in some import orders + monkeypatch.setattr(engine_mod.logger, "info_rank0", lambda msg, *a, **k: lines.append(msg)) + engine_mod.Engine._emit_prefetch_stats(engine) + return "\n".join(lines), taken + + +def test_engine_reports_one_prefetch_window(monkeypatch): + from freetoken.engine.engine import MOE_STATS_INTERVAL + + window = torch.tensor([[0] * 6, [30, 12, 4, 10, 10, 40], [40, 20, 6, 10, 10, 30]]) + out, taken = _emit_prefetch(window, monkeypatch) + assert taken == [1] + assert f"MoE prefetch measure ({MOE_STATS_INTERVAL} decode steps): issued/layer=3.50, useful/layer=1.60" in out + assert "per layer" not in out + out, _ = _emit_prefetch(window, monkeypatch, debug=True) + assert "L1=1.20/3.00/4.00, L2=2.00/4.00/3.00" in out and "L0" not in out + assert _emit_prefetch(torch.zeros((3, 6), dtype=torch.int64), monkeypatch)[0] == "" + + +def test_engine_reports_the_on_mode_window(monkeypatch): + from freetoken.engine.engine import MOE_STATS_INTERVAL + + # issued, useful, resident_hits, calls, rows, misses, late, copied + window = torch.tensor([[0] * 8, [30, 12, 0, 10, 10, 18, 2, 10], [40, 20, 0, 10, 10, 20, 0, 6]]) + out, _ = _emit_prefetch(window, monkeypatch, mode="on") + assert f"MoE prefetch on ({MOE_STATS_INTERVAL} decode steps): issued/layer=3.50, useful/layer=1.60" in out + assert "misses/layer=1.90, coverage=0.457, late=0.125 of 16 copies" in out and "resident_hits" not in out + + +def test_summary_rates(): + from freetoken.moe.prefetch import format_per_layer, format_summary, summarize + + # layer 0 never predicts; layers 1 and 2 over 10 bs-1 steps + counters = torch.tensor([[0] * 6, [30, 12, 4, 10, 10, 40], [40, 20, 6, 10, 10, 30]]) + stats = summarize(counters) + assert stats["layer_calls"] == 20 and stats["tokens"] == 10 + assert stats["issued_per_layer"] == 3.5 and stats["useful_per_layer"] == 1.6 + assert stats["precision"] == 32 / 70 and stats["coverage"] == 32 / 70 + assert stats["useful_per_token"] == 3.2 and stats["resident_hits_per_layer"] == 0.5 + assert "useful/token=3.2" in format_summary(stats) + assert format_per_layer(counters) == "L1=1.20/3.00/4.00, L2=2.00/4.00/3.00" + + +# FREETOKEN_MOE_PREFETCH_VERIFY (moe/verify.py): the mode switch, where the verifier is built, and +# the engine's report; the checks themselves run on the qwen4_exp stacks in test_skeleton.py + + +def test_verify_mode_switch(): + from freetoken.moe.verify import resolve_mode + + assert [resolve_mode(m) for m in ("0", "", "off", "1", " True ", "full", "META")] == [ + "off", "off", "off", "full", "full", "full", "meta"] + with pytest.raises(ValueError): + resolve_mode("bytes") + + +@requires_cuda +def test_cache_builds_the_verifier_only_for_gpu_decode(monkeypatch): + from freetoken.env import ENV + from freetoken.moe.offload_cache import OffloadMoeCache + + def cache(**kw): + return OffloadMoeCache(num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"), **kw) + + monkeypatch.setattr(ENV.MOE_PREFETCH_VERIFY, "value", "0") + assert cache().verify is None and cache(verify_mode="off").verify is None + full = cache(verify_mode="full", prefetch_mode="on") + assert full.verify.mode == "full" and cache(verify_mode="meta").verify.mode == "meta" + # the verifier holds no reference to its cache + ref = weakref.ref(full) + del full + assert ref() is None + with pytest.raises(ValueError): + cache(verify_mode="full", decode_target="hybrid") + monkeypatch.setattr(ENV.MOE_PREFETCH_VERIFY, "value", "1") + for target in ("hybrid", "cpu"): + assert cache(decode_target=target).verify is None + assert cache().verify.mode == "full" + + +@requires_cuda +def test_engine_reports_the_verify_window_and_session(monkeypatch): + from types import SimpleNamespace + + from freetoken.engine import engine as engine_mod + from freetoken.engine.engine import MOE_STATS_INTERVAL + from freetoken.moe.offload_cache import OffloadMoeCache + from freetoken.moe.verify import BYTES_PRE, CHECKS, NUM_FIELDS, PRE_BAD + + cache = OffloadMoeCache(num_layers=3, num_experts=4, cache_size=4, device=torch.device("cuda"), verify_mode="full") + cache.set_bank_sources({ + "gate_up": [torch.randn(4, 32, 8, dtype=torch.bfloat16).pin_memory() for _ in range(3)], + "down": [torch.randn(4, 8, 16, dtype=torch.bfloat16).pin_memory() for _ in range(3)], + }) + v = cache.verify + lines = [] + for level in ("info_rank0", "warning_rank0"): + monkeypatch.setattr(engine_mod.logger, level, lambda msg, *a, level=level, **k: lines.append((level, msg))) + engine = SimpleNamespace(moe_offload_cache=cache) + + v.counters[:, CHECKS] = 256 + engine_mod.Engine._emit_verify_stats(engine) + assert lines == [("info_rank0", f"MoE verify full ({MOE_STATS_INTERVAL} decode steps): 768 layer calls checked, " + "meta_bad=0, pre_bad=0, post_bad=0, post_meta_bad=0; session bad=0")] + lines.clear() + v.counters[:, CHECKS] = 256 + v.counters[1, PRE_BAD] = 1 + # call 7, layer 1, row 3 of a top-2 call, expert 2, slot 1, bank down, before the GEMM; slot 1 holds + # layer 1's expert 2, usage 9 at lru step 9, demand-copied; bytes 64.. differ in 4 words + record = [7, 1, 3, 2, 1, 1, BYTES_PRE, 6, 1, 9, 9, 1, 64, 4] + assert len(record) == NUM_FIELDS + v.ring[0] = torch.tensor(record) + v.cursor.fill_(1) + v._top_k = 2 + engine_mod.Engine._emit_verify_stats(engine) + assert [level for level, _ in lines] == ["warning_rank0", "warning_rank0"] + assert "pre_bad=1" in lines[0][1] and "(meta/pre/post/post_meta by layer: L1=0/1/0/0); session bad=1" in lines[0][1] + assert lines[1][1] == ( + "MoE verify record: bytes_pre call=7 layer=1 row=3 (token 1, rank 1) expert=2 slot=1 bank=down " + "first_bad_byte=64 bad_words=4 | slot now holds L1/e2, slot_for_id=1, usage=9, lru_step=9, slot in: demand copy") + lines.clear() + engine_mod.Engine._emit_verify_stats(engine) # nothing checked and no new record: silent + assert lines == [] + session = v.report_session() + assert session[0] == ("warning", "MoE verify full (session): 1536 layer calls checked, meta_bad=0, pre_bad=1, " + "post_bad=0, post_meta_bad=0 (meta/pre/post/post_meta by layer: L1=0/1/0/0)") + assert session[1:] == [("warning", "MoE verify record: " + v.format_record(record))] + + +# FREETOKEN_MOE_SLOT_AUDIT (moe/slot_audit.py): the switch, where the auditor is built, and the +# engine's report; the recorders and the audit run on the qwen4_exp stacks in test_skeleton.py + + +def test_slot_audit_interval_switch(monkeypatch): + from freetoken.env import ENV + from freetoken.moe.slot_audit import resolve_interval + + monkeypatch.setattr(ENV.MOE_SLOT_AUDIT, "value", 64) + assert [resolve_interval(v) for v in (None, 0, 1, 8)] == [64, 0, 1, 8] + with pytest.raises(ValueError): + resolve_interval(-1) + + +@requires_cuda +def test_cache_builds_the_auditor_only_for_gpu_decode(monkeypatch): + from freetoken.env import ENV + from freetoken.moe.offload_cache import OffloadMoeCache + + def cache(**kw): + return OffloadMoeCache(num_layers=2, num_experts=4, cache_size=4, device=torch.device("cuda"), **kw) + + monkeypatch.setattr(ENV.MOE_SLOT_AUDIT, "value", 0) + assert cache().audit is None and cache(slot_audit=0).audit is None + on = cache(slot_audit=8, prefetch_mode="on") + assert on.audit.interval == 8 and on.slot_audit == 8 + ref = weakref.ref(on) # the auditor holds no reference to its cache + del on + assert ref() is None + with pytest.raises(ValueError): + cache(slot_audit=8, decode_target="hybrid") + monkeypatch.setattr(ENV.MOE_SLOT_AUDIT, "value", 16) + for target in ("hybrid", "cpu"): + assert cache(decode_target=target).audit is None + assert cache().audit.interval == 16 + monkeypatch.setenv("FREETOKEN_SKIP_FAST_INDEX_COPY", "1") + assert cache().audit is None, "copies that move nothing cannot be audited" + + +@requires_cuda +def test_engine_reports_the_audit_window_and_session(monkeypatch): + from types import SimpleNamespace + + from freetoken.engine import engine as engine_mod + from freetoken.engine.engine import MOE_STATS_INTERVAL + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache(num_layers=3, num_experts=4, cache_size=6, device=torch.device("cuda"), slot_audit=2) + cache.set_bank_sources({ + "gate_up": [torch.randn(4, 32, 8, dtype=torch.bfloat16).pin_memory() for _ in range(3)], + "down": [torch.randn(4, 8, 16, dtype=torch.bfloat16).pin_memory() for _ in range(3)], + }) + cache.reset() + lines = [] + for level in ("info_rank0", "warning_rank0"): + monkeypatch.setattr(engine_mod.logger, level, lambda msg, *a, level=level, **k: lines.append((level, msg))) + engine = SimpleNamespace(moe_offload_cache=cache) + engine_mod.Engine._emit_audit_stats(engine) + assert lines == [], "no audit ran: nothing to say" + + # one decode step of layer 1 routing experts 2 and 0: two demand installs and their copy + ids = torch.tensor([[2, 0]], dtype=torch.int32, device="cuda") + cache.ensure_experts(1, ids) + cache.copy_missing() + audit = cache.audit + assert not audit.end_decode_step(cache) and audit.end_decode_step(cache), "every 2 decode steps" + engine_mod.Engine._emit_audit_stats(engine) + assert lines == [("info_rank0", f"MoE slot audit ({MOE_STATS_INTERVAL} decode steps, every 2): 1 audits, 2 held slots " + "an audit (2 byte-checked), bad slots=0 (bytes=0, map=0, stale slot_for_id=0, owner out " + "of range=0), new=0, repeat=0, events=4; session: 1 audits, 0 bad (slot, owner) pairs")] + lines.clear() + slot = int(ids[0, 0]) + cache.bank_views()[1][slot].view(torch.uint8).view(-1)[8:12].bitwise_not_() + audit.scan(cache) + engine_mod.Engine._emit_audit_stats(engine) + assert [level for level, _ in lines] == ["warning_rank0", "warning_rank0"] + assert "1 audits, 2 held slots an audit (2 byte-checked), bad slots=1 (bytes=1," in lines[0][1] + assert lines[1][1].startswith( + f"MoE slot audit record: audit=1 step=0 lru=1 slot={slot} holds L1/e2 (slot_for_id={slot}, usage=1) | bytes differ in " + "down (first at down+8, 1 bad words) | diagnosis: demand install and demand copy of L1/e2 (#"), lines[1][1] + assert "then the bytes changed with no recorded write" in lines[1][1] + # the kernel ranks misses by id: e0 installs first (#0), e2 second (#1); the copy records after both + assert lines[1][1].endswith("history (1 map / 1 byte events in all): #1 demand install L1/e2 @step 0/lru 1 (then held " + "L1/e2, usage 1); #3 demand copy L1/e2 @step 0/lru 1 (slot held L1/e2, usage 1)"), lines[1][1] + lines.clear() + engine_mod.Engine._emit_audit_stats(engine) + assert lines == [] + session = audit.report_session() + assert session[0][0] == "warning" and session[0][1].startswith("MoE slot audit (session, every 2 decode steps): 2 audits") + assert session[1:] == [("warning", "MoE slot audit record: " + audit.format_record(audit.ring[0].tolist()))] diff --git a/tests/moe/test_scored_ensure.py b/tests/moe/test_scored_ensure.py index db734180b6..6bfa84a5ed 100644 --- a/tests/moe/test_scored_ensure.py +++ b/tests/moe/test_scored_ensure.py @@ -6,6 +6,7 @@ (experts.npy, logit_idx.npy, logit_val.npy) to also replay real decode rows. """ import os +from types import SimpleNamespace from pathlib import Path import numpy as np @@ -620,3 +621,321 @@ def test_an_all_pinned_call_never_stores_past_the_slot_arrays(policy): for name, buf in guarded.items(): assert torch.all(buf[s:] == 7), f"{name} written past slot {s - 1}" assert int(cache.num_indices) == 1 and 0 <= int(cache.evict_slots[0]) < s + + +# (h) FREETOKEN_MOE_PREFETCH=on: prefetch installs and low-priority keys ---------------------- + +_PF_L, _PF_E, _PF_S, _PF_K, _PF_W = 4, 64, 100, 6, 16 + + +def _pf_cache(policy, cache_size=_PF_S, num_layers=_PF_L, num_experts=_PF_E): + return _cache(policy, cache_size, num_layers, num_experts, prefetch_mode="on") + + +def _gpu_prefetch(cache, layer, sel, stats_row, width=_PF_W, budget=None): + """prefetch_ensure on the current stream; returns (src_ids, dst_slots), the plan it wrote.""" + from freetoken.moe.offload_kernels import prefetch_ensure_experts + + dev = cache.id_of_slot.device + query = torch.tensor(list(sel) + [-1] * (width - len(sel)), dtype=torch.int32, device=dev) + dst = torch.full((width,), -7, dtype=torch.int32, device=dev) + src = torch.full((width,), -7, dtype=torch.int32, device=dev) + num = torch.full((1,), 99, dtype=torch.int64, device=dev) + ready = torch.ones(1, dtype=torch.int32, device=dev) + prefetch_ensure_experts(cache, layer, query, dst, src, num, stats_row, ready, budget=budget) + n = int(num) + assert int(ready) == 0, "prefetch_ensure clears the copy's ready flag" + assert dst[n:].tolist() == [-7] * (width - n) + return src[:n].cpu().numpy().astype(np.int64), dst[:n].cpu().numpy().astype(np.int64), num + + +def _lowpri_slots(cache): + return set(((cache.id_of_slot >= 0) & (cache.usage == 0)).nonzero().view(-1).tolist()) + + +@_cuda +@pytest.mark.parametrize("policy", ["lru", "kd", "kdfb", "rule"]) +@pytest.mark.parametrize("rows_per_step", [1, 2]) +def test_prefetch_mode_matches_cpu_reference(policy, rows_per_step): + """prefetch_ensure and the low-priority demand ensure against the CPU reference victim for victim, + with adversarial candidates (next routed ids, resident ids, duplicates, padding, junk) and a + materialized prefill in between; the per-layer prefetch stats match the reference's counts.""" + from freetoken.moe.offload_kernels import ensure_experts_scored + from freetoken.moe.prefetch import CALLS, COPIED, ISSUED, LATE, MISSES, RESIDENT_HITS, ROWS, USEFUL + + rng = np.random.default_rng(31 + rows_per_step) + rows = _zipf_rows(40, rows_per_step, seed=33, num_layers=_PF_L, num_experts=_PF_E, top_k=_PF_K) + logits = _random_logits(rows, seed=34, num_experts=_PF_E) if policy == "rule" else None + cache, ref = _pf_cache(policy), _ref(policy, _PF_S, _PF_L, _PF_E) + dev = cache.id_of_slot.device + stats = torch.zeros((_PF_L, 8), dtype=torch.int64, device=dev) + want = np.zeros((_PF_L, 8), np.int64) + ready = torch.zeros(1, dtype=torch.int32, device=dev) + for r in range(rows.shape[0]): + if r == 20: + cache.materialize_layer(2) + ref.materialize(2) + assert not {s for s in _lowpri_slots(cache) if s < _PF_E}, "materialize installs at usage step" + for layer in range(_PF_L): + ids = rows[r, layer] + plan = None + if layer and rng.random() < 0.9: + routed = list(dict.fromkeys(ids.ravel().tolist())) + held = [e for e in range(_PF_E) if ref.slot_of_id[layer * _PF_E + e] >= 0] + sel = rng.choice(routed, size=min(3, len(routed)), replace=False).tolist() + sel += rng.choice(_PF_E, size=4).tolist() + held[:2] + sel[:1] + [-1] + rng.shuffle(sel) + pinned = set(np.flatnonzero((ref.usage == ref.step) & (ref.id_of_slot >= 0)).tolist()) + budget = [None, 0, 1, 2, 4][int(rng.integers(5))] + src, dst, num = _gpu_prefetch(cache, layer, sel, stats[layer], budget=budget) + want_src, want_dst = ref.prefetch(layer, sel, budget) + np.testing.assert_array_equal(src, want_src) + np.testing.assert_array_equal(dst, want_dst) + assert not pinned & set(dst.tolist()), "a prefetch evicted a slot of the last demand call" + assert set(dst.tolist()) <= _lowpri_slots(cache) + want[layer, ISSUED] += len(want_src) + _assert_same_tables(cache, ref) + plan = num + ready.fill_(int(rng.random() < 0.5)) # a copy that has (1) or has not (0) landed + q = torch.from_numpy(np.ascontiguousarray(ids, dtype=np.int32)).to(dev).view(-1) + lg = None if logits is None else torch.from_numpy(np.ascontiguousarray(logits[r, layer])).to(dev, torch.bfloat16) + lowpri_before = _lowpri_slots(cache) + ensure_experts_scored( + cache, layer, q, bump_tok=layer == 0, update_state=True, router_logits=lg, lowpri=True, + pf_count=None if plan is None else (stats[layer], ready, plan, rows_per_step), + ) + out, src, _ = ref.ensure(layer, ids, bump_tok=layer == 0, logits=None if lg is None else lg.float().cpu().numpy(), + thr=NEAR_MISS_THR, lowpri=True) + np.testing.assert_array_equal(q.cpu().numpy(), out) + _assert_same_tables(cache, ref) + # a demand hit makes a prefetched slot an ordinary resident, tagged with its owner's state + converted = lowpri_before & set(out.tolist()) + assert all(int(cache.usage[s]) == int(cache.step) for s in converted) + if policy != "lru": + assert all(int(cache.evict_slot_owner[s]) == int(cache.id_of_slot[s]) for s in converted) + if plan is not None: + n = int(plan) + want[layer, [USEFUL, CALLS, ROWS, MISSES]] += [ref.useful, 1, rows_per_step, src.size] + want[layer, LATE] += int(n > 0 and int(ready) == 0) + want[layer, COPIED] += int(n > 0) + got = stats.cpu().numpy() + np.testing.assert_array_equal(got, want) + assert got[:, USEFUL].sum() > 0 and got[:, LATE].sum() > 0 and got[:, RESIDENT_HITS].sum() == 0 + assert got[:, COPIED].sum() > got[:, LATE].sum() + + +@_cuda +@pytest.mark.parametrize("policy", ["lru", "rule"]) +@pytest.mark.parametrize("rows_per_step", [1, 2, 8]) +def test_prefetch_mode_matches_cpu_reference_at_production_geometry(policy, rows_per_step): + """The same victim-for-victim check at the deployment's shape: 48 layers x 512 experts, 1480 + slots (a 2048-wide, 8-warp victim scan), top-10, up to 8 rows (every graph size that forks) + and 32-wide prefetch queries, with small-prefill ensures (pinned since their first chunk) and a + materialized layer in between. No prefetch ever evicts a slot of the last demand call.""" + from freetoken.moe.offload_kernels import ensure_experts_scored + + num_layers, num_experts, size, width = L, E, 1480, 32 + rng = np.random.default_rng(71 + rows_per_step) + rows = _zipf_rows(4, rows_per_step, seed=72, num_layers=num_layers, num_experts=num_experts, top_k=TOP_K) + logits = _random_logits(rows, seed=73, num_experts=num_experts) if policy == "rule" else None + cache, ref = _pf_cache(policy, size, num_layers, num_experts), _ref(policy, size, num_layers, num_experts) + dev = cache.id_of_slot.device + stats = torch.zeros((num_layers, 8), dtype=torch.int64, device=dev) + ready = torch.zeros(1, dtype=torch.int32, device=dev) + installs = 0 + for r in range(rows.shape[0]): + if r == 2: + cache.materialize_layer(5) + ref.materialize(5) + for layer in range(num_layers): + ids = rows[r, layer] + plan = None + if layer: + routed = list(dict.fromkeys(ids.ravel().tolist())) + held = [e for e in range(num_experts) if ref.slot_of_id[layer * num_experts + e] >= 0] + sel = rng.choice(routed, size=min(6, len(routed)), replace=False).tolist() + sel += rng.choice(num_experts, size=16).tolist() + held[:4] + sel[:2] + [-1, -1] + sel = sel[:width] + rng.shuffle(sel) + pinned = set(np.flatnonzero((ref.usage == ref.step) & (ref.id_of_slot >= 0)).tolist()) + src, dst, plan = _gpu_prefetch(cache, layer, sel, stats[layer], width=width) + want_src, want_dst = ref.prefetch(layer, sel) + np.testing.assert_array_equal(src, want_src) + np.testing.assert_array_equal(dst, want_dst) + assert not pinned & set(dst.tolist()), "a prefetch evicted a slot of the last demand call" + installs += len(dst) + _assert_same_tables(cache, ref) + q = torch.from_numpy(np.ascontiguousarray(ids, dtype=np.int32)).to(dev).view(-1) + lg = None if logits is None else torch.from_numpy(np.ascontiguousarray(logits[r, layer])).to(dev, torch.bfloat16) + ensure_experts_scored( + cache, layer, q, bump_tok=layer == 0, update_state=True, router_logits=lg, lowpri=True, + pf_count=None if plan is None else (stats[layer], ready, plan, rows_per_step), + ) + out, _, _ = ref.ensure(layer, ids, bump_tok=layer == 0, logits=None if lg is None else lg.float().cpu().numpy(), + thr=NEAR_MISS_THR, lowpri=True) + np.testing.assert_array_equal(q.cpu().numpy(), out) + _assert_same_tables(cache, ref) + # a short prefill of one layer between decode steps: chunked, pinned since its first chunk + layer = int(rng.integers(num_layers)) + uniq = np.unique(rng.choice(num_experts, size=60, replace=False)) + pin = cache.step.clone() if policy != "lru" else None + ref_pin = ref.step if policy != "lru" else None + for start in range(0, uniq.size, 32): + part = torch.from_numpy(uniq[start : start + 32].astype(np.int32)).to(dev) + cache.ensure_experts(layer, part, update_state=False, pin_since=pin) + out, _, _ = ref.ensure(layer, uniq[start : start + 32], update_state=False, pin_since=ref_pin, lowpri=True) + np.testing.assert_array_equal(part.cpu().numpy(), out) + _assert_same_tables(cache, ref) + assert installs > 0 and int((ref.id_of_slot >= 0).sum()) == size, "the cache must be full and evicting" + + +def _plant(cache, ref, slot, flat_id, usage): + for t_id, t_slot in ((cache.id_of_slot, cache.slot_for_id.view(-1)), (ref.id_of_slot, ref.slot_of_id)): + t_id[slot] = flat_id + t_slot[flat_id] = slot + cache.usage[slot] = usage + ref.usage[slot] = usage + + +@_cuda +@pytest.mark.parametrize("policy", ["lru", "kd", "kdfb", "rule"]) +def test_lowpri_band_sits_between_empty_slots_and_every_resident(policy): + """An unconsumed prefetch is evicted after every empty slot and before any resident, even the + coldest resident the score caps allow, whatever the slot order.""" + from freetoken.moe import scored_ensure as se + from freetoken.moe.offload_kernels import ensure_experts_scored + + num_layers, num_experts, size = 2, 16, 16 + cache, ref = _pf_cache(policy, size, num_layers, num_experts), _ref(policy, size, num_layers, num_experts) + dev = cache.id_of_slot.device + cache.step.fill_(5) + ref.step = 5 + # slots 0-12: layer-0 residents at the oldest usage 1 with never-seen state (the lowest possible + # score); slot 13 empty; slots 14, 15 unconsumed prefetches + for slot in range(13): + _plant(cache, ref, slot, slot, 1) + _plant(cache, ref, 14, 13, 0) + _plant(cache, ref, 15, 14, 0) + if policy != "lru": + cache.evict_slot_owner.fill_(-1) + victims = [] + for e in (15, 16 + 0, 16 + 1, 16 + 2): # four misses, one per call + layer, expert = divmod(e, num_experts) + q = torch.tensor([expert], dtype=torch.int32, device=dev) + ensure_experts_scored(cache, layer, q, bump_tok=False, update_state=True, lowpri=True) + ref.ensure(layer, [expert], lowpri=True) + victims.append(int(cache.evict_slots[0])) + _assert_same_tables(cache, ref) + assert victims[:3] == [13, 14, 15] and victims[3] < 13 + # the band's keys are below 2 << SLOT_BITS; the caps' worst resident score keeps score + 2^31 >= 2 + beta_step, w_q4, _, _ = se.score_params(se.MAX_BETA, se.MAX_W, 1, 1) + worst = -((se.K_MAX << se.Q) + 1 * beta_step) + ((w_q4 * se.LC_FLOOR) >> 4) + assert worst + se.SCORE_BIAS >= 2 + + +@_cuda +@pytest.mark.parametrize("policy", ["lru", "rule"]) +def test_prefetch_budget_keeps_the_first_non_resident_candidates(policy): + """The install budget counts only non-resident candidates, in query order: resident ids, padding + and duplicates ahead of them take none of it, and ISSUED counts the installs.""" + num_experts = 16 + cache, ref = _pf_cache(policy, num_experts, 2, num_experts), _ref(policy, num_experts, 2, num_experts) + stats = torch.zeros(8, dtype=torch.int64, device="cuda") + cache.ensure_experts(0, torch.tensor([0, 1], dtype=torch.int32, device="cuda")) # step 0 pins every slot + ref.ensure(0, np.arange(2), bump_tok=True, lowpri=True) + src, _, _ = _gpu_prefetch(cache, 1, [4, 6], stats) + ref.prefetch(1, [4, 6]) + assert src.tolist() == [4, 6] + sel = [4, -1, 6, 9, 9, 3, 7, 1] + src, dst, _ = _gpu_prefetch(cache, 1, sel, stats, budget=2) + assert src.tolist() == [9, 3] and int(stats[0]) == 4 + np.testing.assert_array_equal(np.stack(ref.prefetch(1, sel, 2)), np.stack([src, dst])) + _assert_same_tables(cache, ref) + src, _, _ = _gpu_prefetch(cache, 1, [5, 8], stats, budget=0) + assert src.size == 0 and int(stats[0]) == 4 + + +@_cuda +@pytest.mark.parametrize("policy", ["lru", "rule"]) +def test_prefetch_installs_only_into_evictable_slots(policy): + """With fewer evictable slots than candidates a prefetch installs the first candidates only, and + with every slot pinned by the last demand call it installs nothing (no all-pinned fallback).""" + num_experts = 16 + cache, ref = _pf_cache(policy, num_experts, 2, num_experts), _ref(policy, num_experts, 2, num_experts) + stats = torch.zeros(8, dtype=torch.int64, device="cuda") + ids = np.arange(12) + cache.ensure_experts(0, torch.from_numpy(ids.astype(np.int32)).cuda()) + ref.ensure(0, ids, bump_tok=True, lowpri=True) + pinned = cache.id_of_slot.clone() + src, dst, _ = _gpu_prefetch(cache, 1, [9, 3, 7, 1, 5, 11, 2], stats) + assert src.tolist() == [9, 3, 7, 1] and sorted(dst.tolist()) == [12, 13, 14, 15] + np.testing.assert_array_equal(np.stack(ref.prefetch(1, [9, 3, 7, 1, 5, 11, 2])), np.stack([src, dst])) + assert torch.equal(cache.id_of_slot[:12], pinned[:12]) and int(stats[0]) == 4 + all16 = np.arange(16) + cache.ensure_experts(0, torch.from_numpy(all16.astype(np.int32)).cuda()) + before = [t.clone() for t in (cache.id_of_slot, cache.slot_for_id, cache.usage)] + src, dst, num = _gpu_prefetch(cache, 1, [0, 1, 2], stats) + assert int(num) == 0 and src.size == 0 and int(stats[0]) == 4 + for a, b in zip(before, (cache.id_of_slot, cache.slot_for_id, cache.usage)): + assert torch.equal(a, b) + + +@_cuda +@pytest.mark.parametrize("policy", ["kd", "kdfb", "rule"]) +def test_prefetch_scores_victims_from_the_previous_layers_position(policy): + """prefetch_ensure(L) runs right after layer L-1's demand call and scores victims from there: + layer L's residents are next up (d = 1), so one used last token outlives an older resident of + a later layer instead of looking a whole pass away.""" + num_layers, num_experts, size = 8, 16, 16 + cache, ref = _pf_cache(policy, size, num_layers, num_experts), _ref(policy, size, num_layers, num_experts) + cache.step.fill_(7) + ref.step = 7 + cache.evict_tok.fill_(10) + ref.tok = 10 + cache.evict_slot_owner.fill_(-1) + # slots 0-13: layer 1's routed experts (pinned), slot 14: a layer-2 expert used last token, + # slot 15: a layer-5 expert last used three tokens ago; all with the same (never-counted) use count + planted = [(slot, 16 + slot, 7, 10) for slot in range(14)] + [(14, 2 * 16 + 3, 6, 9), (15, 5 * 16 + 4, 5, 7)] + for slot, flat, usage, last in planted: + _plant(cache, ref, slot, flat, usage) + cache.evict_last_tok[flat] = last + ref.last_tok[flat] = last + stats = torch.zeros(8, dtype=torch.int64, device="cuda") + src, dst, _ = _gpu_prefetch(cache, 2, [9], stats) + assert src.tolist() == [9] and dst.tolist() == [15], "the prefetch evicted the target layer's fresh resident" + np.testing.assert_array_equal(np.stack(ref.prefetch(2, [9])), np.stack([src, dst])) + _assert_same_tables(cache, ref) + + +@_cuda +@pytest.mark.parametrize("rows_per_step", [1, 2]) +def test_lowpri_lru_without_prefetches_is_flashlib(rows_per_step): + """Prefetch on routes LRU through the vendored kernel with the low-priority key; while no slot is + a prefetch it evicts exactly what flashlib's lru_ensure does.""" + num_layers, num_experts, size = 8, 64, 200 + on = _pf_cache("lru", size, num_layers, num_experts) + off = _cache("lru", size, num_layers, num_experts) + rows = _zipf_rows(30, rows_per_step, seed=41, num_layers=num_layers, num_experts=num_experts) + ids = torch.from_numpy(rows.astype(np.int32)).cuda() + for r in range(rows.shape[0]): + for layer in range(num_layers): + a, b = ids[r, layer].clone(), ids[r, layer].clone() + on.ensure_experts(layer, a) + off.ensure_experts(layer, b) + assert torch.equal(a, b) + assert torch.equal(on.evict_slots[: int(on.num_indices)], off.evict_slots[: int(off.num_indices)]) + for name in ("slot_for_id", "id_of_slot", "usage", "step"): + assert torch.equal(getattr(on, name), getattr(off, name)), name + + +def test_lru_policy_leaves_flashlib_only_when_prefetch_is_on(monkeypatch): + import freetoken.moe.offload_kernels as ok + + calls = [] + monkeypatch.setattr(ok, "lru_ensure", lambda *a, **kw: calls.append("flashlib")) + monkeypatch.setattr(ok, "ensure_experts_scored", lambda *a, **kw: calls.append(("scored", kw["lowpri"]))) + for mode in ("off", "measure", "on"): + cache = _cache("lru", 12, 2, 8, device="cpu", prefetch_mode=mode) + cache.prefetch = SimpleNamespace(mode=mode, count_args=lambda layer: None) if mode != "off" else None + cache.ensure_experts(1, torch.zeros(2, dtype=torch.int32)) + assert calls == ["flashlib", "flashlib", ("scored", True)] From 524f438ed912ef46917558d7da0e78bb2cccb2a1 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 1 Oct 2026 23:18:22 -0400 Subject: [PATCH 30/94] test(engine): give the kv quant config fake the rope table size #573's max_seq_len reads Assisted-by: Claude Opus 5.5 --- tests/engine/test_kv_quant_config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/engine/test_kv_quant_config.py b/tests/engine/test_kv_quant_config.py index ba4196be09..1f460166d1 100644 --- a/tests/engine/test_kv_quant_config.py +++ b/tests/engine/test_kv_quant_config.py @@ -39,7 +39,7 @@ def _model_config(kind): has_swa_attention=False, has_linear_attention=False, num_layers=4, - rotary_config=SimpleNamespace(max_position=1024), + rotary_config=SimpleNamespace(max_position=1024, table_positions=1024), ) specs = { "full": (_spec("full", AttnType.FULL),), From 206576a515f131b8031594c8a8ad09e9f256dc8f Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 06:55:23 -0400 Subject: [PATCH 31/94] fix(engine): hand the rendezvous listen fd to TCPStore instead of keeping its socket TCPStore owns master_listen_fd and closes it with the store; #461 also kept the Python socket on the engine, whose finalizer then closed whatever file reused that fd number after destroy_process_group at shutdown. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit b54cced6af48bffcce88c280124692d54c69f792) --- python/freetoken/engine/engine.py | 8 ++--- tests/engine/test_distributed_store.py | 49 +++++++++++++++++++++++++- 2 files changed, 52 insertions(+), 5 deletions(-) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index d450e2a187..7a0b49c172 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -673,12 +673,12 @@ def _make_distributed_store(self, config: EngineConfig) -> torch.distributed.Sto sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) sock.bind(("127.0.0.1", config.distributed_port)) sock.listen(1024) - # TCPStore uses this fd directly rather than duplicating it, so it must outlive the - # store; kept alive on self, closed only when the process (and the store) exits. - self._distributed_listen_socket = sock + self._distributed_listen_addr = sock.getsockname() + # TCPStore takes ownership of the fd and closes it with the store; detach so the + # Python socket never closes that number again after it has been reused. return torch.distributed.TCPStore( "127.0.0.1", config.distributed_port, config.tp_info.size, - is_master=True, timeout=timeout, multi_tenant=True, master_listen_fd=sock.fileno(), + is_master=True, timeout=timeout, multi_tenant=True, master_listen_fd=sock.detach(), ) def _init_communication(self, config: EngineConfig) -> torch.distributed.ProcessGroup: diff --git a/tests/engine/test_distributed_store.py b/tests/engine/test_distributed_store.py index f70be12668..4fc2473762 100644 --- a/tests/engine/test_distributed_store.py +++ b/tests/engine/test_distributed_store.py @@ -8,7 +8,9 @@ from __future__ import annotations +import gc import multiprocessing as mp +import os from datetime import timedelta from types import SimpleNamespace @@ -38,7 +40,7 @@ def _run_rank(rank: int, port: int, result_q: mp.Queue) -> None: ) total = torch.tensor([float(rank)]) dist.all_reduce(total) - bind_host = engine._distributed_listen_socket.getsockname()[0] if rank == 0 else None + bind_host = engine._distributed_listen_addr[0] if rank == 0 else None result_q.put((rank, total.item(), bind_host)) dist.destroy_process_group() @@ -64,3 +66,48 @@ def test_two_ranks_rendezvous_with_loopback_only_store(): # the master's listening socket was pre-bound to loopback, not the wildcard # address the C10d TCPStore server binds to on its own assert results[0][1] == "127.0.0.1" + + +def _teardown_rank(port: int, result_q: mp.Queue) -> None: + listen_fds = [] + real_store = dist.TCPStore + + def recording_store(*args, **kwargs): + listen_fds.append(kwargs.get("master_listen_fd")) + return real_store(*args, **kwargs) + + dist.TCPStore = recording_store + engine = Engine.__new__(Engine) + config = SimpleNamespace( + tp_info=DistributedInfo(0, 1), distributed_timeout=TIMEOUT.total_seconds(), + distributed_port=port, + ) + store = engine._make_distributed_store(config) + dist.init_process_group(backend="gloo", rank=0, world_size=1, timeout=TIMEOUT, store=store) + dist.destroy_process_group() + del store + gc.collect() + # the listening fd number is reused by an unrelated file, as a later open() would + fd = listen_fds[0] + r, w = os.pipe() + os.dup2(r, fd) + del engine + gc.collect() + try: + os.fstat(fd) + result_q.put("open") + except OSError as exc: + result_q.put(f"closed: {exc}") + + +def test_the_engine_does_not_close_the_listen_fd_the_store_owns(): + # TCPStore takes ownership of master_listen_fd and closes it itself; an engine that also + # kept the Python socket closed whatever file reused that number when it was freed + ctx = mp.get_context("spawn") + result_q = ctx.Queue() + proc = ctx.Process(target=_teardown_rank, args=(29513, result_q)) + proc.start() + outcome = result_q.get(timeout=30) + proc.join(timeout=30) + assert proc.exitcode == 0 + assert outcome == "open" From 155185c74598be20c28fb64cba7a2e224446093b Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 06:52:14 -0400 Subject: [PATCH 32/94] fix(scheduler): trim the stop-match EOS only at the delivered output limit #464 used `not req.can_decode`, but since #546 overlap runs device_len a step ahead of the host, so a non-terminal EOS was dropped from the matched text one step early. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 1814a3acc1848366a59d16e40d25c3d66a3be301) --- python/freetoken/scheduler/scheduler.py | 4 +- tests/scheduler/test_match_stop_str.py | 71 +++++++++++++++++++++++++ 2 files changed, 74 insertions(+), 1 deletion(-) create mode 100644 tests/scheduler/test_match_stop_str.py diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index 01f1678761..895f012c75 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -478,9 +478,11 @@ def _match_stop_str(self, req: Req) -> str | None: prompt_len = req.max_device_len - req.output_len end = len(req.input_ids) # The frontend omits a terminal EOS, including at the output limit with ignore_eos. + # The limit is the delivered length, as in hit_length: under overlap device_len (and + # so can_decode) runs a step ahead of the host. if ( end > prompt_len - and not req.can_decode + and end >= req.max_device_len and int(req.input_ids[-1]) in self.eos_token_ids ): end -= 1 diff --git a/tests/scheduler/test_match_stop_str.py b/tests/scheduler/test_match_stop_str.py new file mode 100644 index 0000000000..351509b759 --- /dev/null +++ b/tests/scheduler/test_match_stop_str.py @@ -0,0 +1,71 @@ +"""Scheduler-side stop-string matching (Scheduler._match_stop_str) against the same +incrementally decoded text the frontend streams.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import torch + +from freetoken.core import Req, SamplingParams +from freetoken.scheduler.scheduler import Scheduler + +EOS = 0 + + +class _ByteTok: + """Byte-level ids like Qwen's BPE fallback: a 4-byte emoji is four tokens. 0 is EOS.""" + + def decode(self, ids, skip_special_tokens=False): + return b"".join(b"" if t == EOS else bytes([t]) for t in ids).decode( + "utf-8", errors="replace" + ) + + +def _req(prompt: list[int], output_len: int, **sp) -> Req: + return Req( + input_ids=torch.tensor(prompt, dtype=torch.int32), + table_idx=0, + cached_len=0, + output_len=output_len, + uid=1, + sampling_params=SamplingParams(**sp), + cache_handle=None, + ) + + +def _match(req: Req) -> str | None: + stub = SimpleNamespace(tokenizer=_ByteTok(), eos_token_ids={EOS}) + return Scheduler._match_stop_str(stub, req) + + +def _feed(req: Req, ids: list[int]) -> list[str | None]: + seen = [] + for t in ids: + req.append_host(torch.tensor([t], dtype=torch.int32)) + req.complete_one() + seen.append(_match(req)) + return seen + + +def test_a_stop_spanning_more_tokens_than_characters_is_matched(): + dna = list("\N{DNA DOUBLE HELIX}".encode()) # one character, four byte tokens + req = _req([65, 66], 32, stop_strs=["\N{DNA DOUBLE HELIX}"]) + text = list(b"Here is DNA ") + dna + list(b" and more") + seen = _feed(req, text) + first = next(i for i, m in enumerate(seen) if m is not None) + assert first == len(b"Here is DNA ") + len(dna) - 1 + + +def test_an_eos_short_of_the_delivered_limit_stays_in_the_matched_text(): + # Overlap scheduling advances device_len ahead of the host: the budget is spent on the + # device (can_decode is False) while the host holds one token less than the limit. That + # EOS is not terminal, so the frontend decodes it and the scheduler must too. + req = _req([65], 3, stop_strs=["zz"], ignore_eos=True) + for t in (65, EOS): + req.append_host(torch.tensor([t], dtype=torch.int32)) + while req.can_decode: + req.complete_one() + assert req.input_ids.numel() < req.max_device_len + assert _match(req) is None + assert req.stop_decode_status.decoded_ids == [65, EOS] From c9c883e167ccc78d9e8b2f0a25d9fbb6e5200462 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 06:52:27 -0400 Subject: [PATCH 33/94] test(tokenizer): cover a detokenize batch that carries one uid more than once Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 577427a2cecf047ac64616854a51aa338197336b) --- tests/tokenizer/test_detokenize_extras.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/tokenizer/test_detokenize_extras.py b/tests/tokenizer/test_detokenize_extras.py index 5617a10537..4717617b9a 100644 --- a/tests/tokenizer/test_detokenize_extras.py +++ b/tests/tokenizer/test_detokenize_extras.py @@ -79,3 +79,19 @@ def test_reasoning_tokens_count_up_to_the_end_tag(): assert reasoning == 3 # b, c and the tag itself _, none = _run(manager, 2, [1, 2, EOS]) assert none == 0 + + +def test_a_batch_with_one_uid_twice_advances_its_state_in_order(): + # each message must see the decode state the previous one left, and the reasoning + # counts must stay aligned with the texts + manager = DetokenizeManager(_Tok(), frozenset({EOS}), think_end_id=THINK_END) + msgs = [ + DetokenizeMsg(uid=1, next_token=1, finished=False), + DetokenizeMsg(uid=1, next_token=THINK_END, finished=False), + DetokenizeMsg(uid=2, next_token=3, finished=False), + DetokenizeMsg(uid=1, next_token=2, finished=True), + ] + texts, counts = manager.detokenize_with_meta(msgs) + assert texts == ["b", "", "d", "c"] + assert counts == [0, 0, 0, 2] + assert manager.detokenize([DetokenizeMsg(uid=2, next_token=4, finished=True)]) == ["e"] From c6e804ba96e94480599d35f1654256ae421371bf Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 06:53:13 -0400 Subject: [PATCH 34/94] fix(scheduler): match stop strings in the text the request's skip_special_tokens decodes #464's scheduler-side DecodeStatus decoded with special tokens while #393 lets a request skip them on the frontend, so a stop spanning a special token matched different text. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit fc2392b5367e841b329454bddf33514f808417e1) --- python/freetoken/scheduler/scheduler.py | 7 ++++++- python/freetoken/tokenizer/detokenize.py | 7 +++++-- tests/scheduler/test_match_stop_str.py | 25 +++++++++++++++++++----- 3 files changed, 31 insertions(+), 8 deletions(-) diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index 895f012c75..cb6085d344 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -490,7 +490,12 @@ def _match_stop_str(self, req: Req) -> str | None: return None if req.stop_decode_status is None: req.stop_decode_status = DecodeStatus( - decoded_ids=[], decoded_str="", read_offset=0, surr_offset=0, sent_offset=0 + decoded_ids=[], + decoded_str="", + read_offset=0, + surr_offset=0, + sent_offset=0, + skip_special=req.sampling_params.skip_special_tokens, ) state = req.stop_decode_status state.decoded_ids.extend(req.input_ids[prompt_len + len(state.decoded_ids) : end].tolist()) diff --git a/python/freetoken/tokenizer/detokenize.py b/python/freetoken/tokenizer/detokenize.py index b4c0b04733..2eb3029df0 100644 --- a/python/freetoken/tokenizer/detokenize.py +++ b/python/freetoken/tokenizer/detokenize.py @@ -84,9 +84,12 @@ def update(self, read_str: str, surr_str: str) -> str: return self.decoded_str + find_printable_text(new_text) def decode(self, tokenizer: PreTrainedTokenizerBase) -> str: + skip = self.skip_special return self.update( - tokenizer.decode(self.decoded_ids[self.surr_offset :]), - tokenizer.decode(self.decoded_ids[self.surr_offset : self.read_offset]), + tokenizer.decode(self.decoded_ids[self.surr_offset :], skip_special_tokens=skip), + tokenizer.decode( + self.decoded_ids[self.surr_offset : self.read_offset], skip_special_tokens=skip + ), ) diff --git a/tests/scheduler/test_match_stop_str.py b/tests/scheduler/test_match_stop_str.py index 351509b759..e9a5c07da5 100644 --- a/tests/scheduler/test_match_stop_str.py +++ b/tests/scheduler/test_match_stop_str.py @@ -10,16 +10,22 @@ from freetoken.core import Req, SamplingParams from freetoken.scheduler.scheduler import Scheduler -EOS = 0 +EOS, SPECIAL = 0, 1 class _ByteTok: - """Byte-level ids like Qwen's BPE fallback: a 4-byte emoji is four tokens. 0 is EOS.""" + """Byte-level ids like Qwen's BPE fallback: a 4-byte emoji is four tokens. 0 is EOS, + 1 a special marker that skip_special_tokens drops.""" def decode(self, ids, skip_special_tokens=False): - return b"".join(b"" if t == EOS else bytes([t]) for t in ids).decode( - "utf-8", errors="replace" - ) + def piece(t): + if t == EOS: + return b"" + if t == SPECIAL: + return b"" if skip_special_tokens else b"" + return bytes([t]) + + return b"".join(piece(t) for t in ids).decode("utf-8", errors="replace") def _req(prompt: list[int], output_len: int, **sp) -> Req: @@ -69,3 +75,12 @@ def test_an_eos_short_of_the_delivered_limit_stays_in_the_matched_text(): assert req.input_ids.numel() < req.max_device_len assert _match(req) is None assert req.stop_decode_status.decoded_ids == [65, EOS] + + +def test_stops_match_the_text_decoded_with_the_requests_skip_special_tokens(): + # the frontend streams "ab" for a request that skips special tokens; the scheduler has + # to match the stop in that same text, not in "ab" + req = _req([65], 8, stop_strs=["ab"], skip_special_tokens=True) + assert _feed(req, [ord("a"), SPECIAL, ord("b")]) == [None, None, "ab"] + req = _req([65], 8, stop_strs=["ab"]) + assert _feed(req, [ord("a"), SPECIAL, ord("b")]) == [None, None, None] From 4fa89f848cf69bc5c0a71d476f6fdf5617032bd3 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 05:43:14 -0400 Subject: [PATCH 35/94] perf(ple): skip the device n-gram hash for the disk table, which reads only its shape DiskRowTable hashes on the host in its fill; its lookup and prefetch read only the row_ids shape, so every prefill chunk and captured decode step ran the hash for nothing. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 52a18015ffff759a9f2b7733954bd9998b900f32) --- python/freetoken/models/qwen4_exp/ple.py | 17 ++++++-- python/freetoken/models/qwen4_exp/ple_disk.py | 3 ++ tests/models/qwen4_exp/test_ple.py | 40 +++++++++++++++++++ 3 files changed, 57 insertions(+), 3 deletions(-) diff --git a/python/freetoken/models/qwen4_exp/ple.py b/python/freetoken/models/qwen4_exp/ple.py index 44973bc0a6..22bda9e6c8 100644 --- a/python/freetoken/models/qwen4_exp/ple.py +++ b/python/freetoken/models/qwen4_exp/ple.py @@ -581,8 +581,19 @@ def row_ids_reference(self, meta: PLEMetadata) -> torch.Tensor: blocks.append(head_ids + self.ngram_heads_offsets[start:end]) return torch.cat(blocks, dim=-1) + def lookup_row_ids(self, meta: PLEMetadata) -> torch.Tensor: + """``row_ids`` as the attached table's ``prefetch``/``lookup`` take them. A table that + hashes on the host (``reads_row_ids = False``) reads only their shape: skip the hash.""" + if not getattr(self.table, "reads_row_ids", True): + return torch.empty( + (meta.input_ids.numel(), self.num_heads), + dtype=torch.int64, + device=meta.input_ids.device, + ) + return self.row_ids(meta) + def forward(self, meta: PLEMetadata, out: torch.Tensor | None = None) -> torch.Tensor: - return self.table.lookup(self.row_ids(meta), out) + return self.table.lookup(self.lookup_row_ids(meta), out) class _DepthwiseConv1d(BaseOP): @@ -680,7 +691,7 @@ def start_prefetch(self, batch: Batch, meta: PLEMetadata | None = None) -> None: """Hash this forward's n-grams and start the table gather on the side stream.""" if meta is None: meta = build_ple_metadata(batch, self.args, batch.input_ids.device) - row_ids = self.ple_embedding.row_ids(meta) + row_ids = self.ple_embedding.lookup_row_ids(meta) self._pending = (meta, row_ids) self.ple_embedding.table.prefetch(row_ids) @@ -701,7 +712,7 @@ def forward( elif pending is not None and pending[0] is meta: row_ids = pending[1] if row_ids is None: - row_ids = self.ple_embedding.row_ids(meta) + row_ids = self.ple_embedding.lookup_row_ids(meta) embeddings = self.ple_embedding.table.lookup(row_ids).to(R.dtype) key = self.norm_key.forward(self.key_proj.forward(embeddings)) diff --git a/python/freetoken/models/qwen4_exp/ple_disk.py b/python/freetoken/models/qwen4_exp/ple_disk.py index d2eaa99df8..56cfcb922b 100644 --- a/python/freetoken/models/qwen4_exp/ple_disk.py +++ b/python/freetoken/models/qwen4_exp/ple_disk.py @@ -152,6 +152,9 @@ def cancel(self) -> None: class DiskRowTable: """``PLETableBackend`` whose rows are read from disk per fill (--ple-backend disk).""" + # the fill hashes on the host; lookup/prefetch read only the row_ids shape + reads_row_ids = False + def __init__( self, source: PleRowSource, diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index 02abaee1f4..954d898945 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -339,6 +339,46 @@ def test_the_fused_hash_refuses_a_geometry_it_cannot_address(ctx_len, heads_per_ ) + +def test_a_host_hashing_table_skips_the_device_hash(monkeypatch): + """The disk table hashes on the host and reads only the row_ids shape, so neither the + prefetch nor the lookup path runs the device hash for it.""" + from freetoken.models.qwen4_exp.ple import ZeroTable + from freetoken.models.qwen4_exp.ple_disk import DiskRowTable + + assert DiskRowTable.reads_row_ids is False + config = _config() + args = config.qwen4_args + seen = [] + + class HostHashedTable(ZeroTable): + reads_row_ids = False + + def lookup(self, row_ids, out=None): + seen.append(("lookup", tuple(row_ids.shape), row_ids.dtype)) + return super().lookup(row_ids, out) + + def prefetch(self, row_ids): + seen.append(("prefetch", tuple(row_ids.shape), row_ids.dtype)) + + layer = _make_layer(config, table=HostHashedTable(_padded_vocab(args), args.ngram_head_dim)) + embedding = layer.ple_embedding + + def no_hash(*a, **k): + raise AssertionError("the device hash ran for a table that never reads it") + + monkeypatch.setattr(embedding, "row_ids", no_hash) + sequences, contexts = [[3, 4, EOS, 5], [2, 7]], [[EOS, EOS], [21, 22]] + meta = _meta(sequences, contexts) + total = sum(len(s) for s in sequences) + states = torch.zeros(len(sequences), args.ple_state_width, args.ple_conv_state_len) + + layer.start_prefetch(None, meta) + _forward(layer, torch.randn(total, layer.hc_count * layer.hidden_size), meta, states) + embedding.forward(meta) + rows = (total, embedding.num_heads) + assert seen == [("prefetch", rows, torch.int64), ("lookup", rows, torch.int64), ("lookup", rows, torch.int64)] + # -------------------------------------------------------------------------------------- # table backends # -------------------------------------------------------------------------------------- From 5221ce0bf0c1648cf2c4e4e45f31e320466026ec Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 10:24:20 -0400 Subject: [PATCH 36/94] fix(ple): give each graph fill its own readback event and keep every pending fill A shared event re-recorded by the next step before the filler reached synchronize() made fill N wait on graph N, which waits on fill N; and a fill that failed after the next step submitted was dropped, so its graph ran on stale rows without an error. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 79988ccdf596e54b37817971628de84c802f3a58) --- python/freetoken/models/qwen4_exp/ple_disk.py | 34 ++++--- .../qwen4_exp/test_ple_disk_fill_order.py | 95 ++++++++++++++++++- 2 files changed, 117 insertions(+), 12 deletions(-) diff --git a/python/freetoken/models/qwen4_exp/ple_disk.py b/python/freetoken/models/qwen4_exp/ple_disk.py index 56cfcb922b..05b18d5ce7 100644 --- a/python/freetoken/models/qwen4_exp/ple_disk.py +++ b/python/freetoken/models/qwen4_exp/ple_disk.py @@ -8,6 +8,7 @@ import os import queue import threading +from collections import deque from concurrent.futures import Future from contextlib import contextmanager from dataclasses import dataclass @@ -146,6 +147,9 @@ def cancel(self) -> None: self._future.result() except BaseException: # noqa: BLE001 -- the launch error is the one to surface pass + pending = self._table._pending_fills() + if self._future in pending: + pending.remove(self._future) # its outcome was consumed here self._table._flag.zero_() @@ -216,11 +220,10 @@ def __init__( from freetoken.kernel.row_store import probe_wait_sync self._wait_sync = probe_wait_sync(os.getenv(_SYNC_ENV, "auto"), self._device) - # one flag for all graphs: the readback event orders a fill after the previous graph, so signals never overlap + # one flag for all graphs: each fill's readback event orders it after the previous graph, so signals never overlap self._flag = alloc_pinned_tensor(1, dtype=torch.int64) self._flag.zero_() self._token_readback = alloc_pinned_tensor(max_graph_rows, dtype=torch.int32) - self._readback_event = torch.cuda.Event() sync = "wait-sync" if self._wait_sync else "launch-gating" logger.info_rank0(f"PLE disk backend: {self._store.io_backend()}, {sync}") @@ -240,12 +243,17 @@ def close(self) -> None: def _current_stream(self): return torch.cuda.current_stream(self._device) + def _pending_fills(self) -> deque: + return self.__dict__.setdefault("_pending_fill_queue", deque()) + def _raise_failed_fill(self) -> None: - pending = self.__dict__.pop("_pending_fill", None) - if pending is not None and pending.done() and pending.exception() is not None: - raise pending.exception() - if pending is not None and not pending.done(): - self._pending_fill = pending + """Surface a finished graph fill's error on the engine thread, oldest first; a failed + fill signalled its WAIT, so that graph ran on stale rows.""" + pending = self._pending_fills() + while pending and pending[0].done(): + exc = pending.popleft().exception() + if exc is not None: + raise exc def _fill_now(self, runs: Sequence[torch.Tensor], *, graph: bool) -> None: """A fill that completes before its launch, on the store's thread (after any queued graph fill).""" @@ -288,12 +296,15 @@ def host_fill_batch(self, batch: Batch, use_graph: bool) -> _PendingFill | None: if use_graph and self._wait_sync: bs = batch.padded_size self._token_readback[:bs].copy_(batch.input_ids, non_blocking=True) - self._readback_event.record(self._current_stream()) + # one event per fill: a shared one re-recorded by the next step before this job + # reached synchronize() would wait on a graph parked on this job's own WAIT + readback = torch.cuda.Event() + readback.record(self._current_stream()) cancelled = threading.Event() def _job() -> None: try: - self._readback_event.synchronize() + readback.synchronize() tokens = self._token_readback[:bs].to(torch.int64).tolist() if cancelled.is_set(): return @@ -306,8 +317,9 @@ def _job() -> None: signal(self._flag) raise - self._pending_fill = self._filler().submit(_job) - return _PendingFill(self, self._pending_fill, cancelled) + future = self._filler().submit(_job) + self._pending_fills().append(future) + return _PendingFill(self, future, cancelled) # launch-gating: this D2H is the step's readback and orders the fill after sampling tokens = batch.input_ids.to("cpu").to(torch.int64).tolist() runs = [torch.tensor([*ctx, t], dtype=torch.int64) for ctx, t in zip(contexts, tokens)] diff --git a/tests/models/qwen4_exp/test_ple_disk_fill_order.py b/tests/models/qwen4_exp/test_ple_disk_fill_order.py index e7fe58f7e0..4d2bbfd71b 100644 --- a/tests/models/qwen4_exp/test_ple_disk_fill_order.py +++ b/tests/models/qwen4_exp/test_ple_disk_fill_order.py @@ -8,6 +8,7 @@ from __future__ import annotations +import queue import threading import time from types import SimpleNamespace @@ -55,13 +56,13 @@ def synchronize(self): def _table(monkeypatch) -> DiskRowTable: # host-only: no CUDA stream behind the (faked) readback event monkeypatch.setattr(ple_disk.torch.cuda, "current_stream", lambda *a, **k: None) + monkeypatch.setattr(ple_disk.torch.cuda, "Event", _Event) t = DiskRowTable.__new__(DiskRowTable) t._device = None t._wait_sync = True t._flag = torch.zeros(1, dtype=torch.int64) t._store = _Store(t._flag) t._token_readback = torch.zeros(8, dtype=torch.int32) - t._readback_event = _Event() t._graph_pinned = torch.zeros(64, dtype=torch.uint8) t._eager_pinned = [torch.zeros(64, dtype=torch.uint8), torch.zeros(64, dtype=torch.uint8)] t._eager_read = [_Event(), _Event()] @@ -134,3 +135,95 @@ def test_a_failed_launch_leaves_no_stale_signal(monkeypatch): # no WAIT consumed it, so a leftover 1 would let the next step read stale rows assert int(t._flag[0]) == 0 _close(t) + + +class _Gpu: + """One in-order stream: each launched graph parks on WAIT(flag), then RESETs it.""" + + def __init__(self, flag: torch.Tensor): + self.flag = flag + self.enqueued = self.completed = 0 + self.stalled = False + self.cv = threading.Condition() + self._launches: queue.SimpleQueue = queue.SimpleQueue() + threading.Thread(target=self._run, daemon=True).start() + + def launch(self) -> None: + self.enqueued += 1 + self._launches.put(True) + + def close(self) -> None: + self._launches.put(False) + + def _run(self) -> None: + while self._launches.get(): + self.stalled |= not _blocking_launch(self.flag) + with self.cv: + self.completed += 1 + self.cv.notify_all() + + def wait_completed(self, n: int, timeout: float = 5.0) -> bool: + with self.cv: + return self.cv.wait_for(lambda: self.completed >= n, timeout) + + +class _StreamEvent: + """cudaEvent semantics: synchronize() waits for the work enqueued before the LAST record.""" + + def __init__(self, gpu: _Gpu): + self.gpu, self.target = gpu, 0 + + def record(self, stream=None): + self.target = self.gpu.enqueued + + def synchronize(self): + if not self.gpu.wait_completed(self.target, timeout=1.0): + raise TimeoutError("the readback waits on a graph parked on this fill's own WAIT") + + +def _two_steps_behind_a_slow_filler(t: DiskRowTable, gpu: _Gpu) -> None: + """Steps N and N+1 both launch before the filler thread reaches fill N.""" + gate = threading.Event() + t._filler().submit(gate.wait) + for prev, new in (([1, 2, 3], 7), ([1, 2, 3, 7], 8)): + with t.forward_host_ctx(_decode_batch(prev, new), use_graph=True): + gpu.launch() + gate.set() + assert gpu.wait_completed(2), "decode deadlocked" + t._filler().submit(lambda: None).result() # every submitted fill has finished + gpu.close() + + +def test_each_fill_waits_on_its_own_steps_readback(monkeypatch): + """A readback event shared across steps is re-recorded by step N+1 behind graph N, so + fill N would wait on the graph that is parked on fill N's own WAIT.""" + t = _table(monkeypatch) + gpu = _Gpu(t._flag) + monkeypatch.setattr(ple_disk.torch.cuda, "Event", lambda: _StreamEvent(gpu)) + _two_steps_behind_a_slow_filler(t, gpu) + t._raise_failed_fill() + _close(t) + assert not gpu.stalled + # the fake D2H readback is not stream-ordered, so compare only the contexts + assert [c[1][:2] for c in t._store.calls if c[0] == "stage"] == [[2, 3], [3, 7]] + + +def test_a_failed_fill_is_raised_after_the_next_step_was_submitted(monkeypatch): + """Fill N fails (its graph runs on stale rows) after step N+1 already submitted: the error + must still reach the engine thread at the next launch.""" + t = _table(monkeypatch) + gpu = _Gpu(t._flag) + monkeypatch.setattr(ple_disk.torch.cuda, "Event", lambda: _StreamEvent(gpu)) + flush = t._store.flush + + def fail_step_n(signal_addr): + if t._store.calls[-1][1][:2] == [2, 3]: # the run staged just before is step N's + raise OSError("PLE row read failed") + flush(signal_addr) + + t._store.flush = fail_step_n + _two_steps_behind_a_slow_filler(t, gpu) + req = SimpleNamespace(input_ids=torch.tensor([1, 2, 3, 4]), cached_len=1, device_len=4) + with pytest.raises(OSError, match="PLE row read failed"): + t.host_fill_batch(SimpleNamespace(is_decode=False, padded_reqs=[req]), use_graph=False) + _close(t) From 7bb53597694d7da20adb509f248836ef782fde97 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 15:01:03 -0400 Subject: [PATCH 37/94] fix(ple): recycle the fill readback events instead of freeing them on the filler The filler thread dropped each fill's last event reference, and cuEventDestroy there blocks on the driver lock that a blocked cuGraphLaunch holds (host-func nodes in the decode graph make the launch wait for the GPU, parked on the next fill's WAIT), so decode deadlocked with FREETOKEN_CPU_MOE_FLAG_SYNC=0. Events now return to an engine-thread pool once their fill is done. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 198574d266608fd6744f0ad90f01461ecc6b1c2d) --- python/freetoken/models/qwen4_exp/ple_disk.py | 21 ++++++++++++++-- .../qwen4_exp/test_ple_disk_fill_order.py | 25 +++++++++++++++++++ 2 files changed, 44 insertions(+), 2 deletions(-) diff --git a/python/freetoken/models/qwen4_exp/ple_disk.py b/python/freetoken/models/qwen4_exp/ple_disk.py index 05b18d5ce7..a0c71f7046 100644 --- a/python/freetoken/models/qwen4_exp/ple_disk.py +++ b/python/freetoken/models/qwen4_exp/ple_disk.py @@ -150,6 +150,7 @@ def cancel(self) -> None: pending = self._table._pending_fills() if self._future in pending: pending.remove(self._future) # its outcome was consumed here + self._table._retire(self._future) self._table._flag.zero_() @@ -246,12 +247,24 @@ def _current_stream(self): def _pending_fills(self) -> deque: return self.__dict__.setdefault("_pending_fill_queue", deque()) + def _readback_pool(self) -> list: + return self.__dict__.setdefault("_readback_events", []) + + def _retire(self, future: Future) -> None: + """Return a finished fill's readback event for reuse (engine thread only).""" + event = getattr(future, "readback", None) + if event is not None: + self._readback_pool().append(event) + future.readback = None + def _raise_failed_fill(self) -> None: """Surface a finished graph fill's error on the engine thread, oldest first; a failed fill signalled its WAIT, so that graph ran on stale rows.""" pending = self._pending_fills() while pending and pending[0].done(): - exc = pending.popleft().exception() + future = pending.popleft() + self._retire(future) + exc = future.exception() if exc is not None: raise exc @@ -298,7 +311,10 @@ def host_fill_batch(self, batch: Batch, use_graph: bool) -> _PendingFill | None: self._token_readback[:bs].copy_(batch.input_ids, non_blocking=True) # one event per fill: a shared one re-recorded by the next step before this job # reached synchronize() would wait on a graph parked on this job's own WAIT - readback = torch.cuda.Event() + # recycled, never destroyed while serving: cuEventDestroy from the filler blocks on + # the driver lock a blocked cuGraphLaunch holds, and the fill it waits for never runs + pool = self._readback_pool() + readback = pool.pop() if pool else torch.cuda.Event() readback.record(self._current_stream()) cancelled = threading.Event() @@ -318,6 +334,7 @@ def _job() -> None: raise future = self._filler().submit(_job) + future.readback = readback self._pending_fills().append(future) return _PendingFill(self, future, cancelled) # launch-gating: this D2H is the step's readback and orders the fill after sampling diff --git a/tests/models/qwen4_exp/test_ple_disk_fill_order.py b/tests/models/qwen4_exp/test_ple_disk_fill_order.py index 4d2bbfd71b..9d979b64eb 100644 --- a/tests/models/qwen4_exp/test_ple_disk_fill_order.py +++ b/tests/models/qwen4_exp/test_ple_disk_fill_order.py @@ -227,3 +227,28 @@ def fail_step_n(signal_addr): with pytest.raises(OSError, match="PLE row read failed"): t.host_fill_batch(SimpleNamespace(is_decode=False, padded_reqs=[req]), use_graph=False) _close(t) + + +def test_the_readback_event_is_never_freed_on_the_filler_thread(monkeypatch): + """cuEventDestroy on the filler blocks on the driver lock a blocked cuGraphLaunch holds (host + nodes in the graph make the launch wait for the GPU, which is parked on the next fill's WAIT).""" + import gc + + freed_on: list[str] = [] + + class _TrackedEvent(_Event): + def __del__(self): + freed_on.append(threading.current_thread().name) + + t = _table(monkeypatch) + monkeypatch.setattr(ple_disk.torch.cuda, "Event", _TrackedEvent) + for prev, new in (([1, 2, 3], 7), ([1, 2, 3, 7], 8)): + with t.forward_host_ctx(_decode_batch(prev, new), use_graph=True): + assert _blocking_launch(t._flag) + time.sleep(0.05) # let the filler drop its work items + gc.collect() + t._raise_failed_fill() + gc.collect() + assert "ple-filler" not in freed_on, freed_on + assert len(t._readback_pool()) >= 1 # finished fills hand their events back for reuse + _close(t) From 121c0c251c08f4a6c062fdf3d52053b683f54f62 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 05:25:06 -0400 Subject: [PATCH 38/94] fix(kvcache): refuse the host kv-page tier on a pool whose index slab it cannot restore Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit eb462bd73c2fa5eb96b38abe56fb841d82d68a32) --- python/freetoken/kvcache/host_prefix_tier.py | 8 +++++++- tests/kvcache/test_host_tier_nvfp4_guard.py | 7 +++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/python/freetoken/kvcache/host_prefix_tier.py b/python/freetoken/kvcache/host_prefix_tier.py index 39c96dfd6d..5cee24058a 100644 --- a/python/freetoken/kvcache/host_prefix_tier.py +++ b/python/freetoken/kvcache/host_prefix_tier.py @@ -84,8 +84,14 @@ def __init__(self, state_pool, kv_pool=None, *, gdn_slots_host: int = 64, from freetoken.kvcache.kv_host_offload import _reject_block_scaled_pool _reject_block_scaled_pool(kv_pool, "FT_PREFIX_HOST") + if getattr(kv_pool, "_cmp_k_buffer", None) is not None: + # a QSA pool's indexer slab is keyed by the same pages and is not mirrored here + raise ValueError( + "FT_PREFIX_HOST cannot restore a sparse-attention pool's compressed index " + "slab; keep FT_GDN_HOST_TIER alone for this model" + ) codes = kv_pool._kv_buffer # [2, L, P, page, H, D] - scales = kv_pool._scale_buffer # [2, L, P*page, H] | None (fp8) + scales = getattr(kv_pool, "_scale_buffer", None) # [2, L, P*page, H] | None (fp8) _, layers, _, page_size, heads, dim = codes.shape page_bytes = 2 * layers * page_size * heads * dim * codes.element_size() if scales is not None: diff --git a/tests/kvcache/test_host_tier_nvfp4_guard.py b/tests/kvcache/test_host_tier_nvfp4_guard.py index a793fb294a..50227267a6 100644 --- a/tests/kvcache/test_host_tier_nvfp4_guard.py +++ b/tests/kvcache/test_host_tier_nvfp4_guard.py @@ -42,3 +42,10 @@ def test_prefix_tier_gdn_snapshots_do_not_touch_the_kv_pool(): tier = HostPrefixTier(_state_pool(), kv_pool=_kv_pool(block_scaled=True), gdn_slots_host=2, kv_budget_bytes=0, device="cpu") assert tier.num_slots == 2 and not tier.has_kv + + +def test_prefix_tier_refuses_kv_pages_on_a_qsa_pool(): + # the KV-page store mirrors codes and scales, not the QSA indexer's compressed-key slab + pool = SimpleNamespace(**vars(_kv_pool(block_scaled=False)), _cmp_k_buffer=torch.zeros(1)) + with pytest.raises(ValueError, match="FT_PREFIX_HOST .* index slab"): + HostPrefixTier(_state_pool(), kv_pool=pool, gdn_slots_host=2, kv_budget_bytes=1 << 16, device="cpu") From 500b8faa467d7d880776fefa3bfc0cc7f2281719 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 05:51:57 -0400 Subject: [PATCH 39/94] perf(kvcache): keep one gpu gdn snapshot slot when the host snapshot tier is on Evicted snapshots spill to RAM and return on a hit, so the floor of four resident ones only costs VRAM: three 110 MiB Qwen3.8-Flash-Next slots go back to the expert cache. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 6751cf86b302dd539fdf81c1b32aaa8820650a5b) --- python/freetoken/kvcache/linear_state_pool.py | 6 +++++- tests/scheduler/test_hybrid_cache_manager.py | 11 +++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/python/freetoken/kvcache/linear_state_pool.py b/python/freetoken/kvcache/linear_state_pool.py index ff7539410c..d89556899d 100644 --- a/python/freetoken/kvcache/linear_state_pool.py +++ b/python/freetoken/kvcache/linear_state_pool.py @@ -1,6 +1,7 @@ from __future__ import annotations import math +import os import torch from freetoken.distributed import get_tp_info @@ -285,7 +286,10 @@ def _linear_pool_num_slots(config) -> int: if config.cache_type != "hybrid_radix": return mr + 1 # live + dummy/padding ratio = config.linear_state_cache_ratio - n_cache = max(4, int(ratio * mr)) + # With the host snapshot tier (FT_GDN_HOST_TIER=1) evicted snapshots spill to RAM and come + # back on a hit, so the GPU cache need not keep four of them resident. + floor = 1 if os.environ.get("FT_GDN_HOST_TIER", "0") == "1" else 4 + n_cache = max(floor, int(ratio * mr)) return 4 * mr + n_cache + 1 # live + 2 ping-pong + locked committed snapshot + cache + padding diff --git a/tests/scheduler/test_hybrid_cache_manager.py b/tests/scheduler/test_hybrid_cache_manager.py index 07385d7046..165689ab3b 100644 --- a/tests/scheduler/test_hybrid_cache_manager.py +++ b/tests/scheduler/test_hybrid_cache_manager.py @@ -241,5 +241,16 @@ def test_pool_sizing_covers_4mr_floor(): assert _linear_pool_num_slots(c) >= 4 * mr + 1, (mr, _linear_pool_num_slots(c)) +def test_host_snapshot_tier_lowers_the_gpu_snapshot_floor(monkeypatch): + """With FT_GDN_HOST_TIER evicted snapshots live in RAM, so one GPU snapshot slot is enough.""" + from types import SimpleNamespace + from freetoken.kvcache.linear_state_pool import _linear_pool_num_slots + c = SimpleNamespace(max_running_req=1, cache_type="hybrid_radix", linear_state_cache_ratio=0.0) + monkeypatch.delenv("FT_GDN_HOST_TIER", raising=False) + assert _linear_pool_num_slots(c) == 4 + 4 + 1 + monkeypatch.setenv("FT_GDN_HOST_TIER", "1") + assert _linear_pool_num_slots(c) == 4 + 1 + 1 + + if __name__ == "__main__": raise SystemExit(pytest.main([__file__])) From c22561e41fe846ee304817990e53f828f51c8062 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 10:43:02 -0400 Subject: [PATCH 40/94] perf(scheduler): probe the host tier for the exact tombstone prefix Revival only accepts an exact entry, but the longest-prefix lookup re-hashed every shorter page-aligned prefix on a miss and refreshed the LRU age of a shorter entry it then rejected. _rehydrate_snapshot on a 262144-token key, CPU, median of 3: no host entry 1852 ms -> 4.0 ms entry 4096 tokens in 1494 ms -> 2.9 ms entry 8192 below 91 ms -> 3.5 ms exact hit 1.4 ms -> 1.6 ms Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit b349af0cac8ee9334b49023460bc5d7afa9e92d1) --- python/freetoken/kvcache/host_prefix_tier.py | 10 +++++ python/freetoken/scheduler/cache.py | 8 ++-- .../kvcache/test_host_tier_rehydrate_exact.py | 39 ++++++++++++------- 3 files changed, 38 insertions(+), 19 deletions(-) diff --git a/python/freetoken/kvcache/host_prefix_tier.py b/python/freetoken/kvcache/host_prefix_tier.py index 5cee24058a..4295aa8e4b 100644 --- a/python/freetoken/kvcache/host_prefix_tier.py +++ b/python/freetoken/kvcache/host_prefix_tier.py @@ -318,6 +318,16 @@ def lookup(self, input_ids: torch.Tensor, page_size: int = 64) -> HostMatch | No n -= page_size return None + def lookup_exact(self, key_tokens: torch.Tensor) -> HostMatch | None: + """The entry for exactly this prefix: one hash, where lookup re-hashes every shorter + page-aligned prefix on a miss.""" + key = self._key_of(key_tokens) + hit = self._index.get(key) + if hit is None: + return None + self._index.move_to_end(key) + return HostMatch(key_tokens.numel(), hit.kv_slots, hit.gdn_slot) + def evict_host(self, key_tokens: torch.Tensor) -> None: """Remove uma entrada do índice e libera os slots host (KV + GDN).""" ent = self._index.pop(self._key_of(key_tokens), None) diff --git a/python/freetoken/scheduler/cache.py b/python/freetoken/scheduler/cache.py index 4e6b91df88..715dfdc21a 100644 --- a/python/freetoken/scheduler/cache.py +++ b/python/freetoken/scheduler/cache.py @@ -625,10 +625,10 @@ def _rehydrate_snapshot(self, node) -> int | None: keep walking (the match truncates as before).""" tier = self.host_tier key = self.prefix_cache._collect_key(node) - hit = tier.lookup(key, self.page_size) - # lookup is longest-prefix; a shorter entry's state would resume this node's full - # depth with the tokens in between missing from the recurrence - if hit is None or hit.gdn_slot is None or hit.host_len != len(key): + # exact only: a shorter entry's state would resume this node's full depth with the + # tokens in between missing from the recurrence + hit = tier.lookup_exact(key) + if hit is None or hit.gdn_slot is None: return None pool = self.linear_state_pool if pool.num_free_slots < 1: diff --git a/tests/kvcache/test_host_tier_rehydrate_exact.py b/tests/kvcache/test_host_tier_rehydrate_exact.py index f67402041f..d6f4d09215 100644 --- a/tests/kvcache/test_host_tier_rehydrate_exact.py +++ b/tests/kvcache/test_host_tier_rehydrate_exact.py @@ -1,21 +1,19 @@ """The GDN snapshot tier (upstream #525) may only revive a tombstone from a snapshot of -exactly that node's prefix: the host index lookup is longest-prefix, and a shorter -snapshot would resume the recurrence with the tokens in between silently missing.""" +exactly that node's prefix: a shorter snapshot would resume the recurrence with the tokens +in between silently missing.""" from types import SimpleNamespace import torch -from freetoken.kvcache.host_prefix_tier import HostMatch +from freetoken.kvcache.host_prefix_tier import HostPrefixTier from freetoken.scheduler.cache import CacheManager -def _manager(host_len: int, key_len: int = 128): - reads = [] - tier = SimpleNamespace( - lookup=lambda key, page_size: HostMatch(host_len, None, 7), - read_gdn=lambda host_slot, gpu_slot: reads.append((host_slot, gpu_slot)), - ) +def _manager(key_len: int = 128): + state = SimpleNamespace(conv_states=torch.zeros(1, 4, 1, 1), + recurrent_states=torch.zeros(1, 4, 1, 1), slot_states={}) + tier = HostPrefixTier(state, gdn_slots_host=4, device="cpu") pool = SimpleNamespace(num_free_slots=4, alloc=lambda n: [3]) manager = SimpleNamespace( host_tier=tier, @@ -24,16 +22,27 @@ def _manager(host_len: int, key_len: int = 128): linear_state_pool=pool, ensure_mamba_slots=lambda n: None, ) - return manager, reads + return manager, tier, state + + +def _snapshot(tier, state, key, value): + state.conv_states[:, 1] = value + tier.put_prefix(key, None, tier.write_gdn(1)) -def test_a_shorter_host_snapshot_does_not_revive_a_deeper_node(): - manager, reads = _manager(host_len=64) +def test_a_shorter_host_snapshot_neither_revives_nor_refreshes_a_deeper_node(): + manager, tier, state = _manager(key_len=128) + _snapshot(tier, state, torch.arange(64), 1.0) + _snapshot(tier, state, torch.arange(1000, 1064), 2.0) assert CacheManager._rehydrate_snapshot(manager, node=object()) is None - assert reads == [] + # the rejected entry keeps its LRU age + assert [len(e.key_tokens) for e in tier._index.values()] == [64, 64] + assert next(iter(tier._index.values())).key_tokens[0] == 0 + assert torch.all(state.conv_states[:, 3] == 0) def test_an_exact_host_snapshot_revives_the_node(): - manager, reads = _manager(host_len=128) + manager, tier, state = _manager(key_len=128) + _snapshot(tier, state, torch.arange(128), 5.0) assert CacheManager._rehydrate_snapshot(manager, node=object()) == 3 - assert reads == [(7, 3)] + assert torch.all(state.conv_states[:, 3] == 5.0) From 18a082e4cece424496572a04f061b80e24d8f36e Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 10:46:01 -0400 Subject: [PATCH 41/94] fix(scheduler): re-probe the host entry after making room to rehydrate it The evictions that make room for a revived snapshot or a host prefix write their victims to the host tier, which can drop the hit's own entry and hand its GDN and KV slots to the victim, so the stale hit restored another prefix's state. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 44f25067541a6fb84161effe4dd2d7ee8e3a0306) --- python/freetoken/scheduler/cache.py | 12 ++- .../kvcache/test_host_tier_rehydrate_exact.py | 79 ++++++++++++++++++- 2 files changed, 87 insertions(+), 4 deletions(-) diff --git a/python/freetoken/scheduler/cache.py b/python/freetoken/scheduler/cache.py index 715dfdc21a..ff4467328e 100644 --- a/python/freetoken/scheduler/cache.py +++ b/python/freetoken/scheduler/cache.py @@ -632,9 +632,11 @@ def _rehydrate_snapshot(self, node) -> int | None: return None pool = self.linear_state_pool if pool.num_free_slots < 1: - # make room: evicting another snapshot writes IT to the host tier first + # make room: evicting another snapshot writes IT to the host tier first, which may + # drop this entry and hand its host slot to the victim, so probe again self.ensure_mamba_slots(1) - if pool.num_free_slots < 1: + hit = tier.lookup_exact(key) + if pool.num_free_slots < 1 or hit is None or hit.gdn_slot is None: return None slot = pool.alloc(1)[0] tier.read_gdn(hit.gdn_slot, slot) @@ -651,7 +653,11 @@ def _rehydrate_host_prefix(self, ids: torch.Tensor, hit) -> None: pool = self.linear_state_pool if hit.gdn_slot is not None: self.ensure_mamba_slots(1) - if pool.num_free_slots < 1: + # _allocate and ensure_mamba_slots write their victims to the tier, which may drop + # this entry and hand its host slots to them: use what the key maps to now + hit = tier.lookup_exact(ids[: hit.host_len]) + if (pool.num_free_slots < 1 or hit is None or hit.kv_pages is None + or hit.gdn_slot is None): self._free(kv_new) # sem slot GDN não dá pra reidratar; devolve as páginas return slot = pool.alloc(1)[0] diff --git a/tests/kvcache/test_host_tier_rehydrate_exact.py b/tests/kvcache/test_host_tier_rehydrate_exact.py index d6f4d09215..77dae0fd08 100644 --- a/tests/kvcache/test_host_tier_rehydrate_exact.py +++ b/tests/kvcache/test_host_tier_rehydrate_exact.py @@ -1,12 +1,16 @@ """The GDN snapshot tier (upstream #525) may only revive a tombstone from a snapshot of exactly that node's prefix: a shorter snapshot would resume the recurrence with the tokens -in between silently missing.""" +in between silently missing. Making room for the revival evicts other prefixes into the +tier, which may recycle the entry's own host slots, so it must not read stale ones.""" from types import SimpleNamespace import torch +import freetoken.core as core from freetoken.kvcache.host_prefix_tier import HostPrefixTier +from freetoken.kvcache.linear_state_pool import LinearStatePool +from freetoken.models.config import LinearGatedDeltaGroupConfig from freetoken.scheduler.cache import CacheManager @@ -46,3 +50,76 @@ def test_an_exact_host_snapshot_revives_the_node(): _snapshot(tier, state, torch.arange(128), 5.0) assert CacheManager._rehydrate_snapshot(manager, node=object()) == 3 assert torch.all(state.conv_states[:, 3] == 5.0) + + +def _cache_manager(monkeypatch, num_pages, page_size, state_slots, **env): + for k, v in env.items(): + monkeypatch.setenv(k, v) + g = LinearGatedDeltaGroupConfig( + name="linear", layer_ids=(0,), num_key_heads=2, num_value_heads=4, + key_head_dim=16, value_head_dim=16, conv_kernel_dim=4, output_gate="silu", + ) + pool = LinearStatePool(group=g, num_slots=state_slots, dtype=torch.bfloat16, + device=torch.device("cpu"), tp_size=1) + page_table = torch.zeros(2, 64, dtype=torch.int32) + cm = CacheManager(num_pages, page_size, page_table, "hybrid_radix", linear_state_pool=pool) + return cm, pool + + +def _put(cm, pool, ids, state, stamp, kv_pool=None, kv_value=None): + kv = cm._page_to_token(cm._allocate(len(ids) // cm.page_size)) + if kv_pool is not None: + kv_pool._kv_buffer[:, :, kv[:: cm.page_size] // cm.page_size] = kv_value + slot = pool.alloc(1)[0] + pool.conv_states[:, slot] = state + cm.prefix_cache.insert(ids, kv, slot) + cm.prefix_cache.match_prefix(ids).node.timestamp = stamp + + +def _match(cm, ids): + tail = torch.tensor([7, 9], dtype=torch.int32) + return cm.match_req(SimpleNamespace(input_ids=torch.cat([ids, tail]), input_len=len(ids) + 2)) + + +def test_making_room_to_revive_does_not_restore_the_victim_into_the_host_slot(monkeypatch): + # one host slot: evicting B for the revived slot drops A's entry and writes B there + cm, pool = _cache_manager(monkeypatch, 8, 64, 3, FT_GDN_HOST_TIER="1", FT_GDN_HOST_SLOTS="1") + a = torch.arange(128, dtype=torch.int32) + b = torch.arange(1000, 1128, dtype=torch.int32) + _put(cm, pool, a, 1.0, stamp=1) + _put(cm, pool, b, 2.0, stamp=2) + cm.ensure_mamba_slots(1) # A -> host tier, A's node becomes a tombstone + pool.conv_states[:, pool.alloc(1)[0]] = 9.0 # someone else takes the freed slot + + mr = _match(cm, a) + restored = None if mr.mamba_value is None else pool.conv_states[:, mr.mamba_value].unique().tolist() + assert restored in (None, [1.0]), (mr.cuda_handle.cached_len, restored) + assert cm.host_tier.lookup_exact(b) is not None + + +def test_making_room_to_rehydrate_a_host_prefix_does_not_restore_the_victim(monkeypatch): + # a 4-page KV store: writing the 4-page victim V drops the 2-page hit H and reuses its slots + page, n_pages = 4, 6 + kv_pool = SimpleNamespace(_kv_buffer=torch.zeros(2, 1, n_pages, page, 1, 2), + _scale_buffer=None, _block_scale_buffer=None) + monkeypatch.setattr(core, "_GLOBAL_CTX", SimpleNamespace(kv_cache=kv_pool, kv_offloader=None)) + page_bytes = kv_pool._kv_buffer[:, :, 0].numel() * kv_pool._kv_buffer.element_size() + cm, pool = _cache_manager(monkeypatch, n_pages, page, 4, FT_PREFIX_HOST="1", + FT_PREFIX_HOST_GB=repr(4 * page_bytes / (1 << 30)), + FT_GDN_HOST_SLOTS="4") + h = torch.arange(2 * page, dtype=torch.int32) + v = torch.arange(100, 100 + 4 * page, dtype=torch.int32) + _put(cm, pool, h, 1.0, stamp=1, kv_pool=kv_pool, kv_value=1.0) + er = cm.prefix_cache.evict_full(len(h)) # H -> host tier (KV + snapshot), out of the tree + pool.free(er.mamba_slots) + cm._free(er.kv_indices) + _put(cm, pool, v, 2.0, stamp=2, kv_pool=kv_pool, kv_value=2.0) + cm._allocate(len(cm.free_slots)) # a running request holds the rest of the GPU pages + + mr = _match(cm, h) + handle = mr.cuda_handle + if handle.cached_len: + pages = handle.kv_indices[::page] // page + assert pool.conv_states[:, mr.mamba_value].unique().tolist() == [1.0] + assert kv_pool._kv_buffer[:, :, pages].unique().tolist() == [1.0] + assert cm.host_tier.lookup_exact(v) is not None From dedec6d1589b296b25707a282c766ca8663bcb45 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 10:49:19 -0400 Subject: [PATCH 42/94] test(kvcache): cover the engine-stream waits of the prefix tier's kv page copies Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit af5ca25838e9c690ba32a6cae20c72692e61477c) --- tests/kvcache/test_host_tier_live_pool.py | 23 +++++++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/tests/kvcache/test_host_tier_live_pool.py b/tests/kvcache/test_host_tier_live_pool.py index bd12836893..0eea1990e3 100644 --- a/tests/kvcache/test_host_tier_live_pool.py +++ b/tests/kvcache/test_host_tier_live_pool.py @@ -32,7 +32,7 @@ def test_read_after_a_pool_rebuild_lands_in_the_new_tensors(): assert torch.all(pool.recurrent_states[:, 5] == 7.0) -def test_copies_wait_for_the_engine_stream(monkeypatch): +def _fake_copy_stream(monkeypatch): events = [] engine = object() @@ -45,9 +45,28 @@ def synchronize(self): monkeypatch.setattr(tier_module.torch.cuda, "current_stream", lambda device=None: engine) monkeypatch.setattr(tier_module.torch.cuda, "stream", lambda s: contextlib.nullcontext()) + return FakeStream(), events, engine + + +def test_copies_wait_for_the_engine_stream(monkeypatch): + stream, events, engine = _fake_copy_stream(monkeypatch) pool = _pool(3) tier = HostPrefixTier(pool, gdn_slots_host=2, device="cpu") - tier.copy_stream = FakeStream() + tier.copy_stream = stream host = tier.write_gdn(1) tier.read_gdn(host, 2) assert events == [("wait", engine), ("sync",), ("wait", engine), ("sync",)] + + +def test_kv_page_copies_wait_for_the_engine_stream(monkeypatch): + stream, events, engine = _fake_copy_stream(monkeypatch) + kv_pool = SimpleNamespace( + _kv_buffer=torch.zeros(2, 1, 4, 4, 1, 2), # [2, L, P, page, H, D] + _scale_buffer=torch.zeros(2, 1, 16, 1), # [2, L, P*page, H] + _block_scale_buffer=None, + ) + tier = HostPrefixTier(_pool(3), kv_pool, gdn_slots_host=2, kv_budget_bytes=1 << 12, device="cpu") + tier.copy_stream = stream + slots = tier.write_kv(torch.arange(4, 8, dtype=torch.int32)) + tier.read_kv(slots, torch.tensor([2])) + assert events == [("wait", engine), ("sync",), ("wait", engine), ("sync",)] From f7d026bd04271c01c5e0182b3f93d26ef091e6de Mon Sep 17 00:00:00 2001 From: fangyuan-3149 <314945674@qq.com> Date: Sun, 30 Aug 2026 09:35:36 +0800 Subject: [PATCH 43/94] fix(engine): reserve the GDN prefill workspace in the cache budget chunk_gated_delta_rule_fwd_h materializes the per-chunk recurrent states for the whole extend at once -- [NT, H, V, K] with NT = ceil(extend/64) -- plus a transient v_new buffer. Subtract it at the startup sizing and the --moe-cache-auto plan for hybrid-GDN models. (cherry picked from commit 5c3b9141b6b658f4d4ebe59f2a68b2208c6f6771) Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit ec8731a65ff5ca2cce0ff073d273f482a8faa5e6) --- python/freetoken/engine/engine.py | 7 ++++- python/freetoken/kvcache/linear_state_pool.py | 27 +++++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 7a0b49c172..0d43cf2816 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -36,7 +36,8 @@ from freetoken.kvcache.base import CacheRebuildRejected from freetoken.kvcache.cache_status import _supports_swa_ratio from freetoken.kvcache.linear_state_pool import ( - _linear_pool_min_slots, _linear_pool_num_slots, state_pool_bytes, + _linear_pool_min_slots, _linear_pool_num_slots, gdn_prefill_workspace_bytes, + state_pool_bytes, ) logger = init_logger(__name__) @@ -539,6 +540,9 @@ def __init__(self, config: EngineConfig): # off it; the KV pool family owns every geometry-specific formula behind the rest. available_memory = _startup_kv_budget(config.memory_ratio, init_free_memory, new_free) available_memory -= state_pool_bytes(config) + # reserve the GDN prefill transient (chunk recurrent states + v_new); without + # it long prefills OOM inside the GDN kernel when memory_ratio packs the pools + available_memory -= gdn_prefill_workspace_bytes(config) num_gpu_pages = self._pool_cls.solve_num_pages(config, available_memory) if config.kv_host_pages > 0: # KV host offload: the scheduler, page table and radix tree run in a LOGICAL page @@ -794,6 +798,7 @@ def _resolve_auto_moe_cache_size(self, config: EngineConfig, banks, method=None, cache_per_page, fixed_cache_size, page_tokens, min_reserve = self._pool_cls.kv_cost(config) fixed_cache_size += state_pool_bytes(config) # sibling GDN state pool, engine-summed + fixed_cache_size += gdn_prefill_workspace_bytes(config) # GDN prefill transient num_experts = config.model_config.num_experts total_experts = config.model_config.num_moe_layers * num_experts return resolve_moe_cache_auto( diff --git a/python/freetoken/kvcache/linear_state_pool.py b/python/freetoken/kvcache/linear_state_pool.py index d89556899d..16d2fb537e 100644 --- a/python/freetoken/kvcache/linear_state_pool.py +++ b/python/freetoken/kvcache/linear_state_pool.py @@ -261,6 +261,33 @@ def linear_state_bytes_per_req( __all__ = ["LinearStatePool", "linear_state_bytes_per_req"] +_GDN_CHUNK_SIZE = 64 # keep in lockstep with kernel/fla/chunk_delta_h.py CHUNK_SIZE + + +def gdn_prefill_workspace_bytes(config) -> int: + """Peak transient VRAM of one GDN prefill forward. + + ``chunk_gated_delta_rule_fwd_h`` materializes the per-chunk recurrent states for + the WHOLE extend at once -- ``[NT, H, V, K]`` with ``NT = ceil(extend / 64)`` -- + plus the ``v_new`` value buffer; ~192 MiB at an 8192-token extent for + ornith-class geometry. The generic ``(1 - memory_ratio)`` headroom is sized for + CUDA graphs; without reserving this term the GDN kernel OOMs on long prefills + at aggressive memory ratios. 0 for non-GDN models.""" + linear_group = config.model_config.linear_attention_group() + if linear_group is None: + return 0 + extent = min(config.max_extend_tokens, config.max_seq_len) + nt = -(-extent // _GDN_CHUNK_SIZE) + states = ( + nt + * linear_group.num_value_heads + * linear_group.value_head_dim + * linear_group.key_head_dim + ) + v_new = extent * linear_group.num_value_heads * linear_group.value_head_dim + return (states + v_new) * config.dtype.itemsize + + def state_pool_bytes(config, num_slots: int | None = None) -> int: """Total GDN state-pool bytes at ``num_slots`` PHYSICAL slots (default: the startup slot count). The engine adds this to the KV family's fixed cost when budgeting -- From 03bae0975d4c7b9b49f4ae464b59eb662e949dbf Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 05:29:45 -0400 Subject: [PATCH 44/94] fix(engine): reserve the gdn prefill workspace in the rebuild fit-check too The cherry-picked reservation covered startup sizing and the auto plan but not validate_rebuild, so a runtime rebuild could grow the pools back into it. Also price the scheduler's default extend when the config is a bare EngineConfig, and document that h + v_new is a lower bound on the GDN prefill peak. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 9d836ed0b0b6c448975b26b4d86368e087dfe0af) --- python/freetoken/engine/engine.py | 7 ++-- python/freetoken/kvcache/linear_state_pool.py | 13 +++--- tests/engine/test_cache_budget.py | 38 +++++++++++++++++ tests/kvcache/test_linear_state_pool_alloc.py | 42 +++++++++++++++++++ 4 files changed, 91 insertions(+), 9 deletions(-) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 0d43cf2816..ed4525b324 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -540,8 +540,8 @@ def __init__(self, config: EngineConfig): # off it; the KV pool family owns every geometry-specific formula behind the rest. available_memory = _startup_kv_budget(config.memory_ratio, init_free_memory, new_free) available_memory -= state_pool_bytes(config) - # reserve the GDN prefill transient (chunk recurrent states + v_new); without - # it long prefills OOM inside the GDN kernel when memory_ratio packs the pools + # Reserve part of the GDN prefill transient (chunk states + v_new); the rest of the + # prefill activations still live in the (1 - memory_ratio) headroom. available_memory -= gdn_prefill_workspace_bytes(config) num_gpu_pages = self._pool_cls.solve_num_pages(config, available_memory) if config.kv_host_pages > 0: @@ -1241,7 +1241,8 @@ def rebuild_runtime_cache( per_expert_bytes=per_expert_bytes, baseline_free=self._baseline_free, weights_bytes=self._weights_bytes, current_num_pages=self.num_pages, extra_fixed_bytes=( - state_pool_bytes(config, target_mamba) if target_mamba is not None else 0 + (state_pool_bytes(config, target_mamba) if target_mamba is not None else 0) + + gdn_prefill_workspace_bytes(config) ), extra_note=( f", mamba={target_mamba - 1} slots" if target_mamba is not None else "" diff --git a/python/freetoken/kvcache/linear_state_pool.py b/python/freetoken/kvcache/linear_state_pool.py index 16d2fb537e..df5a3396da 100644 --- a/python/freetoken/kvcache/linear_state_pool.py +++ b/python/freetoken/kvcache/linear_state_pool.py @@ -265,18 +265,19 @@ def linear_state_bytes_per_req( def gdn_prefill_workspace_bytes(config) -> int: - """Peak transient VRAM of one GDN prefill forward. + """Bytes of the GDN chunked-prefill buffers that scale with the extend length. ``chunk_gated_delta_rule_fwd_h`` materializes the per-chunk recurrent states for the WHOLE extend at once -- ``[NT, H, V, K]`` with ``NT = ceil(extend / 64)`` -- - plus the ``v_new`` value buffer; ~192 MiB at an 8192-token extent for - ornith-class geometry. The generic ``(1 - memory_ratio)`` headroom is sized for - CUDA graphs; without reserving this term the GDN kernel OOMs on long prefills - at aggressive memory ratios. 0 for non-GDN models.""" + plus the ``v_new`` value buffer (288 MiB at an 8192-token extend for qwen4_exp). + This is a lower bound on the GDN prefill peak, not the peak: w/u/A/o and the + in_proj/conv activations still come out of the ``(1 - memory_ratio)`` headroom. + 0 for non-GDN models.""" linear_group = config.model_config.linear_attention_group() if linear_group is None: return 0 - extent = min(config.max_extend_tokens, config.max_seq_len) + # EngineConfig has no max_extend_tokens (the served config is a SchedulerConfig subclass) + extent = min(getattr(config, "max_extend_tokens", 8192), config.max_seq_len) nt = -(-extent // _GDN_CHUNK_SIZE) states = ( nt diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py index b70227f14f..f5abf0142c 100644 --- a/tests/engine/test_cache_budget.py +++ b/tests/engine/test_cache_budget.py @@ -398,6 +398,44 @@ def slot_limit(self): assert captured["max_slots"] == 5 + +def test_rebuild_fit_check_reserves_the_gdn_prefill_workspace(): + """The runtime-rebuild fit-check prices the GDN prefill workspace like the startup sizing, + so a rebuild cannot grow the pools back into the space reserved for it.""" + from freetoken.engine.engine import Engine + from freetoken.kvcache.base import CacheRebuildRejected + from freetoken.kvcache.linear_state_pool import gdn_prefill_workspace_bytes, state_pool_bytes + from freetoken.models.config import LinearGatedDeltaGroupConfig + + group = LinearGatedDeltaGroupConfig( + name="linear", layer_ids=(0, 1), num_key_heads=2, num_value_heads=4, + key_head_dim=16, value_head_dim=16, conv_kernel_dim=4, output_gate="silu", + ) + config = SimpleNamespace( + model_config=SimpleNamespace(linear_attention_group=lambda: group), + dtype=torch.bfloat16, tp_info=SimpleNamespace(size=1), + max_extend_tokens=1024, max_seq_len=4096, + ) + captured = {} + + def validate_rebuild(config, **kwargs): + captured.update(kwargs) + raise CacheRebuildRejected("stop before teardown") + + engine = SimpleNamespace( + config=config, moe_offload_cache=None, kv_offloader=None, + linear_state_pool=SimpleNamespace(num_slots=5), + kv_cache=SimpleNamespace(validate_rebuild=validate_rebuild), + _baseline_free=1 << 30, _weights_bytes=0, num_pages=64, + ) + engine._target_moe_and_expert_bytes = lambda size: (0, 0) + with pytest.raises(CacheRebuildRejected): + Engine.rebuild_runtime_cache(engine, num_pages=32) + + workspace = gdn_prefill_workspace_bytes(config) + assert workspace > 0 + assert captured["extra_fixed_bytes"] == state_pool_bytes(config, 5) + workspace + # --------------------------------------------------------------------------- # offload-cache sizing guard + auto-resolution (_require_offload_cache_size / _adjust_config), # the floor rule compute_cache_floors documents above. diff --git a/tests/kvcache/test_linear_state_pool_alloc.py b/tests/kvcache/test_linear_state_pool_alloc.py index e9fe405426..e0edd0438e 100644 --- a/tests/kvcache/test_linear_state_pool_alloc.py +++ b/tests/kvcache/test_linear_state_pool_alloc.py @@ -9,6 +9,7 @@ from freetoken.kvcache.linear_state_pool import ( LinearStatePool, + gdn_prefill_workspace_bytes, linear_state_bytes_per_req, state_pool_bytes, ) @@ -189,3 +190,44 @@ def test_slot_state_bytes_for_the_real_geometry(): delta = linear_state_bytes_per_req(group, 1, torch.bfloat16, (spec,)) - \ linear_state_bytes_per_req(group, 1, torch.bfloat16) assert delta == 4 * 2560 * 9 * 2 == 180 * 1024 + + +def _gdn_config(**fields): + # qwen4_exp GDN geometry: 48 value heads of 128x128 + group = LinearGatedDeltaGroupConfig( + name="linear", layer_ids=(0, 1), + num_key_heads=16, num_value_heads=48, + key_head_dim=128, value_head_dim=128, conv_kernel_dim=4, output_gate="silu", + ) + mc = SimpleNamespace(linear_attention_group=lambda: group) + return SimpleNamespace(model_config=mc, dtype=torch.bfloat16, max_seq_len=262144, **fields) + + +def test_gdn_prefill_workspace_for_the_real_geometry(): + # h [8192/64 chunks, 48, 128, 128] bf16 = 192 MiB, plus v_new [8192, 48, 128] = 96 MiB + assert gdn_prefill_workspace_bytes(_gdn_config(max_extend_tokens=8192)) == 288 * 2**20 + assert gdn_prefill_workspace_bytes(_gdn_config(max_extend_tokens=2048)) == 72 * 2**20 + + +def test_gdn_prefill_workspace_without_a_scheduler_extend_budget(): + # a bare EngineConfig has no max_extend_tokens: price the scheduler's 8192 default + assert gdn_prefill_workspace_bytes(_gdn_config()) == 288 * 2**20 + + +def test_gdn_prefill_workspace_is_capped_by_the_sequence_length(): + config = _gdn_config(max_extend_tokens=8192) + config.max_seq_len = 100 # 2 chunks, 100 value rows + assert gdn_prefill_workspace_bytes(config) == (2 * 48 * 128 * 128 + 100 * 48 * 128) * 2 + + +def test_gdn_prefill_workspace_is_zero_without_a_linear_group(): + config = _gdn_config(max_extend_tokens=8192) + config.model_config.linear_attention_group = lambda: None + assert gdn_prefill_workspace_bytes(config) == 0 + + +def test_gdn_prefill_workspace_chunk_matches_the_kernel(): + from freetoken.kernel.fla.chunk_delta_h import CHUNK_SIZE + from freetoken.kvcache.linear_state_pool import _GDN_CHUNK_SIZE + + assert _GDN_CHUNK_SIZE == CHUNK_SIZE From 4b73a931fc71b6fd90792c9ac99931e8d9416545 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 10:25:31 -0400 Subject: [PATCH 45/94] fix(engine): count min_tokens progress from the device length Under overlap scheduling the sampler prepares batch N before batch N-1's token lands in req.input_ids, so the host count lagged by one and kept the stop ids masked for one token past min_tokens. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 8e4891fd54037f19aa42d43422a50cafc38513d9) --- python/freetoken/engine/sample.py | 4 +++- tests/engine/test_logits_processors.py | 26 ++++++++++++++++++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py index bca85c56fa..89d2a6fb88 100644 --- a/python/freetoken/engine/sample.py +++ b/python/freetoken/engine/sample.py @@ -220,7 +220,9 @@ def _plan(self, batch: Batch, params) -> LogitsPlan | None: min_ids: list[list[int]] = [] for local, (req, p) in enumerate(picked): if p.min_tokens > 0 and p.min_tokens_stop_ids: - generated = len(req.input_ids) - (req.max_device_len - req.output_len) + # device_len, not the host ids: under overlap scheduling the previous token + # reaches req.input_ids only after this batch is prepared + generated = req.device_len - (req.max_device_len - req.output_len) if generated < p.min_tokens: min_rows.append(local) min_ids.append([t for t in p.min_tokens_stop_ids if 0 <= t < self.vocab_size]) diff --git a/tests/engine/test_logits_processors.py b/tests/engine/test_logits_processors.py index a9645f4210..c9f498183f 100644 --- a/tests/engine/test_logits_processors.py +++ b/tests/engine/test_logits_processors.py @@ -85,6 +85,31 @@ def test_min_tokens_masks_the_stop_ids_of_the_row_only(): assert torch.isfinite(out[0]).all() +@pytest.mark.parametrize("overlap", [True, False]) +def test_min_tokens_releases_the_stop_ids_at_output_index_min_tokens(overlap): + """Overlap scheduling prepares batch N before batch N-1's token reaches req.input_ids on + the host; the count must not lag by that token and hold EOS one step too long.""" + from freetoken.core import Req + + sp = SamplingParams(min_tokens=3, max_tokens=10) + sp.min_tokens_stop_ids = [0] + req = Req(input_ids=torch.tensor([5, 6], dtype=torch.int32), table_idx=0, cached_len=0, + output_len=10, uid=1, sampling_params=sp, cache_handle=None) + sampler = Sampler(CPU, V) + masked, pending = [], None + for _ in range(6): + plan = sampler.prepare(SimpleNamespace(reqs=[req])).plan + masked.append(plan is not None and plan.min_rows is not None) + req.complete_one() # the forward launch advances the device length + if overlap: + if pending is not None: # batch N-1 drains after batch N launched + req.append_host(pending) + pending = torch.tensor([7], dtype=torch.int32) + else: + req.append_host(torch.tensor([7], dtype=torch.int32)) + assert [i for i, m in enumerate(masked) if m] == [0, 1, 2] + + def test_min_p_drops_tokens_below_the_fraction_of_the_top_probability(): logits = torch.log(torch.tensor([[0.5, 0.3, 0.1, 0.05, 0.03, 0.01, 0.005, 0.005]])) plan = _plan( @@ -104,6 +129,7 @@ def _req( ids = torch.tensor(prompt + generated, dtype=torch.int32) return SimpleNamespace( input_ids=ids, + device_len=len(ids), output_len=max_tokens, max_device_len=len(prompt) + max_tokens, sampling_params=SamplingParams(**sp), From 67fc938119a2c9491ab92f972bde7127428185df Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 1 Oct 2026 23:23:00 -0400 Subject: [PATCH 46/94] test(engine): import pytest for the parametrized min_tokens test Assisted-by: Claude Opus 5.5 --- tests/engine/test_logits_processors.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/engine/test_logits_processors.py b/tests/engine/test_logits_processors.py index c9f498183f..abd2e1f939 100644 --- a/tests/engine/test_logits_processors.py +++ b/tests/engine/test_logits_processors.py @@ -5,6 +5,7 @@ from types import SimpleNamespace +import pytest import torch from freetoken.core import SamplingParams from freetoken.engine.sample import LogitsPlan, Sampler, apply_logits_processors From 47bc6236c4486e17f22c23687f523a3e60440293 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 06:01:42 -0400 Subject: [PATCH 47/94] fix(server): let protocol-level thinking requests win over --default-thinking-mode returned early, so reasoning_effort / reasoning.effort / thinking.type were silently turned off under --default-thinking-mode chat. count_tokens also skipped the default. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 8276ccd49a8fcb9123ea28f113ac39684b725665) --- python/freetoken/server/anthropic_api.py | 16 ++-- python/freetoken/server/args.py | 17 ++-- python/freetoken/server/model_meta.py | 19 +++++ python/freetoken/server/openai_api.py | 33 +------- python/freetoken/server/responses_api.py | 9 +-- tests/server/test_default_thinking_mode.py | 94 +++++++++++++++++++--- 6 files changed, 127 insertions(+), 61 deletions(-) diff --git a/python/freetoken/server/anthropic_api.py b/python/freetoken/server/anthropic_api.py index 4cb4b34822..a7c88b141a 100644 --- a/python/freetoken/server/anthropic_api.py +++ b/python/freetoken/server/anthropic_api.py @@ -166,7 +166,9 @@ async def handle_anthropic_count_tokens(req: AnthropicCountTokensRequest, state: # so it must not fall into the convert/empty-prompt ValueError branch. try: messages, template_tools, _, ctk = convert_anthropic_prompt( - req, reasoning_parser=getattr(state.config, "reasoning_parser", None) + req, + reasoning_parser=getattr(state.config, "reasoning_parser", None), + default_thinking_mode=getattr(state.config, "default_thinking_mode", "auto"), ) except ValueError as exc: return _anthropic_error_response(400, "invalid_request_error", str(exc)) @@ -198,6 +200,7 @@ async def handle_anthropic_count_tokens(req: AnthropicCountTokensRequest, state: def convert_anthropic_prompt( req: AnthropicMessagesRequest | AnthropicCountTokensRequest, reasoning_parser: str | None = None, + default_thinking_mode: str | None = None, ) -> tuple[list[dict[str, Any]], list[dict[str, Any]] | None, list[dict[str, Any]] | None, dict[str, Any]]: """(messages, template_tools, parser_tools, chat_template_kwargs) — the prompt side of the conversion, shared by /v1/messages and /v1/messages/count_tokens so @@ -310,7 +313,7 @@ def convert_anthropic_prompt( # Native extended-thinking toggle -> template kwargs, broadcast in every # spelling the ecosystem's templates read (a bare enable_thinking bool is # inert for templates that read a different knob, e.g. M3's thinking_mode). - from .model_meta import thinking_toggle_kwargs + from .model_meta import apply_default_thinking_mode, thinking_toggle_kwargs ctk: dict[str, Any] = {} if req.thinking: @@ -318,6 +321,10 @@ def convert_anthropic_prompt( ctk = thinking_toggle_kwargs(True) elif req.thinking.get("type") == "disabled": ctk = thinking_toggle_kwargs(False) + else: + # Any thinking block (e.g. "adaptive") is the client's explicit choice; the server + # default only covers requests that leave thinking out. + ctk = apply_default_thinking_mode(ctk, default_thinking_mode) return render_messages(messages), template_tools, parser_tools, ctk @@ -330,12 +337,9 @@ def convert_anthropic_to_genspec( inline_system_policy: str = "auto", default_thinking_mode: str | None = None, ) -> GenSpec: - from .openai_api import apply_default_thinking_mode - messages, template_tools, parser_tools, ctk = convert_anthropic_prompt( - req, reasoning_parser=reasoning_parser + req, reasoning_parser=reasoning_parser, default_thinking_mode=default_thinking_mode ) - ctk = apply_default_thinking_mode(ctk, default_thinking_mode) return GenSpec( messages=messages, sampling_params=resolve_sampling( diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 4600bc68a0..f753c60757 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -52,10 +52,10 @@ class ServerArgs(SchedulerConfig): # responses. None disables it (default for models without a reasoning protocol). reasoning_parser: str | None = None # Server-wide default thinking mode for reasoning-capable models: "auto" keeps the - # current per-request behavior; "chat" forces enable_thinking=False for every request - # that does not explicitly set it in chat_template_kwargs (for OpenAI-compatible - # clients that never send template kwargs, like Vercel AI SDK or llama-swap); - # "thinking" forces thinking on the same way. + # current per-request behavior; "chat" turns thinking off for every request that does + # not choose itself (template kwargs or a protocol-level effort/thinking field), for + # OpenAI-compatible clients that never send template kwargs, like Vercel AI SDK or + # llama-swap; "thinking" turns it on the same way. default_thinking_mode: str = "auto" # "model": fill unspecified request sampling params from generation_config.json # (temperature/top_k/top_p), like sglang. "none": use framework defaults only. @@ -720,10 +720,11 @@ def _infer_reasoning_parser(model_path: str) -> str | None: choices=["auto", "chat", "thinking"], help=( "Server-wide default thinking mode for reasoning-capable models. 'auto' keeps " - "the current per-request behavior; 'chat' forces enable_thinking=False for " - "every request that does not explicitly set it in chat_template_kwargs (for " - "OpenAI-compatible clients that never send template kwargs, like Vercel AI SDK " - "or llama-swap); 'thinking' forces thinking on the same way." + "the current per-request behavior; 'chat' turns thinking off for every request " + "that does not choose itself (chat_template_kwargs, reasoning_effort, " + "reasoning.effort, thinking) -- for OpenAI-compatible clients that never send " + "template kwargs, like Vercel AI SDK or llama-swap; 'thinking' turns it on the " + "same way." ), ) diff --git a/python/freetoken/server/model_meta.py b/python/freetoken/server/model_meta.py index 8ce8009a05..c9187ad629 100644 --- a/python/freetoken/server/model_meta.py +++ b/python/freetoken/server/model_meta.py @@ -119,6 +119,25 @@ def effort_toggle_kwargs( return mapped +def apply_default_thinking_mode( + ctk: dict[str, Any] | None, + default_mode: str | None, +) -> dict[str, Any] | None: + """Merge the server's --default-thinking-mode into a request's template kwargs. + + An explicit per-request value always wins; the server default only fills + in what the request left unset. "auto" (the default) is a no-op. Adapters + apply it after folding their protocol-level toggles (effort_toggle_kwargs), + so those count as explicit too. + """ + if default_mode not in ("chat", "thinking"): + return ctk + ctk = dict(ctk or {}) + if not any(key in ctk for key in _THINKING_KWARG_KEYS): + ctk.update(thinking_toggle_kwargs(default_mode == "thinking")) + return ctk + + def moe_total_experts(config: Any) -> int: """Total routed-expert slots the model has: experts per layer x MoE layers. Matches the engine's own basis (``Engine._resolve_auto_moe_cache_size``), so a residency rate derived diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index b09c480e1d..54c10ac2e6 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -71,34 +71,6 @@ def _thinking_type(req: Any) -> str | None: -def apply_default_thinking_mode( - ctk: dict[str, Any] | None, - default_mode: str | None, -) -> dict[str, Any] | None: - """Merge the server's --default-thinking-mode into a request's template kwargs. - - An explicit per-request value always wins; the server default only fills - in what the request left unset. "auto" (the default) is a no-op. - """ - if not default_mode or default_mode == "auto": - return ctk - if ctk is None: - ctk = {} - else: - ctk = dict(ctk) - has_explicit = ( - "enable_thinking" in ctk - or "thinking" in ctk - or "thinking_mode" in ctk - ) - if not has_explicit: - if default_mode == "chat": - ctk["enable_thinking"] = False - elif default_mode == "thinking": - ctk["enable_thinking"] = True - return ctk - - def chat_request_to_genspec( req: ChatCompletionRequest, model_sampling: dict[str, Any], @@ -106,12 +78,13 @@ def chat_request_to_genspec( default_thinking_mode: str | None = None, ) -> GenSpec: """OpenAI ChatCompletionRequest -> GenSpec (the OpenAI 'to_sampling_params').""" - from .model_meta import effort_toggle_kwargs + from .model_meta import apply_default_thinking_mode, effort_toggle_kwargs - ctk = apply_default_thinking_mode(req.chat_template_kwargs, default_thinking_mode) + ctk = req.chat_template_kwargs thinking_type = _thinking_type(req) if req.reasoning_effort or thinking_type: ctk = effort_toggle_kwargs(req.reasoning_effort, ctk, thinking_type=thinking_type) + ctk = apply_default_thinking_mode(ctk, default_thinking_mode) if req.continue_final_message: ctk = {**ctk, "continue_final_message": True} sampling_params = _resolve_sampling(req, model_sampling, default_max_tokens=default_max_tokens) diff --git a/python/freetoken/server/responses_api.py b/python/freetoken/server/responses_api.py index 9fd063a992..a84c68cb6d 100644 --- a/python/freetoken/server/responses_api.py +++ b/python/freetoken/server/responses_api.py @@ -245,15 +245,12 @@ def convert_responses_to_genspec( else: template_tools, parser_tools = split_tool_lists(raw_tools, selected) - from .model_meta import effort_toggle_kwargs + from .model_meta import apply_default_thinking_mode, effort_toggle_kwargs - from .openai_api import apply_default_thinking_mode - - ctk = apply_default_thinking_mode( - dict(getattr(req, "chat_template_kwargs", None) or {}), default_thinking_mode - ) + ctk = dict(getattr(req, "chat_template_kwargs", None) or {}) if req.reasoning: ctk = effort_toggle_kwargs(req.reasoning.get("effort"), ctk) + ctk = apply_default_thinking_mode(ctk, default_thinking_mode) return GenSpec( messages=render_messages(messages), diff --git a/tests/server/test_default_thinking_mode.py b/tests/server/test_default_thinking_mode.py index ac979e625d..7f9938f60c 100644 --- a/tests/server/test_default_thinking_mode.py +++ b/tests/server/test_default_thinking_mode.py @@ -1,7 +1,18 @@ """Unit tests for the --default-thinking-mode server flag merge logic.""" from __future__ import annotations -from freetoken.server.openai_api import apply_default_thinking_mode +import asyncio +from types import SimpleNamespace + +from freetoken.server import anthropic_api as A +from freetoken.server.anthropic_models import AnthropicCountTokensRequest, AnthropicMessagesRequest +from freetoken.server.api_models import ChatCompletionRequest +from freetoken.server.model_meta import apply_default_thinking_mode +from freetoken.server.openai_api import chat_request_to_genspec +from freetoken.server.responses_api import ResponsesRequest, convert_responses_to_genspec + +OFF = {"enable_thinking": False, "thinking_mode": "disabled"} +ON = {"enable_thinking": True, "thinking_mode": "enabled"} def test_auto_is_noop(): @@ -13,17 +24,14 @@ def test_auto_is_noop(): def test_chat_fills_unset_kwargs(): - assert apply_default_thinking_mode(None, "chat") == {"enable_thinking": False} - assert apply_default_thinking_mode({}, "chat") == {"enable_thinking": False} - assert apply_default_thinking_mode({"other": 1}, "chat") == { - "other": 1, - "enable_thinking": False, - } + assert apply_default_thinking_mode(None, "chat") == OFF + assert apply_default_thinking_mode({}, "chat") == OFF + assert apply_default_thinking_mode({"other": 1}, "chat") == {"other": 1, **OFF} def test_thinking_fills_unset_kwargs(): - assert apply_default_thinking_mode(None, "thinking") == {"enable_thinking": True} - assert apply_default_thinking_mode({}, "thinking") == {"enable_thinking": True} + assert apply_default_thinking_mode(None, "thinking") == ON + assert apply_default_thinking_mode({}, "thinking") == ON def test_explicit_request_value_wins(): @@ -38,14 +46,78 @@ def test_explicit_request_value_wins(): assert apply_default_thinking_mode({"thinking_mode": "thinking"}, "chat") == { "thinking_mode": "thinking" } + # A raw graded-effort request asks for thinking too. + assert apply_default_thinking_mode({"reasoning_effort": "high"}, "chat") == { + "reasoning_effort": "high" + } def test_original_dict_not_mutated(): original = {"other": 1} result = apply_default_thinking_mode(original, "chat") - assert result == {"other": 1, "enable_thinking": False} + assert result == {"other": 1, **OFF} assert original == {"other": 1} def test_unknown_mode_is_noop(): - assert apply_default_thinking_mode({"other": 1}, "bogus") == {"other": 1} \ No newline at end of file + assert apply_default_thinking_mode({"other": 1}, "bogus") == {"other": 1} + + +def _chat(**extra): + body = {"model": "m", "messages": [{"role": "user", "content": "hi"}], **extra} + return chat_request_to_genspec( + ChatCompletionRequest.model_validate(body), {}, default_thinking_mode="chat" + ).chat_template_kwargs + + +def test_chat_default_fills_a_request_without_thinking_controls(): + assert _chat() == OFF + + +def test_chat_default_keeps_protocol_level_thinking_requests(): + assert _chat(reasoning_effort="high")["enable_thinking"] is True + assert _chat(thinking={"type": "enabled"})["enable_thinking"] is True + assert _chat(reasoning_effort="none")["enable_thinking"] is False + + +def test_responses_chat_default_keeps_reasoning_effort(): + def ctk(**extra): + req = ResponsesRequest.model_validate({"model": "m", "input": "hi", **extra}) + return convert_responses_to_genspec(req, {}, default_thinking_mode="chat").chat_template_kwargs + + assert ctk() == OFF + on = ctk(reasoning={"effort": "medium"}) + assert on["enable_thinking"] is True and on["reasoning_effort"] == "medium" + + +def _anthropic(**extra): + return AnthropicMessagesRequest.model_validate( + {"model": "claude-x", "max_tokens": 64, "messages": [{"role": "user", "content": "hi"}], **extra} + ) + + +def test_anthropic_chat_default_applies_only_without_a_thinking_block(): + def ctk(**extra): + return A.convert_anthropic_to_genspec( + _anthropic(**extra), {}, default_thinking_mode="chat" + ).chat_template_kwargs + + assert ctk() == OFF + assert ctk(thinking={"type": "enabled", "budget_tokens": 1024}) == ON + assert ctk(thinking={"type": "adaptive"}) == {} + + +def test_anthropic_count_tokens_renders_the_generation_prompt_under_the_default(monkeypatch): + captured = {} + + async def fake_count(messages, template_tools, ctk, state): + captured["ctk"] = ctk + return 1 + + monkeypatch.setattr(A, "count_prompt_tokens", fake_count) + state = SimpleNamespace(config=SimpleNamespace(reasoning_parser=None, default_thinking_mode="chat")) + body = {"model": "claude-x", "messages": [{"role": "user", "content": "hi"}]} + asyncio.run(A.handle_anthropic_count_tokens(AnthropicCountTokensRequest.model_validate(body), state)) + + generation = A.convert_anthropic_to_genspec(_anthropic(), {}, default_thinking_mode="chat") + assert captured["ctk"] == generation.chat_template_kwargs == OFF From de1dfc95a7130ff8b38b91bb616d7ad827b3abd4 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 05:57:18 -0400 Subject: [PATCH 48/94] fix(server): type allOf-wrapped refs and untyped enum/const tool parameters #435 sends any parameter it cannot type to the loose JSON parser, so an untyped string enum member such as "1" came back as the integer 1, outside the enum. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit ab2693c977c64e425d428468364e2a2743f170e4) --- .../freetoken/server/function_call_parser.py | 38 +++++++++- ..._function_call_parser_schema_resolution.py | 70 +++++++++++++++++++ 2 files changed, 106 insertions(+), 2 deletions(-) diff --git a/python/freetoken/server/function_call_parser.py b/python/freetoken/server/function_call_parser.py index 08c60a2993..88cc65b367 100644 --- a/python/freetoken/server/function_call_parser.py +++ b/python/freetoken/server/function_call_parser.py @@ -382,7 +382,27 @@ def _resolve_local_ref(self, schema: Dict, root_schema: Dict, _seen: Optional[se return self._resolve_local_ref(merged, root_schema, _seen) return merged - def _normalize_param_schema(self, schema: Dict, root_schema: Dict) -> Dict: + @staticmethod + def _literal_schema_type(schema: Dict) -> Optional[str]: + """Type of an untyped ``enum``/``const`` schema, inferred from its literals.""" + if "const" in schema: + values = [schema["const"]] + elif isinstance(schema.get("enum"), list): + values = schema["enum"] + else: + return None + values = [v for v in values if v is not None] + if not values: + return None + if all(isinstance(v, str) for v in values): + return "string" + if all(isinstance(v, bool) for v in values): + return "boolean" + if all(isinstance(v, int) and not isinstance(v, bool) for v in values): + return "integer" + return "loose" + + def _normalize_param_schema(self, schema: Dict, root_schema: Dict, _depth: int = 0) -> Dict: """Resolve the schema parts needed for tool argument typing.""" if not isinstance(schema, dict): return {"type": "string"} @@ -395,6 +415,18 @@ def _normalize_param_schema(self, schema: Dict, root_schema: Dict) -> Dict: return result if isinstance(param_type, str): return schema + all_of = schema.get("allOf") + if ( + isinstance(all_of, list) + and len(all_of) == 1 + and isinstance(all_of[0], dict) + and _depth < 16 # bounds a self-referential allOf chain + ): + # Pydantic v1 / OpenAPI 3.0 wrap a $ref in a one-element allOf to attach + # siblings such as default/description; type it like the bare $ref. + merged = dict(all_of[0]) + merged.update({k: v for k, v in schema.items() if k != "allOf"}) + return self._normalize_param_schema(merged, root_schema, _depth + 1) for keyword in ("oneOf", "anyOf"): branches = schema.get(keyword) if not isinstance(branches, list): @@ -405,6 +437,8 @@ def _normalize_param_schema(self, schema: Dict, root_schema: Dict) -> Dict: continue branch = self._resolve_local_ref(branch, root_schema) branch_type = branch.get("type") + if not isinstance(branch_type, str): + branch_type = self._literal_schema_type(branch) if isinstance(branch_type, str) and branch_type != "null": types.append(branch_type) types = list(dict.fromkeys(types)) @@ -420,7 +454,7 @@ def _normalize_param_schema(self, schema: Dict, root_schema: Dict) -> Dict: result["type"] = "array" return result result = dict(schema) - result["type"] = "loose" + result["type"] = self._literal_schema_type(schema) or "loose" return result def _get_tool_schema(self, func_name: str, tools: List[Tool]) -> Dict: diff --git a/tests/server/test_function_call_parser_schema_resolution.py b/tests/server/test_function_call_parser_schema_resolution.py index fbc0fe270b..468acc3bd2 100644 --- a/tests/server/test_function_call_parser_schema_resolution.py +++ b/tests/server/test_function_call_parser_schema_resolution.py @@ -375,6 +375,67 @@ def _parse_args(parser_name: str, prop_schema: dict, root_extra: dict | None, ra 5, id="ambiguous-union-loose-json", ), + + # An enum/const with no "type" takes its type from the literals, so a + # numeric-looking string member is not coerced out of the enum. + pytest.param( + {"enum": ["1", "2"]}, + None, + "1", + "1", + id="untyped-string-enum", + ), + pytest.param( + {"enum": [1, 2]}, + None, + "2", + 2, + id="untyped-integer-enum", + ), + pytest.param( + {"const": "10"}, + None, + "10", + "10", + id="untyped-string-const", + ), + pytest.param( + { + "anyOf": [ + {"enum": ["1", "2"]}, + {"type": "null"}, + ] + }, + None, + "1", + "1", + id="anyof-untyped-string-enum-nullable", + ), + + # Pydantic v1 / OpenAPI 3.0 wrap a $ref in a one-element allOf to attach + # siblings such as default/description. + pytest.param( + {"allOf": [{"$ref": "#/$defs/Mode"}], "default": "a"}, + { + "$defs": { + "Mode": {"type": "string", "enum": ["a", "true"]}, + } + }, + "true", + "true", + id="allof-ref-string-enum", + ), + pytest.param( + {"allOf": [{"$ref": "#/$defs/Limit"}], "description": "max rows"}, + { + "$defs": { + "Limit": {"type": "integer"}, + } + }, + "5", + 5, + id="allof-ref-integer", + ), ] @@ -454,6 +515,15 @@ def test_schema_aware_parsers_resolve_indirect_schema_types( }, id="mutually-recursive-ref", ), + pytest.param( + {"allOf": [{"$ref": "#/$defs/A"}]}, + { + "$defs": { + "A": {"allOf": [{"$ref": "#/$defs/A"}]}, + } + }, + id="self-referential-allof", + ), ] From 9e677ee5c5135cb5ec52b21770b16005a3c932c7 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 06:03:44 -0400 Subject: [PATCH 49/94] fix(server): keep /v1/models/ listing the model after the retrieve route Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 7ccc433dcc8608b23f33c1a4be6313b5f9307c47) --- python/freetoken/server/openai_api.py | 3 +++ tests/server/test_openai_api.py | 26 ++++++++++++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index 54c10ac2e6..7f94efc250 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -176,6 +176,9 @@ async def v1_model(model_id: str): OpenAI-compatible clients may use this endpoint to validate a model name before sending completion requests. """ + if not model_id: + # The :path converter also matches "/v1/models/", which used to redirect to the list. + return await v1_models() state = get_state() served_model_id = _served_model_name(state) diff --git a/tests/server/test_openai_api.py b/tests/server/test_openai_api.py index 376cc2a7ba..7e6d4ed30e 100644 --- a/tests/server/test_openai_api.py +++ b/tests/server/test_openai_api.py @@ -523,6 +523,32 @@ def test_models_route_prefers_the_last_rebuild_over_the_load_time_pool(): assert card["max_model_len"] == 40000 +def test_model_retrieve_route_returns_the_served_card_or_404(): + state = FakeState([]) + state.config.max_seq_len = 262144 + app = FastAPI() + register_openai_routes(app, lambda: state, lambda: {}) + client = TestClient(app) + + card = client.get("/v1/models/unit-model").json() + assert card["id"] == "unit-model" and card["context_length"] == 262144 + + missing = client.get("/v1/models/other-model") + assert missing.status_code == 404 + assert missing.json()["error"]["code"] == "model_not_found" + + +def test_models_route_with_a_trailing_slash_still_lists(): + state = FakeState([]) + app = FastAPI() + register_openai_routes(app, lambda: state, lambda: {}) + + response = TestClient(app).get("/v1/models/") + + assert response.status_code == 200 + assert response.json()["data"][0]["id"] == "unit-model" + + async def _collect(generator): return [chunk async for chunk in generator] From a1b477863eb3c42b7d214d852a2eb76ae0f43207 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 06:22:47 -0400 Subject: [PATCH 50/94] fix(server): open the streamed thinking block with the signature field The full response and the closing signature_delta already send signature="", but exclude_none dropped it from content_block_start, where the SDK types it as required. Taken from upstream #517 (b0442764b014bd32a68ba29055eb20eda9bfb026). Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 55a8e72f21cb679b585cfe6f7f738aafedeb7369) --- python/freetoken/server/anthropic_api.py | 4 +++- tests/server/test_anthropic_api.py | 3 ++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/python/freetoken/server/anthropic_api.py b/python/freetoken/server/anthropic_api.py index a7c88b141a..545cf4e7cc 100644 --- a/python/freetoken/server/anthropic_api.py +++ b/python/freetoken/server/anthropic_api.py @@ -480,7 +480,9 @@ def _open_thinking() -> str: block_open = "thinking" return _event(AnthropicStreamEvent( type="content_block_start", index=block_index, - content_block=AnthropicContentBlock(type="thinking", thinking=""), + # The SDK types signature as a required string on the opening block, before any + # signature_delta. Local reasoning is unsigned, like the full-response path. + content_block=AnthropicContentBlock(type="thinking", thinking="", signature=""), )) def _stop_block() -> list[str]: diff --git a/tests/server/test_anthropic_api.py b/tests/server/test_anthropic_api.py index 1a09b83a1c..890339522c 100644 --- a/tests/server/test_anthropic_api.py +++ b/tests/server/test_anthropic_api.py @@ -431,7 +431,8 @@ def test_stream_reasoning_as_thinking_block(): "message_stop", ] blocks = [e[1] for e in collected if e[0] == "content_block_start"] - assert blocks[0]["content_block"]["type"] == "thinking" + # signature is a required string in the SDK's thinking block, present from the start. + assert blocks[0]["content_block"] == {"type": "thinking", "thinking": "", "signature": ""} assert blocks[0]["index"] == 0 assert blocks[1]["content_block"]["type"] == "text" assert blocks[1]["index"] == 1 From 6cd147bb3f2daab15b87ce7a706383a0494c22a9 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 06:22:47 -0400 Subject: [PATCH 51/94] test(server): exercise detector SDK and lifecycle contracts Upstream #517 tests without its docs and pyproject dev-extra: the SDK tests importorskip anthropic/openai/httpx so the server suite still collects without them. (cherry picked from commit b0442764b014bd32a68ba29055eb20eda9bfb026, tests only) Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 50f99d167304858d9eaebfbcd3d00aa3c1f06430) --- tests/server/test_accounting.py | 121 ++++++++++++ tests/server/test_function_call_parser.py | 129 +++++++++++++ tests/server/test_protocol_compat.py | 213 ++++++++++++++++++++++ tests/server/test_reasoning_parser.py | 108 +++++++++++ tests/server/test_supervisor.py | 37 ++++ 5 files changed, 608 insertions(+) create mode 100644 tests/server/test_accounting.py create mode 100644 tests/server/test_protocol_compat.py create mode 100644 tests/server/test_reasoning_parser.py diff --git a/tests/server/test_accounting.py b/tests/server/test_accounting.py new file mode 100644 index 0000000000..fe29d414d8 --- /dev/null +++ b/tests/server/test_accounting.py @@ -0,0 +1,121 @@ +"""Failure/retry and concurrency contracts for the terminal accounting barrier.""" +import asyncio +from types import SimpleNamespace + +import pytest + +from freetoken.message import UserReply +from freetoken.server.accounting import AccountingDrainError, prepare_stop_accounting +from freetoken.server.stats import StatsTracker + + +def state_with_request(): + stats = StatsTracker() + stats.on_new_user(7) + return SimpleNamespace(stats=stats, maintenance_state="serving", instance_id="instance", + ready_at=None, config=SimpleNamespace(served_model_name="test")) + + +def terminal(state, tokens=3): + state.stats.observe(UserReply(uid=7, incremental_output="", finished=True, + completion_tokens_delta=tokens, error="aborted")) + + +def test_drain_timeout_then_late_terminal_can_retry_without_losing_tokens(): + async def run(): + state = state_with_request() + aborted = [] + + async def abort(uid): + aborted.append(uid) + state.stats.on_abort(uid) # enqueueing abort is NOT a terminal reply + + state.abort_user = abort + with pytest.raises(AccountingDrainError, match="abort barrier timed out"): + await prepare_stop_accounting(state, drain_timeout_s=0, abort_timeout_s=0) + assert state.maintenance_state == "stopping" + assert state.stats.active == 1 and not hasattr(state, "_sealed_accounting") + terminal(state, tokens=11) # last sampled-token delta arrives AFTER the timeout + sealed = await prepare_stop_accounting(state, drain_timeout_s=0, abort_timeout_s=0) + assert sealed["completion_tokens_total"] == 11 + assert sealed["drain_complete"] is True + assert aborted == [7] + sealed["completion_tokens_total"] = -1 + assert (await prepare_stop_accounting(state))["completion_tokens_total"] == 11 + asyncio.run(run()) + + +def test_concurrent_prepare_stop_serializes_abort_and_returns_independent_snapshots(): + async def run(): + state = state_with_request() + entered, release = asyncio.Event(), asyncio.Event() + aborted = [] + + async def abort(uid): + assert state.maintenance_state == "stopping" + aborted.append(uid) + entered.set() + await release.wait() + terminal(state) + + state.abort_user = abort + first = asyncio.create_task(prepare_stop_accounting(state, drain_timeout_s=0)) + await entered.wait() + second = asyncio.create_task(prepare_stop_accounting(state, drain_timeout_s=0)) + await asyncio.sleep(0) # second caller reaches the shared lock, no wall-clock race + assert not second.done() + release.set() + a, b = await asyncio.gather(first, second) + assert a == b and a is not b + assert aborted == [7] + assert a["completion_tokens_total"] == 3 + asyncio.run(run()) + + +def test_abort_transport_failure_preserves_engine_and_closes_gate(): + state = state_with_request() + + async def abort(uid): + raise ConnectionError("scheduler disconnected") + + state.abort_user = abort + with pytest.raises(AccountingDrainError, match="failed to abort active requests") as error: + asyncio.run(prepare_stop_accounting(state, drain_timeout_s=0, abort_timeout_s=0)) + assert isinstance(error.value.__cause__, ConnectionError) + assert state.maintenance_state == "stopping" and state.stats.active == 1 + assert not hasattr(state, "_sealed_accounting") + + +def test_unidentified_active_request_cannot_be_sealed(): + state = state_with_request() + state.stats = SimpleNamespace(active=1, inflight_uids=()) + with pytest.raises(AccountingDrainError, match="unidentified"): + asyncio.run(prepare_stop_accounting(state, drain_timeout_s=0)) + assert state.maintenance_state == "stopping" + assert not hasattr(state, "_sealed_accounting") + + +@pytest.mark.parametrize("maintenance", ["rebuilding", "unknown"]) +def test_invalid_prepare_transition_preserves_previous_state(maintenance): + state = state_with_request() + state.maintenance_state = maintenance + with pytest.raises(AccountingDrainError): + asyncio.run(prepare_stop_accounting(state, drain_timeout_s=0)) + assert state.maintenance_state == maintenance + assert not hasattr(state, "_sealed_accounting") + + +def test_cancelled_drain_releases_lock_without_reopening_admission(): + async def run(): + state = state_with_request() + task = asyncio.create_task(prepare_stop_accounting(state, drain_timeout_s=10)) + await asyncio.sleep(0) # admission closes before the first drain wait + assert state.maintenance_state == "stopping" + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert not state._accounting_prepare_lock.locked() + assert not hasattr(state, "_sealed_accounting") + terminal(state) + assert (await prepare_stop_accounting(state))["drain_complete"] is True + asyncio.run(run()) diff --git a/tests/server/test_function_call_parser.py b/tests/server/test_function_call_parser.py index bc22296760..23cd12272e 100644 --- a/tests/server/test_function_call_parser.py +++ b/tests/server/test_function_call_parser.py @@ -371,3 +371,132 @@ def test_streaming_support_flags(): # test_streaming_model_matrix.py::test_non_streaming_detector_falls_back_to_buffered_parse). for name in SUPPORTED_TOOL_CALL_PARSERS: assert FunctionCallParser(TOOLS, tool_call_parser=name).supports_streaming() is True + + +# One class per concrete detector makes format coverage visible in pytest output. +# A shared contract avoids accidentally giving a newly added format weaker checks. +class _DetectorContract: + parser_name: str + block: str + truncate_before: str + buffered_recovery = False + + def parser(self): + return FunctionCallParser(OPENCODE_TOOLS, self.parser_name) + + def test_complete_call(self): + result = self.parser().parse_non_stream(self.block) + assert result.normal_text == "" + assert len(result.calls) == 1 + assert result.calls[0].name == "read" + assert json.loads(result.calls[0].parameters) == {"filePath": "/tmp/test_calc.py"} + + def test_every_character_boundary(self): + parser = self.parser() + texts, calls = _feed(parser, self.block) + assert "".join(texts) + parser.finish_stream() == "" + assert [call.name for call in calls if call.name] == ["read"] + assert json.loads("".join(call.parameters for call in calls)) == { + "filePath": "/tmp/test_calc.py" + } + # A completed call must never be duplicated by the end-of-stream recovery. + assert parser.recover_truncated_call() == [] + + def test_truncated_mid_call_recovery(self): + parser = self.parser() + cut = self.block.rindex(self.truncate_before) + texts, calls = _feed(parser, self.block[:cut]) + recovered = parser.recover_truncated_call() + assert bool(recovered) is self.buffered_recovery + # Incremental detectors have already consumed the opener and emitted the + # call identity. Their recovery method MUST NOT invent a duplicate call: + # generation.py recovers missing arguments from their partial-parse ledger. + # Buffered M3 instead recovers a complete call from the retained markup. + all_calls = calls + recovered + assert [call.name for call in all_calls if call.name] == ["read"] + arguments = "".join(call.parameters for call in all_calls) + try: + parsed = json.loads(arguments) + except json.JSONDecodeError: + parsed = json.loads(parser.unstreamed_arguments(0)) + assert parsed == {"filePath": "/tmp/test_calc.py"} + assert parser.recover_truncated_call() == [] + assert "".join(texts) + parser.finish_stream() == "" + + +class TestQwen25Detector(_DetectorContract): + parser_name = 'qwen25' + block = '{"name": "read", "arguments": {"filePath": "/tmp/test_calc.py"}}' + truncate_before = '' + + +class TestMistralDetector(_DetectorContract): + parser_name = 'mistral' + block = '[TOOL_CALLS] [{"name": "read", "arguments": {"filePath": "/tmp/test_calc.py"}}]' + truncate_before = '}]' + + +class TestLlama32Detector(_DetectorContract): + parser_name = 'llama3' + block = '<|python_tag|>{"name": "read", "arguments": {"filePath": "/tmp/test_calc.py"}}' + truncate_before = '}' + + +class TestGlm47Detector(_DetectorContract): + parser_name = 'glm47' + block = 'readfilePath/tmp/test_calc.py' + truncate_before = '' + + +class TestDeepSeekV32Detector(_DetectorContract): + parser_name = 'deepseekv32' + block = '<｜DSML｜function_calls><｜DSML｜invoke name="read"><｜DSML｜parameter name="filePath" string="true">/tmp/test_calc.py' + truncate_before = '' + + +class TestQwen3CoderDetector(_DetectorContract): + parser_name = 'qwen3_coder' + block = '/tmp/test_calc.py' + truncate_before = '' + + +class TestGemma4Detector(_DetectorContract): + parser_name = 'gemma4' + block = '<|tool_call>call:read{filePath:<|"|>/tmp/test_calc.py<|"|>}' + truncate_before = '' + + +class TestMiniMaxDetector(_DetectorContract): + parser_name = 'minimax' + block = '/tmp/test_calc.py' + truncate_before = '' + + +class TestMiniMaxM3Detector(_DetectorContract): + parser_name = 'minimax_m3' + block = ']<]minimax[>[\n]<]minimax[>[]<]minimax[>[/tmp/test_calc.py]<]minimax[>[]<]minimax[>[\n]<]minimax[>[' + truncate_before = ']<]minimax[>[' + buffered_recovery = True + + +class TestGptOssDetector(_DetectorContract): + parser_name = 'gpt_oss' + block = '<|start|>assistant<|channel|>commentary to=functions.read <|constrain|>json<|message|>{"filePath": "/tmp/test_calc.py"}<|end|>' + truncate_before = '<|end|>' + + +class TestMuseGlimmerDetector(_DetectorContract): + parser_name = 'muse_glimmer' + block = '<|start|>assistant to=read<|message|>\n\n/tmp/test_calc.py\n\n<|eot|>' + truncate_before = '' + + +def test_contract_classes_cover_every_concrete_detector(): + from freetoken.server.function_call_parser import BaseFormatDetector + + covered = { + type(cls().parser().detector) + for cls in _DetectorContract.__subclasses__() + } + assert covered == set(BaseFormatDetector.__subclasses__()) + assert covered == set(FunctionCallParser.ToolCallParserEnum.values()) diff --git a/tests/server/test_protocol_compat.py b/tests/server/test_protocol_compat.py new file mode 100644 index 0000000000..8dc22908d6 --- /dev/null +++ b/tests/server/test_protocol_compat.py @@ -0,0 +1,213 @@ +"""Real SDK model/stream validation over the real HTTP adapters and GenSpec path. + +Only the scheduler/tokenizer transport is scripted. In-process ASGI transport +keeps these deterministic and CPU-only; it does not prove TCP chunk timing, +disconnect behavior, tokenizer templates, or a real model's output quality. +""" +import asyncio +import json +from types import SimpleNamespace + +import pytest +from fastapi import FastAPI + +from freetoken.message import TokenizeMsg, UserReply +from freetoken.server.anthropic_api import register_anthropic_routes +from freetoken.server.openai_api import register_openai_routes +from freetoken.server.responses_api import register_responses_routes + +# The SDKs are optional test dependencies: skip what needs a missing one instead of +# failing collection of the whole server suite. The Anthropic half skips per test. +httpx = pytest.importorskip("httpx") +openai = pytest.importorskip("openai") + + +SCHEMA = {"type": "object", "properties": {"city": {"type": "string"}}} +TOOL = {"type": "function", "function": {"name": "get_weather", "parameters": SCHEMA}} +CALL = '{"name":"get_weather","arguments":{"city":"Paris"}}' + + +class ScriptedBackend: + def __init__(self, text, *, reasoning=False): + self.config = SimpleNamespace( + model_path="/models/test", served_model_name="test", max_seq_len=4096, + tool_call_parser="qwen25", reasoning_parser="qwen3" if reasoning else None, + ) + self.text = text + self.sent = [] + self.maintenance_state = "serving" + + def new_user(self): + return 42 + + async def send_one(self, message): + assert isinstance(message, TokenizeMsg) + self.sent.append(message) + + async def wait_for_ack(self, uid): + # One character per backend reply exercises adapter stream aggregation + # as well as the SDK SSE decoder. Usage deltas remain authoritative. + for i, char in enumerate(self.text): + yield UserReply(uid=uid, incremental_output=char, + finished=i == len(self.text) - 1, + prompt_tokens_delta=5 if i == 0 else 0, + completion_tokens_delta=1) + + async def stream_with_cancellation(self, events, request, uid): + async for event in events: + yield event + + +def client(state, vendor="openai"): + app = FastAPI() + for register in (register_openai_routes, register_anthropic_routes, register_responses_routes): + register(app, lambda: state, dict) + # Recent Anthropic SDKs use httpx2 and reject httpx client instances. + # Select the transport matching its public default client class; older SDKs + # in the supported range still use httpx. Neither path bypasses SDK validation. + http = httpx + if vendor == "anthropic": + anthropic = pytest.importorskip("anthropic") + if not issubclass(anthropic.DefaultAsyncHttpxClient, httpx.AsyncClient): + http = pytest.importorskip("httpx2") + transport = http.AsyncClient(transport=http.ASGITransport(app=app)) + cls = openai.AsyncOpenAI if vendor == "openai" else anthropic.AsyncAnthropic + return cls(base_url="http://test/v1" if vendor == "openai" else "http://test", + api_key="test", http_client=transport, max_retries=0, + _strict_response_validation=True) + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("tool", [False, True]) +def test_openai_chat_sdk_parses_reasoning_text_tools_and_usage(stream, tool): + async def run(): + state = ScriptedBackend("consider" + (CALL if tool else "Hello"), reasoning=True) + async with client(state) as sdk: + response = await sdk.chat.completions.create( + model="test", messages=[{"role": "user", "content": "Hi"}], + tools=[TOOL], stream=stream, stream_options={"include_usage": True}, + ) + if stream: + chunks = [chunk async for chunk in response] + choices = [choice for chunk in chunks for choice in chunk.choices] + assert "".join(getattr(c.delta, "reasoning_content", "") or "" for c in choices) == "consider" + assert [c.finish_reason for c in choices if c.finish_reason] == ["tool_calls" if tool else "stop"] + if tool: + calls = [tc for c in choices for tc in c.delta.tool_calls or []] + assert [tc.function.name for tc in calls if tc.function.name] == ["get_weather"] + assert json.loads("".join(tc.function.arguments or "" for tc in calls)) == {"city": "Paris"} + assert all(tc.index == 0 for tc in calls) + else: + assert "".join(c.delta.content or "" for c in choices) == "Hello" + usage = chunks[-1].usage + else: + assert response.choices[0].message.reasoning_content == "consider" + assert response.choices[0].finish_reason == ("tool_calls" if tool else "stop") + if tool: + call = response.choices[0].message.tool_calls[0] + assert call.id and call.function.name == "get_weather" + assert json.loads(call.function.arguments) == {"city": "Paris"} + else: + assert response.choices[0].message.content == "Hello" + usage = response.usage + assert usage.prompt_tokens == 5 + assert usage.completion_tokens == len(state.text) + assert usage.total_tokens == 5 + len(state.text) + assert len(state.sent) == 1 + asyncio.run(run()) + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("tool", [False, True]) +def test_anthropic_sdk_assembles_typed_message(stream, tool): + pytest.importorskip("anthropic") + + async def run(): + state = ScriptedBackend("consider" + (CALL if tool else "Hello"), reasoning=True) + async with client(state, "anthropic") as sdk: + kwargs = dict(model="test", max_tokens=256, + messages=[{"role": "user", "content": "Hi"}], + tools=[{"name": "get_weather", "input_schema": SCHEMA}]) + if stream: + async with sdk.messages.stream(**kwargs) as response: + events = [event async for event in response] + result = await response.get_final_message() + assert events[0].type == "message_start" + assert events[-1].type == "message_stop" + else: + result = await sdk.messages.create(**kwargs) + assert result.role == "assistant" and result.model == "test" + assert result.stop_reason == ("tool_use" if tool else "end_turn") + assert result.content[0].type == "thinking" + assert result.content[0].thinking == "consider" + if tool: + # The established full-response shape includes an empty text + # block before tools; preserving that is part of compatibility. + assert all(b.text == "" for b in result.content if b.type == "text") + block = next(b for b in result.content if b.type == "tool_use") + assert block.id + assert block.name == "get_weather" and block.input == {"city": "Paris"} + else: + block = result.content[1] + assert block.type == "text" and block.text == "Hello" + assert result.usage.input_tokens == 5 + assert result.usage.output_tokens == len(state.text) + assert len(state.sent) == 1 + asyncio.run(run()) + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("tool", [False, True]) +def test_responses_sdk_parses_output_and_events(stream, tool): + async def run(): + state = ScriptedBackend(CALL if tool else "Hello") + async with client(state) as sdk: + response = await sdk.responses.create( + model="test", input="Hi", stream=stream, + tools=[{"type": "function", "name": "get_weather", "parameters": SCHEMA}], + ) + if stream: + events = [event async for event in response] + assert events[0].type == "response.created" + assert events[-1].type == "response.completed" + assert [e.sequence_number for e in events] == list(range(len(events))) + result = events[-1].response + deltas = [e for e in events if e.type == ("response.function_call_arguments.delta" if tool else "response.output_text.delta")] + text = "".join(e.delta for e in deltas) + assert json.loads(text) == {"city": "Paris"} if tool else text == "Hello" + else: + result = response + assert result.status == "completed" + if tool: + call = next(item for item in result.output if item.type == "function_call") + assert call.call_id and call.name == "get_weather" + assert json.loads(call.arguments) == {"city": "Paris"} + else: + assert result.output_text == "Hello" + assert result.usage.input_tokens == 5 + assert result.usage.output_tokens == len(state.text) + assert result.usage.total_tokens == 5 + len(state.text) + assert len(state.sent) == 1 + asyncio.run(run()) + + +@pytest.mark.parametrize("protocol", ["chat", "responses", "anthropic"]) +def test_sdk_classifies_invalid_request_as_400(protocol): + vendor = "anthropic" if protocol == "anthropic" else "openai" + sdk_module = pytest.importorskip("anthropic") if vendor == "anthropic" else openai + + async def run(): + state = ScriptedBackend("unused") + error = sdk_module.BadRequestError + async with client(state, vendor) as sdk: + with pytest.raises(error) as caught: + if protocol == "chat": + await sdk.chat.completions.create(model="test", messages=[{"role": "user", "content": "Hi"}], n=2) + elif protocol == "responses": + await sdk.responses.create(model="test", input="Hi", background=True) + else: + await sdk.messages.create(model="test", max_tokens=0, messages=[{"role": "user", "content": "Hi"}]) + assert caught.value.status_code == 400 + assert "invalid_request_error" in str(caught.value) + assert state.sent == [] + asyncio.run(run()) diff --git a/tests/server/test_reasoning_parser.py b/tests/server/test_reasoning_parser.py new file mode 100644 index 0000000000..05c207d0e1 --- /dev/null +++ b/tests/server/test_reasoning_parser.py @@ -0,0 +1,108 @@ +"""Character-boundary contracts for every reasoning detector and template mode.""" +from types import SimpleNamespace + +import pytest + +from freetoken.server.generation import GenSpec, _make_reasoning_parser +from freetoken.server.reasoning_parser import BaseReasoningParser, ReasoningParser + + +class _ReasoningContract: + name: str + opener: str + closer: str + content_opener = "" + + def parser(self): + return ReasoningParser(self.name, force_reasoning=False) + + def stream(self, text): + parser = self.parser() + parts = [parser.parse_stream_chunk(char) for char in text] + parts.append(parser.flush()) + assert parser.flush() == ("", "") + return tuple("".join(pair[i] for pair in parts) for i in (0, 1)) + + def test_complete_reasoning(self): + text = self.opener + "consider" + self.closer + self.content_opener + "answer" + assert self.parser().parse_non_stream(text) == ("consider", "answer") + + def test_every_character_boundary(self): + text = self.opener + "consider" + self.closer + self.content_opener + "answer" + assert self.stream(text) == ("consider", "answer") + + def test_truncated_reasoning_flush(self): + # Reasoning detectors expose flush(), not recover_truncated_call(). EOS + # inside an open reasoning body must retain it in the reasoning channel. + assert self.stream(self.opener + "partial thought") == ("partial thought", "") + + +class TestDeepSeekV32ReasoningParser(_ReasoningContract): + name = "deepseekv32" + opener = "" + closer = "" + + +class TestThinkReasoningParser(_ReasoningContract): + name = "qwen3" + opener = "" + closer = "" + + +class TestMiniMaxM3ReasoningParser(_ReasoningContract): + name = "minimax_m3" + opener = "" + closer = "" + + +class TestGemmaThoughtReasoningParser(_ReasoningContract): + name = "gemma4" + opener = "<|channel>thought\n" + closer = "" + + +class TestGptOssHarmonyReasoningParser(_ReasoningContract): + name = "gpt_oss" + opener = "<|channel|>analysis<|message|>" + closer = "<|end|>" + content_opener = "<|start|>assistant<|channel|>final<|message|>" + + +class TestMuseGlimmerReasoningParser(_ReasoningContract): + name = "muse_glimmer" + opener = "<|start|>assistant to=self<|message|>" + closer = "<|eom|>" + content_opener = "<|start|>assistant to=user<|message|>" + + +def test_contract_classes_cover_every_reasoning_detector(): + covered = {type(cls().parser().detector) for cls in _ReasoningContract.__subclasses__()} + assert covered == set(BaseReasoningParser.__subclasses__()) + assert covered == set(ReasoningParser.ReasoningParserEnum.values()) + + +@pytest.mark.parametrize("tools", [None, [{"type": "function", "function": {"name": "read"}}]]) +@pytest.mark.parametrize(("name", "kwargs", "forced"), [ + ("qwen3", {}, True), + ("qwen3", {"enable_thinking": True}, True), + ("qwen3", {"enable_thinking": False}, False), + ("glm", {}, True), + ("glm", {"enable_thinking": True}, True), + ("glm", {"enable_thinking": False}, False), + ("gemma4", {}, False), + ("gemma4", {"thinking_mode": "thinking"}, True), + ("gemma4", {"enable_thinking": True}, True), + ("gemma4", {"thinking": True}, True), + ("gemma4", {"thinking_mode": "chat", "enable_thinking": False}, False), + ("minimax_m3", {}, False), + ("minimax_m3", {"thinking_mode": "adaptive"}, False), + ("minimax_m3", {"thinking_mode": "disabled"}, False), + ("minimax_m3", {"thinking_mode": "enabled"}, True), +]) +def test_generation_force_reasoning_matches_template(name, kwargs, forced, tools): + spec = GenSpec(messages=[], sampling_params=None, chat_template_kwargs=kwargs, + template_tools=tools, parser_tools=tools) + parser = _make_reasoning_parser(spec, SimpleNamespace(config=SimpleNamespace(reasoning_parser=name))) + # Observe routing, not an internal flag: tools must not turn a disabled + # template's ordinary response into invisible reasoning. + assert parser.parse_non_stream("probe") == (("probe", "") if forced else ("", "probe")) diff --git a/tests/server/test_supervisor.py b/tests/server/test_supervisor.py index 3568d90505..c4772603b0 100644 --- a/tests/server/test_supervisor.py +++ b/tests/server/test_supervisor.py @@ -1,5 +1,7 @@ from __future__ import annotations +import pytest + import queue from queue import Empty as _Empty @@ -270,3 +272,38 @@ def test_byte_bar_emits_to_installed_sink_then_stops_after_clear(): bar.update(100) bar.close() assert seen == [] + + +@pytest.mark.parametrize("shutdown", [False, True]) +def test_shutdown_flag_is_read_at_later_worker_death(monkeypatch, shutdown): + """Readiness must not snapshot shutdown state: the flag changes on a later poll. + + Keep TP0 alive and kill a different worker, so the first process cannot mask + a dead tokenizer. No real processes or sleep timing enter the assertion. + """ + import queue + from types import SimpleNamespace + from freetoken.server import supervisor + + alive = {"scheduler": True, "tokenizer": True} + processes = [SimpleNamespace(name=name, is_alive=lambda name=name: alive[name]) + for name in alive] + acks = queue.Queue() + acks.put("scheduler ready") + acks.put("tokenizer ready") + handle = BackendHandle(ack_queue=acks, processes=processes, expected_acks=2) + stopping = False + seen = [] + + def next_poll(_): + nonlocal stopping + stopping = shutdown + alive["tokenizer"] = False + + monkeypatch.setattr(supervisor.time, "sleep", next_poll) + supervisor.run_backend_supervisor( + handle, LoadProgress(), on_ready=lambda: seen.append("ready"), + on_failure=lambda message: seen.append(message), + is_shutting_down=lambda: stopping, + ) + assert seen == (["ready"] if shutdown else ["ready", "backend worker tokenizer exited"]) From 7ac93f59806d0a3b2dddea44fbbea23911881d50 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 06:49:51 -0400 Subject: [PATCH 52/94] fix(server): drop /v1/completions from the --enable-metrics-report help, which it does not serve Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit c6b7ce7cc462c99c561a08071ec9cd308fdcc26e) --- python/freetoken/server/args.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index f753c60757..393832258c 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -643,7 +643,7 @@ def _infer_reasoning_parser(model_path: str) -> str | None: default=ServerArgs.enable_metrics_report, help=( "Serve a per-request `metrics` object next to usage on /v1/chat/completions, " - "/v1/completions, /v1/messages and /v1/responses: ttft_ms, prefill_time_ms and " + "/v1/messages and /v1/responses: ttft_ms, prefill_time_ms and " "prefill_tokens_per_second (the prefill span measured by the scheduler itself), " "decode_time_ms and decode_tokens_per_second, cached_prompt_tokens and " "total_time_ms. Streaming responses carry it on the same final chunk as usage, so " From ebe78888fad5c27a45598631e05b0ba5679aad31 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 10:36:38 -0400 Subject: [PATCH 53/94] fix(server): keep developer ahead of system and reject images when merging system messages render_messages hoisted system ahead of a leading developer message, which the tokenizer then merged as "S\n\nD"; it now leaves developer requests to the tokenizer's per-template mapping. Both merge sites joined str() of an image part list into the prompt, data URL included; they now raise the template's 400 instead. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit e60b2c39d78626e8ecc5c21c7f5b1d77664698ba) --- python/freetoken/server/generation.py | 13 ++++++----- python/freetoken/tokenizer/tokenize.py | 13 +++++++++-- tests/server/test_mm_input.py | 32 ++++++++++++++++++++++++++ tests/tokenizer/test_tokenize.py | 3 +++ 4 files changed, 53 insertions(+), 8 deletions(-) diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index b86bcab0f0..0ed933feb8 100644 --- a/python/freetoken/server/generation.py +++ b/python/freetoken/server/generation.py @@ -25,7 +25,7 @@ from freetoken.core import SamplingParams from freetoken.message import TokenizeMsg from freetoken.mm.media import collect_image_refs, fetch_image_bytes, image_reject_reason -from freetoken.tokenizer.tokenize import resolve_thinking_mode +from freetoken.tokenizer.tokenize import join_system_contents, resolve_thinking_mode from freetoken.tokenizer.inline_system import InlineSystemError try: @@ -525,6 +525,11 @@ def render_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: hoist system messages to the front and merge multiples into one to satisfy that constraint.""" rendered = [_render_message(m) for m in messages] + if any(m.get("role") == "developer" for m in rendered): + # The tokenizer maps developer per template and then merges in order + # (_map_developer_role); hoisting system alone here would put it ahead of + # a developer message that came first. + return rendered system_msgs = [m for m in rendered if m.get("role") == "system"] if system_msgs and rendered[0].get("role") != "system": non_system = [m for m in rendered if m.get("role") != "system"] @@ -533,12 +538,8 @@ def render_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: # system messages that are not at loop.first, i.e. after the first). system_msgs = [m for m in rendered if m.get("role") == "system"] if len(system_msgs) > 1: - sep = chr(10) + chr(10) - merged_content = sep.join( - str(m.get("content") or "") for m in system_msgs if m.get("content") - ) non_system = [m for m in rendered if m.get("role") != "system"] - rendered = [{"role": "system", "content": merged_content}] + non_system + rendered = [{"role": "system", "content": join_system_contents(system_msgs)}] + non_system return rendered diff --git a/python/freetoken/tokenizer/tokenize.py b/python/freetoken/tokenizer/tokenize.py index 9648974dab..31b74bdac2 100644 --- a/python/freetoken/tokenizer/tokenize.py +++ b/python/freetoken/tokenizer/tokenize.py @@ -57,6 +57,16 @@ def resolve_thinking_mode(chat_template_kwargs: dict[str, Any] | None, tools: An _EFFORT_PROBE_MESSAGES = [{"role": "user", "content": "ping"}] +def join_system_contents(messages: list[dict[str, Any]]) -> str: + """The contents of several system/developer messages as one system prompt, in order. + Image content (a part list left by text flattening) cannot be joined into text: raise + the error templates give for an image in the system turn.""" + contents = [m.get("content") for m in messages if m.get("content")] + if any(not isinstance(c, str) for c in contents): + raise ValueError("System message cannot contain images") + return "\n\n".join(contents) + + def _map_developer_role( messages: list[dict[str, Any]], chat_template: str | None ) -> list[dict[str, Any]]: @@ -74,8 +84,7 @@ def _map_developer_role( if len(system) == 1 and mapped[0] is system[0]: return mapped rest = [m for m in mapped if m.get("role") != "system"] - content = "\n\n".join(str(m.get("content") or "") for m in system if m.get("content")) - return [{"role": "system", "content": content}] + rest + return [{"role": "system", "content": join_system_contents(system)}] + rest class TokenizeManager: diff --git a/tests/server/test_mm_input.py b/tests/server/test_mm_input.py index 5ac5834bf5..a694a50a84 100644 --- a/tests/server/test_mm_input.py +++ b/tests/server/test_mm_input.py @@ -173,3 +173,35 @@ def test_render_messages_no_system_unchanged(): {"role": "user", "content": "Bye"}, ]) assert [m["role"] for m in msgs] == ["user", "assistant", "user"] + + +def test_render_messages_keeps_a_leading_developer_message_ahead_of_system(): + """[developer, system, user]: the tokenizer maps developer -> system and merges in order, so + render_messages must not hoist system ahead of it first. A template that knows the + developer role gets both messages as sent.""" + from freetoken.tokenizer.tokenize import _map_developer_role + + msgs = [ + {"role": "developer", "content": "D"}, + {"role": "system", "content": "S"}, + {"role": "user", "content": "hi"}, + ] + rendered = render_messages(msgs) + assert [m["role"] for m in rendered] == ["developer", "system", "user"] + mapped = _map_developer_role(rendered, "{% if message.role == 'system' %}") + assert [m["role"] for m in mapped] == ["system", "user"] and mapped[0]["content"] == "D\n\nS" + assert _map_developer_role(rendered, "{% if message.role == 'developer' %}") is rendered + + +def test_render_messages_rejects_an_image_in_a_merged_system_message(): + """Merging system messages joins text; an image part must not be pasted into the prompt as + the repr of its part list (with the whole data URL).""" + with pytest.raises(ValueError, match="System message cannot contain images"): + render_messages([ + {"role": "system", "content": "A"}, + {"role": "system", "content": [ + {"type": "text", "text": "B"}, + {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{PNG}"}}, + ]}, + {"role": "user", "content": "hi"}, + ]) diff --git a/tests/tokenizer/test_tokenize.py b/tests/tokenizer/test_tokenize.py index 390fcd4cce..9e441886b2 100644 --- a/tests/tokenizer/test_tokenize.py +++ b/tests/tokenizer/test_tokenize.py @@ -845,3 +845,6 @@ def test_mapped_developer_messages_are_merged_into_one_leading_system_message(): assert mapped[0]["content"] == "be brief\n\nanswer in German" late = [{"role": "user", "content": "hi"}, {"role": "developer", "content": "be brief"}] assert [m["role"] for m in _map_developer_role(late, template)] == ["system", "user"] + image = [{"type": "text", "text": "B"}, {"type": "image", "freetoken_ref": {"kind": "url", "data": "x"}}] + with pytest.raises(ValueError, match="System message cannot contain images"): + _map_developer_role([{"role": "developer", "content": "A"}, {"role": "system", "content": image}], template) From 0838391ad68b0692c16fe049730e0caa96de2f75 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 1 Oct 2026 23:24:07 -0400 Subject: [PATCH 54/94] test(server): ask for n=0 in the invalid-request sdk test, since n=2 is served here Assisted-by: Claude Opus 5.5 --- tests/server/test_protocol_compat.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/server/test_protocol_compat.py b/tests/server/test_protocol_compat.py index 8dc22908d6..d91582e423 100644 --- a/tests/server/test_protocol_compat.py +++ b/tests/server/test_protocol_compat.py @@ -202,7 +202,7 @@ async def run(): async with client(state, vendor) as sdk: with pytest.raises(error) as caught: if protocol == "chat": - await sdk.chat.completions.create(model="test", messages=[{"role": "user", "content": "Hi"}], n=2) + await sdk.chat.completions.create(model="test", messages=[{"role": "user", "content": "Hi"}], n=0) elif protocol == "responses": await sdk.responses.create(model="test", input="Hi", background=True) else: From 60ef53c2430ac4767b63bee56eefe1d8e65a6a55 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 05:40:55 -0400 Subject: [PATCH 55/94] fix(moe): keep unpinned layers off the small-prefill path and take it in every layer An unpinned layer's copy_missing is the whole-layer pageable copy, which cannot serve the LRU slot remap, so it keeps streaming. The per-layer fallback when a layer routes to more than half the cache is gone: the engine's num_experts slot floor already keeps every chunk resident, and a uniform path cannot leave a later layer in an overlap buffer that was never begun. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 37c2479212b8b7028827218b23d1cee68e2c274c) --- python/freetoken/layers/moe.py | 44 +++++++++--------- tests/moe/test_small_prefill_on_demand.py | 54 +++++++++++++++++++++++ 2 files changed, 78 insertions(+), 20 deletions(-) diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index e1b4b8517a..d7298ab9fc 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -424,22 +424,26 @@ def _prefill_routed( pass through unmapped.""" cache = self.offload_cache assert cache is not None - if 0 < hidden_states.shape[0] <= _SMALL_PREFILL_TOKENS and cache.decode_target in ("gpu", "hybrid"): + if ( + 0 < hidden_states.shape[0] <= _SMALL_PREFILL_TOKENS + and cache.decode_target in ("gpu", "hybrid") + and not cache.is_unpinned_layer(self.layer_id) + ): # A short extension (an agent turn over a cached prefix) touches a fraction of each - # layer's experts; fetching those beats streaming all of them (~68 GB here). The same - # token count holds in every layer, so the overlap double buffer is never begun. - slots = self._ensure_unique(cache, topk_ids) - if slots is not None: - return self._expert_gemm( - cache, - hidden_states, - topk_weights, - slots, - views=cache.bank_views(), - n=None, - alphas=cache.alphas_for_slots(self.layer_id), - is_prefill=False, - ) + # layer's experts; fetching those beats streaming all of them (~68 GB here). An + # unpinned layer's copy_missing presumes position == expert id, so it streams below; + # unpinned layers only exist with the prefill overlap off, so no layer can enter + # the overlap double buffer after an earlier layer skipped begin_prefill. + return self._expert_gemm( + cache, + hidden_states, + topk_weights, + self._ensure_unique(cache, topk_ids), + views=cache.bank_views(), + n=None, + alphas=cache.alphas_for_slots(self.layer_id), + is_prefill=False, + ) if cache.prefill_overlap: views = self._wait_prefill_overlap(cache) out = self._expert_gemm( @@ -469,9 +473,8 @@ def _prefill_routed( cache.release_prefill_layer(self.layer_id) return out - def _ensure_unique(self, cache: OffloadMoeCache, topk_ids: torch.Tensor) -> torch.Tensor | None: - """Make this layer's routed experts resident and return ``topk_ids`` mapped to slots, or - None (take the whole-layer path) when they cannot all be resident at once. + def _ensure_unique(self, cache: OffloadMoeCache, topk_ids: torch.Tensor) -> torch.Tensor: + """Make this layer's routed experts resident and return ``topk_ids`` mapped to slots. The LRU kernel is sized for decode: its block is next_pow2(ids) x next_pow2(slots), so a prefill's tokens x top_k ids overflow Triton's element limit, and even a block under it @@ -479,8 +482,9 @@ def _ensure_unique(self, cache: OffloadMoeCache, topk_ids: torch.Tensor) -> torc decode-width chunks; each chunk's misses are copied before the next ensure reuses the plan.""" uniq, inverse = torch.unique(topk_ids.reshape(-1), return_inverse=True) n = uniq.numel() - if n > cache.cache_size // 2: # leave room so a chunk never evicts the one before it - return None + # Each ensure stamps a newer LRU step, so a later chunk evicts the earlier chunks' slots + # only once every other slot is gone; the engine's num_experts slot floor rules that out. + assert n <= cache.cache_size, f"{n} routed experts exceed the {cache.cache_size}-slot cache" uniq = uniq.to(torch.int32) # scored eviction ranks an earlier chunk's fresh installs coldest, so pin them explicitly pin_since = cache.step.clone() if cache.cache_policy_id else None diff --git a/tests/moe/test_small_prefill_on_demand.py b/tests/moe/test_small_prefill_on_demand.py index 70cec00019..717f835d9c 100644 --- a/tests/moe/test_small_prefill_on_demand.py +++ b/tests/moe/test_small_prefill_on_demand.py @@ -56,6 +56,60 @@ def fake_gemm(cache_, hs, w, ids, *, views, n, alphas, is_prefill): assert "ensure" not in calls and "begin_prefill" not in calls[1:] and got["n"] == layer.num_experts +def _record_paths(monkeypatch, layer, cache): + """Stub the cache movement and the GEMM; log each call and what the GEMM was handed.""" + calls, got = [], {} + + def ensure(layer_id, ids): + calls.append(("ensure", ids.tolist())) + ids.add_(100) # the LRU kernel rewrites its query to slot ids in place + + monkeypatch.setattr(cache, "ensure_experts", ensure) + monkeypatch.setattr(cache, "materialize_layer", lambda layer_id: calls.append(("materialize",))) + monkeypatch.setattr(cache, "copy_missing", lambda: calls.append(("copy",))) + + def fake_gemm(cache_, hs, w, ids, *, views, n, alphas, is_prefill): + got.update(ids=ids.tolist(), n=n, is_prefill=is_prefill) + return hs + + monkeypatch.setattr(layer, "_expert_gemm", fake_gemm) + return calls, got + + +def test_small_prefill_leaves_an_unpinned_layer_on_the_whole_layer_path(monkeypatch): + # an unpinned (locked) layer's copy_missing is the whole-layer pageable copy, which presumes + # position == expert id and cannot serve an LRU slot remap + layer, cache = _layer_and_cache(torch.device("cpu")) + monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", 8) + cache._unpinned_layers = frozenset({layer.layer_id}) + calls, got = _record_paths(monkeypatch, layer, cache) + layer._prefill_routed(torch.randn(3, 16), torch.full((3, 2), 0.5), torch.zeros(3, 2, dtype=torch.int32)) + assert calls == [("materialize",), ("copy",)] and got["n"] == layer.num_experts + + +def test_small_prefill_path_does_not_depend_on_how_many_experts_a_layer_routes_to(monkeypatch): + # 6 distinct experts in an 8-slot cache: every layer of a short prefill still takes the + # on-demand path, so no layer can fall back into an overlap buffer that was never begun + layer, cache = _layer_and_cache(torch.device("cpu")) + monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", 8) + monkeypatch.setattr(moe_mod, "_ENSURE_CHUNK_IDS", 4) + calls, got = _record_paths(monkeypatch, layer, cache) + ids = torch.tensor([[5, 0], [1, 7], [2, 5], [4, 0]], dtype=torch.int32) + layer._prefill_routed(torch.randn(4, 16), torch.full((4, 2), 0.5), ids.clone()) + assert calls == [("ensure", [0, 1, 2, 4]), ("copy",), ("ensure", [5, 7]), ("copy",)] + assert got == {"ids": (ids + 100).tolist(), "n": None, "is_prefill": False} + + +def test_small_prefill_refuses_more_experts_than_the_cache_holds(monkeypatch): + layer, cache = _layer_and_cache(torch.device("cpu")) + monkeypatch.setattr(moe_mod, "_SMALL_PREFILL_TOKENS", 8) + _record_paths(monkeypatch, layer, cache) + cache.cache_size = 4 # below the engine's num_experts floor + ids = torch.tensor([[5, 0], [1, 7], [2, 5], [4, 0]], dtype=torch.int32) + with pytest.raises(AssertionError, match="exceed"): + layer._prefill_routed(torch.randn(4, 16), torch.full((4, 2), 0.5), ids) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") def test_small_prefill_matches_the_whole_layer_path(monkeypatch): from freetoken.core import get_global_ctx # noqa: F401 (layer reads the ctx only in forward()) From 1cb0b998c50d261e50431e9688739630d3aecf8e Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 1 Oct 2026 23:24:22 -0400 Subject: [PATCH 56/94] test(moe): let the small-prefill ensure fake take the eviction policy kwargs Assisted-by: Claude Opus 5.5 --- tests/moe/test_small_prefill_on_demand.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/moe/test_small_prefill_on_demand.py b/tests/moe/test_small_prefill_on_demand.py index 717f835d9c..8b7b9d802c 100644 --- a/tests/moe/test_small_prefill_on_demand.py +++ b/tests/moe/test_small_prefill_on_demand.py @@ -60,7 +60,7 @@ def _record_paths(monkeypatch, layer, cache): """Stub the cache movement and the GEMM; log each call and what the GEMM was handed.""" calls, got = [], {} - def ensure(layer_id, ids): + def ensure(layer_id, ids, **_): calls.append(("ensure", ids.tolist())) ids.add_(100) # the LRU kernel rewrites its query to slot ids in place From e1af69593ba6f66627bd4dae628b32f07497d5bc Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 4 Sep 2026 17:31:34 -0400 Subject: [PATCH 57/94] test(qwen4_exp): compare chunked prefill to one-shot at bf16 tolerance `torch.equal` fails for all three cut points on torch 2.11.0+cu130 / flashinfer 0.6.18 / triton 3.6.0 (sm_89): cuBLAS picks a different kernel for M=5000 and M=5000-cut, so the indexer's q and k round differently, a few rows flip one block at the top-k margin, and the layer output differs by up to 2.6e-3 (a few bf16 ulps). Both paths are deterministic run to run. Use the tolerance of the neighbouring oracle tests. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt (cherry picked from commit 3849e3caeebe4f1376708fe677cc0ecb106fb78c) Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit e700e8fb7339eb958507401c39d81b2623f56ca7) --- tests/models/qwen4_exp/test_qsa_backend.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/models/qwen4_exp/test_qsa_backend.py b/tests/models/qwen4_exp/test_qsa_backend.py index 9302fad172..3827b62416 100644 --- a/tests/models/qwen4_exp/test_qsa_backend.py +++ b/tests/models/qwen4_exp/test_qsa_backend.py @@ -4,7 +4,8 @@ tokens every complete block is selected, so QSA IS dense attention: the selection must be exactly the causal prefix and the layer output must match ``TorchDenseQSAReference`` (fp32) and a flashinfer dense run over the same pool; -(b) chunked prefill at unaligned cut points equals one-shot prefill (the dual-source compress); +(b) chunked prefill at unaligned cut points matches one-shot prefill to bf16 tolerance (the + dual-source compress); (c) a captured decode replay equals the eager decode step; (d) an fp8 KV pool (``--kv-cache-dtype fp8``) keeps block selection bit-identical to the 16-bit run -- only the selected K/V rows are read back as e4m3 codes -- and the layer @@ -136,7 +137,11 @@ def test_flashinfer_dense_matches_the_sparse_path(): @requires_cuda @pytest.mark.parametrize("cut", [1001, 4096, 4097], ids=["unaligned", "page-boundary", "boundary+1"]) def test_chunked_prefill_matches_one_shot(cut: int): - """Cut points that are not multiples of index_ratio exercise the dual-source compress.""" + """Cut points that are not multiples of index_ratio exercise the dual-source compress. + + Not bit for bit: the projections see a different M (5000 vs 5000 - cut rows) and cuBLAS + may pick a different kernel per shape, so the indexer's q and k round differently and a + few rows flip one block at the top-k margin. Same tolerance as the oracle tests above.""" config = parsed_config() fixture = Fixture(config, num_pages=512) attn = fixture.layer(QSA_LAYER) @@ -148,7 +153,7 @@ def test_chunked_prefill_matches_one_shot(cut: int): attn.forward(x[:cut], fixture.batch([head], "prefill")) tail = fixture.req(1, cut, length) got = attn.forward(x[cut:], fixture.batch([tail], "prefill")) - assert torch.equal(got, one_shot[cut:]) + torch.testing.assert_close(got.float(), one_shot[cut:].float(), rtol=2e-2, atol=2e-2) @requires_cuda From 5df3b9a9dd485ed965c395134149b4f0804730a8 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 09:23:20 -0400 Subject: [PATCH 58/94] test(qwen4_exp): compare the chunked prefill's compressed slab rows to the one-shot run At bf16 output tolerance a broken dual-source compress only swaps a block in a few rows and passes; the slab keys are unit-scale, so a wrong straddling-group key stands out. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit cedf179739d2f17e8dc8ee27152f35c5584c24c4) --- tests/models/qwen4_exp/test_qsa_backend.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/tests/models/qwen4_exp/test_qsa_backend.py b/tests/models/qwen4_exp/test_qsa_backend.py index 3827b62416..8e836f5632 100644 --- a/tests/models/qwen4_exp/test_qsa_backend.py +++ b/tests/models/qwen4_exp/test_qsa_backend.py @@ -155,6 +155,18 @@ def test_chunked_prefill_matches_one_shot(cut: int): got = attn.forward(x[cut:], fixture.batch([tail], "prefill")) torch.testing.assert_close(got.float(), one_shot[cut:].float(), rtol=2e-2, atol=2e-2) + # A wrong key for the group straddling the cut only swaps one block in a few rows, which + # the output tolerance hides; the unit-scale slab rows show it as an O(1) error. + ratio = config.qwen4_args.index_ratio + slab = fixture.pool.cmp_k_cache(0) + group_starts = slice(0, length - length % ratio, ratio) + one_shot_rows, chunked_rows = ( + (fixture.page_table[table_idx, group_starts] // ratio).long() for table_idx in (0, 1) + ) + torch.testing.assert_close( + slab[chunked_rows].float(), slab[one_shot_rows].float(), rtol=2e-2, atol=2e-2 + ) + @requires_cuda def test_decode_graph_replay_matches_eager(): From 219d3dc2b88e430fd2584c907d7bcfb86542c316 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 05:20:22 -0400 Subject: [PATCH 59/94] feat(engine): keep the token embedding in pinned host ram (--embed-weights host) The embedding is a row lookup, so a single-GPU engine can leave the table in page-locked host RAM and gather each step's rows over UVA with the PLE row kernel (5 KiB per decode token, graph-safe): 1.2 GiB of VRAM back for Qwen3.8-Flash-Next. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 9f2b7a2e33424e2c0e06b240f1fb7a408c850212) --- python/freetoken/engine/config.py | 2 + python/freetoken/engine/engine.py | 35 +++++++++++++++++ python/freetoken/kernel/triton/ple.py | 8 +++- python/freetoken/layers/embedding.py | 31 ++++++++++++--- python/freetoken/server/args.py | 10 +++++ tests/kernels/test_host_embedding.py | 55 +++++++++++++++++++++++++++ 6 files changed, 134 insertions(+), 7 deletions(-) create mode 100644 tests/kernels/test_host_embedding.py diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index f64829f311..3167093aca 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -38,6 +38,8 @@ class EngineConfig: # DeepSeek-V4.1 Decoder SWA Bounded Replay: "bounded" (default) runs the decoder layers on each # prompt's last window; "exact" runs them on every token. swa_decoder_replay: str = "bounded" + # --embed-weights: "host" keeps the token embedding in pinned host RAM, gathered over PCIe. + embed_weights: str = "gpu" # Expert-bank host load (--expert-load): auto|serial|parallel. "auto" reads scattered # experts in parallel but falls back to serial when free RAM can't cover the banks + the # parallel reader's extra (non-reclaimable) whole-shard buffer; "serial" forces the diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index ed4525b324..b6763d7881 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -107,6 +107,39 @@ def _sgl_flash_attn_available() -> bool: return True +def _place_embeddings_on_host(model) -> None: + """--embed-weights host: every input embedding table moves to pinned host RAM, except one an lm_head is tied to.""" + from freetoken.layers.base import BaseOP + from freetoken.layers.embedding import ParallelLMHead, VocabParallelEmbedding + + found: list[VocabParallelEmbedding] = [] + tied: set[int] = set() + stack, seen = [model], set() + while stack: + op = stack.pop() + if id(op) in seen: + continue + seen.add(id(op)) + if isinstance(op, ParallelLMHead): + if op.tied_embedding is not None: + tied.add(id(op.tied_embedding)) + continue + if isinstance(op, VocabParallelEmbedding): + found.append(op) + continue + for value in vars(op).values(): + items = value if isinstance(value, (list, tuple)) else (value,) + stack.extend(v for v in items if isinstance(v, BaseOP)) + moved = [e for e in found if id(e) not in tied] + if not moved: + raise ValueError("--embed-weights host: the model has no untied input embedding to move") + for emb in moved: + emb.place_on_host() + torch.cuda.empty_cache() + size = sum(e.weight.numel() * e.weight.element_size() for e in moved) + logger.info_rank0(f"Token embedding in pinned host RAM ({mem_GB(size)}), gathered over PCIe") + + def _startup_kv_budget(memory_ratio: float, init_free_memory: int, new_free_memory: int) -> int: """Bytes available to the KV pool at startup: ratio-scaled pre-load free memory minus what the resident model consumed. Kept as a pure function so the composition with the @@ -490,6 +523,8 @@ def __init__(self, config: EngineConfig): ) # before the residency snapshot, so streamed blocks are not charged as resident weights self.model.place_encoder_weights(config.mm.encoder_weights) + if config.embed_weights == "host": + _place_embeddings_on_host(self.model) post_weights_free = self._sync_get_memory()[0] self._weights_bytes = self._baseline_free - post_weights_free # Pool-budget baseline for the desktop cache sliders: free VRAM after the weights are diff --git a/python/freetoken/kernel/triton/ple.py b/python/freetoken/kernel/triton/ple.py index 2314e94e72..cded101a2e 100644 --- a/python/freetoken/kernel/triton/ple.py +++ b/python/freetoken/kernel/triton/ple.py @@ -21,6 +21,7 @@ # Latency-bound over PCIe, so keep the block small and let many of them be in flight. _NUM_WARPS = 1 +_TL_TABLE_DTYPES = {torch.bfloat16: tl.bfloat16, torch.float16: tl.float16} @triton.jit @@ -33,6 +34,7 @@ def _ple_gather_kernel( EMB_DIM: tl.constexpr, IS_FP8: tl.constexpr, BLOCK_D: tl.constexpr, + TABLE_DTYPE: tl.constexpr, ): row = tl.program_id(0) idx = tl.load(ids_ptr + row).to(tl.int64) @@ -50,7 +52,7 @@ def _ple_gather_kernel( base = table_ptr.to(tl.int64).to(tl.pointer_type(tl.uint8)) values = e4m3_u8_to_f32(tl.load(base + idx * EMB_DIM + offsets, mask=mask, other=0)) else: - base = table_ptr.to(tl.int64).to(tl.pointer_type(tl.bfloat16)) + base = table_ptr.to(tl.int64).to(tl.pointer_type(TABLE_DTYPE)) values = tl.load(base + idx * EMB_DIM + offsets, mask=mask, other=0.0).to(tl.float32) values = tl.where(in_range, values * scale, 0.0) tl.store( @@ -68,11 +70,12 @@ def ple_gather_rows( out: torch.Tensor, scale: float = 1.0, is_fp8: bool = True, + table_dtype: torch.dtype = torch.bfloat16, ) -> torch.Tensor: """Gather ``row_ids`` from the host-resident table at ``table_ptr`` into ``out``. ``row_ids`` is a flat device int tensor; ``out`` is ``[row_ids.numel(), embed_dim]`` - bf16 on the same device. ``table_ptr`` is the address the GPU must dereference + on the same device. A non-fp8 table holds ``table_dtype`` elements. ``table_ptr`` is the address the GPU must dereference (``kernel/pinned.device_ptr``), not necessarily the host ``data_ptr``. """ n = row_ids.numel() @@ -87,6 +90,7 @@ def ple_gather_rows( EMB_DIM=embed_dim, IS_FP8=is_fp8, BLOCK_D=triton.next_power_of_2(embed_dim), + TABLE_DTYPE=_TL_TABLE_DTYPES[table_dtype], num_warps=_NUM_WARPS, ) return out diff --git a/python/freetoken/layers/embedding.py b/python/freetoken/layers/embedding.py index 6939156cd6..fc5deaac58 100644 --- a/python/freetoken/layers/embedding.py +++ b/python/freetoken/layers/embedding.py @@ -36,16 +36,37 @@ def __init__( self._embed_scale = embed_scale self._embed_scale_t: torch.Tensor | None = None self._comm = DistributedCommunicator() + self._host_ptr: int | None = None + + def place_on_host(self) -> None: + """Move the table into pinned host RAM; lookups then read their rows over PCIe (UVA).""" + from freetoken.kernel.pinned import alloc_pinned_tensor, device_ptr + + if self.tp_size > 1: + raise ValueError("a host-resident embedding serves a single GPU only") + host = alloc_pinned_tensor(*self.weight.shape, dtype=self.weight.dtype) + host.copy_(self.weight) + self.weight = host + self._host_ptr = device_ptr(host) @nvtx_annotate("Embedding") def forward(self, x: torch.Tensor) -> torch.Tensor: from freetoken.kernel import indexing - y = indexing( - weights=self.weight, - indices=x, - vocab_range=self.vocab_range if self.tp_size > 1 else None, - ) + if self._host_ptr is not None: + from freetoken.kernel.triton.ple import ple_gather_rows + + rows, dim = self.weight.shape + y = torch.empty(x.numel(), dim, dtype=self.weight.dtype, device=x.device) + ple_gather_rows(self._host_ptr, rows, dim, x.reshape(-1), y, is_fp8=False, + table_dtype=self.weight.dtype) + y = y.view(*x.shape, dim) + else: + y = indexing( + weights=self.weight, + indices=x, + vocab_range=self.vocab_range if self.tp_size > 1 else None, + ) if self.tp_size > 1: y = self._comm.all_reduce(y) diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 393832258c..fc2c925b78 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -781,6 +781,16 @@ def _infer_reasoning_parser(model_path: str) -> str | None: ), ) + parser.add_argument( + "--embed-weights", + default=ServerArgs.embed_weights, + choices=["gpu", "host"], + help=( + "Where the token embedding table lives. 'host' keeps it in pinned host RAM and gathers " + "the rows each step needs over PCIe, freeing its VRAM (single GPU only)." + ), + ) + parser.add_argument( "--nvfp4-backend", action=_DeprecatedAlias, diff --git a/tests/kernels/test_host_embedding.py b/tests/kernels/test_host_embedding.py new file mode 100644 index 0000000000..76c8c612c3 --- /dev/null +++ b/tests/kernels/test_host_embedding.py @@ -0,0 +1,55 @@ +"""--embed-weights host: the UVA row gather from pinned host RAM must equal the GPU table lookup.""" + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + + +@pytest.fixture(autouse=True) +def _tp(): + from freetoken.distributed import set_tp_info, try_get_tp_info + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + + +def _embedding(dtype: torch.dtype, rows: int = 1000, dim: int = 2560): + from freetoken.layers.embedding import VocabParallelEmbedding + + torch.manual_seed(0) + emb = VocabParallelEmbedding(rows, dim) + emb.weight = torch.randn(rows, dim, device="cuda", dtype=dtype) + return emb + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_host_rows_equal_the_gpu_lookup(dtype): + emb = _embedding(dtype) + ids = torch.randint(0, 1000, (37,), device="cuda", dtype=torch.int32) + want = emb.forward(ids) + emb.place_on_host() + assert not emb.weight.is_cuda and emb.weight.is_pinned() + got = emb.forward(ids) + assert got.is_cuda and got.dtype is dtype + assert torch.equal(got, want) + + +def test_host_lookup_replays_in_a_cuda_graph_with_new_ids(): + emb = _embedding(torch.bfloat16) + table = emb.weight.clone() + emb.place_on_host() + ids = torch.zeros(8, device="cuda", dtype=torch.int64) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + emb.forward(ids) # compile outside the capture + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = emb.forward(ids) + for seed in range(3): + ids.copy_(torch.randint(0, 1000, (8,), generator=torch.Generator().manual_seed(seed))) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, table[ids]) From 798ae0123af74b9ada99d1b122e8b687b0747e1c Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 10:51:34 -0400 Subject: [PATCH 60/94] fix(layers): gather fp32 host embeddings and refuse other dtypes at placement Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit ccfb2b8148c8c108897f37b7f337e99cbb5a0134) --- python/freetoken/kernel/triton/ple.py | 2 +- python/freetoken/layers/embedding.py | 3 +++ tests/kernels/test_host_embedding.py | 17 +++++++++++++---- 3 files changed, 17 insertions(+), 5 deletions(-) diff --git a/python/freetoken/kernel/triton/ple.py b/python/freetoken/kernel/triton/ple.py index cded101a2e..8fa25acdfd 100644 --- a/python/freetoken/kernel/triton/ple.py +++ b/python/freetoken/kernel/triton/ple.py @@ -21,7 +21,7 @@ # Latency-bound over PCIe, so keep the block small and let many of them be in flight. _NUM_WARPS = 1 -_TL_TABLE_DTYPES = {torch.bfloat16: tl.bfloat16, torch.float16: tl.float16} +_TL_TABLE_DTYPES = {torch.bfloat16: tl.bfloat16, torch.float16: tl.float16, torch.float32: tl.float32} @triton.jit diff --git a/python/freetoken/layers/embedding.py b/python/freetoken/layers/embedding.py index fc5deaac58..860c220010 100644 --- a/python/freetoken/layers/embedding.py +++ b/python/freetoken/layers/embedding.py @@ -41,9 +41,12 @@ def __init__( def place_on_host(self) -> None: """Move the table into pinned host RAM; lookups then read their rows over PCIe (UVA).""" from freetoken.kernel.pinned import alloc_pinned_tensor, device_ptr + from freetoken.kernel.triton.ple import _TL_TABLE_DTYPES if self.tp_size > 1: raise ValueError("a host-resident embedding serves a single GPU only") + if self.weight.dtype not in _TL_TABLE_DTYPES: + raise ValueError(f"a host-resident embedding cannot gather a {self.weight.dtype} table") host = alloc_pinned_tensor(*self.weight.shape, dtype=self.weight.dtype) host.copy_(self.weight) self.weight = host diff --git a/tests/kernels/test_host_embedding.py b/tests/kernels/test_host_embedding.py index 76c8c612c3..866687a9f1 100644 --- a/tests/kernels/test_host_embedding.py +++ b/tests/kernels/test_host_embedding.py @@ -3,7 +3,7 @@ import pytest import torch -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +needs_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") @pytest.fixture(autouse=True) @@ -14,16 +14,17 @@ def _tp(): set_tp_info(rank=0, size=1) -def _embedding(dtype: torch.dtype, rows: int = 1000, dim: int = 2560): +def _embedding(dtype: torch.dtype, rows: int = 1000, dim: int = 2560, device: str = "cuda"): from freetoken.layers.embedding import VocabParallelEmbedding torch.manual_seed(0) emb = VocabParallelEmbedding(rows, dim) - emb.weight = torch.randn(rows, dim, device="cuda", dtype=dtype) + emb.weight = torch.randn(rows, dim, device=device, dtype=dtype) return emb -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@needs_cuda +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) def test_host_rows_equal_the_gpu_lookup(dtype): emb = _embedding(dtype) ids = torch.randint(0, 1000, (37,), device="cuda", dtype=torch.int32) @@ -35,6 +36,14 @@ def test_host_rows_equal_the_gpu_lookup(dtype): assert torch.equal(got, want) +def test_a_table_the_gather_cannot_read_is_refused_before_pinning(): + emb = _embedding(torch.float64, rows=4, dim=8, device="cpu") + with pytest.raises(ValueError, match="cannot gather a torch.float64 table"): + emb.place_on_host() + assert emb._host_ptr is None + + +@needs_cuda def test_host_lookup_replays_in_a_cuda_graph_with_new_ids(): emb = _embedding(torch.bfloat16) table = emb.weight.clone() From a7c2afbcf750ddd86a3c5d597a268bba4bb0ce68 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 10:54:43 -0400 Subject: [PATCH 61/94] fix(kernel): compute the host gather output offset in int64 No regression test: overflowing it needs an output of 2^31+ elements (4 GiB+ in bf16). Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit cc5af7adc5d99c17c0edc3f7947040ac35a48445) --- python/freetoken/kernel/triton/ple.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/freetoken/kernel/triton/ple.py b/python/freetoken/kernel/triton/ple.py index 8fa25acdfd..466fce11a0 100644 --- a/python/freetoken/kernel/triton/ple.py +++ b/python/freetoken/kernel/triton/ple.py @@ -36,7 +36,8 @@ def _ple_gather_kernel( BLOCK_D: tl.constexpr, TABLE_DTYPE: tl.constexpr, ): - row = tl.program_id(0) + # the output offset row * EMB_DIM can pass 2^31 for a long prefill at a wide hidden size + row = tl.program_id(0).to(tl.int64) idx = tl.load(ids_ptr + row).to(tl.int64) in_range = (idx >= 0) & (idx < num_rows) idx = tl.where(in_range, idx, 0) From adac17cd93f659132652499ef6e4438dce346aa7 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 10:56:09 -0400 Subject: [PATCH 62/94] fix(engine): charge the host embedding's pinned bytes to the pin budget Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit a8b2f0a6924c9362b014c200dd702115b0c33d89) --- python/freetoken/engine/engine.py | 16 +++++++++------- tests/engine/test_cache_budget.py | 16 ++++++++++++++++ 2 files changed, 25 insertions(+), 7 deletions(-) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index b6763d7881..bc780377ef 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -107,8 +107,10 @@ def _sgl_flash_attn_available() -> bool: return True -def _place_embeddings_on_host(model) -> None: - """--embed-weights host: every input embedding table moves to pinned host RAM, except one an lm_head is tied to.""" +def _place_embeddings_on_host(model) -> int: + """--embed-weights host: every input embedding table moves to pinned host RAM, except one an lm_head is tied to. + + Returns the pinned bytes.""" from freetoken.layers.base import BaseOP from freetoken.layers.embedding import ParallelLMHead, VocabParallelEmbedding @@ -138,6 +140,7 @@ def _place_embeddings_on_host(model) -> None: torch.cuda.empty_cache() size = sum(e.weight.numel() * e.weight.element_size() for e in moved) logger.info_rank0(f"Token embedding in pinned host RAM ({mem_GB(size)}), gathered over PCIe") + return size def _startup_kv_budget(memory_ratio: float, init_free_memory: int, new_free_memory: int) -> int: @@ -523,8 +526,7 @@ def __init__(self, config: EngineConfig): ) # before the residency snapshot, so streamed blocks are not charged as resident weights self.model.place_encoder_weights(config.mm.encoder_weights) - if config.embed_weights == "host": - _place_embeddings_on_host(self.model) + embed_host_bytes = _place_embeddings_on_host(self.model) if config.embed_weights == "host" else 0 post_weights_free = self._sync_get_memory()[0] self._weights_bytes = self._baseline_free - post_weights_free # Pool-budget baseline for the desktop cache sliders: free VRAM after the weights are @@ -538,11 +540,11 @@ def __init__(self, config: EngineConfig): self.cpu_moe_executor = None # Host-side auxiliary stores (qwen4_exp's pinned PLE table): after the weights so a # load failure is not masked, before the MoE offload cache so the bank residency - # planning sees the pin quota the table already spent. - self._host_tables_bytes = 0 + # planning sees the pin quota the table (and a host embedding) already spent. + self._host_tables_bytes = embed_host_bytes if hasattr(self.model, "load_host_tables"): with _weight_load_context(): - self._host_tables_bytes = int(self.model.load_host_tables(config) or 0) + self._host_tables_bytes += int(self.model.load_host_tables(config) or 0) if is_offload_moe_strategy(config.moe_strategy): self._init_offload_moe_cache(config) if hasattr(self.model, "prepare_for_runtime"): diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py index f5abf0142c..a8fd045fc7 100644 --- a/tests/engine/test_cache_budget.py +++ b/tests/engine/test_cache_budget.py @@ -673,3 +673,19 @@ def test_moe_pool_flags_and_decode_batch_clamp(): ) _clamp_decode_batch_to_pools(config, [12, 16, 9]) assert (config.max_running_req, config.cuda_graph_max_bs, config.cuda_graph_bs) == (2, 2, [1, 2]) + + +def test_host_embedding_placement_reports_the_bytes_it_pinned(monkeypatch): + # the engine charges these bytes to the pin budget before it plans the expert banks + from freetoken.distributed import set_tp_info, try_get_tp_info + from freetoken.engine.engine import _place_embeddings_on_host + from freetoken.layers.embedding import ParallelLMHead, VocabParallelEmbedding + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + monkeypatch.setattr(VocabParallelEmbedding, "place_on_host", lambda self: None) + embed, tied = VocabParallelEmbedding(100, 64), VocabParallelEmbedding(50, 64) + embed.weight = torch.empty(100, 64, dtype=torch.bfloat16) + head = ParallelLMHead(50, 64, tie_word_embeddings=True, tied_embedding=tied) + model = SimpleNamespace(embed_tokens=embed, per_layer_embed=tied, lm_head=head) + assert _place_embeddings_on_host(model) == 100 * 64 * 2 From 1f1b2a235f9752619fe0becfde33237462941915 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 10:57:23 -0400 Subject: [PATCH 63/94] fix(engine): refuse --embed-weights host for a tied embedding or tp > 1 at config time Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 579e03a7aae077ac209cbf22ed4de6cda5bb390f) --- python/freetoken/engine/config.py | 10 ++++++++ tests/engine/test_embed_weights.py | 39 ++++++++++++++++++++++++++++++ 2 files changed, 49 insertions(+) create mode 100644 tests/engine/test_embed_weights.py diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 3167093aca..c3643e46c5 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -169,8 +169,18 @@ def model_config(self) -> ModelConfig: quant = checkpoint_quant_config(self.model_path, hf_config, spec) set_quant_config(quant) model_config = _load_attr(spec.module, spec.parse_config)(hf_config) + self._check_embed_weights(model_config) return replace(model_config, quant=quant) + def _check_embed_weights(self, model_config) -> None: + # the parsed flag, not the hf key: gemma4 ties by default and gguf infers it from the tensors + if self.embed_weights != "host": + return + if model_config.tie_word_embeddings: + raise ValueError("--embed-weights host needs an untied embedding (a tied lm_head reads the table on the GPU)") + if self.tp_info.size > 1: + raise ValueError("--embed-weights host serves a single GPU only (--tensor-parallel-size 1)") + @property def max_seq_len(self) -> int: if self.max_seq_len_override is not None: diff --git a/tests/engine/test_embed_weights.py b/tests/engine/test_embed_weights.py new file mode 100644 index 0000000000..b8cb0fdc26 --- /dev/null +++ b/tests/engine/test_embed_weights.py @@ -0,0 +1,39 @@ +"""--embed-weights host: a setup the host gather cannot serve is refused while the config is still a dataclass, not after the weight load.""" + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.distributed import DistributedInfo +from freetoken.engine.config import EngineConfig + + +def _config(monkeypatch, *, tie: bool, tp: int = 1, embed_weights: str = "host"): + import freetoken.engine.config as engine_config + + hf = SimpleNamespace( + architectures=["Qwen3ForCausalLM"], model_type="qwen3", num_hidden_layers=2, num_attention_heads=4, + num_key_value_heads=2, hidden_size=64, vocab_size=128, intermediate_size=128, rms_norm_eps=1e-6, + max_position_embeddings=1024, hidden_act="silu", tie_word_embeddings=tie, + ) + monkeypatch.setattr(engine_config, "cached_load_hf_config", lambda path: hf) + monkeypatch.setattr(engine_config, "checkpoint_quant_config", lambda *args: None) + return EngineConfig( + model_path="/fake", tp_info=DistributedInfo(rank=0, size=tp), dtype=torch.bfloat16, embed_weights=embed_weights + ) + + +def test_a_tied_embedding_is_refused(monkeypatch): + with pytest.raises(ValueError, match="needs an untied embedding"): + _config(monkeypatch, tie=True).model_config + + +def test_tensor_parallel_is_refused(monkeypatch): + with pytest.raises(ValueError, match="single GPU only"): + _config(monkeypatch, tie=False, tp=2).model_config + + +def test_an_untied_single_gpu_model_and_the_gpu_default_pass(monkeypatch): + assert not _config(monkeypatch, tie=False).model_config.tie_word_embeddings + assert _config(monkeypatch, tie=True, tp=2, embed_weights="gpu").model_config.tie_word_embeddings From 009f5fd2ebec4db8fe59341768438300bc4096a1 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 1 Oct 2026 23:26:39 -0400 Subject: [PATCH 64/94] test(engine): fake cached_load_hf_config with the overrides argument #573 passes Assisted-by: Claude Opus 5.5 --- tests/engine/test_embed_weights.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/engine/test_embed_weights.py b/tests/engine/test_embed_weights.py index b8cb0fdc26..d3aa102f94 100644 --- a/tests/engine/test_embed_weights.py +++ b/tests/engine/test_embed_weights.py @@ -17,7 +17,7 @@ def _config(monkeypatch, *, tie: bool, tp: int = 1, embed_weights: str = "host") num_key_value_heads=2, hidden_size=64, vocab_size=128, intermediate_size=128, rms_norm_eps=1e-6, max_position_embeddings=1024, hidden_act="silu", tie_word_embeddings=tie, ) - monkeypatch.setattr(engine_config, "cached_load_hf_config", lambda path: hf) + monkeypatch.setattr(engine_config, "cached_load_hf_config", lambda path, *_: hf) monkeypatch.setattr(engine_config, "checkpoint_quant_config", lambda *args: None) return EngineConfig( model_path="/fake", tp_info=DistributedInfo(rank=0, size=tp), dtype=torch.bfloat16, embed_weights=embed_weights From bf6c6ee45b4c971f9615dc8bf7a77cd93a6854c9 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 1 Oct 2026 23:27:21 -0400 Subject: [PATCH 65/94] test(ple): drive overlapped graph fills against an in-order fake stream so a post-launch fence fails as a hang Upstream #588's _graph_consumed fence is left out of the merge: our graph fill runs on the filler thread before its own launch passes its WAIT, so that fence deadlocks the second step. Assisted-by: Claude Opus 5.5 --- .../qwen4_exp/test_ple_disk_no_graph_fence.py | 136 ++++++++++++++++++ 1 file changed, 136 insertions(+) create mode 100644 tests/models/qwen4_exp/test_ple_disk_no_graph_fence.py diff --git a/tests/models/qwen4_exp/test_ple_disk_no_graph_fence.py b/tests/models/qwen4_exp/test_ple_disk_no_graph_fence.py new file mode 100644 index 0000000000..cb9bc32356 --- /dev/null +++ b/tests/models/qwen4_exp/test_ple_disk_no_graph_fence.py @@ -0,0 +1,136 @@ +"""The wait-sync graph fill must never wait on an event recorded after its own launch. + +Upstream #588 fences ``_graph_pinned`` with ``_graph_consumed``, recorded once the launch returns and +synchronized by the next graph fill. Upstream runs the fill after the launch returns; here (fork +#65/#67) the fill runs on the filler thread BEFORE its own graph can pass its WAIT, so that fence +makes fill k wait on graph k, which waits on fill k: decode hangs on the second overlapped step. +The previous graph's read of ``_graph_pinned`` is already ordered before the next fill by the +per-fill readback event. This test drives the overlapped decode loop against an asynchronous, +in-order fake stream (launch returns at once; events record a stream position) and fails if any +step stalls, including when the fence is reintroduced. +""" + +from __future__ import annotations + +import threading +import time +from types import SimpleNamespace + +import torch + +import freetoken.models.qwen4_exp.ple_disk as ple_disk +from freetoken.models.qwen4_exp.ple_disk import DiskRowTable + +EOS = 9 + + +class _Stream: + """One in-order stream: a graph item WAITs on the flag, then RESETs it; launch is async.""" + + def __init__(self, flag: torch.Tensor): + self.flag = flag + self.cv = threading.Condition() + self.queue: list[str] = [] + self.enqueued = self.done = 0 + self.stop = False + threading.Thread(target=self._run, daemon=True).start() + + def enqueue(self, kind: str) -> None: + with self.cv: + self.queue.append(kind) + self.enqueued += 1 + self.cv.notify_all() + + def _run(self) -> None: + while True: + with self.cv: + self.cv.wait_for(lambda: self.queue or self.stop) + if self.stop: + return + kind = self.queue.pop(0) + if kind == "graph": + deadline = time.monotonic() + 10.0 + while int(self.flag[0]) != 1 and time.monotonic() < deadline: + time.sleep(0.0005) + self.flag[0] = 0 + with self.cv: + self.done += 1 + self.cv.notify_all() + + def wait_done(self, n: int, timeout: float) -> bool: + with self.cv: + return self.cv.wait_for(lambda: self.done >= n, timeout) + + def close(self) -> None: + with self.cv: + self.stop = True + self.cv.notify_all() + + +class _Event: + """cudaEvent semantics: synchronize() waits for the work enqueued before the LAST record().""" + + def __init__(self, stream: _Stream): + self.stream, self.pos = stream, 0 + + def record(self, _stream=None) -> None: + self.pos = self.stream.enqueued + + def synchronize(self) -> None: + if not self.stream.wait_done(self.pos, 2.0): + raise TimeoutError(f"a fill waited on stream position {self.pos}, which never completed") + + +class _Store: + def __init__(self, flag: torch.Tensor): + self.flag = flag + + def stage(self, tokens_addr, n, staging_addr): + pass + + def flush(self, signal_addr): + if signal_addr: + self.flag[0] = 1 + + +def test_overlapped_graph_fills_never_wait_on_their_own_launch(monkeypatch): + flag = torch.zeros(1, dtype=torch.int64) + stream = _Stream(flag) + monkeypatch.setattr(ple_disk.torch.cuda, "current_stream", lambda *a, **k: None) + monkeypatch.setattr(ple_disk.torch.cuda, "Event", lambda *a, **k: _Event(stream)) + t = DiskRowTable.__new__(DiskRowTable) + t._device = None + t._wait_sync = True + t._flag = flag + t._store = _Store(flag) + t._token_readback = torch.zeros(8, dtype=torch.int32) + t._graph_pinned = torch.zeros(64, dtype=torch.uint8) + t._eager_pinned = [torch.zeros(64, dtype=torch.uint8), torch.zeros(64, dtype=torch.uint8)] + t._eager_read = [_Event(stream), _Event(stream)] + t._eager_slot = 0 + # if upstream's post-launch fence is ever reintroduced, model it with real stream semantics + # so this test fails by hanging the way production would, not on a missing attribute + t._graph_consumed = _Event(stream) + t._token_bytes = 1 + t.eos_token_id = EOS + t.image_token_id = None + + tokens = [1, 2, 3] + steps = 12 + try: + for step in range(steps): + new = 10 + step + req = SimpleNamespace(input_ids=torch.tensor(tokens + [new]), device_len=len(tokens) + 1) + batch = SimpleNamespace(is_decode=True, reqs=[req], padded_size=1, + input_ids=torch.tensor([new], dtype=torch.int32)) + with t.forward_host_ctx(batch, use_graph=True): + stream.enqueue("graph") # cuGraphLaunch returns before the GPU passes the WAIT + tokens.append(new) + # the overlap loop keeps the engine at most one launch ahead of the GPU + assert stream.wait_done(step, 4.0), f"decode hung: graph {step - 1} never completed" + t._raise_failed_fill() + assert stream.wait_done(steps, 4.0), "decode hung on the last step" + t._raise_failed_fill() + finally: + t.close() + stream.close() From cee58f6d4c1c4a7d8ab6557bbe67324964f19e37 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 7 Oct 2026 17:33:31 -0400 Subject: [PATCH 66/94] fix(kernel): keep one e4m3_f32_to_u8, #408's copy duplicates the one #603 added to main Assisted-by: Claude Opus 5.5 --- python/freetoken/kernel/triton/e4m3_compat.py | 27 ------------------- 1 file changed, 27 deletions(-) diff --git a/python/freetoken/kernel/triton/e4m3_compat.py b/python/freetoken/kernel/triton/e4m3_compat.py index 6c8e568b25..f0a4464a5a 100644 --- a/python/freetoken/kernel/triton/e4m3_compat.py +++ b/python/freetoken/kernel/triton/e4m3_compat.py @@ -211,33 +211,6 @@ def round_e4m3(x): return tl.where(tl.abs(x) >= 0.015625, y_norm, y_sub) -@jit -def e4m3_f32_to_u8(x): - """Encode an fp32 value that ALREADY lies on the e4m3 grid -- the output of - :func:`round_e4m3`, clamped to +-448 -- into its e4m3 byte code. This is the - encoder the pre-sm_89 emulated path needs to STORE fp8-sized data (the fp8 - type itself is unavailable there, so the bytes live in a uint8 buffer that - :func:`e4m3_u8_to_f32` decodes back). - - Normal range: read the (unbiased) exponent and the now-zero-padded fp32 - mantissa back out of the fp32 header. Subnormal range (|x| < 2^-6, grid step - 2^-9): the value is an exact multiple of 2^-9, so ``|x| * 512`` IS the mantissa - field -- the sign bit has to be carried in by hand, since that branch never - looks at the header. The same 0.015625 boundary as :func:`round_e4m3` keeps the - two consistent: ``e4m3_u8_to_f32(e4m3_f32_to_u8(round_e4m3(x)))`` is x's - single-rounded value for every input, and no code it emits is a NaN pattern (the - caller's +-448 clamp caps the code at 0x7E). ``-0.0`` encodes as 0x00 after - round_e4m3 (which documents returning +0.0 for it). - """ - u = x.to(tl.uint32, bitcast=True) - sign = ((u >> 31) & 1).to(tl.int32) - exp = ((u >> 23) & 0xFF).to(tl.int32) - 127 - mant = ((u >> 20) & 7).to(tl.int32) - normal = (sign << 7) | ((exp + 7) << 3) | mant - sub = (sign << 7) | (tl.abs(x) * 512.0).to(tl.int32) - return tl.where(tl.abs(x) >= 0.015625, normal, sub).to(tl.uint8) - - @jit def kv_load_e4m3_tile_f32(ptrs, mask): """Load a tile of KV e4m3 codes and widen it to fp32. From 4ec2dbd2723b88430930289802445d756634a418 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 7 Oct 2026 17:38:06 -0400 Subject: [PATCH 67/94] fix(server): leave late Anthropic system messages to the inline-system policy instead of hoisting them #487's hoist ran before #608's fold on /v1/messages and moved a late system reminder into the system head, so every reminder missed the prefix cache; the OpenAI and Responses paths keep the hoist the qwen templates need. Assisted-by: Claude Opus 5.5 --- python/freetoken/server/anthropic_api.py | 3 ++- python/freetoken/server/generation.py | 7 ++++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/python/freetoken/server/anthropic_api.py b/python/freetoken/server/anthropic_api.py index 545cf4e7cc..35d902b9d6 100644 --- a/python/freetoken/server/anthropic_api.py +++ b/python/freetoken/server/anthropic_api.py @@ -326,7 +326,8 @@ def convert_anthropic_prompt( # default only covers requests that leave thinking out. ctk = apply_default_thinking_mode(ctk, default_thinking_mode) - return render_messages(messages), template_tools, parser_tools, ctk + # late system reminders stay in place: the tokenizer folds them instead (prefix reuse) + return render_messages(messages, hoist_system=False), template_tools, parser_tools, ctk def convert_anthropic_to_genspec( diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index 0ed933feb8..e5a215658b 100644 --- a/python/freetoken/server/generation.py +++ b/python/freetoken/server/generation.py @@ -516,16 +516,17 @@ def pick(value, key, framework): return params -def render_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: +def render_messages(messages: list[dict[str, Any]], *, hoist_system: bool = True) -> list[dict[str, Any]]: """Normalize OpenAI-shaped message dicts for the chat template: flatten text content parts to a string and decode tool-call arguments from JSON. Raises ValueError on a non-text content part (text-only server). Shared by all adapters. Some chat templates (e.g. Qwen3.6) require the system message at index 0; hoist system messages to the front and merge multiples into one to satisfy - that constraint.""" + that constraint. ``hoist_system=False`` leaves late system messages in place + for the tokenizer's inline-system policy (the Anthropic path).""" rendered = [_render_message(m) for m in messages] - if any(m.get("role") == "developer" for m in rendered): + if not hoist_system or any(m.get("role") == "developer" for m in rendered): # The tokenizer maps developer per template and then merges in order # (_map_developer_role); hoisting system alone here would put it ahead of # a developer message that came first. From 4c7e83cbf1c92a31afd33143366b08e2f72e3a9d Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 7 Oct 2026 17:39:01 -0400 Subject: [PATCH 68/94] test(server): give main's DeepSeek-V4.1 detector its contract class and the count fake the inline-system kwarg Assisted-by: Claude Opus 5.5 --- tests/server/test_default_thinking_mode.py | 2 +- tests/server/test_function_call_parser.py | 13 ++++++++++++- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/tests/server/test_default_thinking_mode.py b/tests/server/test_default_thinking_mode.py index 7f9938f60c..56854641f2 100644 --- a/tests/server/test_default_thinking_mode.py +++ b/tests/server/test_default_thinking_mode.py @@ -110,7 +110,7 @@ def ctk(**extra): def test_anthropic_count_tokens_renders_the_generation_prompt_under_the_default(monkeypatch): captured = {} - async def fake_count(messages, template_tools, ctk, state): + async def fake_count(messages, template_tools, ctk, state, **_): captured["ctk"] = ctk return 1 diff --git a/tests/server/test_function_call_parser.py b/tests/server/test_function_call_parser.py index 23cd12272e..8af5bc88b6 100644 --- a/tests/server/test_function_call_parser.py +++ b/tests/server/test_function_call_parser.py @@ -454,6 +454,12 @@ class TestDeepSeekV32Detector(_DetectorContract): truncate_before = '' +class TestDeepSeekV41Detector(_DetectorContract): + parser_name = 'deepseekv41' + block = '<｜DSML｜ calls><｜DSML｜ invoke name="read"><｜DSML｜ parameter name="filePath" string="true">/tmp/test_calc.py' + truncate_before = '' + + class TestQwen3CoderDetector(_DetectorContract): parser_name = 'qwen3_coder' block = '/tmp/test_calc.py' @@ -498,5 +504,10 @@ def test_contract_classes_cover_every_concrete_detector(): type(cls().parser().detector) for cls in _DetectorContract.__subclasses__() } - assert covered == set(BaseFormatDetector.__subclasses__()) + concrete, todo = set(), list(BaseFormatDetector.__subclasses__()) + while todo: # a detector may extend another (DeepSeekV41Detector < DeepSeekV32Detector) + cls = todo.pop() + concrete.add(cls) + todo.extend(cls.__subclasses__()) + assert covered == concrete assert covered == set(FunctionCallParser.ToolCallParserEnum.values()) From 52609c55902fba980f41441fac241332db5eb9e4 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 7 Oct 2026 17:45:44 -0400 Subject: [PATCH 69/94] feat(qwen4_exp): read llm-compressor (compressed-tensors) NVFP4 routed experts The expert spec now takes its tensor names and global-scale direction from the checkpoint's quant config, as qwen3_5_moe does, so RedHatAI/Qwen3.8-Flash-Next-NVFP4 loads and converts; the nvidia modelopt checkpoint keeps the same spec. Assisted-by: Claude Opus 5.5 --- python/freetoken/models/qwen4_exp/weight.py | 35 +++++++----- tests/models/qwen4_exp/test_weight.py | 59 +++++++++++++++++++++ tests/models/qwen4_exp/test_weight_ckpt.py | 5 +- 3 files changed, 83 insertions(+), 16 deletions(-) diff --git a/python/freetoken/models/qwen4_exp/weight.py b/python/freetoken/models/qwen4_exp/weight.py index 91c8c1bd17..26f47af508 100644 --- a/python/freetoken/models/qwen4_exp/weight.py +++ b/python/freetoken/models/qwen4_exp/weight.py @@ -28,7 +28,7 @@ from freetoken.models.nvfp4_banks import ( Nvfp4ExpertSourceSpec, ) -from freetoken.layers.quantization import get_quant_config +from freetoken.layers.quantization import QuantKind, get_quant_config from freetoken.models.register import get_model_spec from freetoken.moe.host_banks import HostBank, read_range_into from freetoken.utils import cached_load_hf_config, div_even, download_hf_weight, init_logger @@ -37,20 +37,16 @@ logger = init_logger(__name__) -# Routed NVFP4 experts (nvidia modelopt layout): per-expert, un-fused. Matched against the RAW -# weight_map key in nvfp4_banks. The ``model.language_model.`` anchor excludes the MTP head's -# stacked ``mtp.layers.N.mlp.experts.*`` tensors. -_EXPERT_KEY_RE = re.compile( +# Routed NVFP4 experts: per-expert, un-fused, under the checkpoint's dialect names (nvidia modelopt +# or llm-compressor). Matched against the RAW weight_map key in nvfp4_banks. The +# ``model.language_model.`` anchor excludes the MTP head's stacked ``mtp.layers.N.mlp.experts.*``. +_EXPERT_KEY_RE = ( r"^model\.language_model\.layers\.(?P\d+)\.mlp\.experts\.(?P\d+)\." - r"(?Pgate_proj|up_proj|down_proj)\.(?Pweight|weight_scale|weight_scale_2)$" + r"(?Pgate_proj|up_proj|down_proj)\.(?P{kinds})$" ) +# role -> the expert bank reader's canonical (modelopt) tensor kind +_BANK_KINDS = {"weight": "weight", "weight_scale": "weight_scale", "weight_global": "weight_scale_2"} _EXPERT_RE = re.compile(r"\.mlp\.experts\.\d+\.") -_NVFP4_SOURCE_SPEC = Nvfp4ExpertSourceSpec( - key_pattern=_EXPERT_KEY_RE, - proj_to_role={"gate_proj": "gate", "up_proj": "up", "down_proj": "down"}, - layer_to_bank=lambda layer, config: layer, # every layer is MoE - desc="Qwen3.8-Flash-Next NVFP4 experts", -) # Per-tensor modelopt quant scales; consumed with their ``.weight`` (experts) or unused. _SCALE_SUFFIXES = (".weight_scale", ".weight_scale_2", ".input_scale") @@ -543,8 +539,19 @@ def load_ple_table(model_path: str, qwen4_args, *, pin: bool = True, # ====================================================================================== -def nvfp4_expert_spec(model_path: str, config): - return _NVFP4_SOURCE_SPEC +def nvfp4_expert_spec(model_path: str, config) -> Nvfp4ExpertSourceSpec: + """The per-expert NVFP4 layout under the checkpoint's dialect names (modelopt or llm-compressor).""" + quant = get_quant_config() + stored = quant.stored_tensors(QuantKind.NVFP4) + kind_map = {stored[role].name: kind for role, kind in _BANK_KINDS.items()} + return Nvfp4ExpertSourceSpec( + key_pattern=re.compile(_EXPERT_KEY_RE.format(kinds="|".join(map(re.escape, kind_map)))), + proj_to_role={"gate_proj": "gate", "up_proj": "up", "down_proj": "down"}, + layer_to_bank=lambda layer, config: layer, # every layer is MoE + desc=f"Qwen3.8-Flash-Next NVFP4 experts ({quant.dialect})", + kind_map=kind_map, + global_reciprocal=stored["weight_global"].reciprocal, + ) __all__ = [ diff --git a/tests/models/qwen4_exp/test_weight.py b/tests/models/qwen4_exp/test_weight.py index 1a9a100edb..1167eb1b4b 100644 --- a/tests/models/qwen4_exp/test_weight.py +++ b/tests/models/qwen4_exp/test_weight.py @@ -518,3 +518,62 @@ def test_checkpoint_disagreeing_with_its_quant_config_is_rejected(tmp_path, quan (tmp_path / "config.json").write_text(json.dumps(_config_json(quantization_config))) with pytest.raises(ValueError, match=match): _load(str(tmp_path)) + + +_CT_NVFP4_EXPERTS = { + "quant_method": "compressed-tensors", "format": "nvfp4-pack-quantized", "quantization_status": "compressed", + "config_groups": {"group_0": { + "format": "nvfp4-pack-quantized", "targets": [r"re:.*mlp\.experts\..*(gate|up|down)_proj$"], + "weights": {"num_bits": 4, "type": "float", "symmetric": True, "group_size": 16, + "strategy": "tensor_group", "dynamic": False, "scale_dtype": "torch.float8_e4m3fn"}, + "input_activations": {"num_bits": 4, "type": "float", "symmetric": True, "group_size": 16, + "strategy": "tensor_group", "dynamic": "local", "scale_dtype": "torch.float8_e4m3fn"}, + }}, + "ignore": ["lm_head"], +} + + +def test_compressed_tensors_experts_load_like_their_modelopt_twin(tmp_path, monkeypatch): + """An llm-compressor export (RedHatAI) names the routed experts weight_packed / weight_global_scale, + stores the quant-side global (the reciprocal of modelopt's weight_scale_2) and adds an unused + input_global_scale: its pieces must equal the modelopt twin's bit for bit.""" + from freetoken.layers.quantization import QuantConfig, set_quant_config + from freetoken.layers.quantization import configs as qconfigs + from freetoken.models.nvfp4_banks import iter_nvfp4_expert_pieces + from freetoken.models.qwen4_exp.weight import nvfp4_expert_spec + + monkeypatch.setattr(qconfigs, "_QUANT_CONFIG", qconfigs._QUANT_CONFIG) # restored at teardown + layers, experts, hid, inter = 2, 3, 32, 16 + g = torch.Generator().manual_seed(0) + ct, mo = {}, {} + for layer in range(layers): + for e in range(experts): + for proj, (rows, cols) in {"gate_proj": (inter, hid), "up_proj": (inter, hid), "down_proj": (hid, inter)}.items(): + base = f"{LM}.layers.{layer}.mlp.experts.{e}.{proj}" + packed = torch.randint(0, 256, (rows, cols // 2), dtype=torch.uint8, generator=g) + scale = (torch.rand(rows, cols // 16, generator=g) * 4 + 0.1).to(torch.float8_e4m3fn) + glob = torch.rand(1, generator=g) * 1000 + 1 + ct |= {f"{base}.weight_packed": packed, f"{base}.weight_scale": scale, + f"{base}.weight_global_scale": glob, f"{base}.input_global_scale": torch.rand(1, generator=g)} + mo |= {f"{base}.weight": packed, f"{base}.weight_scale": scale, + f"{base}.weight_scale_2": 1.0 / glob, f"{base}.input_scale": torch.rand(1, generator=g)} + config = SimpleNamespace(num_experts=experts, num_moe_layers=layers, moe_intermediate_size=inter) + + def load(name, tensors, quant): + folder = tmp_path / name + folder.mkdir() + save_file(tensors, str(folder / "model.safetensors")) + set_quant_config(QuantConfig.from_hf({"quantization_config": quant})) + spec = nvfp4_expert_spec(str(folder), config) + pieces = iter_nvfp4_expert_pieces(str(folder), config, spec, drop_page_cache=lambda path: None, primary=False) + return spec, {(lay, e): piece for lay, e, _, piece in pieces} + + ct_spec, ct_pieces = load("ct", ct, _CT_NVFP4_EXPERTS) + mo_spec, mo_pieces = load("mo", mo, {"quant_method": "modelopt", "quant_algo": "NVFP4", "ignore": ["lm_head"]}) + assert ct_spec.global_reciprocal and not mo_spec.global_reciprocal + assert sorted(ct_pieces) == sorted(mo_pieces) == [(lay, e) for lay in range(layers) for e in range(experts)] + for key, piece in ct_pieces.items(): + assert sorted(piece) == sorted(mo_pieces[key]) + for role, tensor in piece.items(): + twin = mo_pieces[key][role] + assert tensor.dtype == twin.dtype and torch.equal(tensor.view(torch.uint8), twin.view(torch.uint8)), (key, role) diff --git a/tests/models/qwen4_exp/test_weight_ckpt.py b/tests/models/qwen4_exp/test_weight_ckpt.py index 605c40f18b..38d32bd1af 100644 --- a/tests/models/qwen4_exp/test_weight_ckpt.py +++ b/tests/models/qwen4_exp/test_weight_ckpt.py @@ -22,10 +22,10 @@ from freetoken.models.nvfp4_banks import iter_nvfp4_expert_pieces from freetoken.models.qwen4_exp.config import parse_config from freetoken.models.qwen4_exp.weight import ( - _NVFP4_SOURCE_SPEC, _ZERO_CENTERED_NORM_SUFFIXES, iter_weights, load_ple_table, + nvfp4_expert_spec, ) from freetoken.moe.expert_banks import build_expert_banks from freetoken.moe.host_banks import HostResidency @@ -271,8 +271,9 @@ def layer0_expert_banks(): """The real NVFP4 source-bank loader, restricted to layer 0 (1.32 GiB instead of 63.5).""" if not torch.cuda.is_available(): pytest.skip("expert bank pinning needs CUDA") + install_quant_config(MODEL_PATH) spec = dataclasses.replace( - _NVFP4_SOURCE_SPEC, layer_to_bank=lambda layer, config: 0 if layer == 0 else None + nvfp4_expert_spec(MODEL_PATH, None), layer_to_bank=lambda layer, config: 0 if layer == 0 else None ) config = SimpleNamespace(num_experts=E, hidden_size=H, moe_intermediate_size=I, num_moe_layers=1) From d65e5dbac2543870d6d3e756738c0fd1c66a62f1 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 7 Oct 2026 18:46:57 -0400 Subject: [PATCH 70/94] test(moe): let the prefetch counter test's copy spies pass #610's blocks_per_bank through Assisted-by: Claude Opus 5.5 --- tests/models/qwen4_exp/test_skeleton.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/models/qwen4_exp/test_skeleton.py b/tests/models/qwen4_exp/test_skeleton.py index 2b47748f53..0803cbca28 100644 --- a/tests/models/qwen4_exp/test_skeleton.py +++ b/tests/models/qwen4_exp/test_skeleton.py @@ -1117,13 +1117,13 @@ def spy_ensure(layer, ids, **kw): calls.append(("ensure", layer, ids.clone(), None if logits is None or policy != "rule" else logits.clone())) return real_ensure(layer, ids, **kw) - def spy_slim(*args): + def spy_slim(*args, **kw): streams["slim"].add(torch.cuda.current_stream().cuda_stream) - return real_slim(*args) + return real_slim(*args, **kw) - def spy_multi(*args): + def spy_multi(*args, **kw): streams["demand"].add(torch.cuda.current_stream().cuda_stream) - return real_multi(*args) + return real_multi(*args, **kw) def install(): monkeypatch.setattr(cache, "prefetch_ensure", spy_pf) From 0f282b01cab62fded7ffde7e5c164e10bf25cc17 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Tue, 13 Oct 2026 09:43:17 -0400 Subject: [PATCH 71/94] test(engine): give #650's rebuild engine fake the kv offloader and gdn group the fit check reads The carried host kv tier (FlashML-org/FreeToken#499) rejects a pages rebuild through engine.kv_offloader, and the fork's gdn workspace reserve asks the model config for its linear attention group; the SimpleNamespace shell had neither. Assisted-by: Claude Opus 5.5 Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Bu6LgoxLR4wETqb7RPR2vt --- tests/engine/test_rebuild_runtime_cache_ordering.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/engine/test_rebuild_runtime_cache_ordering.py b/tests/engine/test_rebuild_runtime_cache_ordering.py index 229431b7a4..aae7d49cf0 100644 --- a/tests/engine/test_rebuild_runtime_cache_ordering.py +++ b/tests/engine/test_rebuild_runtime_cache_ordering.py @@ -125,12 +125,14 @@ def _build_engine(monkeypatch, ledger: _Ledger, *, moe: int, kv: int, budget: in config = types.SimpleNamespace( page_size=16, max_seq_len=4096, max_running_req=4, memory_ratio=1.0, cuda_graph_max_bs=4, swa_num_pages_override=None, - model_config=types.SimpleNamespace(vocab_size=128, model_is_mrope=False), + model_config=types.SimpleNamespace( + vocab_size=128, model_is_mrope=False, linear_attention_group=lambda: None + ), ) page_table = torch.zeros((config.max_running_req + 1, 256), dtype=torch.int32) eng = types.SimpleNamespace( config=config, moe_offload_cache=_FakeMoeCache(ledger, moe), kv_cache=_FakeKvPool(ledger, kv), - linear_state_pool=None, num_pages=kv, device=torch.device("cpu"), stream=None, model=None, + linear_state_pool=None, kv_offloader=None, num_pages=kv, device=torch.device("cpu"), stream=None, model=None, _baseline_free=budget, _weights_bytes=0, page_table=page_table, max_seq_len=4096, ctx=types.SimpleNamespace(page_table=page_table), dummy_req=types.SimpleNamespace(table_idx=0), From 67a920458cbd61cf8a0e4667ca6713eb7a08ddc0 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Tue, 13 Oct 2026 09:43:17 -0400 Subject: [PATCH 72/94] test(server): expect out-of-range penalties to fail request validation Main (FlashML-org/FreeToken#527) bounds presence/frequency penalties to [-2, 2] on the request models, so they never reach the handler's 400 path. Assisted-by: Claude Opus 5.5 Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Bu6LgoxLR4wETqb7RPR2vt --- tests/server/test_openai_extras.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/server/test_openai_extras.py b/tests/server/test_openai_extras.py index 115112b9d0..19140d8b72 100644 --- a/tests/server/test_openai_extras.py +++ b/tests/server/test_openai_extras.py @@ -8,6 +8,7 @@ import json from types import SimpleNamespace +import pydantic import pytest from freetoken.message import TokenizeMsg, UserReply from freetoken.server.api_models import ( @@ -250,8 +251,6 @@ def test_sampling_extras_reach_sampling_params(): "field", [ {"min_p": 1.5}, - {"presence_penalty": 3.0}, - {"frequency_penalty": -2.5}, {"repetition_penalty": 0.0}, {"logit_bias": {"x": 1.0}}, {"min_tokens": -1}, @@ -263,6 +262,13 @@ def test_bad_sampling_extras_are_400(field): assert body.status_code == 400, field +@pytest.mark.parametrize("field", [{"presence_penalty": 3.0}, {"frequency_penalty": -2.5}]) +def test_out_of_range_penalties_fail_request_validation(field): + # main's request models bound both penalties to [-2, 2] (FastAPI answers 422) + with pytest.raises(pydantic.ValidationError): + chat(**field) + + def test_continue_final_message_needs_an_assistant_tail_and_reaches_the_template(): spec = chat_request_to_genspec( chat( From 956715011308a66802b0cd839acfc481e48ff77d Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Tue, 13 Oct 2026 11:28:51 -0400 Subject: [PATCH 73/94] fix(kernel): decode nvfp4 qsa prefill tiles narrow at head_dim 256 The wide prefill profile (BLOCK_N 64, 2 warps) that the nvfp4 path inherits from the bf16 tuning has load_nvfp4 decode a 64 x 256 fp32 tile for K and V on 64 threads. Triton 3.8 compiles it to 40 registers with ~1.4 KB of spills: 439 ms instead of 45 ms at 4096 rows on sm_120, and a ~556 MiB driver local-memory reservation outside the torch allocator that OOMs this deployment on its first prefill. BLOCK_N 16 at 4 warps does not spill: 12.3 ms on Triton 3.8 (33 ms on 3.6), same output to bf16 rounding. Upstream: FlashML-org/FreeToken#675, root cause posted on FlashML-org/FreeToken#408. Assisted-by: Claude Opus 5.5 Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Bu6LgoxLR4wETqb7RPR2vt --- python/freetoken/kernel/triton/qsa/attend.py | 6 +++ tests/kernels/test_qsa_nvfp4.py | 55 ++++++++++++++++++++ 2 files changed, 61 insertions(+) diff --git a/python/freetoken/kernel/triton/qsa/attend.py b/python/freetoken/kernel/triton/qsa/attend.py index f6e223385b..9414ca01c3 100644 --- a/python/freetoken/kernel/triton/qsa/attend.py +++ b/python/freetoken/kernel/triton/qsa/attend.py @@ -407,6 +407,12 @@ def qsa_sparse_paged_attention( else: block_n, target_splits, partial_warps = 64, 1, 2 + if nvfp4 and block_n == 64 and head_dim >= 256: + # The wide bf16 profile has load_nvfp4 decode a BLOCK_N x HEAD_DIM fp32 tile for K and V + # across 64 threads (256 values per thread at HEAD_DIM 256): Triton 3.8 compiles that to + # 40 registers with ~1.4 KB of spills, ~10x slower and a ~550 MiB local-memory reservation + # outside the torch allocator. Narrow tiles at 4 warps do not spill and are faster on 3.6 too. + block_n, partial_warps = 16, 4 num_tiles = triton.cdiv(logical_indices.shape[1], block_n) # Avoid empty splits when the selection width is smaller than the profile. max_useful_splits = 1 << (num_tiles.bit_length() - 1) diff --git a/tests/kernels/test_qsa_nvfp4.py b/tests/kernels/test_qsa_nvfp4.py index 591c6f1d4e..ec96e34563 100644 --- a/tests/kernels/test_qsa_nvfp4.py +++ b/tests/kernels/test_qsa_nvfp4.py @@ -139,3 +139,58 @@ def test_qsa_nvfp4_splitk_reads_value_rows(): ref = _decode(codes, row, block).view(3, PAGE, 1, HEAD_DIM) want = qsa_sparse_paged_attention(q, ref, ref, *args) assert torch.equal(got, want) + + +def test_qsa_nvfp4_wide_head_prefill_uses_narrow_tiles(monkeypatch): + """At HEAD_DIM 256 the prefill profile must not decode 64-token fp32 tiles on 2 warps: + Triton 3.8 spills that variant (~10x slower, ~550 MiB of driver local memory).""" + from freetoken.kernel.triton.qsa import attend + + head_dim, kv_heads, query_heads, rows, pages = 256, 2, 24, 32, 3 + seen = [] + kernel = attend._qsa_sparse_paged_gqa_splitk_kernel + + class Spy: + def __getitem__(self, grid): + launch = kernel[grid] + + def run(*args, **kwargs): + seen.append((kwargs["BLOCK_N"], kwargs["num_warps"])) + return launch(*args, **kwargs) + return run + + torch.manual_seed(29) + slots = pages * PAGE + k = torch.randn(slots, kv_heads * head_dim, device="cuda", dtype=torch.bfloat16) + v = torch.randn_like(k) * 0.25 + codes = [torch.zeros((slots, kv_heads, head_dim // 2), dtype=torch.uint8, device="cuda") for _ in range(2)] + rows_s = [torch.zeros((slots, kv_heads), dtype=torch.float32, device="cuda") for _ in range(2)] + blocks = [torch.zeros((slots, kv_heads, head_dim // 16), dtype=torch.uint8, device="cuda") for _ in range(2)] + quantize_nvfp4_to_cache( + k, v, torch.arange(slots, dtype=torch.int32, device="cuda"), codes[0], codes[1], + rows_s[0], rows_s[1], blocks[0], blocks[1], + ) + q = torch.randn(rows, query_heads, head_dim, device="cuda", dtype=torch.bfloat16) + indices, block_table, token_to_req = _layout(rows, topk=128) + packed = (pages, PAGE, kv_heads, head_dim // 2) + monkeypatch.setattr(attend, "_qsa_sparse_paged_gqa_splitk_kernel", Spy()) + got = qsa_sparse_paged_attention( + q, codes[0].view(packed), codes[1].view(packed), indices, block_table, token_to_req, + k_scale=rows_s[0], v_scale=rows_s[1], kv_quant="nvfp4", + k_block_scale=blocks[0], v_block_scale=blocks[1], + ) + assert seen == [(16, 4)] + + def decode(c, r, b): + lo, hi = (c & 15).to(torch.long), (c >> 4).to(torch.long) + code = torch.stack((lo, hi), dim=-1).reshape(*c.shape[:-1], head_dim) + mag = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], device=c.device)[code & 7] + val = torch.where(code < 8, mag, -mag) + blk = b.view(torch.float8_e4m3fn).to(torch.float32).repeat_interleave(16, dim=-1) + return (val * blk * r.unsqueeze(-1)).to(torch.bfloat16).view(pages, PAGE, kv_heads, head_dim) + + want = qsa_sparse_paged_attention( + q, decode(codes[0], rows_s[0], blocks[0]), decode(codes[1], rows_s[1], blocks[1]), + indices, block_table, token_to_req, + ) + torch.testing.assert_close(got.float(), want.float(), atol=2e-3, rtol=2e-2) From e30702584b16bf6905bd01f9f4a0452851ff1edb Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 05:16:35 -0400 Subject: [PATCH 74/94] fix(kernel): size triton's autotune flush buffer to the device l2 A kernel autotunes on its first launch of each key, which can be the first request after the cache planner has handed out nearly all VRAM; triton's flat 256 MB flush buffer was an OOM there on an 11 GB card. Four times the L2 still evicts it between timed runs. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit ffbf001acec04f4f882d6cccbc1d0d4cbbec633f) (cherry picked from commit e097977e525302a6f773ace34ca553666e593876) --- python/freetoken/engine/engine.py | 3 +++ .../freetoken/kernel/triton/autotune_cache.py | 26 ++++++++++++++++++- 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 138a0d1e37..d26141a1eb 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -512,6 +512,9 @@ def __init__(self, config: EngineConfig): logger.info_rank0(f"Free memory before loading model: {mem_GB(init_free_memory)}") # ======================= Model initialization ======================== + from freetoken.kernel.triton.autotune_cache import bound_autotune_flush_buffer + + bound_autotune_flush_buffer(self.device.index) set_rope_device(self.device) with torch.device("meta"), torch_dtype(config.dtype): self.model = create_model(config.model_config) diff --git a/python/freetoken/kernel/triton/autotune_cache.py b/python/freetoken/kernel/triton/autotune_cache.py index ea95768368..2e57fadde9 100644 --- a/python/freetoken/kernel/triton/autotune_cache.py +++ b/python/freetoken/kernel/triton/autotune_cache.py @@ -34,4 +34,28 @@ {"cache_results": _CACHE_RESULTS} if _SUPPORTS_AUTOTUNE_CACHE else {} ) -__all__ = ["autotune_cache_kwargs"] + + +def bound_autotune_flush_buffer(device_index: int) -> None: + """Size triton's L2-flush buffer for autotune benchmarks to the device instead of a flat 256 MB. + + A kernel autotunes on its first launch of each key, which can be the first request, after the + cache planner has handed out nearly all VRAM; a 256 MB transient there is an OOM on a small + card. Four times the L2 still evicts everything between timed runs. + """ + import torch + + try: + from triton.backends.nvidia.driver import CudaDriver + except ImportError: + return + l2 = int(getattr(torch.cuda.get_device_properties(device_index), "L2_cache_size", 0) or 0) + nbytes = min(256 << 20, max(8 << 20, 4 * l2)) + + def get_empty_cache_for_benchmark(self): + return torch.empty(nbytes // 4, dtype=torch.int, device="cuda") + + CudaDriver.get_empty_cache_for_benchmark = get_empty_cache_for_benchmark + + +__all__ = ["autotune_cache_kwargs", "bound_autotune_flush_buffer"] From 4b263dd0e125851474ec85c7135b3b9da3cd70f4 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Tue, 13 Oct 2026 10:02:35 -0400 Subject: [PATCH 75/94] test(utils): drive the async coalesced-pull test with asyncio.run Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 32007ecb48aee2dc619fc9ceb89ad9b5dd70ef14) --- tests/utils/test_mp_coalesced.py | 31 +++++++++++++++++-------------- 1 file changed, 17 insertions(+), 14 deletions(-) diff --git a/tests/utils/test_mp_coalesced.py b/tests/utils/test_mp_coalesced.py index e599db7265..c64ee1c5c3 100644 --- a/tests/utils/test_mp_coalesced.py +++ b/tests/utils/test_mp_coalesced.py @@ -92,7 +92,7 @@ def test_sub_queue_coalesced(self): sub.stop() -async def test_async_pull_coalesced(): +def test_async_pull_coalesced(): import sys from freetoken.utils.mp import ZmqAsyncPullQueue as _AsyncPull @@ -101,16 +101,19 @@ async def test_async_pull_coalesced(): # zmq.asyncio needs a Selector loop on Windows; the project runs on Linux CI. if sys.platform == "win32": pytest.skip("zmq.asyncio requires SelectorEventLoop on Windows (prod runs on Linux)") - push = ZmqAsyncPushQueue(_addr(5604), create=True, encoder=lambda o: o) - pull = ZmqAsyncPullQueue(_addr(5604), create=False, decoder=_identity_decoder) - try: - await push.put({"n": 1}) - await push.put({"n": 2}) - await asyncio.sleep(0.1) - - # Even if both messages arrive in a single frame, two gets must succeed. - assert await pull.get() == {"n": 1} - assert await pull.get() == {"n": 2} - finally: - push.stop() - pull.stop() \ No newline at end of file + async def scenario(): + push = ZmqAsyncPushQueue(_addr(5604), create=True, encoder=lambda o: o) + pull = ZmqAsyncPullQueue(_addr(5604), create=False, decoder=_identity_decoder) + try: + await push.put({"n": 1}) + await push.put({"n": 2}) + await asyncio.sleep(0.1) + + # Even if both messages arrive in a single frame, two gets must succeed. + assert await pull.get() == {"n": 1} + assert await pull.get() == {"n": 2} + finally: + push.stop() + pull.stop() + + asyncio.run(scenario()) From 5ca859097f72fd8c2485d25fa3bbfdc14b7051f4 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 23 Sep 2026 20:35:40 -0400 Subject: [PATCH 76/94] fix(engine): keep multimodal placeholder ids out of the repetition-penalty prompt mask Fork #47's sampler plan scattered the raw prompt ids into a [rows, vocab+1] mask; image tokens carry pseudo-ids >= 1,000,000, so repetition_penalty with an image indexed out of bounds -- a device-side assert that takes the server down. They now land in the scratch column. Also keeps presence/frequency out of the LogitsPlan: with #527 device counts and #393 plan both applying them, a request with a penalty plus any plan processor was penalized twice. (cherry picked from commit 66053ded) Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit b9883902735cadb2232e3ae4d9dd18e02481a189) --- python/freetoken/engine/sample.py | 23 ++++++----------------- tests/engine/test_logits_processors.py | 23 ++++++++++++++++++++--- tests/engine/test_sample_logprobs.py | 4 ++-- 3 files changed, 28 insertions(+), 22 deletions(-) diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py index 89d2a6fb88..8e05f80694 100644 --- a/python/freetoken/engine/sample.py +++ b/python/freetoken/engine/sample.py @@ -172,35 +172,24 @@ def _plan(self, batch: Batch, params) -> LogitsPlan | None: ) pad = self.vocab_size - want_hist = any( - p.presence_penalty != 0.0 or p.frequency_penalty != 0.0 or p.repetition_penalty != 1.0 - for _, p in picked - ) - if want_hist: + # presence/frequency are not plan work: they ride the device counts (Sampler.prepare's penalties) + if any(p.repetition_penalty != 1.0 for _, p in picked): hist: list[list[int]] = [] prompts: list[list[int]] = [] - want_prompt = False for req, p in picked: prompt_len = req.max_device_len - req.output_len ids = req.input_ids - penalized = p.presence_penalty != 0.0 or p.frequency_penalty != 0.0 or p.repetition_penalty != 1.0 - hist.append(ids[prompt_len:].tolist() if penalized else []) if p.repetition_penalty != 1.0: - want_prompt = True + hist.append(ids[prompt_len:].tolist()) # multimodal placeholders are pseudo-ids >= vocab_size (MM_PAD_SHIFT_VALUE); # send them to the scratch column, not past the scatter target prompts.append(ids[:prompt_len].clamp(max=pad).tolist()) else: + hist.append([]) prompts.append([]) plan.hist_ids = _padded_rows(hist, pad, self.device) - if want_prompt: - plan.prompt_ids = _padded_rows(prompts, pad, self.device) - if any(p.presence_penalty != 0.0 for _, p in picked): - plan.presence = make_device_tensor([[p.presence_penalty] for _, p in picked], torch.float32, self.device) - if any(p.frequency_penalty != 0.0 for _, p in picked): - plan.frequency = make_device_tensor([[p.frequency_penalty] for _, p in picked], torch.float32, self.device) - if any(p.repetition_penalty != 1.0 for _, p in picked): - plan.repetition = make_device_tensor([[p.repetition_penalty] for _, p in picked], torch.float32, self.device) + plan.prompt_ids = _padded_rows(prompts, pad, self.device) + plan.repetition = make_device_tensor([[p.repetition_penalty] for _, p in picked], torch.float32, self.device) bias_rows: list[int] = [] bias_ids: list[int] = [] diff --git a/tests/engine/test_logits_processors.py b/tests/engine/test_logits_processors.py index abd2e1f939..237ce9f602 100644 --- a/tests/engine/test_logits_processors.py +++ b/tests/engine/test_logits_processors.py @@ -172,9 +172,8 @@ def test_prepare_builds_no_plan_without_processors_and_a_plan_with_them(): logits = torch.zeros((4, V)) out = apply_logits_processors(logits, plan, V) - assert ( - out[1, 3].item() == -1.0 - ) # presence on the repeated 3 (repetition on 0 is a no-op) + # presence is the sampler's device counts, not the plan; repetition on 0 is a no-op + assert plan.presence is None and out[1, 3].item() == 0.0 assert out[1, 5].item() == 2.0 assert out[2, 0].item() == float("-inf") and out[2, 7].item() == float("-inf") assert torch.isfinite(out[3]).all() @@ -182,6 +181,24 @@ def test_prepare_builds_no_plan_without_processors_and_a_plan_with_them(): assert sampler.sample(logits, args)[1].item() == 5 +def test_presence_and_frequency_ride_the_device_counts_not_the_plan(): + # the plan must not apply them a second time on top of the sampler's own counts + sampler = Sampler(CPU, V) + only = SimpleNamespace(reqs=[_req([1], [], 4, presence_penalty=1.0, frequency_penalty=0.5)]) + args = sampler.prepare(only) + assert args.plan is None and [row for row, *_ in args.penalties] == [0] + + req = _req([1], [], 4, presence_penalty=1.0, logit_bias=[[5, 0.25]]) + args = sampler.prepare(SimpleNamespace(reqs=[req])) + assert args.plan is not None and [row for row, *_ in args.penalties] == [0] + req.output_token_counts[3] = 1 + logits = torch.zeros((1, V)) + logits[0, 3] = 1.5 + # once: 3 -> 0.5 beats 5 -> 0.25; twice would drop 3 to -0.5 and pick 5 + assert sampler.sample(logits, args).item() == 3 + assert req.output_token_counts[3].item() == 2 + + def test_repetition_penalty_ignores_multimodal_placeholder_ids_in_the_prompt(): # image tokens carry pseudo-ids >= MM_PAD_SHIFT_VALUE; scattering them into the # [rows, V+1] prompt mask indexed out of bounds (a device-side assert on CUDA) diff --git a/tests/engine/test_sample_logprobs.py b/tests/engine/test_sample_logprobs.py index 6bd204604d..fae5b7c25b 100644 --- a/tests/engine/test_sample_logprobs.py +++ b/tests/engine/test_sample_logprobs.py @@ -137,8 +137,8 @@ def test_compute_logprobs_with_prepare_keeps_max_top_logprobs_clamped(monkeypatc batch = SimpleNamespace( reqs=[ - SimpleNamespace(sampling_params=SimpleNamespace(needs_logits_processing=False, logprobs=True, top_logprobs=17, is_greedy=True)), - SimpleNamespace(sampling_params=SimpleNamespace(needs_logits_processing=False, logprobs=False, top_logprobs=0, is_greedy=True)), + SimpleNamespace(sampling_params=SimpleNamespace(needs_logits_processing=False, presence_penalty=0.0, frequency_penalty=0.0, logprobs=True, top_logprobs=17, is_greedy=True)), + SimpleNamespace(sampling_params=SimpleNamespace(needs_logits_processing=False, presence_penalty=0.0, frequency_penalty=0.0, logprobs=False, top_logprobs=0, is_greedy=True)), ] ) args = sampler.prepare(batch) From 7de35d9c2a31e08b3e98f1f3a532cab27b85ef27 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 06:43:17 -0400 Subject: [PATCH 77/94] fix(engine): apply the repetition penalty from device state, not a per-step prompt upload (up to 262K ids) and a synchronous pageable copy, which stalled the overlap loop and saw the output one token late. The prompt mask is now built once per request, the output side reuses #527's device counts, and the order is vLLM's (repetition, then presence and frequency). Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 343d7421dfebf72244b9e9d16261dba073b48cbc) (cherry picked from commit 3df811dee9bcd4af287e3cc94abb689d531e58bd) --- python/freetoken/core.py | 11 +- python/freetoken/engine/sample.py | 188 ++++++++++++++----------- tests/engine/test_logits_processors.py | 83 +++++++---- tests/engine/test_sample_logprobs.py | 5 +- 4 files changed, 172 insertions(+), 115 deletions(-) diff --git a/python/freetoken/core.py b/python/freetoken/core.py index ab5b81b0af..a7811ed712 100644 --- a/python/freetoken/core.py +++ b/python/freetoken/core.py @@ -30,7 +30,8 @@ class SamplingParams: # min_p: drop tokens whose probability is below min_p x the top probability. min_p: float = 0.0 # OpenAI presence/frequency penalties over the generated tokens; repetition_penalty is the - # HF/vLLM multiplicative penalty over prompt + generated tokens (1.0 = off). + # HF/vLLM multiplicative penalty over prompt + generated tokens (1.0 = off). The sampler + # keeps their state on the device (Req.output_token_counts, Req.prompt_token_mask). presence_penalty: float = 0.0 frequency_penalty: float = 0.0 repetition_penalty: float = 1.0 @@ -60,11 +61,10 @@ def is_greedy(self) -> bool: @property def needs_logits_processing(self) -> bool: + # the penalties are not LogitsPlan work: the sampler applies them from its own + # device-side state (Sampler.prepare's penalties). return bool( self.min_p > 0.0 - or self.presence_penalty != 0.0 - or self.frequency_penalty != 0.0 - or self.repetition_penalty != 1.0 or self.logit_bias or self.min_tokens > 0 ) @@ -106,7 +106,10 @@ class Req: # handler must not free resources under an in-flight forward; it sets this flag and # _process_last_data frees the request when the batch drains (after copy_done.synchronize). aborted: bool = False + output_token_counts: torch.Tensor | None = field(default=None, init=False, repr=False) stop_decode_status: DecodeStatus | None = None + # [vocab] bool on the device: the prompt's ids, for repetition_penalty. Built once. + prompt_token_mask: torch.Tensor | None = field(default=None, init=False, repr=False) def __post_init__(self) -> None: assert self.input_ids.is_cpu diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py index 8e05f80694..9d12b5f571 100644 --- a/python/freetoken/engine/sample.py +++ b/python/freetoken/engine/sample.py @@ -1,7 +1,7 @@ from __future__ import annotations -from dataclasses import dataclass -from typing import TYPE_CHECKING, List +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, List, NamedTuple import torch from freetoken.utils import is_sm90_supported, nvtx_annotate @@ -13,41 +13,49 @@ @dataclass class LogitsPlan: """Inputs of the logits processors for the rows of one batch that asked for any - (presence / frequency / repetition penalties, logit_bias, min_tokens, min_p). - - Built on the host in ``Sampler.prepare`` from the requests' SamplingParams and - token histories, applied on the device by ``apply_logits_processors`` before the - sampling kernel. ``None`` when no request in the batch uses a processor, so the - default path stays exactly as it was. Row indices inside the plan are LOCAL (0..r-1, - in the order of ``rows``); ``rows`` maps them back to the batch. + (logit_bias, min_tokens, min_p). The penalties are not here: the sampler applies + them from per-request device state (``RowPenalty``). + + Built on the host in ``Sampler.prepare`` from the requests' SamplingParams, applied + on the device by ``apply_logits_processors`` before the sampling kernel. ``None`` + when no request in the batch uses a processor, so the default path stays exactly + as it was. Row indices inside the plan are LOCAL (0..r-1, in the order of ``rows``); + ``rows`` maps them back to the batch. """ rows: torch.Tensor # [r] int64 batch rows temps: torch.Tensor # [r, 1] float32 temperatures (min_p works on probabilities) - presence: torch.Tensor | None = None # [r, 1] float32 - frequency: torch.Tensor | None = None # [r, 1] float32 - repetition: torch.Tensor | None = None # [r, 1] float32 (1.0 = off) - # Generated ids per row, right-padded with vocab_size (a scratch column that is - # dropped), for the presence / frequency counts and the repetition mask. - hist_ids: torch.Tensor | None = None # [r, L] int64 - # Prompt ids per row, same padding: repetition_penalty also covers the prompt. - prompt_ids: torch.Tensor | None = None # [r, P] int64 bias_rows: torch.Tensor | None = None # [b] int64 local rows bias_ids: torch.Tensor | None = None # [b] int64 bias_vals: torch.Tensor | None = None # [b] float32 - # Rows still under min_tokens and the ids (EOS + stop_token_ids, padded) they must - # not sample yet. + # Rows still under min_tokens and the ids (EOS + stop_token_ids, padded with + # vocab_size, a scratch column that is dropped) they must not sample yet. min_rows: torch.Tensor | None = None # [m] int64 local rows min_ids: torch.Tensor | None = None # [m, S] int64 min_p: torch.Tensor | None = None # [r, 1] float32 (0 = off for that row) +class RowPenalty(NamedTuple): + """One batch row's presence / frequency / repetition penalty. ``counts`` are the + request's generated-token counts ([V] int32), updated on the device after each draw so + overlap scheduling never sees them late; ``prompt_mask`` ([V] bool) marks the prompt's + ids for the repetition penalty and is None when that is off.""" + + row: int + counts: torch.Tensor + presence: float + frequency: float + repetition: float = 1.0 + prompt_mask: torch.Tensor | None = None + + @dataclass class BatchSamplingArgs: temperatures: torch.Tensor | None top_k: torch.Tensor | None = None top_p: torch.Tensor | None = None greedy_mask: torch.Tensor | None = None + penalties: list[RowPenalty] = field(default_factory=list) plan: LogitsPlan | None = None logprob_rows: torch.Tensor | None = None max_top_logprobs: int = 0 @@ -60,47 +68,39 @@ def make_device_tensor(data: List, dtype: torch.dtype, device: torch.device) -> def _padded_rows(rows: list[list[int]], pad: int, device: torch.device) -> torch.Tensor: - """[len(rows), max_len] int64 with ``pad`` on the right. A synchronous copy: the - histories can be long (prompt + output) and a pinned buffer must not be freed under - an in-flight non_blocking copy.""" + """[len(rows), max_len] int64 with ``pad`` on the right.""" width = max(1, max((len(r) for r in rows), default=1)) - out = torch.full((len(rows), width), pad, dtype=torch.int64) - for i, r in enumerate(rows): - if r: - out[i, : len(r)] = torch.as_tensor(r, dtype=torch.int64) - return out.to(device) + return make_device_tensor([r + [pad] * (width - len(r)) for r in rows], torch.int64, device) + + +def apply_penalties(logits: torch.Tensor, penalties: list[RowPenalty]) -> torch.Tensor: + """A float32 copy of ``logits`` with each row's penalties applied in vLLM's order: + repetition (over prompt + generated ids), then frequency and presence (generated).""" + logits = logits.float().clone() + for row, counts, presence, frequency, repetition, prompt_mask in penalties: + if repetition != 1.0: + x = logits[row] + seen = counts > 0 + if prompt_mask is not None: + seen = seen | prompt_mask + # HF / vLLM semantics: divide positive logits, multiply negative ones. + penalized = torch.where(x > 0, x / repetition, x * repetition) + logits[row] = torch.where(seen, penalized, x) + if presence or frequency: + logits[row] -= frequency * counts + presence * (counts > 0) + return logits def apply_logits_processors( logits: torch.Tensor, plan: LogitsPlan, vocab_size: int ) -> torch.Tensor: - """A float32 copy of ``logits`` with the plan applied to its rows, in vLLM's order: - repetition -> presence -> frequency penalties, logit_bias, the min_tokens mask, then - min_p. Rows outside the plan are copied unchanged.""" + """A float32 copy of ``logits`` with the plan applied to its rows: logit_bias, the + min_tokens mask, then min_p. Rows outside the plan are copied unchanged.""" out = logits.to(torch.float32, copy=True) sub = out[plan.rows] # [r, V] (advanced indexing copies) - r = sub.shape[0] device = sub.device neg_inf = float("-inf") - if plan.hist_ids is not None: - counts = torch.zeros((r, vocab_size + 1), dtype=torch.float32, device=device) - counts.scatter_add_(1, plan.hist_ids, torch.ones_like(plan.hist_ids, dtype=torch.float32)) - counts = counts[:, :vocab_size] - if plan.repetition is not None: - seen = counts > 0 - if plan.prompt_ids is not None: - in_prompt = torch.zeros((r, vocab_size + 1), dtype=torch.bool, device=device) - in_prompt.scatter_(1, plan.prompt_ids, torch.ones_like(plan.prompt_ids, dtype=torch.bool)) - seen = seen | in_prompt[:, :vocab_size] - # HF / vLLM semantics: divide positive logits, multiply negative ones. - penalized = torch.where(sub > 0, sub / plan.repetition, sub * plan.repetition) - sub = torch.where(seen, penalized, sub) - if plan.presence is not None: - sub = sub - plan.presence * (counts > 0).to(sub.dtype) - if plan.frequency is not None: - sub = sub - plan.frequency * counts - if plan.bias_ids is not None: sub.index_put_((plan.bias_rows, plan.bias_ids), plan.bias_vals, accumulate=True) @@ -155,11 +155,7 @@ class Sampler: vocab_size: int def _plan(self, batch: Batch, params) -> LogitsPlan | None: - """The LogitsPlan for this batch, or None when no request asks for a processor. - Token histories come from the host-side ``req.input_ids``; under overlap - scheduling the previous step's token may not be appended yet, so a penalty sees - the history one token late. That is the accepted cost of keeping the counts off - the device.""" + """The LogitsPlan for this batch, or None when no request asks for a processor.""" MIN_T = 1e-6 rows = [i for i, p in enumerate(params) if p.needs_logits_processing] if not rows: @@ -172,25 +168,6 @@ def _plan(self, batch: Batch, params) -> LogitsPlan | None: ) pad = self.vocab_size - # presence/frequency are not plan work: they ride the device counts (Sampler.prepare's penalties) - if any(p.repetition_penalty != 1.0 for _, p in picked): - hist: list[list[int]] = [] - prompts: list[list[int]] = [] - for req, p in picked: - prompt_len = req.max_device_len - req.output_len - ids = req.input_ids - if p.repetition_penalty != 1.0: - hist.append(ids[prompt_len:].tolist()) - # multimodal placeholders are pseudo-ids >= vocab_size (MM_PAD_SHIFT_VALUE); - # send them to the scratch column, not past the scatter target - prompts.append(ids[:prompt_len].clamp(max=pad).tolist()) - else: - hist.append([]) - prompts.append([]) - plan.hist_ids = _padded_rows(hist, pad, self.device) - plan.prompt_ids = _padded_rows(prompts, pad, self.device) - plan.repetition = make_device_tensor([[p.repetition_penalty] for _, p in picked], torch.float32, self.device) - bias_rows: list[int] = [] bias_ids: list[int] = [] bias_vals: list[float] = [] @@ -223,9 +200,48 @@ def _plan(self, batch: Batch, params) -> LogitsPlan | None: plan.min_p = make_device_tensor([[p.min_p] for _, p in picked], torch.float32, self.device) return plan + def _prompt_mask(self, req) -> torch.Tensor: + """[V] bool marking the request's prompt ids, uploaded once per request.""" + prompt_len = req.max_device_len - req.output_len + # multimodal placeholders are pseudo-ids >= vocab_size (MM_PAD_SHIFT_VALUE); + # send them to the scratch column, not past the scatter target + ids = req.input_ids[:prompt_len].to(torch.int64).clamp(max=self.vocab_size) + if self.device.type == "cuda": + ids = ids.pin_memory().to(self.device, non_blocking=True) + mask = torch.zeros(self.vocab_size + 1, dtype=torch.bool, device=self.device) + mask[ids] = True + return mask[: self.vocab_size] + def prepare(self, batch: Batch) -> BatchSamplingArgs: params = [r.sampling_params for r in batch.reqs] is_greedy = [p.is_greedy for p in params] + penalties = [] + for row, req in enumerate(batch.reqs): + p = req.sampling_params + repetition = p.repetition_penalty + if not (p.presence_penalty or p.frequency_penalty or repetition != 1.0): + continue + if not req.can_decode: + continue + if req.output_token_counts is None: + req.output_token_counts = torch.zeros( + self.vocab_size, dtype=torch.int32, device=self.device + ) + prompt_mask = None + if repetition != 1.0: + if req.prompt_token_mask is None: + req.prompt_token_mask = self._prompt_mask(req) + prompt_mask = req.prompt_token_mask + penalties.append( + RowPenalty( + row, + req.output_token_counts, + p.presence_penalty, + p.frequency_penalty, + repetition, + prompt_mask, + ) + ) plan = self._plan(batch, params) want_logprobs = [p.logprobs for p in params] logprob_rows = ( @@ -239,6 +255,7 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs: if all(is_greedy): return BatchSamplingArgs( temperatures=None, + penalties=penalties, plan=plan, logprob_rows=logprob_rows, max_top_logprobs=max_top_logprobs, @@ -270,6 +287,7 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs: top_k=top_k, top_p=top_p, greedy_mask=greedy_mask, + penalties=penalties, plan=plan, logprob_rows=logprob_rows, max_top_logprobs=max_top_logprobs, @@ -278,16 +296,26 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs: @nvtx_annotate("Sampler") def sample(self, logits: torch.Tensor, args: BatchSamplingArgs) -> torch.Tensor: with torch.cuda.nvtx.range("Sampler"): + if args.penalties: + logits = apply_penalties(logits, args.penalties) + # after the penalties: min_p must see the final distribution if args.plan is not None: logits = apply_logits_processors(logits, args.plan, self.vocab_size) if args.temperatures is None: # greedy sampling - return torch.argmax(logits, dim=-1) - tokens = sample_impl(logits.float(), args.temperatures, args.top_k, args.top_p) - if args.greedy_mask is not None: - # Mixed batches still run probability sampling for all rows, but - # greedy rows must follow argmax's deterministic tie-breaking. - greedy_tokens = torch.argmax(logits, dim=-1).to(tokens.dtype) - tokens = torch.where(args.greedy_mask, greedy_tokens, tokens) + tokens = torch.argmax(logits, dim=-1) + else: + tokens = sample_impl(logits.float(), args.temperatures, args.top_k, args.top_p) + if args.greedy_mask is not None: + # Mixed batches still run probability sampling for all rows, but + # greedy rows must follow argmax's deterministic tie-breaking. + greedy_tokens = torch.argmax(logits, dim=-1).to(tokens.dtype) + tokens = torch.where(args.greedy_mask, greedy_tokens, tokens) + # Update on the sampling stream: overlapped scheduling can prepare the next + # batch before the previous token reaches Req.input_ids on the CPU. + for row, counts, *_ in args.penalties: + counts.scatter_add_( + 0, tokens[row : row + 1].long(), counts.new_ones(1) + ) return tokens def compute_logprobs( diff --git a/tests/engine/test_logits_processors.py b/tests/engine/test_logits_processors.py index 237ce9f602..1829f35f32 100644 --- a/tests/engine/test_logits_processors.py +++ b/tests/engine/test_logits_processors.py @@ -8,7 +8,13 @@ import pytest import torch from freetoken.core import SamplingParams -from freetoken.engine.sample import LogitsPlan, Sampler, apply_logits_processors +from freetoken.engine.sample import ( + LogitsPlan, + RowPenalty, + Sampler, + apply_logits_processors, + apply_penalties, +) V = 8 CPU = torch.device("cpu") @@ -32,33 +38,17 @@ def _ids(rows: list[list[int]]) -> torch.Tensor: return out -def test_presence_and_frequency_penalties_count_generated_tokens(): - logits = torch.zeros((2, V)) - plan = _plan( - [1], - hist_ids=_ids([[3, 3, 5]]), - presence=torch.tensor([[0.5]]), - frequency=torch.tensor([[0.25]]), - ) - out = apply_logits_processors(logits, plan, V) - assert torch.equal(out[0], logits[0]) # row outside the plan untouched - assert out[1, 3].item() == -(0.5 + 2 * 0.25) - assert out[1, 5].item() == -(0.5 + 0.25) - assert out[1, 0].item() == 0.0 - - def test_repetition_penalty_covers_prompt_and_output_with_hf_semantics(): logits = torch.tensor([[2.0, -2.0, 2.0, -2.0, 1.0, 1.0, 1.0, 1.0]]) - plan = _plan( - [0], - hist_ids=_ids([[0, 1]]), # generated: 0 (positive), 1 (negative) - prompt_ids=_ids([[2, 3]]), # prompt: 2 (positive), 3 (negative) - repetition=torch.tensor([[2.0]]), - ) - out = apply_logits_processors(logits, plan, V) + counts = torch.zeros(V, dtype=torch.int32) + counts[[0, 1]] = 1 # generated: 0 (positive), 1 (negative) + prompt = torch.zeros(V, dtype=torch.bool) + prompt[[2, 3]] = True # prompt: 2 (positive), 3 (negative) + out = apply_penalties(logits, [RowPenalty(0, counts, 0.0, 0.0, 2.0, prompt)]) assert out[0, 0].item() == 1.0 and out[0, 2].item() == 1.0 # divided assert out[0, 1].item() == -4.0 and out[0, 3].item() == -4.0 # multiplied assert out[0, 4].item() == 1.0 # unseen + assert logits[0, 0].item() == 2.0 # the caller's logits stay untouched def test_logit_bias_adds_per_row(): @@ -134,6 +124,9 @@ def _req( output_len=max_tokens, max_device_len=len(prompt) + max_tokens, sampling_params=SamplingParams(**sp), + can_decode=True, + output_token_counts=None, + prompt_token_mask=None, ) @@ -164,16 +157,15 @@ def test_prepare_builds_no_plan_without_processors_and_a_plan_with_them(): plan = args.plan assert plan is not None assert plan.rows.tolist() == [1, 2, 3] - assert plan.hist_ids[0].tolist()[:2] == [3, 3] # row 1's generated tokens - assert plan.prompt_ids[0].tolist()[:2] == [1, 2] + assert [pen.row for pen in args.penalties] == [1] + assert batch.reqs[1].prompt_token_mask.nonzero().flatten().tolist() == [1, 2] assert plan.bias_rows.tolist() == [0] and plan.bias_ids.tolist() == [5] assert plan.min_rows.tolist() == [1] # local row of the third request only assert plan.min_ids[0].tolist() == [0, 7] logits = torch.zeros((4, V)) out = apply_logits_processors(logits, plan, V) - # presence is the sampler's device counts, not the plan; repetition on 0 is a no-op - assert plan.presence is None and out[1, 3].item() == 0.0 + assert out[1, 3].item() == 0.0 # the penalties are not plan work assert out[1, 5].item() == 2.0 assert out[2, 0].item() == float("-inf") and out[2, 7].item() == float("-inf") assert torch.isfinite(out[3]).all() @@ -181,7 +173,7 @@ def test_prepare_builds_no_plan_without_processors_and_a_plan_with_them(): assert sampler.sample(logits, args)[1].item() == 5 -def test_presence_and_frequency_ride_the_device_counts_not_the_plan(): +def test_penalties_ride_the_device_state_not_the_plan(): # the plan must not apply them a second time on top of the sampler's own counts sampler = Sampler(CPU, V) only = SimpleNamespace(reqs=[_req([1], [], 4, presence_penalty=1.0, frequency_penalty=0.5)]) @@ -205,5 +197,38 @@ def test_repetition_penalty_ignores_multimodal_placeholder_ids_in_the_prompt(): sampler = Sampler(CPU, V) batch = SimpleNamespace(reqs=[_req([1, 1_000_000, 2], [3], 4, repetition_penalty=2.0)]) args = sampler.prepare(batch) - out = apply_logits_processors(torch.ones(1, V), args.plan, V) + out = apply_penalties(torch.ones(1, V), args.penalties) assert out[0, 1] == 0.5 and out[0, 2] == 0.5 and out[0, 4] == 1.0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_min_p_with_temperature_draws_only_the_kept_tokens(): + # min_p leaves -inf over nearly the whole row; the sampling softmax must not turn that + # into NaN (which the Pascal fallback then samples as a uniform row) + V = 32768 + sampler = Sampler(torch.device("cuda"), V) + req = _req([1, 2], [], 4, temperature=0.8, min_p=0.1) + args = sampler.prepare(SimpleNamespace(reqs=[req])) + logits = torch.full((1, V), -10.0, device="cuda") + kept = [100, 20000, 31000] + logits[0, kept] = torch.tensor([5.0, 4.5, 4.0], device="cuda") + for _ in range(32): + assert sampler.sample(logits, args).item() in kept + + +def test_repetition_penalty_reads_the_prompt_once_and_counts_output_on_the_device(): + # the prompt (up to the whole context) is uploaded once per request, not every step, + # and the generated side comes from the counts sample() keeps, not the host history + sampler = Sampler(CPU, V) + req = _req([1, 2], [], 4, repetition_penalty=2.0) + batch = SimpleNamespace(reqs=[req]) + args = sampler.prepare(batch) + mask = req.prompt_token_mask + logits = torch.tensor([[0.0, 1.0, 1.0, 4.0, 0.5, 0.0, 0.0, 0.0]]) + assert sampler.sample(logits, args).item() == 3 + assert req.output_token_counts.tolist() == [0, 0, 0, 1, 0, 0, 0, 0] + req.input_ids = torch.tensor([9, 9, 3], dtype=torch.int32) # a stale host view is not read + args = sampler.prepare(batch) + assert req.prompt_token_mask is mask + out = apply_penalties(logits, args.penalties) + assert out[0].tolist() == [0.0, 0.5, 0.5, 2.0, 0.5, 0.0, 0.0, 0.0] diff --git a/tests/engine/test_sample_logprobs.py b/tests/engine/test_sample_logprobs.py index fae5b7c25b..88fffcc2ed 100644 --- a/tests/engine/test_sample_logprobs.py +++ b/tests/engine/test_sample_logprobs.py @@ -2,6 +2,7 @@ import torch +from freetoken.core import SamplingParams from freetoken.engine.sample import BatchSamplingArgs, Sampler @@ -137,8 +138,8 @@ def test_compute_logprobs_with_prepare_keeps_max_top_logprobs_clamped(monkeypatc batch = SimpleNamespace( reqs=[ - SimpleNamespace(sampling_params=SimpleNamespace(needs_logits_processing=False, presence_penalty=0.0, frequency_penalty=0.0, logprobs=True, top_logprobs=17, is_greedy=True)), - SimpleNamespace(sampling_params=SimpleNamespace(needs_logits_processing=False, presence_penalty=0.0, frequency_penalty=0.0, logprobs=False, top_logprobs=0, is_greedy=True)), + SimpleNamespace(sampling_params=SamplingParams(logprobs=True, top_logprobs=17)), + SimpleNamespace(sampling_params=SamplingParams(logprobs=False, top_logprobs=0)), ] ) args = sampler.prepare(batch) From be43d10e64da2bc0359bcbf9acf2280368a58e2c Mon Sep 17 00:00:00 2001 From: Siarhei Berdachuk Date: Mon, 5 Oct 2026 19:11:56 +0300 Subject: [PATCH 78/94] fix(server): clarify unsupported response_format error text After soft-accepting json_object/json_schema, the old "no constrained decoding" message was misleading for remaining unknown types. (cherry picked from commit 51e17e6c3380af25d34a301859a8991224f30c2f) (cherry picked from commit e2df92f2ade5cd0ed4f5f1bd421eef96ded793cb) --- python/freetoken/server/openai_api.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index 7f94efc250..79dff72963 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -305,8 +305,10 @@ async def handle_chat_completion( if logprobs_conflict is not None: return create_error_response(logprobs_conflict, param="logprobs") if _response_format_unsupported(req.response_format): + rtype = (req.response_format or {}).get("type") return create_error_response( - "response_format json_object/json_schema is not supported (no constrained decoding)", + f"response_format.type {rtype!r} is not supported " + "(supported: text, json_object, json_schema)", param="response_format", ) if not 1 <= req.n <= MAX_N: @@ -1072,7 +1074,11 @@ def _completion_unsupported_reason(req: CompletionRequest) -> str | None: if logprobs_error is not None: return logprobs_error if _response_format_unsupported(req.response_format): - return "response_format json_object/json_schema is not supported (no constrained decoding)" + rtype = (req.response_format or {}).get("type") + return ( + f"response_format.type {rtype!r} is not supported " + "(supported: text, json_object, json_schema)" + ) return None From bf58e805debc8f4078a20eedcbffb3f203114bb6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?L=C3=A9on=20Avic=20Simmons?= Date: Sat, 3 Oct 2026 09:45:33 -0400 Subject: [PATCH 79/94] fix(mm): classify image fetch errors and re-check redirect targets (cherry picked from commit 8952c053603397040c0bdf58c9d3e020b82e319e) (cherry picked from commit 5eeb4802f3dde07fdc1b4b1c74a9b33d1726e264) --- docs/cli.md | 2 +- python/freetoken/mm/media.py | 124 ++++++++++++++++++++++++++-------- tests/server/test_mm_input.py | 121 +++++++++++++++++++++++++++++++++ 3 files changed, 219 insertions(+), 28 deletions(-) diff --git a/docs/cli.md b/docs/cli.md index 5369395734..0a301a7777 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -205,7 +205,7 @@ so a client can gate its attachment controls without reading the checkpoint conf | `--image-min-tokens`, `--image-max-tokens` | processor defaults | Per-image token budget: the image processor resizes every image to take between these many tokens, converted to the family's own units by its processor. A family with fixed budgets honors the maximum only and refuses one below its smallest budget at start-up | | `--mm-processor-kwargs` | none | JSON object of extra keyword arguments for the checkpoint's image processor call, for knobs the token budget does not cover; applied after the budget, so an explicit key wins | | `--mm-embed-cache-device` | cpu | Where encoded image embeddings live between prefill chunks. `cpu` keeps them out of the VRAM budget; `cuda` skips the copy back | -| `--allowed-media-domains` | any | Comma-separated hostname allowlist for image URLs; requests for other domains are rejected with a 400. Empty allows any domain | +| `--allowed-media-domains` | any | Comma-separated hostname allowlist for image URLs; requests for other domains are rejected with a 400. Empty allows any domain. Redirects are re-checked against the list and may not land on a private or loopback address | | `--allowed-local-media-path` | off | Directory `file://` image refs may be read from; unset rejects local files | ## ft shell diff --git a/python/freetoken/mm/media.py b/python/freetoken/mm/media.py index 5c3c4ef789..a51e5bc98b 100644 --- a/python/freetoken/mm/media.py +++ b/python/freetoken/mm/media.py @@ -1,6 +1,8 @@ """Client-supplied image handling: ref collection, the accept gate, and byte fetching. Raises plain ValueError on bad input; the server layer maps it to its wire error type. +The message names the class of failure only: no local path, allowlist root, errno or +library text reaches the client (that detail goes to the server log). """ from __future__ import annotations @@ -8,10 +10,15 @@ import base64 from typing import TYPE_CHECKING, Any +from freetoken.utils import init_logger + if TYPE_CHECKING: from freetoken.server.args import ServerArgs +logger = init_logger(__name__) + _MAX_IMAGE_BYTES = 32 << 20 +_MAX_REDIRECTS = 5 def collect_image_refs(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: @@ -56,6 +63,57 @@ def _check_media_domain(url: str, config: ServerArgs) -> None: ) +def _check_redirect_target(url: str, config: ServerArgs) -> None: + """A redirect hop is the remote server's choice, not the client's URL: it must pass the + allowlist like the original, and may never land on a loopback / private / link-local + address (an allowed public host 302-ing into the operator's network is the SSRF shape).""" + import ipaddress + from urllib.parse import urlparse + + if not url.startswith(("http://", "https://")): + raise ValueError("the URL redirects to a non-http(s) location") + try: + _check_media_domain(url, config) + except ValueError: + raise ValueError("the URL redirects outside the allowed domains") from None + host = (urlparse(url).hostname or "").rstrip(".").lower() + if host == "localhost" or host.endswith(".localhost"): + raise ValueError("the URL redirects to a private or loopback address") + try: + ip = ipaddress.ip_address(host) + except ValueError: + return # a name; the allowlist above is the gate for names (no DNS resolution here) + if isinstance(ip, ipaddress.IPv6Address) and ip.ipv4_mapped is not None: + ip = ip.ipv4_mapped + if not ip.is_global or ip.is_multicast: # covers loopback, private, link-local, CGNAT, unspecified + raise ValueError("the URL redirects to a private or loopback address") + + +async def _fetch_remote_media(url: str, config: ServerArgs) -> bytes: + """GET an http(s) ref, following redirects one hop at a time so every target is re-checked.""" + import httpx + + _check_media_domain(url, config) + async with httpx.AsyncClient(timeout=30.0, follow_redirects=False) as client: + for _ in range(_MAX_REDIRECTS + 1): + async with client.stream("GET", url) as resp: + if resp.next_request is not None: # a 3xx with a Location: httpx built the hop + url = str(resp.next_request.url) + _check_redirect_target(url, config) + continue + resp.raise_for_status() + declared = resp.headers.get("content-length", "") + if declared.isdigit() and int(declared) > _MAX_IMAGE_BYTES: + raise ValueError(f"image exceeds {_MAX_IMAGE_BYTES} bytes") + buf = bytearray() + async for chunk in resp.aiter_bytes(): + buf += chunk + if len(buf) > _MAX_IMAGE_BYTES: + raise ValueError(f"image exceeds {_MAX_IMAGE_BYTES} bytes") + return bytes(buf) + raise ValueError(f"the URL redirects more than {_MAX_REDIRECTS} times") + + def _load_local_media(url: str, config: ServerArgs) -> bytes: """Read a file:// ref; the resolved path must be a strict subpath of --allowed-local-media-path.""" from pathlib import Path @@ -70,41 +128,53 @@ def _load_local_media(url: str, config: ServerArgs) -> bytes: # resolve() follows symlinks, so a link inside the root escaping it is rejected too resolved = filepath.resolve() if Path(root).resolve() not in resolved.parents: - raise ValueError( - f"the file path {filepath} must be a subpath of " - f"--allowed-local-media-path {root}" - ) + raise ValueError("the file path must be a subpath of --allowed-local-media-path") return resolved.read_bytes() +async def _fetch_one(ref: dict[str, Any], config: ServerArgs) -> bytes: + """One ref to bytes. Our own ValueErrors are already client-safe and pass through; anything + a library raises is reduced to its failure class and logged with the detail.""" + data = ref.get("data") or "" + try: + if ref.get("kind") == "b64": + return base64.b64decode(data) + if data.startswith("data:"): + return base64.b64decode(data.split(",", 1)[1]) + except Exception as exc: # noqa: BLE001 -- binascii.Error / IndexError, input-driven + raise ValueError("image data is not valid base64") from exc + if data.startswith(("http://", "https://")): + import httpx + + try: + return await _fetch_remote_media(data, config) + except ValueError: + raise + except httpx.HTTPStatusError as exc: + raise ValueError(f"the image URL returned HTTP {exc.response.status_code}") from exc + except httpx.TimeoutException as exc: + raise ValueError("timed out fetching the image URL") from exc + except Exception as exc: # noqa: BLE001 -- transport / protocol / bad URL + logger.warning("image URL fetch failed: %r", exc) + raise ValueError("the image URL is unreachable") from exc + if data.startswith("file://"): + try: + return _load_local_media(data, config) + except ValueError: + raise + except OSError as exc: + logger.warning("local image read failed: %r", exc) + raise ValueError("the local file is missing or unreadable") from exc + raise ValueError("unsupported image source (expect http(s)/file url or base64)") + + async def fetch_image_bytes(refs: list[dict[str, Any]], config: ServerArgs) -> list[bytes]: """Resolve collected image refs (URLs / base64) to raw bytes, in order.""" out: list[bytes] = [] for ref in refs: - data = ref.get("data") or "" try: - if ref.get("kind") == "b64": - out.append(base64.b64decode(data)) - elif data.startswith("data:"): - out.append(base64.b64decode(data.split(",", 1)[1])) - elif data.startswith(("http://", "https://")): - import httpx - - _check_media_domain(data, config) - # redirect targets are not re-checked against the allowlist - async with httpx.AsyncClient(timeout=30.0, follow_redirects=True) as client: - resp = await client.get(data) - resp.raise_for_status() - if len(resp.content) > _MAX_IMAGE_BYTES: - raise ValueError(f"image exceeds {_MAX_IMAGE_BYTES} bytes") - out.append(resp.content) - elif data.startswith("file://"): - out.append(_load_local_media(data, config)) - else: - raise ValueError( - "unsupported image source (expect http(s)/file url or base64)" - ) - except Exception as exc: # noqa: BLE001 -- input-driven, client-classifiable + out.append(await _fetch_one(ref, config)) + except ValueError as exc: raise ValueError(f"could not load image: {exc}") from exc return out diff --git a/tests/server/test_mm_input.py b/tests/server/test_mm_input.py index a694a50a84..33561acdcd 100644 --- a/tests/server/test_mm_input.py +++ b/tests/server/test_mm_input.py @@ -10,6 +10,7 @@ import pytest +from freetoken.mm.media import _MAX_IMAGE_BYTES as _MAX from freetoken.mm.media import collect_image_refs, fetch_image_bytes, image_reject_reason from freetoken.server.generation import GenerationError, render_messages from freetoken.server.stats import derive_model_card @@ -205,3 +206,123 @@ def test_render_messages_rejects_an_image_in_a_merged_system_message(): ]}, {"role": "user", "content": "hi"}, ]) +# ---- error bodies and redirect policy (no network: httpx.MockTransport stands in for the remote) ---- + + +def _mock_remote(monkeypatch, routes): + """Serve ``routes`` ({url: (status, headers, body) | Exception}) through the client media.py builds; + returns the list of URLs the transport actually saw.""" + import httpx + + seen = [] + + def handler(request): + seen.append(str(request.url)) + route = routes[str(request.url)] + if isinstance(route, Exception): + raise route + status, headers, body = route + return httpx.Response(status, headers=headers, content=body, request=request) + + real = httpx.AsyncClient + monkeypatch.setattr(httpx, "AsyncClient", lambda **kw: real(transport=httpx.MockTransport(handler), **kw)) + return seen + + +def _fetch(url, config): + return asyncio.run(fetch_image_bytes([{"kind": "url", "data": url}], config)) + + +def test_local_media_error_bodies_carry_no_path_root_or_errno(tmp_path): + root = tmp_path / "media" + root.mkdir() + config = _config(allowed_local_media_path=str(root)) + outside = tmp_path / "secret.png" + outside.write_bytes(b"x") + (root / "dir").mkdir() + + cases = { + f"file://{outside}": "subpath", # escapes the root + f"file://{root / 'missing.png'}": "missing or unreadable", # ENOENT + f"file://{root / 'dir'}": "missing or unreadable", # EISDIR + } + for url, expect in cases.items(): + with pytest.raises(ValueError, match=expect) as ei: + _fetch(url, config) + body = str(ei.value) + assert str(tmp_path) not in body and "Errno" not in body, body + assert "secret.png" not in body and "missing.png" not in body, body + + +def test_remote_fetch_failures_are_classified_without_library_text(monkeypatch): + import httpx + + big = _MAX + 1 + seen = _mock_remote(monkeypatch, { + "https://cdn.example.com/404.png": (404, {}, b""), + "https://cdn.example.com/down.png": httpx.ConnectError("[Errno -2] Name or service not known"), + "https://cdn.example.com/slow.png": httpx.ReadTimeout("timed out"), + "https://cdn.example.com/declared.png": (200, {"content-length": str(big)}, b""), + "https://cdn.example.com/streamed.png": (200, {}, b"\0" * big), + }) + config = _config(allowed_media_domains="cdn.example.com") + for name, expect in [("404", "HTTP 404"), ("down", "unreachable"), ("slow", "timed out"), + ("declared", "exceeds"), ("streamed", "exceeds")]: + with pytest.raises(ValueError, match=expect) as ei: + _fetch(f"https://cdn.example.com/{name}.png", config) + assert "Errno" not in str(ei.value) + assert len(seen) == 5 + + +def test_redirects_are_rechecked_against_the_allowlist(monkeypatch): + seen = _mock_remote(monkeypatch, { + "https://cdn.example.com/a.png": (302, {"location": "https://evil.com/a.png"}, b""), + "https://cdn.example.com/b.png": (302, {"location": "/real.png"}, b""), + "https://cdn.example.com/real.png": (200, {}, b"fakepng"), + "https://evil.com/a.png": (200, {}, b"pwned"), + }) + config = _config(allowed_media_domains="cdn.example.com") + with pytest.raises(ValueError, match="redirects outside the allowed domains") as ei: + _fetch("https://cdn.example.com/a.png", config) + assert "evil.com" not in str(ei.value) + # a same-host relative redirect is followed and served + assert _fetch("https://cdn.example.com/b.png", config) == [b"fakepng"] + # the disallowed hop was never requested + assert seen == ["https://cdn.example.com/a.png", "https://cdn.example.com/b.png", "https://cdn.example.com/real.png"] + + +@pytest.mark.parametrize("target", [ + "http://127.0.0.1/x", "http://10.0.0.1/x", "http://169.254.169.254/latest/meta-data", + "http://[::1]/x", "http://[::ffff:10.0.0.1]/x", "http://localhost/x", "http://0.0.0.0/x", + "http://100.64.0.1/x", +]) +def test_redirects_to_private_or_loopback_addresses_are_refused(monkeypatch, target): + seen = _mock_remote(monkeypatch, { + "https://public.example.com/a.png": (302, {"location": target}, b""), + target: (200, {}, b"internal"), + }) + # even with an empty allowlist (any domain admitted) a redirect may not reach inside + with pytest.raises(ValueError, match="private or loopback"): + _fetch("https://public.example.com/a.png", _config()) + assert seen == ["https://public.example.com/a.png"] + + +def test_redirect_hop_count_and_scheme_are_bounded(monkeypatch): + chain = {f"https://h.example.com/{i}.png": (302, {"location": f"https://h.example.com/{i + 1}.png"}, b"") + for i in range(7)} + chain["https://h.example.com/7.png"] = (200, {}, b"fakepng") + chain["https://h.example.com/file.png"] = (302, {"location": "file:///etc/passwd"}, b"") + seen = _mock_remote(monkeypatch, chain) + with pytest.raises(ValueError, match="redirects more than"): + _fetch("https://h.example.com/0.png", _config()) + assert len(seen) == 6 # the original plus five hops; the sixth hop is never requested + with pytest.raises(ValueError, match="non-http"): + _fetch("https://h.example.com/file.png", _config()) + # a short chain that stays public is fine + assert _fetch("https://h.example.com/5.png", _config()) == [b"fakepng"] + + +def test_bad_base64_is_named_without_decoder_text(): + for ref in ({"kind": "b64", "data": "not base64!!"}, {"kind": "url", "data": "data:image/png;base64"}): + with pytest.raises(ValueError, match="not valid base64"): + asyncio.run(fetch_image_bytes([ref], _config())) From 4c83febb3fe98ea3ff2520768befe21e398aca0f Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 06:06:12 -0400 Subject: [PATCH 80/94] test(engine): cover presence/frequency penalties on the greedy and sampled paths Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 30b818114ce353da17bf6c3c9933e6c752eb7c11) (cherry picked from commit dc7cfc13fa5eaef06e105530b55e752529c9d0c6) --- tests/engine/test_sample.py | 84 +++++++++++++++++++++++++++++++++++++ 1 file changed, 84 insertions(+) create mode 100644 tests/engine/test_sample.py diff --git a/tests/engine/test_sample.py b/tests/engine/test_sample.py new file mode 100644 index 0000000000..c1cc925f00 --- /dev/null +++ b/tests/engine/test_sample.py @@ -0,0 +1,84 @@ +"""Presence/frequency penalties in engine.sample.Sampler (greedy and sampled rows).""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.core import SamplingParams +from freetoken.engine import sample as S + + +def _batch(*params: SamplingParams, can_decode: bool = True): + reqs = [ + SimpleNamespace(sampling_params=p, can_decode=can_decode, output_token_counts=None) + for p in params + ] + return SimpleNamespace(reqs=reqs) + + +def test_no_penalty_allocates_nothing(): + sampler = S.Sampler(torch.device("cpu"), 5) + batch = _batch(SamplingParams(temperature=0.0)) + args = sampler.prepare(batch) + assert args.penalties == [] and batch.reqs[0].output_token_counts is None + + +def test_frequency_penalty_counts_generated_tokens_on_the_greedy_path(): + sampler = S.Sampler(torch.device("cpu"), 5) + batch = _batch(SamplingParams(temperature=0.0, frequency_penalty=10.0)) + logits = torch.tensor([[0.0, 5.0, 1.0, 0.0, 0.0]]) + + assert sampler.sample(logits, sampler.prepare(batch)).tolist() == [1] + # Token 1 now carries -10, so the runner-up wins; the caller's logits stay untouched. + assert sampler.sample(logits, sampler.prepare(batch)).tolist() == [2] + assert batch.reqs[0].output_token_counts.tolist() == [0, 1, 1, 0, 0] + assert logits[0, 1] == 5.0 + + +def test_intermediate_prefill_chunks_are_not_counted(): + sampler = S.Sampler(torch.device("cpu"), 5) + batch = _batch(SamplingParams(temperature=0.0, presence_penalty=1.0), can_decode=False) + assert sampler.prepare(batch).penalties == [] + + +def test_penalized_logits_reach_the_probability_sampler(monkeypatch): + """Sampled rows (the torch top-k/top-p fallback on Pascal) see the penalized logits.""" + seen = {} + + def fake_sample_impl(logits, temperatures, top_k, top_p): + seen["logits"] = logits.clone() + return torch.argmax(logits, dim=-1) + + monkeypatch.setattr(S, "sample_impl", fake_sample_impl) + monkeypatch.setattr(S, "make_device_tensor", lambda data, dtype, device: torch.tensor(data, dtype=dtype)) + sampler = S.Sampler(torch.device("cpu"), 4) + batch = _batch( + SamplingParams(temperature=1.0, top_k=2, presence_penalty=3.0), + SamplingParams(temperature=1.0, top_k=2), + ) + batch.reqs[0].output_token_counts = torch.tensor([0, 2, 0, 0], dtype=torch.int32) + logits = torch.tensor([[0.0, 4.0, 2.0, 0.0], [0.0, 4.0, 2.0, 0.0]]) + + tokens = sampler.sample(logits, sampler.prepare(batch)) + + assert seen["logits"][0].tolist() == [0.0, 1.0, 2.0, 0.0] # presence: -3 once, not per count + assert seen["logits"][1].tolist() == [0.0, 4.0, 2.0, 0.0] + assert tokens.tolist() == [2, 1] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_penalty_through_the_pre_volta_torch_sampling_fallback(monkeypatch): + import freetoken.kernel.backend as backend + import freetoken.utils.arch as arch + + monkeypatch.setattr(backend, "is_flashinfer_installed", lambda: False) + monkeypatch.setattr(arch, "is_sm70_supported", lambda: False) + sampler = S.Sampler(torch.device("cuda"), 4) + batch = _batch(SamplingParams(temperature=1.0, top_k=1, frequency_penalty=5.0)) + logits = torch.tensor([[0.0, 4.0, 2.0, 0.0]], device="cuda") + + assert sampler.sample(logits, sampler.prepare(batch)).tolist() == [1] + assert sampler.sample(logits, sampler.prepare(batch)).tolist() == [2] From 923fd4a125499ee4921799044867478a808e8954 Mon Sep 17 00:00:00 2001 From: "Eva (port)" Date: Sat, 3 Oct 2026 01:21:44 +0300 Subject: [PATCH 81/94] perf(model): fuse hc combine with the next block's RMSNorm via hc_combine_norm the vendored combine+norm kernel was exported but unwired; mix() takes a pre-computed rn, the model loop carries the fused pair across the layer boundary, skipped before PLE layers where ple(hidden) lands between. (cherry picked from commit 7636dd76cbe64989cbdec09e5c09a4c67be89542) (cherry picked from commit 0e847f8cd849e4a75cab781a7846af972afbbc58) --- python/freetoken/models/qwen4_exp/hc.py | 30 +++++++++++++++++----- python/freetoken/models/qwen4_exp/model.py | 30 ++++++++++++++++------ 2 files changed, 45 insertions(+), 15 deletions(-) diff --git a/python/freetoken/models/qwen4_exp/hc.py b/python/freetoken/models/qwen4_exp/hc.py index f13808dccd..1574cd473f 100644 --- a/python/freetoken/models/qwen4_exp/hc.py +++ b/python/freetoken/models/qwen4_exp/hc.py @@ -17,6 +17,7 @@ from freetoken.kernel.triton.hc import ( grouped_gemma_rmsnorm, hc_combine, + hc_combine_norm, hc_gate_mix, hc_silu, ) @@ -119,14 +120,16 @@ def _down(self, rn: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor | None]: # both slices keep unit inner stride, so the kernels read them without a copy return down[:, : self.lowrank], down[:, self.lowrank : self.lowrank + self.hc_count] - def _mix_kernel(self, R: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor | None]: - rn = grouped_gemma_rmsnorm(R, self.hc_norm.weight, self.hc_norm.eps, self.hc_count) + def _mix_kernel(self, R: torch.Tensor, rn: torch.Tensor | None = None) -> Tuple[torch.Tensor, torch.Tensor | None]: + if rn is None: + rn = grouped_gemma_rmsnorm(R, self.hc_norm.weight, self.hc_norm.eps, self.hc_count) lora, s = self._down(rn) gate = self.input_mix_weight_up.forward(hc_silu(lora, self.hc_count)) return hc_gate_mix(rn, gate, self.hc_count), s - def _mix_torch(self, R: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor | None]: - rn = grouped_plus_one_rms_norm(R, self.hc_norm.weight, self.hc_norm.eps, self.hc_count) + def _mix_torch(self, R: torch.Tensor, rn: torch.Tensor | None = None) -> Tuple[torch.Tensor, torch.Tensor | None]: + if rn is None: + rn = grouped_plus_one_rms_norm(R, self.hc_norm.weight, self.hc_norm.eps, self.hc_count) lora, s = self._down(rn) lora = F.silu(lora.float() / self.hc_count) gate = self.input_mix_weight_up.forward(lora.to(R.dtype)) @@ -140,9 +143,12 @@ def _combine_torch(self, R: torch.Tensor, y: torch.Tensor, s: torch.Tensor) -> t out = out + y.float().unsqueeze(-2) * inject.unsqueeze(-1) return out.flatten(-2).to(R.dtype) - def mix(self, R: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor | None]: - """Return the block input ``x [T, hidden]`` and the inject logits ``s [T, hc_count]`` (None if no combine).""" - return self._mix_kernel(R) if R.is_cuda else self._mix_torch(R) + def mix(self, R: torch.Tensor, rn: torch.Tensor | None = None) -> Tuple[torch.Tensor, torch.Tensor | None]: + """Return the block input ``x [T, hidden]`` and the inject logits ``s [T, hc_count]`` (None if no combine). + + ``rn`` is the pre-normalized R from a preceding :meth:`combine_norm` (same norm weight as + this block's ``hc_norm``), letting the caller fuse the two kernels; None recomputes it.""" + return self._mix_kernel(R, rn) if R.is_cuda else self._mix_torch(R, rn) def combine(self, R: torch.Tensor, y: torch.Tensor, s: torch.Tensor) -> torch.Tensor: """Inject the block output ``y [T, hidden]`` back into every stream of ``R``.""" @@ -150,5 +156,15 @@ def combine(self, R: torch.Tensor, y: torch.Tensor, s: torch.Tensor) -> torch.Te return hc_combine(R, y, s, self.hc_count) return self._combine_torch(R, y, s) + def combine_norm( + self, R: torch.Tensor, y: torch.Tensor, s: torch.Tensor, next_block: "GatedResidual" + ) -> Tuple[torch.Tensor, torch.Tensor | None]: + """combine() fused with ``next_block``'s hc_norm; returns (R', Rn'). CPU keeps the split path (Rn' = None).""" + if not R.is_cuda: + return self.combine(R, y, s), None + return hc_combine_norm( + R, y, s, next_block.hc_norm.weight, next_block.hc_norm.eps, self.hc_count + ) + __all__ = ["GatedResidual", "GroupedPlusOneRMSNorm", "grouped_plus_one_rms_norm"] diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py index 7f4f826807..6ce2da985b 100644 --- a/python/freetoken/models/qwen4_exp/model.py +++ b/python/freetoken/models/qwen4_exp/model.py @@ -75,17 +75,19 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = "") -> N ) @nvtx_annotate("Layer_{}", layer_id_field="_layer_id") - def forward(self, hidden: torch.Tensor, batch: Batch) -> torch.Tensor: + def forward(self, hidden: torch.Tensor, batch: Batch, rn: torch.Tensor | None = None): + """Returns (R, mlp_out, inject) un-combined: the caller fuses the mlp combine with the + next consumer's norm (next layer's attn mix, or the final mixer).""" if self.ple is not None: hidden = hidden + self.ple.forward(hidden, batch) - block_input, inject = self.attn_hyper_connection.mix(hidden) + block_input, inject = self.attn_hyper_connection.mix(hidden, rn=rn) if self._is_linear: block_output = self.linear_attn.forward(block_input) else: block_output = self.self_attn.forward(block_input, batch) - hidden = self.attn_hyper_connection.combine(hidden, block_output, inject) - block_input, inject = self.mlp_hyper_connection.mix(hidden) - return self.mlp_hyper_connection.combine(hidden, self.mlp.forward(block_input), inject) + hidden, rn_mlp = self.attn_hyper_connection.combine_norm(hidden, block_output, inject, self.mlp_hyper_connection) + block_input, inject = self.mlp_hyper_connection.mix(hidden, rn=rn_mlp) + return hidden, self.mlp.forward(block_input), inject class Qwen4ExpModel(BaseOP): @@ -121,13 +123,25 @@ def forward(self, input_ids: torch.Tensor, batch: Batch) -> torch.Tensor: meta = build_ple_metadata(batch, self._ple[0].args, input_ids.device) for ple in self._ple: # gather the pinned-host PLE rows while the early layers run ple.start_prefetch(batch, meta) - for layer in self.layers.op_list: - hidden = layer.forward(hidden, batch) + # The last block of each layer combines through the layer boundary, fusing the + # combine with the next consumer's hc_norm -- except before a PLE layer, whose + # `hidden + ple(hidden)` happens after the boundary and changes what the norm sees. + layers = self.layers.op_list + rn = None + for i, layer in enumerate(layers): + hidden, mlp_out, inject = layer.forward(hidden, batch, rn=rn) + nxt = layers[i + 1] if i + 1 < len(layers) else None + if nxt is not None and nxt.ple is not None: + hidden = layer.mlp_hyper_connection.combine(hidden, mlp_out, inject) + rn = None + else: + consumer = nxt.attn_hyper_connection if nxt is not None else self.hyper_connection_mixer + hidden, rn = layer.mlp_hyper_connection.combine_norm(hidden, mlp_out, inject, consumer) if meta is not None: # single writer: the layers only read the context, so a second PLE layer's # prefetch sees the un-rolled window commit_ngram_context(meta, getattr(batch, "fla_metadata", None)) - return self.hyper_connection_mixer.mix(hidden)[0] + return self.hyper_connection_mixer.mix(hidden, rn=rn)[0] class Qwen4ExpForCausalLM(BaseLLMModel): From ada9f8d7fcd9eed98be80d710c05470d5561ff0e Mon Sep 17 00:00:00 2001 From: "Eva (port)" Date: Sat, 3 Oct 2026 01:21:44 +0300 Subject: [PATCH 82/94] fix(kernel): bit-exact rmsnorm reduction tree in hc_combine_norm the fused kernel summed squares over a two-axis padded tile; the standalone grouped RMSNorm reduces one flat next_pow2 vector - different fp32 tree, ulp-level drift caught by the greedy gate. Reshape to the identical tree. (cherry picked from commit c1c5e76f5d98e08f0ff525fe71415bac028af3bc) (cherry picked from commit 8f8d7fbd7f0940e6ea4f0fa1b5ddfee4892efe0f) --- python/freetoken/kernel/triton/hc.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/python/freetoken/kernel/triton/hc.py b/python/freetoken/kernel/triton/hc.py index 45d71bd731..acae534c45 100644 --- a/python/freetoken/kernel/triton/hc.py +++ b/python/freetoken/kernel/triton/hc.py @@ -326,9 +326,11 @@ def _hc_combine_norm_kernel( tl.store(out_ptr + row * stride_out + offs, out, mask=mask_inner) out = out.to(tl.float32) - # Keep the two-axis reduction: flattening the padded tile is ~40% slower - # at decode sizes. - sum_sq = tl.sum(tl.sum(out * out, axis=1), axis=0) + # The standalone grouped RMSNorm reduces one padded [next_pow2(HC_DIM)] + # vector; reshape keeps the identical tree (same element + # order, same masked zeros), so the fused statistic is bit-exact vs the split. + flat = tl.reshape(out * out, [NUM_TILES_PAD * BLOCK_SIZE]) + sum_sq = tl.sum(flat) rrms = tl.rsqrt(sum_sq / HC_DIM + EPS) if launch_pdl: From bad406cc627c803d776034dfed9206ffe168220a Mon Sep 17 00:00:00 2001 From: "Eva (port)" Date: Sat, 3 Oct 2026 01:22:10 +0300 Subject: [PATCH 83/94] test(kernel): bitwise pin for hc_combine_norm against the split path (cherry picked from commit 7cba55ea8ba0d3576523c75ef0fcacb3d3e817bb) (cherry picked from commit 081534bb5574c7074ff37a2ecb9dca5a645073f1) --- tests/kernels/test_hc_combine_norm.py | 48 +++++++++++++++++++++++++++ 1 file changed, 48 insertions(+) create mode 100644 tests/kernels/test_hc_combine_norm.py diff --git a/tests/kernels/test_hc_combine_norm.py b/tests/kernels/test_hc_combine_norm.py new file mode 100644 index 0000000000..7f46755479 --- /dev/null +++ b/tests/kernels/test_hc_combine_norm.py @@ -0,0 +1,48 @@ +"""Pins for the fused hc combine + grouped RMSNorm path (hc_combine_norm). + +hc_combine_norm replaces the standalone pair (hc_combine -> grouped_gemma_rmsnorm) +in the Qwen4-exp decoder loop. The swap is only sound if it is BITWISE identical: +the hyper-connection streams feed every later layer, and a bf16 flip here compounds +into different greedy tokens by the end of the stack. We assert torch.equal, not +assert_close: an approx-tolerance test would not have caught the regression this +path shipped with at first attempt -- the fused kernel reduced squares over the +[8, block] tile view (a two-axis tree) while the standalone kernel reduces the +flat row tree; fp32 ulps that differ between trees flip bf16 rounding. Both paths +on the same data must return the same bytes: fails before the reduction-tree fix +in _hc_combine_norm_kernel, passes after. +""" + +from __future__ import annotations + +import pytest +import torch + +from freetoken.kernel.triton.hc import hc_combine, grouped_gemma_rmsnorm, hc_combine_norm + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + +HC, H = 4, 2560 # model geometry: the reduction tree specializes on the block layout +DIM = HC * H +EPS = 1e-6 + + +def _inputs(seed: int, T: int): + g = torch.Generator(device="cuda").manual_seed(seed) + bf = lambda shape: torch.randn(shape, generator=g, device="cuda", dtype=torch.float32).to(torch.bfloat16) + return ( + bf((T, DIM)), # residual streams R + bf((T, H)), # block output y + bf((T, HC)), # inject scores s + (torch.randn(DIM, generator=g, device="cuda", dtype=torch.float32) * 0.1).to(torch.bfloat16), # norm weight + ) + + +@pytest.mark.parametrize("T", [1, 8]) +@pytest.mark.parametrize("seed", range(6)) +def test_combine_norm_matches_split_bitwise(seed: int, T: int) -> None: + R, y, s, w = _inputs(seed, T) + r2_split = hc_combine(R, y, s, HC) + rn_split = grouped_gemma_rmsnorm(r2_split, w, EPS, HC) + r2_fused, rn_fused = hc_combine_norm(R, y, s, w, EPS, HC) + assert torch.equal(r2_fused, r2_split) + assert torch.equal(rn_fused, rn_split) From 51debf6cc1490b99f0024558a881e04ba09920e9 Mon Sep 17 00:00:00 2001 From: Siarhei Berdachuk Date: Mon, 5 Oct 2026 19:11:42 +0300 Subject: [PATCH 84/94] fix(server): soft-accept OpenAI response_format json_object/json_schema Rejecting response_format with "no constrained decoding" breaks OpenAI-compatible clients that rely on json_object/json_schema even when the model can emit valid JSON from a prompt. Accept those wire types and prepend a best-effort system hint; schema is not token-level enforced. Unknown response_format types still return 400. (cherry picked from commit ca7d704bb84e255d8677b788f295e9cb6ebb06f0) (cherry picked from commit db64d8eff1a7778b6bc093634303d31e1324af76) --- python/freetoken/server/openai_api.py | 44 +++++++++++++++++++++++++-- 1 file changed, 42 insertions(+), 2 deletions(-) diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index 79dff72963..7a345cadfa 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -71,6 +71,43 @@ def _thinking_type(req: Any) -> str | None: +def _json_mode_system_message(response_format: dict[str, Any] | None) -> dict[str, str] | None: + """Best-effort JSON mode hint for OpenAI response_format. + + FreeToken has no constrained/guided decoding yet. OpenAI-compatible clients + still send response_format=json_object|json_schema. Rejecting them with 400 + breaks those clients even though the model can emit valid JSON from a + prompt. Accept the wire shape and prepend a system reminder so the model is + steered toward JSON; the schema is NOT token-level enforced. + """ + if not response_format: + return None + rtype = response_format.get("type") + if rtype == "json_object": + return { + "role": "system", + "content": ( + "Respond with a single valid JSON object only. " + "Do not wrap it in markdown fences or add commentary." + ), + } + if rtype == "json_schema": + wrapper = response_format.get("json_schema") or {} + name = wrapper.get("name") or "response" + schema = wrapper.get("schema", wrapper) + return { + "role": "system", + "content": ( + f"Respond with a single JSON object that validates against the " + f"JSON Schema below (schema name: '{name}'). Put required fields " + "at the top level of the object — do not wrap them under an extra " + "key named after the schema. Do not wrap the answer in markdown " + f"fences or add commentary.\nJSON Schema:\n{json.dumps(schema, ensure_ascii=False)}" + ), + } + return None + + def chat_request_to_genspec( req: ChatCompletionRequest, model_sampling: dict[str, Any], @@ -1063,8 +1100,11 @@ def _usage( def _response_format_unsupported(response_format: dict[str, Any] | None) -> bool: - # We have no constrained/guided decoding; only plain text ('text' or unset) is honored. - return response_format is not None and response_format.get("type") not in (None, "text") + # Soft-accept OpenAI json_object/json_schema (prompt-steered; no constrained decoding). + # Unknown types still 400. See _json_mode_system_message. + if response_format is None: + return False + return response_format.get("type") not in (None, "text", "json_object", "json_schema") def _completion_unsupported_reason(req: CompletionRequest) -> str | None: From d3e8e3563a818e82351e8728a0bb2648d7b13db7 Mon Sep 17 00:00:00 2001 From: Siarhei Berdachuk Date: Mon, 5 Oct 2026 19:18:06 +0300 Subject: [PATCH 85/94] refactor(server): drop response_format system-hint injection Soft-accepting json_object/json_schema is enough when clients already request JSON in the prompt. The injected schema wrapper was redundant and could distort the output shape. (cherry picked from commit 65907c8bc2c1c11a4af11dff328f7cc293806433) (cherry picked from commit c82625f6eed239a9edcc6e76e935d3b51359f425) --- python/freetoken/server/openai_api.py | 45 ++------------------------- 1 file changed, 3 insertions(+), 42 deletions(-) diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index 7a345cadfa..0fc8113bc8 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -68,46 +68,6 @@ def _thinking_type(req: Any) -> str | None: if value in ("enabled", "disabled"): return value return None - - - -def _json_mode_system_message(response_format: dict[str, Any] | None) -> dict[str, str] | None: - """Best-effort JSON mode hint for OpenAI response_format. - - FreeToken has no constrained/guided decoding yet. OpenAI-compatible clients - still send response_format=json_object|json_schema. Rejecting them with 400 - breaks those clients even though the model can emit valid JSON from a - prompt. Accept the wire shape and prepend a system reminder so the model is - steered toward JSON; the schema is NOT token-level enforced. - """ - if not response_format: - return None - rtype = response_format.get("type") - if rtype == "json_object": - return { - "role": "system", - "content": ( - "Respond with a single valid JSON object only. " - "Do not wrap it in markdown fences or add commentary." - ), - } - if rtype == "json_schema": - wrapper = response_format.get("json_schema") or {} - name = wrapper.get("name") or "response" - schema = wrapper.get("schema", wrapper) - return { - "role": "system", - "content": ( - f"Respond with a single JSON object that validates against the " - f"JSON Schema below (schema name: '{name}'). Put required fields " - "at the top level of the object — do not wrap them under an extra " - "key named after the schema. Do not wrap the answer in markdown " - f"fences or add commentary.\nJSON Schema:\n{json.dumps(schema, ensure_ascii=False)}" - ), - } - return None - - def chat_request_to_genspec( req: ChatCompletionRequest, model_sampling: dict[str, Any], @@ -1100,8 +1060,9 @@ def _usage( def _response_format_unsupported(response_format: dict[str, Any] | None) -> bool: - # Soft-accept OpenAI json_object/json_schema (prompt-steered; no constrained decoding). - # Unknown types still 400. See _json_mode_system_message. + # Soft-accept OpenAI json_object/json_schema. FreeToken has no constrained + # decoding; clients that already ask for JSON in the prompt still need the + # wire field accepted (HTTP 200) instead of a hard 400. Unknown types 400. if response_format is None: return False return response_format.get("type") not in (None, "text", "json_object", "json_schema") From 56561e27dea4596cac9706a58a59d7fded3f5e6a Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Tue, 13 Oct 2026 13:14:39 -0400 Subject: [PATCH 86/94] test(engine): drop the pre-volta sampling penalty test, whose torch fallback this branch does not carry Assisted-by: Claude Opus 5.5 Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Bu6LgoxLR4wETqb7RPR2vt --- tests/engine/test_sample.py | 15 --------------- 1 file changed, 15 deletions(-) diff --git a/tests/engine/test_sample.py b/tests/engine/test_sample.py index c1cc925f00..28e7195b58 100644 --- a/tests/engine/test_sample.py +++ b/tests/engine/test_sample.py @@ -67,18 +67,3 @@ def fake_sample_impl(logits, temperatures, top_k, top_p): assert seen["logits"][0].tolist() == [0.0, 1.0, 2.0, 0.0] # presence: -3 once, not per count assert seen["logits"][1].tolist() == [0.0, 4.0, 2.0, 0.0] assert tokens.tolist() == [2, 1] - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") -def test_penalty_through_the_pre_volta_torch_sampling_fallback(monkeypatch): - import freetoken.kernel.backend as backend - import freetoken.utils.arch as arch - - monkeypatch.setattr(backend, "is_flashinfer_installed", lambda: False) - monkeypatch.setattr(arch, "is_sm70_supported", lambda: False) - sampler = S.Sampler(torch.device("cuda"), 4) - batch = _batch(SamplingParams(temperature=1.0, top_k=1, frequency_penalty=5.0)) - logits = torch.tensor([[0.0, 4.0, 2.0, 0.0]], device="cuda") - - assert sampler.sample(logits, sampler.prepare(batch)).tolist() == [1] - assert sampler.sample(logits, sampler.prepare(batch)).tolist() == [2] From f42330e921ae20e919bc469b74ad8495dae7e16b Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Tue, 13 Oct 2026 13:14:39 -0400 Subject: [PATCH 87/94] fix(moe): align on triton past sgl_kernel's 1024-expert limit sgl_kernel's moe_align_block_size (0.4.5 and 0.4.9) reads out of bounds once num_experts + 1 exceeds 1024: an illegal memory access that poisons the CUDA context. A prefill over the offload slot cache (FlashML-org/FreeToken#676) aligns against the slot count, 2100 here, so the first short prefill crashed the server; the bf16 and NVFP4 prefill paths both go through it. The Triton align has no such limit and matches sgl_kernel's output below it. Assisted-by: Claude Opus 5.5 Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Bu6LgoxLR4wETqb7RPR2vt --- python/freetoken/moe/fused.py | 8 +++++++- tests/moe/test_nvfp4_backends.py | 19 +++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/python/freetoken/moe/fused.py b/python/freetoken/moe/fused.py index 5c34fb33ca..2f8af47c89 100644 --- a/python/freetoken/moe/fused.py +++ b/python/freetoken/moe/fused.py @@ -44,6 +44,9 @@ def fused_topk( return fused_topk_softmax(gating_output, topk, renormalize, num_token_non_padded) +_SGL_ALIGN_MAX_EXPERTS = 1024 + + def moe_align_block_size( topk_ids: torch.Tensor, block_size: int, num_experts: int ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: @@ -97,7 +100,10 @@ def moe_align_block_size( return moe_align_block_size_triton(topk_ids, block_size, num_experts) - if not is_sgl_kernel_installed(): + # sgl_kernel's align kernel takes num_experts + 1 counters in one block and reads out of + # bounds past 1024 of them (0.4.5 and 0.4.9). A prefill over an offload slot cache passes the + # slot count here (#676), which is thousands; the Triton kernel has no such limit. + if not is_sgl_kernel_installed() or num_experts + 1 > _SGL_ALIGN_MAX_EXPERTS: from freetoken.kernel.triton.moe_align import ( moe_align_block_size as triton_moe_align_block_size, ) diff --git a/tests/moe/test_nvfp4_backends.py b/tests/moe/test_nvfp4_backends.py index f3b3253c6e..8dd08cd71e 100644 --- a/tests/moe/test_nvfp4_backends.py +++ b/tests/moe/test_nvfp4_backends.py @@ -639,3 +639,22 @@ def test_b12x_pack_keeps_per_layer_banks_and_flat_alphas(): assert len(cache.bank_sources["gate_up"]) == L assert sum(t.shape[0] for t in cache.bank_sources["gate_up"]) == total assert cache.gate_up_alpha.shape == (total,) + + +@cuda +@pytest.mark.parametrize("num_experts", [1023, 2100]) +def test_moe_align_handles_a_slot_cache_wider_than_1024(num_experts): + """A prefill over the offload slot cache aligns against its slot count (#676): sgl_kernel's + align reads out of bounds past 1024 counters, so a 2100-slot cache must still align.""" + from freetoken.moe.fused import moe_align_block_size + + torch.manual_seed(0) + ids = torch.randint(0, num_experts, (300, 10), device="cuda", dtype=torch.int32) + sorted_ids, expert_ids, padded = moe_align_block_size(ids, 16, num_experts) + n = int(padded) + flat = ids.flatten() + real = sorted_ids[:n][sorted_ids[:n] < flat.numel()] + assert sorted(real.tolist()) == list(range(flat.numel())) + block_of = torch.arange(n, device="cuda") // 16 + rows = sorted_ids[:n] < flat.numel() + assert torch.equal(flat[sorted_ids[:n][rows].long()], expert_ids[block_of[rows]]) From f0919098f89283dae6b11fd33ca3419925371659 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 15 Oct 2026 08:40:45 -0400 Subject: [PATCH 88/94] fix(moe): refuse the global-slot moe features when the cache has several slot pools Upstream #689 splits the offload expert cache into slot pools whose slot ids are pool-local. The scored eviction policies, expert prefetch, its verify and slot-audit debug modes and the #676 on-demand prefill GEMM all address the cache by global slot over one typed bank, so with several pools they would install or read experts in another pool's rows. Refuse the combination at construction; one pool (the default, and this deployment) behaves as before. Assisted-by: Claude Opus 5.5 Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Bu6LgoxLR4wETqb7RPR2vt --- python/freetoken/moe/offload_cache.py | 24 ++++++++++++++++++++++++ tests/moe/test_offload.py | 23 +++++++++++++++++++++++ 2 files changed, 47 insertions(+) diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index d7b03e4733..a12ab593a4 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -371,6 +371,30 @@ def __post_init__(self) -> None: self.prefetch = self._init_prefetch() self.verify = self._init_verify() self.audit = self._init_audit() + self._refuse_single_pool_features() + + def _refuse_single_pool_features(self) -> None: + """The scored policies, prefetch, its debug checks and the on-demand prefill GEMM address + the cache by global slot id over one typed bank, so with several slot pools (#689) they + would install and read pool-local slots in the wrong rows: refuse that combination.""" + if len(self.pool_sizes) == 1: + return + used = [ + name + for name, on in ( + (f"--moe-cache-policy {self.cache_policy}", self.cache_policy != "lru"), + ("FREETOKEN_MOE_PREFETCH", self.prefetch_mode not in (None, "off")), + ("FREETOKEN_MOE_PREFETCH_VERIFY", self.verify_mode not in (None, "off")), + ("FREETOKEN_MOE_SLOT_AUDIT", bool(self.slot_audit)), + ("--moe-prefill-demand-tokens", bool(self.prefill_demand_tokens)), + ) + if on + ] + if used: + raise ValueError( + f"{', '.join(used)} need one MoE slot pool, but this cache has {len(self.pool_sizes)} " + "(--moe-cache-scope per-layer, or experts in several quant formats)" + ) def _init_audit(self): from freetoken.kernel.fast_index_copy import _skip_fast_index_copy_enabled diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index 2c33443b36..05ab1d700b 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -1312,3 +1312,26 @@ def boom(addr, nbytes): with hb.PinPipeline() as pins: pins(1, {"gate_up": hb.HostBank((4,), torch.uint8)}) assert plan2.actual == {1: hb.HostResidency.PAGEABLE.value} + + +@pytest.mark.parametrize( + "kwargs, flag", + [ + ({"cache_policy": "rule"}, "--moe-cache-policy rule"), + ({"prefetch_mode": "on"}, "FREETOKEN_MOE_PREFETCH"), + ({"verify_mode": "meta"}, "FREETOKEN_MOE_PREFETCH_VERIFY"), + ({"slot_audit": 8}, "FREETOKEN_MOE_SLOT_AUDIT"), + ({"prefill_demand_tokens": 64}, "--moe-prefill-demand-tokens"), + ], +) +def test_offload_cache_refuses_global_slot_features_with_several_pools(kwargs, flag): + # these address slots globally over one typed bank: with pool-local slot ids they would + # install and read experts in another pool's rows + from freetoken.moe.offload_cache import OffloadMoeCache + + _init_tp() + common = dict(num_layers=2, num_experts=4, cache_size=16, device=torch.device("cpu"), **kwargs) + with pytest.raises(ValueError, match=f"{flag}.*one MoE slot pool"): + OffloadMoeCache(pool_of_layer=[0, 1], pool_sizes=[8, 8], **common) + OffloadMoeCache(**common) + OffloadMoeCache(num_layers=2, num_experts=4, cache_size=16, device=torch.device("cpu"), pool_of_layer=[0, 1], pool_sizes=[8, 8]) From 475e4d7b65465a9f9ea5d982d81ff1881e570f89 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 15 Oct 2026 08:41:49 -0400 Subject: [PATCH 89/94] test(moe): restore the #676 merge's on-demand test updates the replay dropped The 5080deploy merge of #676 updated these tests inside the merge commit (the on-demand path runs the prefill GEMM over the slot cache, and the layer hands ensure_experts unique ids). Re-merging #676's head on the new base does not carry edits made inside a merge, so they are restored here. Assisted-by: Claude Opus 5.5 Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Bu6LgoxLR4wETqb7RPR2vt --- tests/moe/test_offload.py | 9 +++++---- tests/moe/test_small_prefill_on_demand.py | 4 ++-- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index 05ab1d700b..818dd37096 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -138,9 +138,10 @@ def test_offload_moe_layer_short_prefill_loads_its_experts_through_the_slot_cach ) monkeypatch.setattr(cache, "materialize_layer", lambda layer_id: calls.setdefault("materialized", layer_id)) - def ensure(layer_id, expert_ids, *, counts_as_decode=True): + def ensure(layer_id, expert_ids, *, counts_as_decode=True, **kwargs): + # the layer hands the unique routed ids (decode-width chunks); rewrite each to a slot calls["ensured"] = (layer_id, expert_ids.tolist(), counts_as_decode) - expert_ids.copy_(torch.tensor([[5, 3]], dtype=torch.int32)) + expert_ids.add_(3) monkeypatch.setattr(cache, "ensure_experts", ensure) monkeypatch.setattr(cache, "copy_missing", lambda: calls.setdefault("copied", True)) @@ -157,11 +158,11 @@ def fake_fused(hidden_states, w1, w2, got_topk_weights, got_topk_ids, activation assert out is hidden_states assert "materialized" not in calls - assert calls["ensured"] == (0, [[2, 1]], False) + assert calls["ensured"] == (0, [1, 2], False) assert calls["copied"] is True # the grouped GEMM reads the whole slot cache at the slots the admission wrote assert calls["w1"].shape[0] == cache.cache_size - assert calls["topk_ids"].tolist() == [[5, 3]] + assert calls["topk_ids"].tolist() == [[5, 4]] calls.clear() layer.prefill_forward(torch.randn(2, 8), torch.randn(2, 4)) diff --git a/tests/moe/test_small_prefill_on_demand.py b/tests/moe/test_small_prefill_on_demand.py index 8b7b9d802c..f7971c67d5 100644 --- a/tests/moe/test_small_prefill_on_demand.py +++ b/tests/moe/test_small_prefill_on_demand.py @@ -51,7 +51,7 @@ def fake_gemm(cache_, hs, w, ids, *, views, n, alphas, is_prefill): hs = torch.randn(tokens, 16) layer._prefill_routed(hs, torch.full((tokens, 2), 0.5), torch.zeros(tokens, 2, dtype=torch.int32)) if on_demand: - assert calls == ["ensure", "copy"] and got == {"n": None, "is_prefill": False} + assert calls == ["ensure", "copy"] and got == {"n": cache.cache_size, "is_prefill": True} else: assert "ensure" not in calls and "begin_prefill" not in calls[1:] and got["n"] == layer.num_experts @@ -97,7 +97,7 @@ def test_small_prefill_path_does_not_depend_on_how_many_experts_a_layer_routes_t ids = torch.tensor([[5, 0], [1, 7], [2, 5], [4, 0]], dtype=torch.int32) layer._prefill_routed(torch.randn(4, 16), torch.full((4, 2), 0.5), ids.clone()) assert calls == [("ensure", [0, 1, 2, 4]), ("copy",), ("ensure", [5, 7]), ("copy",)] - assert got == {"ids": (ids + 100).tolist(), "n": None, "is_prefill": False} + assert got == {"ids": (ids + 100).tolist(), "n": cache.cache_size, "is_prefill": True} def test_small_prefill_refuses_more_experts_than_the_cache_holds(monkeypatch): From ac05b40341195103b3908eed6d499bb385a014c7 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 15 Oct 2026 08:44:36 -0400 Subject: [PATCH 90/94] fix(moe): keep bank_views for the one-pool slot cache Upstream #689 replaced OffloadMoeCache.bank_views with per-pool and staging views. The #676 on-demand prefill GEMM runs over the whole slot cache (n=cache_size) and still calls bank_views, so the first short prefill raised AttributeError. Restore it for a cache with one pool and one bank layout, where a slot id is its bank row; the pool guard already refuses on-demand prefill with several pools. tests/moe/test_small_prefill_on_demand.py covers the call. Assisted-by: Claude Opus 5.5 Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Bu6LgoxLR4wETqb7RPR2vt --- python/freetoken/moe/offload_cache.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index a12ab593a4..f7eb6c659e 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -555,6 +555,16 @@ def pool_state(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: def pool_slots(self, layer_id: int) -> int: return self.pool_sizes[self.pool_of_layer[layer_id]] + def bank_views(self, n: int | None = None) -> tuple[torch.Tensor, ...]: + """Per-bank views of the whole ``[S]`` slot cache, or of its first ``n`` slots, for a cache + with one pool and one bank layout, where a slot id is its bank row (the on-demand prefill + GEMM over the slot cache and the fork's slot checks).""" + assert self.banks, "set_bank_sources must register the banks first" + assert len(self.pool_sizes) == 1 and not self._raw_banks, "bank_views needs one slot pool and one bank layout" + if n is None: + return tuple(cache for _, cache in self.banks) + return tuple(cache[:n] for _, cache in self.banks) + def pool_bank_views(self, layer_id: int) -> tuple[torch.Tensor, ...]: """Per-bank slot views of the pool that serves ``layer_id``, indexed by its local slot ids.""" assert self.banks, "set_bank_sources must register the banks first" From 993fd3e12b247c867006f00278888d56dacc717c Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 15 Oct 2026 08:44:37 -0400 Subject: [PATCH 91/94] test: give the rebuild and scored-ensure fakes the slot-pool api of #689 #650's engine fake and the decode-forward fake predate upstream #689: the rebuild now splits the MoE target into pool sizes (decode_target, _moe_pools, _moe_pool_bytes, validate_rebuild/rebuild take the sizes) and decode reads pool_bank_views. Assisted-by: Claude Opus 5.5 Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Bu6LgoxLR4wETqb7RPR2vt --- tests/engine/test_cache_budget.py | 2 +- tests/engine/test_rebuild_runtime_cache_ordering.py | 12 ++++++++++-- tests/moe/test_scored_ensure.py | 2 +- 3 files changed, 12 insertions(+), 4 deletions(-) diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py index a8fd045fc7..e997d2ba90 100644 --- a/tests/engine/test_cache_budget.py +++ b/tests/engine/test_cache_budget.py @@ -428,7 +428,7 @@ def validate_rebuild(config, **kwargs): kv_cache=SimpleNamespace(validate_rebuild=validate_rebuild), _baseline_free=1 << 30, _weights_bytes=0, num_pages=64, ) - engine._target_moe_and_expert_bytes = lambda size: (0, 0) + engine._target_moe_and_expert_bytes = lambda size, sizes=None: (0, 0) with pytest.raises(CacheRebuildRejected): Engine.rebuild_runtime_cache(engine, num_pages=32) diff --git a/tests/engine/test_rebuild_runtime_cache_ordering.py b/tests/engine/test_rebuild_runtime_cache_ordering.py index aae7d49cf0..ab01f8a18f 100644 --- a/tests/engine/test_rebuild_runtime_cache_ordering.py +++ b/tests/engine/test_rebuild_runtime_cache_ordering.py @@ -54,14 +54,15 @@ def __init__(self, ledger: _Ledger, size: int): # one uint8 per expert row -> expert_bytes_per_slot == 1 unit self.bank_sources = {"w13": [torch.zeros((self.num_experts, 1), dtype=torch.uint8)]} self.cache_size = size + self.decode_target = "gpu" self._held = size ledger.alloc("moe", size) - def validate_rebuild(self, cache_size: int) -> None: + def validate_rebuild(self, cache_size: int, pool_sizes: list[int] | None = None) -> None: if cache_size < self.num_experts: raise ValueError(f"cache_size {cache_size} < num_experts {self.num_experts}") - def rebuild(self, cache_size: int) -> None: + def rebuild(self, cache_size: int, pool_sizes: list[int] | None = None) -> None: self.validate_rebuild(cache_size) self.ledger.free("moe", self._held) self._held = 0 @@ -140,6 +141,13 @@ def _build_engine(monkeypatch, ledger: _Ledger, *, moe: int, kv: int, budget: in graph_runner=_FakeGraphRunner(), _sync_get_memory=lambda: (0, 0), ) + # one MoE slot pool (#689): the rebuild splits the target into its pool sizes first + from freetoken.moe.pool_policy import make_pools + + eng._moe_pools = lambda batch: make_pools([[0]], 8, 1, None) + eng._moe_expert_counts = None + eng._moe_pool_groups, eng._moe_layer_slot_bytes = [[0]], [1] # one uint8 per expert row + eng._moe_pool_bytes = Engine._moe_pool_bytes.__get__(eng) eng._target_moe_and_expert_bytes = Engine._target_moe_and_expert_bytes.__get__(eng) eng._resize_kv_pool = Engine._resize_kv_pool.__get__(eng) eng._refresh_seq_state = Engine._refresh_seq_state.__get__(eng) diff --git a/tests/moe/test_scored_ensure.py b/tests/moe/test_scored_ensure.py index 6bfa84a5ed..21318c257b 100644 --- a/tests/moe/test_scored_ensure.py +++ b/tests/moe/test_scored_ensure.py @@ -570,7 +570,7 @@ def test_decode_forward_hands_the_router_logits_to_ensure(monkeypatch, policy): ) monkeypatch.setattr(cache, "ensure_experts", lambda lid, ids, **kw: seen.update(kw)) monkeypatch.setattr(cache, "copy_missing", lambda: None) - monkeypatch.setattr(cache, "bank_views", lambda: ()) + monkeypatch.setattr(cache, "pool_bank_views", lambda layer_id: ()) monkeypatch.setattr(layer, "_expert_gemm", lambda *a, **kw: None) router_logits = torch.randn(2, 8) layer.decode_forward(torch.randn(2, 8), router_logits) From 3694e64fc568af64ccaf8e505758bd71e5a803f3 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Thu, 15 Oct 2026 08:58:40 -0400 Subject: [PATCH 92/94] fix(moe): index the per-layout copy row bytes in the verify and audit gathers Upstream #689 keeps one fused-copy row-byte descriptor per bank layout (_copy_feat_bytes[layout]). The prefetch verify and slot-audit gathers still passed the whole list, so every enabled check failed at launch (TypeError: expected DLTensor*, got ffi.Array). Pass the layer's layout entry. The 83 failing test_skeleton/test_prefetch audit and verify tests pass again (385 passed). Assisted-by: Claude Opus 5.5 Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Bu6LgoxLR4wETqb7RPR2vt --- python/freetoken/moe/slot_audit.py | 2 +- python/freetoken/moe/verify.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/python/freetoken/moe/slot_audit.py b/python/freetoken/moe/slot_audit.py index 981239758f..12d9f40ad9 100644 --- a/python/freetoken/moe/slot_audit.py +++ b/python/freetoken/moe/slot_audit.py @@ -493,7 +493,7 @@ def _gather(self, cache, layer: int, rows: torch.Tensor, num: torch.Tensor) -> N from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit fast_index_copy_multi_jit( - self.scratch_ptrs, cache._copy_src_ptrs[layer], cache._copy_feat_bytes, self.iota, rows, num, + self.scratch_ptrs, cache._copy_src_ptrs[layer], cache._copy_feat_bytes[cache.layer_layout[layer]], self.iota, rows, num, ) return from freetoken.kernel import fast_index_copy_jit diff --git a/python/freetoken/moe/verify.py b/python/freetoken/moe/verify.py index 1f096c9c3a..cc71e19d9d 100644 --- a/python/freetoken/moe/verify.py +++ b/python/freetoken/moe/verify.py @@ -401,7 +401,7 @@ def _gather(self, cache, layer: int) -> None: from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit fast_index_copy_multi_jit( - self.scratch_ptrs, cache._copy_src_ptrs[layer], cache._copy_feat_bytes, self.iota, self.u_ids, self.u_num, + self.scratch_ptrs, cache._copy_src_ptrs[layer], cache._copy_feat_bytes[cache.layer_layout[layer]], self.iota, self.u_ids, self.u_num, ) return from freetoken.kernel import fast_index_copy_jit From 3c1fc7bf4d44e7de270d45d7547dce97a08dfb4f Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 10:56:34 -0400 Subject: [PATCH 93/94] test(kernel): cover the device-sized autotune flush buffer Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 99c07575e037d41cdc0c16a7a993aee5566f70ba) --- tests/kernels/test_autotune_cache.py | 29 ++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) create mode 100644 tests/kernels/test_autotune_cache.py diff --git a/tests/kernels/test_autotune_cache.py b/tests/kernels/test_autotune_cache.py new file mode 100644 index 0000000000..3416117f79 --- /dev/null +++ b/tests/kernels/test_autotune_cache.py @@ -0,0 +1,29 @@ +"""The autotune L2-flush buffer is sized to the device instead of triton's flat 256 MB.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +MB = 1 << 20 + + +@pytest.mark.parametrize("l2_bytes, flush_bytes", [ + (2816 * 1024, 11 * MB), # GTX 1080 Ti: 4x its 2.75 MB L2 + (0, 8 * MB), # no L2 size reported: the floor + (128 * MB, 256 * MB), # a large L2: triton's own 256 MB is the cap +]) +def test_flush_buffer_is_four_times_the_l2_within_bounds(monkeypatch, l2_bytes, flush_bytes): + from triton.backends.nvidia.driver import CudaDriver + + from freetoken.kernel.triton.autotune_cache import bound_autotune_flush_buffer + + monkeypatch.setattr(CudaDriver, "get_empty_cache_for_benchmark", CudaDriver.get_empty_cache_for_benchmark) + monkeypatch.setattr(torch.cuda, "get_device_properties", lambda index: SimpleNamespace(L2_cache_size=l2_bytes)) + bound_autotune_flush_buffer(0) + bound_autotune_flush_buffer(0) # the engine may call it again; the patch must not stack + monkeypatch.setattr(torch, "empty", lambda n, dtype, device: SimpleNamespace(nbytes=n * 4, dtype=dtype, device=device)) + buf = CudaDriver.get_empty_cache_for_benchmark(None) + assert (buf.nbytes, buf.dtype, buf.device) == (flush_bytes, torch.int, "cuda") From d5907fc670a06bd0ee7c589bb369214d9dcd2969 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Sat, 26 Sep 2026 00:28:53 -0400 Subject: [PATCH 94/94] fix(server): accept top_k <= 0 as "disabled", as vLLM and the OpenAI-compatible servers do Upstream #223 made resolve_sampling reject top_k=0 with HTTP 400 ("top_k must be -1 (disabled) or >= 1"). Clients send top_k=0 to mean "off": Open WebUI with Ollama-style parameters and some SDKs do, and vLLM and the other OpenAI-compatible servers accept any top_k <= 0 as "no top-k filter". A checkpoint whose generation_config.json carries top_k=0 made every request that omits top_k fail the same way under --sampling-defaults=model. resolve_sampling is the one place that validates top_k (the chat, legacy completions, Responses and Anthropic converters and #393's n>1 fan-out all go through it; the pydantic request models put no bound on it), so it now maps any top_k <= 0, from the request or from the checkpoint defaults, to the engine's -1. The sampler already read top_k < 1 as "off". The NaN/inf/negative temperature and top_p (0, 1] checks stay. #223's test that expected ("top_k", 0) -> 400 now asserts that 0 and -5 are accepted and reach the engine as -1; tests/server/test_sampling_validation.py covers the request and checkpoint cases and that the other range checks still raise. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RUamGLmkVPkWmPPBEUqQLP --- python/freetoken/server/generation.py | 7 +++-- tests/server/test_openai_api.py | 27 ++++++++++++++++++- tests/server/test_sampling_validation.py | 33 ++++++++++++++++++++++++ 3 files changed, 64 insertions(+), 3 deletions(-) diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index 89c898a0e9..00d94abebf 100644 --- a/python/freetoken/server/generation.py +++ b/python/freetoken/server/generation.py @@ -476,8 +476,11 @@ def pick(value, key, framework): if not math.isfinite(resolved_top_p) or not 0 < resolved_top_p <= 1: raise ValueError(f"top_p must be in (0, 1], got {resolved_top_p}") resolved_top_k = pick(top_k, "top_k", -1) - if resolved_top_k != -1 and resolved_top_k < 1: - raise ValueError(f"top_k must be -1 (disabled) or >= 1, got {resolved_top_k}") + # Any top_k <= 0 means "no top-k filter", as in vLLM and the other OpenAI-compatible + # servers: Ollama-style clients (Open WebUI) and some SDKs send top_k=0 for "off", and a + # checkpoint's generation_config.json can carry it too. Map it to the engine's -1. + if resolved_top_k <= 0: + resolved_top_k = -1 resolved_min_p = float(pick(min_p, "min_p", 0.0)) resolved_rep = float(pick(repetition_penalty, "repetition_penalty", 1.0)) if not 0.0 <= resolved_min_p <= 1.0: diff --git a/tests/server/test_openai_api.py b/tests/server/test_openai_api.py index 86781df84d..fdd2b11159 100644 --- a/tests/server/test_openai_api.py +++ b/tests/server/test_openai_api.py @@ -853,7 +853,6 @@ async def collect(): [ ("temperature", -0.5), ("top_p", 0.0), - ("top_k", 0), ], ) def test_chat_completion_rejects_invalid_sampling(field, value): @@ -906,3 +905,29 @@ async def chat_completion(req: ChatCompletionRequest): assert response.status_code == 200 assert state.sent is not None + + +@pytest.mark.parametrize("top_k", [0, -5]) +def test_chat_completion_treats_non_positive_top_k_as_disabled(top_k): + """Ollama-style clients (Open WebUI) and some SDKs send top_k=0 to mean "off"; vLLM and + the other OpenAI-compatible servers accept any top_k <= 0 as disabled, so it must not 400.""" + app = FastAPI() + state = FakeState([]) + + @app.post("/v1/chat/completions") + async def chat_completion(req: ChatCompletionRequest): + return await handle_chat_completion(req, request=None, state=state, model_sampling={}) + + response = TestClient(app).post( + "/v1/chat/completions", + json={ + "model": "unit-model", + "messages": [{"role": "user", "content": "hi"}], + "stream": True, + "top_k": top_k, + }, + ) + + assert response.status_code == 200 + assert state.sent is not None + assert state.sent.sampling_params.top_k == -1 diff --git a/tests/server/test_sampling_validation.py b/tests/server/test_sampling_validation.py index 3b2e1d6473..b13db921be 100644 --- a/tests/server/test_sampling_validation.py +++ b/tests/server/test_sampling_validation.py @@ -21,3 +21,36 @@ def test_repetition_penalty_must_be_positive_and_finite(value): def test_a_normal_repetition_penalty_is_kept(): assert _resolve(repetition_penalty=1.1).repetition_penalty == pytest.approx(1.1) + + +def _resolve_with(**overrides): + kwargs = dict( + temperature=None, top_k=None, top_p=None, max_tokens=8, ignore_eos=False, + model_sampling={}, + ) + kwargs.update(overrides) + return resolve_sampling(**kwargs) + + +@pytest.mark.parametrize("value", [0, -1, -5]) +def test_a_non_positive_top_k_means_disabled(value): + # vLLM's rule: any top_k <= 0 is "no top-k filter"; clients send 0 for "off" + assert _resolve_with(top_k=value).top_k == -1 + + +def test_a_checkpoint_default_top_k_of_zero_means_disabled(): + assert _resolve_with(model_sampling={"top_k": 0}).top_k == -1 + + +def test_a_positive_top_k_is_kept(): + assert _resolve_with(top_k=20).top_k == 20 + + +@pytest.mark.parametrize( + ("field", "value"), + [("temperature", float("nan")), ("temperature", float("inf")), ("temperature", -0.5), + ("top_p", 0.0), ("top_p", 1.5), ("top_p", float("nan"))], +) +def test_the_other_sampling_ranges_are_still_rejected(field, value): + with pytest.raises(ValueError, match=field): + _resolve_with(**{field: value})