Cobordism /cobordism-lm-head warm-path speedup

Date: 2026-05-04 Author: read-only investigation, post wave-22 cache deploy Status: PROPOSAL — no source edits made Targets: worker-cobordism-g4-s{0..7}.taylorbuley.workers.dev Budget claim: ~1.5 bule (1 day, narrow blast radius — port one Q6_K row routine + matmul to wasm SIMD path that already exists for Q4_K)

1. Executive summary

The wave-22 RAM cache eliminated the 3 s R2 round-trip but left 1.74 s of warm per-shard wall in the /cobordism-lm-head hot path. That residual is not fetch-bound, not network-bound, and not Q6_K-SIMD-bound on Rust — it is pure JavaScript scalar dequant + scalar dot product. The cobordism worker does not invoke the Rust/wasm kernel at all: it imports dequantizeRow from its own apps/worker-cobordism/src/dequantize.ts, which is a hand-rolled JS port (line 11-19 docstring even says: "Why JS instead of the Rust/WASM kernel: ... For qwen-coder-7b that tensor is ~306 MiB at Q4_K, well over the 128 MiB Cloudflare Workers ceiling"). The matmul is a literal for h in 0..hiddenDim: acc += lastRow[h] * row[h] JS loop in apps/worker-cobordism/src/index.ts:917-921.

The fix is to thread the already-resident 38 MiB quant slice (now in isolate-RAM via the wave-22 cache) through the already-existing wasm32 SIMD mat_vec_q4k/mat_vec_q6k kernels in open-source/gnosis/distributed-inference/src/math.rs. The original "too-big-for-128 MiB" justification was correct in 2026-04 but is now obsolete: the wave-22 cache already pays the 38 MiB resident cost in JS-side ArrayBuffer memory, so paying it again as a wasm-linear-memory copy is the only real cost (one-time per cache-fill, not per request).

Predicted lift: lm-head warm wall 1.74 s → 0.20-0.35 s (5-9× speedup), moving cobordism's per-token share from 8.5 % → ~1 %, lifting end-to-end TPS by +5-7 % (0.0285 → ~0.0301-0.0305).

2. Per-step breakdown of the 1.74 s warm wall

Per apps/worker-cobordism/src/index.ts:898-922, after the cache hit the worker executes per shard, per token:

for i in 0..32768 (vocab slice):
    rowBytes = sliceBytes.subarray(rowOffset, rowOffset + 3528)   // ~free
    row      = dequantizeRow(rowBytes, "Q6_K")                    // JS scalar
    acc      = 0
    for h in 0..5376:                                             // JS scalar
        acc += lastRow[h] * row[h]
    partialLogits[i] = acc

Element-op count

352 M elementwise ops in 1 743 ms → ~202 MOPS sustained. That is the canonical V8 scalar-JS floor for tight FMA-style loops on a Workers isolate (CF Workers do not enable --turbo-fast-math; loops with Float32Array gather + scalar mul/add typically clock 150-300 MOPS).

Where the 1.74 s splits

There is no per-step instrumentation in the handler — the X-Compute-Time-Ms header reports the aggregate compute window after the cache is consumed. We can divide the cost by op count, since dequant and dot are roughly the same elementwise cost in JS:

(The softcap loop Math.tanh(x/30) * 30 over 32 768 floats is real but sub-25 ms in V8.)

The partialLogits.buffer → Uint8Array zero-copy view at line 938 is genuinely free; serialization to the wire is the 131 072-byte Response.body return, which is just a buffer move.

