Skip to content

feat(k3): map prepared weights read-only - #965

Open
bherald wants to merge 3 commits into
JustVugg:devfrom
bherald:feat/k3-mmap-prepared-weights
Open

feat(k3): map prepared weights read-only#965
bherald wants to merge 3 commits into
JustVugg:devfrom
bherald:feat/k3-mmap-prepared-weights

Conversation

@bherald

@bherald bherald commented Aug 11, 2026

Copy link
Copy Markdown
Contributor

Summary

Add an opt-in, CPU-only K3_MMAP=1 path for Kimi K3 containers whose dense,
MLA, shared-expert, and head tensors are already stored in their final U8 data

  • F32 scale representation.
  • add a small cross-platform read-only mapping primitive that carries the
    aligned OS view and Windows mapping handle required for exact cleanup;
  • map prepared int8 and int4-g64 tensor bytes plus their F32 scale sidecars
    instead of copying them into private heap;
  • leave the existing loader and all defaults unchanged;
  • refuse Vulkan, CUDA, load-time conversion, incompatible scale formats, and
    mapping failures instead of silently falling back; and
  • test unaligned tensor offsets, byte-exact mapped-versus-copied CPU math, the
    backend refusal contract, and Windows file/directory cleanup.

This is the focused st_map + K3 vertical slice discussed in #826. It does not
implement TRUNK_RESIDENT_LAYERS, prefetching, page locking, or an OS working-
set policy.

Field result

On a Ryzen 9 7845HX / 64 GB Windows host using the official Kimi K3 text model,
int4/MLA8/Head8, a 1-GiB expert cache, 1K context, K3_VK=0, and K3_CUDA=0,
the mapped path reported 1,906 prepared views / 33.8 GB file-backed. Mapping
reduced peak private memory from the prior heap-copy path's roughly 35.36 GB to
about 7.32 GB.

Mapping alone did not bound the active working set. The host therefore applied
its own verified hard working-set maximum to only the exact K3 engine; that
limit is external to this PR. Three isolated 26-token smoke repeats at 16 GiB
completed in 72.945-73.892 seconds with 2,126-2,578 MiB pagefile growth and
zero final K3 engines.

A follow-up exposed an important limit. This runtime uses K3_CHUNK=32; prompts
above 32 total chat tokens require another full 93-layer prefill traversal. A
40-token synthetic prompt repeatedly crossed the unchanged 3,072-MiB pagefile
guard even with a 19-GiB host cap. Five concise 31-32-token synthetic cases did
complete safely at 19 GiB in 86.491-92.282 seconds with 1,759-1,881 MiB
pagefile growth and zero final K3 engines, but only four returned the expected
token. The NoFallback case returned an unexpected token. These results show
that prepared read-only mapping makes host-managed residency possible inside a
narrow single-prefill envelope; they are not a production-readiness,
throughput, broad-quality, or full-model-equivalence claim.

Validation

  • make -C c check
  • make -C c test-asan
  • Windows native suite and 484 Python/API tests passed with 49 expected
    Windows skips
  • Two Windows builds were byte-identical
  • Exact copied-versus-mapped int8 and int4 CPU math tests passed
  • Three new bounded full-model repeats returned the same one-token content
    hash and clean final process state
  • Five bounded single-prefill quality cases completed safely; four returned
    the expected token and the mismatch failed closed
  • CUDA changes were not made; mapped mode explicitly refuses enabled CUDA
    or Vulkan backends
  • Performance claims include the host, exact settings, and repeatable
    measurements above

Compatibility

  • The default CPU build remains dependency-free
  • No model files, generated binaries, receipts, or benchmark artifacts are
    included
  • K3_MMAP defaults off and changes no existing loader behavior

Closes no issue; contributes the mapped-residency slice toward #826 while
leaving the broader TRUNK_RESIDENT_LAYERS design open.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant