Skip to content

fix(st): publish the lazy shard-mapping table without a data race - #1549

Merged
JustVugg merged 3 commits into
JustVugg:devfrom
basil-k-aji-dev:fix/st-shard-table-race
Sep 17, 2026
Merged

JustVugg merged 3 commits into
JustVugg:devfrom
basil-k-aji-dev:fix/st-shard-table-race

Conversation

@basil-k-aji-dev

Copy link
Copy Markdown
Contributor

Closes #1519.

The race

st_shard_mapped() wrote g_st_shard_base / g_st_shard_len / g_st_shard_tried with no lock, no atomics and no barrier. It is reached concurrently: glm53.c:1587's

#pragma omp parallel for schedule(dynamic, 1)
for (int r = 0; r < reads; r++)
    expert_read(m, index, union_ids[base + i], &cache->s[slot_of[i]]);

calls expert_read → st_map_shard_range → st_shard_mapped from every worker at once, and qwen38_core.h:1102-1104 does the same. expert_read already knows it runs in parallel — it guards m->miss with #pragma omp atomic — but the table it reaches through was unguarded.

Two distinct failures follow:

  1. A leaked mapping. Two threads both read g_st_shard_tried[fd] == 0, both map the file. One base pointer overwrites the other and that mapping is never unmapped — and shard mappings are deliberately kept for the process lifetime, so it is never reclaimed.
  2. A shard that rejects every range. The base was published before the length, so a third thread could take the non-NULL base and still read g_st_shard_len[fd] as 0. st_map_shard_range then fails off > g_st_shard_len[fd] - nbytes for every range on that shard, silently demoting it to pread for the rest of the run.

Reproduction

TSan will not run on my host (aarch64, 47-bit VMA; TSan supports 39/42/48), so I reproduced the consequences directly instead. The harness copies the three globals and the lookup body verbatim out of st.h and drives them from the same #pragma omp parallel for schedule(dynamic, 1) shape, counting how many times the map actually happens:

############ BEFORE (code as on main) ############
rounds=200  mmaps=201 (expected 200)  len-seen-as-0=1
rounds=200  mmaps=203 (expected 200)  len-seen-as-0=0
rounds=200  mmaps=205 (expected 200)  len-seen-as-0=2

############ AFTER (this patch) ############
rounds=200  mmaps=200 (expected 200)  len-seen-as-0=0
rounds=200  mmaps=200 (expected 200)  len-seen-as-0=0
rounds=200  mmaps=200 (expected 200)  len-seen-as-0=0

Both failure modes appear on main and neither survives the patch. Happy to add the harness to the tree as a test if you'd like it — I left it out because it is a copy of the function body rather than a test of the real one.

The fix

  • Claim each fd exactly once with a compare-exchange on g_st_shard_tried.
  • Store the length before release-storing the base, and read the base with acquire, so a non-NULL base always carries its length.
  • The thread that loses the claim does not block: NULL is the documented "use pread" answer, so it takes the pread path for that one call rather than parking an OpenMP worker behind an mmap.

Why lock-free and not the mutex used elsewhere

qwen38_core.h solves the identical "lazy table published from a parallel region" problem with pthread_mutex_t + double-checked locking + acquire/release — after, by its own comment, a SIGSEGV "in about one run in twelve on a CUDA warmstart". Copying that idiom was my first instinct.

It does not work here: st.h is also compiled into c/tools/check_glm53_container.c, whose documented build line in its own header is

gcc -O2 -std=gnu11 -Ic -o check_container c/tools/check_glm53_container.c -lm

with neither -pthread nor -fopenmp. Adding #include <pthread.h> to st.h would silently invalidate that. The __atomic_* builtins are already used ~49 times in the tree and need no new dependency.

Verification

Check Result
gcc -O2 -std=gnu11 -Ic ... -lm (the tool's documented line, no pthread/OpenMP) builds clean
gcc -fopenmp -pthread -c glm53.c compiles, 188 KB object
qwen38.c, olmoe.c, qwen36.c syntax check all OK
st.h standalone with -Wall -Wextra no new warnings

Behaviour is unchanged on the fast path: once a base is published, the lookup is a single acquire load, as before.

Note, not fixed here

st_map_experts_enabled() has the same shape (static int on = -1; initialised on first call) and is technically a data race under the memory model, though a benign one — every thread computes the same value from the same getenv. I left it alone to keep this diff to the reported bug; say the word and I'll fold it in.

JustVugg and others added 3 commits September 15, 2026 14:39
…OTICE

The LICENSE file still carried the Apache 2.0 template placeholders
(Copyright [yyyy] [name of copyright owner]), so the notice was never
actually applied to this work. Fill it in, add the NOTICE file Apache 2.0
expects for the copyright line, and point the four READMEs at both.
license: name the copyright holder and add NOTICE
st_shard_mapped() wrote g_st_shard_base/_len/_tried with no lock, no atomics
and no barrier, and glm53.c reaches it from inside a `#pragma omp parallel
for` (the batch loop at glm53.c:1587 -> expert_read -> st_map_shard_range).
qwen38_core.h reaches it from a parallel region too.

Two threads that both observe g_st_shard_tried[fd] == 0 both map the file, so
one mapping is leaked for the process lifetime; and because the base pointer
was published before the length, a third thread could take the non-NULL base
and still read g_st_shard_len[fd] as 0, at which point st_map_shard_range
rejects every range for that shard.

Claim each fd once with a compare-exchange on g_st_shard_tried, store the
length before release-storing the base, and read the base with acquire. A
thread that loses the claim does not block: NULL is the documented "use
pread" answer, so it takes the pread path for that one call instead of
parking an OpenMP worker behind an mmap.

Lock-free rather than the pthread mutex qwen38_core.h uses for its lazy HITS
table, because st.h is also compiled into c/tools/check_glm53_container.c,
whose documented build line is plain `gcc -O2 -std=gnu11 -Ic ... -lm` with
neither -pthread nor -fopenmp.

Latent today because COLI_MAP_EXPERTS is opt-in; JustVugg#1350 proposes flipping that
default, which would make it live for everyone.

Closes JustVugg#1519.
Edo771977 added a commit to Edo771977/colibri that referenced this pull request Sep 16, 2026
@JustVugg
JustVugg changed the base branch from main to dev September 17, 2026 07:06
@JustVugg
JustVugg merged commit 3326b5f into JustVugg:dev Sep 17, 2026
28 checks passed
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.

st.h: the lazy shard-mapping table is written without synchronisation, and glm53 reaches it from an OpenMP region

2 participants