Skip to content

Detach tokenizer storage from model mmap #942

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
riccardomenegazzo wants to merge 1 commit into
base: main
Choose a base branch
from
+220 −22
Open

Detach tokenizer storage from model mmap #942

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 13 additions & 2 deletions Makefile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ else
NATIVE_CPU_FLAG ?= -march=native
endif
SAMPLING_TEST := tests/test_sampling
VOCAB_STORAGE_TEST := tests/test_vocab_storage
GLM53_KDA_TEST := tests/test_glm53_kda
GLM53_KDA_ROCM_TEST := tests/test_glm53_kda_rocm

Expand Down Expand Up @@ -199,6 +200,7 @@ rocm: strix-halo
test-rocm:
$(MAKE) -B ds4_test ds4_agent_test ds4-eval q4k-dot-test mxfp4-dot-test \
tests/test_layer_pack tests/test_engine_mgpu_placement tests/test_gpu_args \
$(VOCAB_STORAGE_TEST) \
ds4 ds4-server ds4-bench ds4-agent \
CORE_OBJS="ds4.o ds4_image.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_rocm.o ds4_rocm_compat.o ds4_rocm_unavailable.o ds4_layer_pack.o $(ROCM_MMQ_OBJS)" \
CFLAGS="$(CFLAGS) $(ROCM_HOST_CFLAGS) -DDS4_ROCM_BUILD" \
Expand All @@ -211,6 +213,7 @@ test-rocm:
./tests/test_engine_mgpu_placement
./tests/test_gpu_args
./tests/test_gpu_args_cli.sh
./tests/test_vocab_storage

ds4: ds4_cli.o ds4_help.o linenoise.o ds4_gpu_args.o $(CORE_OBJS)
$(DS4_LINK) -o $@ $^ $(DS4_LINK_LIBS)
Expand Down Expand Up @@ -480,6 +483,12 @@ tests/test_sampling.o: tests/test_sampling.c ds4.h
tests/test_sampling: tests/test_sampling.o ds4_cpu_test_hooks.o ds4_image.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o
$(CC) $(CFLAGS) -o $@ $^ $(LDLIBS)

tests/test_vocab_storage.o: tests/test_vocab_storage.c ds4.h
$(CC) $(CFLAGS) -DDS4_TEST_HOOKS -I. -c -o $@ $<

tests/test_vocab_storage: tests/test_vocab_storage.o ds4_cpu_test_hooks.o ds4_image.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o
$(CC) $(CFLAGS) -o $@ $^ $(LDLIBS)

ifneq ($(UNAME_S),Darwin)
tests/test_gpu_xdev.o: tests/test_gpu_xdev.c ds4_gpu.h ds4_gpu_mgpu.h
$(CC) $(CFLAGS) -I. -I$(CUDA_HOME)/include -c -o $@ $<
Expand Down Expand Up @@ -555,7 +564,8 @@ endif

test: ds4_test ds4_agent_test ds4-eval q4k-dot-test mxfp4-dot-test \
tests/test_layer_pack tests/test_engine_mgpu_placement tests/test_gpu_args \
tests/test_deepseek4_vision_image $(SAMPLING_TEST) ds4 ds4-server ds4-bench ds4-agent
tests/test_deepseek4_vision_image $(SAMPLING_TEST) $(VOCAB_STORAGE_TEST) \
ds4 ds4-server ds4-bench ds4-agent
./ds4-eval --self-test-extractors
./ds4_agent_test
./ds4_test
Expand All @@ -564,6 +574,7 @@ test: ds4_test ds4_agent_test ds4-eval q4k-dot-test mxfp4-dot-test \
./tests/test_gpu_args
./tests/test_gpu_args_cli.sh
./tests/test_sampling
./tests/test_vocab_storage
./tests/test_deepseek4_vision_image

dspark-acceptance: ds4
Expand Down Expand Up @@ -600,4 +611,4 @@ mxfp4-dot-test: tests/test_mxfp4_dot.c
./tests/test_mxfp4_dot