3. The wasm32 SIMD gap (and what's actually missing)

There are two layers to the gap, in order of leverage:

3a. Cobordism worker bypasses Rust kernel entirely (PRIMARY GAP)

apps/worker-cobordism/src/dequantize.ts:1-54 documents the deliberate choice to use a JS port of the dequant routines, with a justification that no longer holds:

"The Rust dequant in open-source/gnosis/distributed-inference/src/math.rs expects the full quantized tensor to be loaded into WASM linear memory ... For qwen-coder-7b that tensor is ~306 MiB at Q4_K, well over the 128 MiB Cloudflare Workers ceiling."

That was true when each /cobordism-lm-head call held only one row at a time. As of wave-22 the worker is already holding 38 MiB resident in JS RAM per isolate (lmHeadCache.quantBytes), well below the 128 MiB ceiling. Copying that 38 MiB once per cache fill into wasm-linear memory is a one-time ~30 ms cost, after which every subsequent /cobordism-lm-head call is a single mat_vec_q6k(out, lastRow, slice, 32768, 5376) wasm call.

The Rust kernel exists, is wasm-bindgen-exported (wasm_bindings.rs), and is the same one driving the trisplit layer workers' Q4_K matmuls.

3b. wasm32 Q6_K SIMD branch is currently scalar (SECONDARY GAP)

Even if we threaded the matmul through wasm, mat_vec_q6k_row for #[cfg(target_arch = "wasm32")] at math.rs:1257-1297 is scalar. There is a BUG-FIX 2026-04-27 comment explaining that the original SIMD attempt had a "4× SIMD splat-and-sum-all-lanes bug" identical to the aarch64 one and was reverted.

Compare to mat_vec_q4k_row at math.rs:549-629, which has a complete working wasm32 SIMD branch using f32x4_splat, v128_load, f32x4_mul, f32x4_add — a 1:1 port of the NEON path that the comment at lines 531-548 explicitly documents was contributed during wave-22.

So the Q6_K SIMD port is missing because it failed once and was reverted, but the Q4_K port is the working template. The two block layouts differ (Q4_K = 144 B / Q6_K = 210 B) but the SIMD pattern (per-block super-scale broadcast + per-nibble unpack + f32x4 MAC accumulate + horizontal reduce) is mechanically identical.

A correct Q6_K wasm SIMD reimplementation at the same fidelity as Q4_K should land 2-3× speedup over the current scalar Rust path, and the scalar Rust path is itself ~3× faster than scalar JS (V8 vs LLVM codegen). Combined: ~6-9× over the current JS hot path.

4. Proposed fix (3 sentences)

In apps/worker-cobordism/src/index.ts handleCobordismLmHead, replace the per-row JS dequant + JS dot loop (lines 898-922) with a single mat_vec_q6k_dispatch wasm-bindgen call that consumes the cache's sliceBytes (already resident, 38 MiB) and writes directly into partialLogits. Hold a per-isolate wasm-linear-memory mirror of sliceBytes keyed by the same cacheKey, populated once per cache fill (adds ~30 ms one-time, then zero per request). Re-port the wasm32 SIMD branch of mat_vec_q6k_row in open-source/gnosis/distributed-inference/src/math.rs using the mat_vec_q4k_row wasm32 implementation at lines 549-629 as a 1:1 template, and gate the new path behind a unit test that compares byte-identical output against the existing scalar Rust path on a fixed seed.

5. Predicted TPS impact

Working from the wave-22 measured baseline (1.74 s warm hit, 4.75 s cold miss, ~75 % isolate-warm hit rate after settle):

Per-shard wall (parallel-fanned across 8)

The "dequant-once" line is a free secondary win available after wasm threading: instead of dequantizing all 32 768 rows fresh on every token (the model is autoregressive — same lm-head matrix, different residual each step), keep an F32 mirror after the first dequant and amortise the 176 M-elem dequant cost across N generated tokens. For an 80-token generation that is an 80× amortisation of the dequant share.

End-to-end TPS

Plug into the wave-23 post-everything formula (per_token = embed_wall + 60 × per_layer_warm + lm_head_wall):

Cobordism share of per-token cost drops from 8.5 % → ~0.9 % at the SIMD-plus-mirror tier. The mesh becomes truly layer-chain bound (>97 % of wall is the trisplit critical path), and the next attack surface becomes Death #4 (split-b parallel fan) or frame-coalescing.

Sensitivity to isolate hit rate

If the warm-hit rate stays at 75 % (per wave-22 bench), the steady-state average lm-head wall is 0.75 × 0.15 + 0.25 × cold_miss. Cold miss is fetch-bound (3 s R2 read) and the wasm port does not improve it. Effective average:

0.75 × 0.150 + 0.25 × 4.752 = 0.112 + 1.188 = 1.300 s

vs current 0.75 × 1.743 + 0.25 × 4.752 = 2.495 s from the wave-22 doc. Net steady-state lift = 1.195 s saved per token. On a 35 s baseline that is +3.4 % TPS at p75 hit rate, growing toward +5-7 % as warmup saturates the isolate pool.

6. Bule cost

6a. Caveat on the dequant-once F32 mirror

The proposal §4 mentions an F32 mirror; on closer math the full mirror is 32 768 × 5 376 × 4 = 704 MiB — far over the 128 MiB Workers ceiling and not viable as a flat array. Two viable substitutes:

1. Skip it. The wasm SIMD Q6_K dequant alone is fast enough; the matmul is the dominant share post-port and the dequant-on-the-fly pattern is fine.

2. Per-row LRU. Cache the most-recently dequantized rows (e.g. 4096 rows × 5376 elems × 4 B = ~88 MiB) keyed by vocab id. Useful only if token-id locality is high — for next-token sampling against a conversational prompt this is empirically true (top-K sampling visits ~50-200 unique vocab ids per session). Defer to a follow-up wave; not part of the 1.5 bule estimate.

The TPS table above shows numbers without the dequant-once mirror (the "Wasm SIMD Q6_K port" row, 0.29 s wall, +1.04× TPS). That is the honest steady-state lift of this proposal as scoped.

7. Risk register

8. Provenance

• Cobordism handler: ./apps/worker-cobordism/src/index.ts:731-956
• JS dequant module: ./apps/worker-cobordism/src/dequantize.ts
• Rust math kernel: ./open-source/gnosis/distributed-inference/src/math.rs:1171-1343 (Q6_K branches), :549-629 (Q4_K wasm SIMD template)
• Wave-21 cobordism bench: ./docs/cobordism-inference-bench-result.md
• Wave-22 cache deploy: ./docs/cobordism-ram-cache-deploy-result.md
• Wave-23 post-everything bench: ./docs/post-everything-tps-bench-result.md
• Wasm exports: ./open-source/gnosis/distributed-inference/src/wasm_bindings.rs