Skip to content

Detach tokenizer storage from model mmap - #942

Open
riccardomenegazzo wants to merge 1 commit into
antirez:mainfrom
riccardomenegazzo:fix/tokenizer-owned-storage
Open

Detach tokenizer storage from model mmap#942
riccardomenegazzo wants to merge 1 commit into
antirez:mainfrom
riccardomenegazzo:fix/tokenizer-owned-storage

Conversation

@riccardomenegazzo

Copy link
Copy Markdown

Related to #922.

DS4 currently stores vocabulary tokens and BPE merge keys as borrowed ds4_str spans into the GGUF mmap. This leaves tokenization and token decoding able to fault model storage long after model loading has completed.

This change detaches only tokenizer string metadata from the model mapping:

  • computes the exact token and merge byte count with overflow checks;
  • allocates one compact arena owned by ds4_vocab;
  • copies token strings and merge keys into that arena;
  • builds token_to_id and merge_rank exclusively from owned spans;
  • releases the arena in vocab_free().

Tensor data remains mmap-backed. Metal, CUDA, ROCm, distributed inference and SSD expert streaming retain their existing storage semantics.

Once the vocabulary is loaded, converting a token ID back to text should not require waking a 96 GiB model disk.

Regression test

tests/test_vocab_storage.c creates synthetic GGUF-style token and merge arrays inside an anonymous mapping, loads the tokenizer, verifies that every retained span is detached, and then removes the source mapping with munmap().

After the mapping no longer exists, the test exercises token lookup, BPE merge lookup and emission, token byte access, and the actual ds4_token_text() literal-special path from #922.

Validation

Run on Linux x86_64:

  • make -B ds4.o
  • make -j2 cpu
  • ./tests/test_vocab_storage
  • ./tests/test_engine_mgpu_placement - 109/109 checks passed
  • ./tests/test_sampling - passed
  • ASan + UBSan regression build — passed

The original 96 GiB external-SSD reproduction and Metal/CUDA/ROCm runtime tests were not available in my environment. The patch does not modify backend or tensor execution code.

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