clean:
rm -f ds4 ds4-server ds4-bench ds4-eval ds4-agent ds4_cpu ds4_native ds4_server_test ds4_test ds4_agent_test gguf-tools/quality-testing/score_official gguf-tools/quality-testing/score_official.o speed-bench/metal_decode_schedule_bench speed-bench/metal_prefill_variant_bench speed-bench/*.o tests/test_q4k_dot tests/test_mxfp4_dot tests/test_mxfp4_metal tests/test_mxfp4_rocm tests/test_mxfp4_cuda tests/test_metal_session_batch tests/test_glm53_kda tests/test_glm53_kda_rocm tests/test_glm53_vision_engine tests/test_glm53_vision_prompt tests/test_deepseek4_vision_image tests/test_gpu_xdev tests/test_gpu_model_cache tests/test_gpu_lookup_cache_strict tests/test_engine_mgpu_refusal tests/test_engine_mgpu_runtime tests/test_engine_correctness tests/test_sampling tests/test_cuda_session_batch tests/test_cuda_mixed_batch tests/*.o *.o tests/cuda_long_context_smoke tests/cuda_long_context_smoke.o
rm -f ds4 ds4-server ds4-bench ds4-eval ds4-agent ds4_cpu ds4_native ds4_server_test ds4_test ds4_agent_test gguf-tools/quality-testing/score_official gguf-tools/quality-testing/score_official.o speed-bench/metal_decode_schedule_bench speed-bench/metal_prefill_variant_bench speed-bench/*.o tests/test_q4k_dot tests/test_mxfp4_dot tests/test_mxfp4_metal tests/test_mxfp4_rocm tests/test_mxfp4_cuda tests/test_metal_session_batch tests/test_glm53_kda tests/test_glm53_kda_rocm tests/test_glm53_vision_engine tests/test_glm53_vision_prompt tests/test_deepseek4_vision_image tests/test_gpu_xdev tests/test_gpu_model_cache tests/test_gpu_lookup_cache_strict tests/test_engine_mgpu_refusal tests/test_engine_mgpu_runtime tests/test_engine_correctness tests/test_sampling tests/test_vocab_storage tests/test_cuda_session_batch tests/test_cuda_mixed_batch tests/*.o *.o tests/cuda_long_context_smoke tests/cuda_long_context_smoke.o
214 changes: 194 additions & 20 deletions ds4.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38182,10 +38182,10 @@ static void embed_prompt(
* =========================================================================
*
* DeepSeek V4 Flash stores a GPT-2 style byte-level BPE tokenizer in GGUF.
* The implementation below is intentionally small. It loads token strings
* and merge ranks from the mmaped file, builds two open-addressed hash tables,
* and applies BPE to user text. Chat special tokens are inserted directly by
* ID; user text goes through BPE.
* The implementation below is intentionally small. It copies token strings
* and merge keys into one compact owned arena, builds two open-addressed hash
* tables, and applies BPE to user text. Tensor data remains mmap-backed. Chat
* special tokens are inserted directly by ID; user text goes through BPE.
*/

typedef struct {
Expand Down Expand Up @@ -38288,6 +38288,7 @@ bool ds4_tokens_starts_with(const ds4_tokens *tokens, const ds4_tokens *prefix)

struct ds4_vocab {
ds4_str *token;
char *string_storage; /* Owns token strings and hash table keys. */
int n_vocab;
int bos_id;
int eos_id;
Expand Down Expand Up @@ -39199,40 +39200,89 @@ static int vocab_lookup_optional(const ds4_vocab *vocab, const char *text) {
return token;
}

/* Load token strings, special token ids, and merge ranks from GGUF metadata. */

static void vocab_load(ds4_vocab *vocab, const ds4_model *model) {
memset(vocab, 0, sizeof(*vocab));
static size_t vocab_array_storage_bytes(const ds4_model *model,
ds4_array_ref array) {
size_t total = 0;
ds4_cursor c = cursor_at(model, array.data_pos);
for (uint64_t i = 0; i < array.len; i++) {
ds4_str s;
if (!cursor_string(&c, &s)) ds4_die(c.error);
if (s.len > SIZE_MAX - total) {
ds4_die("GGUF tokenizer string data is too large");
}
total += (size_t)s.len;
}
return total;
}

ds4_array_ref tokens;
ds4_array_ref merges;
if (!model_get_array(model, "tokenizer.ggml.tokens", &tokens) ||
tokens.type != GGUF_VALUE_STRING ||
tokens.len > INT32_MAX) {
ds4_die("GGUF tokenizer token table is missing or invalid");
static ds4_str vocab_copy_string(ds4_cursor *c, char **dst,
size_t *remaining) {
ds4_str borrowed;
if (!cursor_string(c, &borrowed)) ds4_die(c->error);
if (borrowed.len > *remaining) {
ds4_die("GGUF tokenizer string data changed while loading");
}
if (!model_get_array(model, "tokenizer.ggml.merges", &merges) ||
merges.type != GGUF_VALUE_STRING) {
ds4_die("GGUF tokenizer merge table is missing or invalid");
memcpy(*dst, borrowed.ptr, (size_t)borrowed.len);
ds4_str owned = { *dst, borrowed.len };
*dst += (size_t)borrowed.len;
*remaining -= (size_t)borrowed.len;
return owned;
}

static void vocab_load_string_storage(ds4_vocab *vocab,
const ds4_model *model,
ds4_array_ref tokens,
ds4_array_ref merges) {
const size_t token_bytes = vocab_array_storage_bytes(model, tokens);
const size_t merge_bytes = vocab_array_storage_bytes(model, merges);
if (merge_bytes > SIZE_MAX - token_bytes) {
ds4_die("GGUF tokenizer string data is too large");
}

const size_t storage_bytes = token_bytes + merge_bytes;
vocab->string_storage = xmalloc(storage_bytes ? storage_bytes : 1);
char *dst = vocab->string_storage;
size_t remaining = storage_bytes;

vocab->n_vocab = (int)tokens.len;
vocab->token = xcalloc((size_t)vocab->n_vocab, sizeof(vocab->token[0]));
table_init(&vocab->token_to_id, tokens.len);

ds4_cursor c = cursor_at(model, tokens.data_pos);
for (int i = 0; i < vocab->n_vocab; i++) {
if (!cursor_string(&c, &vocab->token[i])) ds4_die(c.error);
vocab->token[i] = vocab_copy_string(&c, &dst, &remaining);
table_put(&vocab->token_to_id, vocab->token[i], i);
}

table_init(&vocab->merge_rank, merges.len);
c = cursor_at(model, merges.data_pos);
for (uint64_t i = 0; i < merges.len; i++) {
ds4_str merge;
if (!cursor_string(&c, &merge)) ds4_die(c.error);
ds4_str merge = vocab_copy_string(&c, &dst, &remaining);
table_put(&vocab->merge_rank, merge, (int)i);
}
if (remaining != 0) {
ds4_die("GGUF tokenizer string data changed while loading");
}
}

/* Load token strings, special token ids, and merge ranks from GGUF metadata. */

static void vocab_load(ds4_vocab *vocab, const ds4_model *model) {
memset(vocab, 0, sizeof(*vocab));

ds4_array_ref tokens;
ds4_array_ref merges;
if (!model_get_array(model, "tokenizer.ggml.tokens", &tokens) ||
tokens.type != GGUF_VALUE_STRING ||
tokens.len > INT32_MAX) {
ds4_die("GGUF tokenizer token table is missing or invalid");
}
if (!model_get_array(model, "tokenizer.ggml.merges", &merges) ||
merges.type != GGUF_VALUE_STRING) {
ds4_die("GGUF tokenizer merge table is missing or invalid");
}

vocab_load_string_storage(vocab, model, tokens, merges);

if (DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) {
if (!model_get_token_id(model, "tokenizer.ggml.bos_token_id", &vocab->bos_id)) {
Expand Down Expand Up @@ -39284,6 +39334,7 @@ static void vocab_free(ds4_vocab *vocab) {
free(vocab->token);
table_free(&vocab->token_to_id);
table_free(&vocab->merge_rank);
free(vocab->string_storage);
memset(vocab, 0, sizeof(*vocab));
}

Expand Down Expand Up @@ -39621,6 +39672,129 @@ static bool vocab_token_is_literal_special(ds4_str s) {
return false;
}

#ifdef DS4_TEST_HOOKS
static size_t ds4_test_write_gguf_string(uint8_t *dst, size_t pos,
const char *s) {
const uint64_t len = (uint64_t)strlen(s);
memcpy(dst + pos, &len, sizeof(len));
pos += sizeof(len);
memcpy(dst + pos, s, (size_t)len);
return pos + (size_t)len;
}

static bool ds4_test_vocab_keys_are_detached(const str_i32_table *table,
uintptr_t map_begin,
uintptr_t map_end) {
for (uint64_t i = 0; i < table->cap; i++) {
if (!table->entry[i].used) continue;
const uintptr_t ptr = (uintptr_t)table->entry[i].key.ptr;
if (ptr >= map_begin && ptr < map_end) return false;
}
return true;
}

int ds4_test_vocab_storage_detached(void) {
static const char *tokens[] = {
"a",
"b",
"ab",
"<\xef\xbd\x9c" "special" "\xef\xbd\x9c>",
};
static const char *merges[] = { "a b" };

size_t map_size = 0;
for (size_t i = 0; i < sizeof(tokens) / sizeof(tokens[0]); i++) {
map_size += sizeof(uint64_t) + strlen(tokens[i]);
}
for (size_t i = 0; i < sizeof(merges) / sizeof(merges[0]); i++) {
map_size += sizeof(uint64_t) + strlen(merges[i]);
}

#if defined(MAP_ANONYMOUS)
const int map_flags = MAP_PRIVATE | MAP_ANONYMOUS;
#elif defined(MAP_ANON)
const int map_flags = MAP_PRIVATE | MAP_ANON;
#else
return 0;
#endif
uint8_t *mapping = mmap(NULL, map_size, PROT_READ | PROT_WRITE,
map_flags, -1, 0);
if (mapping == MAP_FAILED) return 0;

size_t pos = 0;
const uint64_t token_pos = pos;
for (size_t i = 0; i < sizeof(tokens) / sizeof(tokens[0]); i++) {
pos = ds4_test_write_gguf_string(mapping, pos, tokens[i]);
}
const uint64_t merge_pos = pos;
for (size_t i = 0; i < sizeof(merges) / sizeof(merges[0]); i++) {
pos = ds4_test_write_gguf_string(mapping, pos, merges[i]);
}

ds4_model model = {
.fd = -1,
.map = mapping,
.size = pos,
};
const ds4_array_ref token_array = {
.type = GGUF_VALUE_STRING,
.len = sizeof(tokens) / sizeof(tokens[0]),
.data_pos = token_pos,
};
const ds4_array_ref merge_array = {
.type = GGUF_VALUE_STRING,
.len = sizeof(merges) / sizeof(merges[0]),
.data_pos = merge_pos,
};
ds4_vocab vocab = {0};
vocab_load_string_storage(&vocab, &model, token_array, merge_array);

const uintptr_t map_begin = (uintptr_t)mapping;
const uintptr_t map_end = map_begin + pos;
bool detached = vocab.string_storage != NULL;
for (int i = 0; detached && i < vocab.n_vocab; i++) {
const uintptr_t ptr = (uintptr_t)vocab.token[i].ptr;
detached = ptr < map_begin || ptr >= map_end;
}
detached = detached &&
ds4_test_vocab_keys_are_detached(&vocab.token_to_id,
map_begin, map_end) &&
ds4_test_vocab_keys_are_detached(&vocab.merge_rank,
map_begin, map_end);

if (munmap(mapping, map_size) != 0 || !detached) {
vocab_free(&vocab);
return 0;
}

int token = -1;
int rank = -1;
bool ok = table_get(&vocab.token_to_id, "ab", 2, &token) &&
token == 2 &&
table_get(&vocab.merge_rank, "a b", 3, &rank) &&
rank == 0 &&
vocab.token[2].len == 2 &&
memcmp(vocab.token[2].ptr, "ab", 2) == 0;

ds4_engine *engine = xcalloc(1, sizeof(*engine));
engine->vocab = vocab;
size_t text_len = 0;
char *text = ds4_token_text(engine, 3, &text_len);
ok = ok && text_len == vocab.token[3].len &&
memcmp(text, vocab.token[3].ptr, text_len) == 0;
free(text);
free(engine);

token_vec encoded = {0};
bpe_emit_piece(&vocab, (ds4_str){ "ab", 2 }, &encoded);
ok = ok && encoded.len == 1 && encoded.v[0] == 2;

token_vec_free(&encoded);
vocab_free(&vocab);
return ok ? 1 : 0;
}
#endif

char *ds4_token_text(ds4_engine *e, int token, size_t *len) {
ds4_vocab *vocab = &e->vocab;
if (token < 0 || token >= vocab->n_vocab) {
Expand Down
1 change: 1 addition & 0 deletions ds4.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -449,6 +449,7 @@ int ds4_sample_logits(const float *logits, int n_vocab, float temperature,
int top_k, float top_p, float min_p, uint64_t *rng);
int ds4_session_sample(ds4_session *s, float temperature, int top_k, float top_p, float min_p, uint64_t *rng);
#ifdef DS4_TEST_HOOKS
int ds4_test_vocab_storage_detached(void);
int ds4_test_sample_logits(const float *logits, uint32_t n_vocab,
float temperature, int top_k,
float top_p, float min_p, uint64_t *rng,
Expand Down
12 changes: 12 additions & 0 deletions tests/test_vocab_storage.c
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
#include "../ds4.h"

#include <stdio.h>

int main(void) {
if (!ds4_test_vocab_storage_detached()) {
fprintf(stderr, "FAIL: tokenizer strings still depend on source mapping\n");
return 1;
}
puts("PASS: tokenizer storage survives source mapping removal");
return 0;
}