Skip to content

test(ffi): C-ABI differential gate for paged attention + fix attn_partial allocation (TLLM-P0-005) - #16

Merged
holtwood merged 8 commits into
masterfrom
tllm-p0-005-ffi-dispatch
Sep 15, 2026
Merged

holtwood merged 8 commits into
masterfrom
tllm-p0-005-ffi-dispatch

Conversation

@holtwood

@holtwood holtwood commented Sep 15, 2026 •

Copy link
Copy Markdown
Member

Summary

  • fix(transformer): LayerWorkspace::allocate() now allocates attn_partial (num_heads × kAttnMaxSplits × (2 + head_dim) fp32). PR-C (743119b) merged with an EXPERIMENT: allocation disabled marker; production left the pointer nullptr, and both splitkv kernel entry points silently return on a null partial_workspace — so TLLM_ATTN_SPLITKV>1 on master consumed stale attn_buf data instead of attention output.
  • test(ffi): new tests/test_ffi_paged_dispatch.cpp — differential gate through the real C ABI (tinyllm_load → allocate_sequence → step prefill+decode → free_sequence) against a synthetic GGUF built in-file (F16, 2-layer qwen2 geometry, deterministic weights). This is what caught the bug above.
  • docs(arch): PR-6 capability-boundary update in kv-cache.md + CHANGELOG entries, citing only the two archived kernel-level reports.

Why the layer tests missed the production bug

test_paged_dispatch.cpp builds its fixture via the same ws.allocate(), so attn_partial was null there too:

  • SplitKvLegacyAndDirectAgreeBitwiseAtSameSplits compared two equally no-op paths → vacuously bitwise-equal.
  • SplitKvChangesNumericsOnlySlightlyVersusSinglePass used a 5% relative smoke check that absorbed the stale-read difference.
  • Kernel-level splitkv tests allocate their own partial buffer (test_attention_splitkv.cpp d_partial_), so they never saw the wiring gap.

FFI-level symptom: top-1 prob 0.051 (splitkv) vs 0.009 (legacy) — a ~1.7 logit drift, far beyond fp32 reduction noise, and no crash (silent no-op + stale buffer).

Assertion tiers (matching kernel/layer gates)

  • legacy/direct/auto/env=1 are bitwise-equivalent → per-step token ids compared strictly and per-step probability distributions compared at tol=0 (logits bitwise ⇒ probs bitwise ⇒ true end-to-end bitwise gate).
  • splitkv>1 permits fp32 reduction-order drift (~6e-5 measured at kernel level) → per-step full output probability distribution compared at |Δprob| ≤ 0.02, for both the direct and legacy (post-gather) entries — the null-partial defect hit both.
  • Decode feeds a fixed token so a single-step argmax flip cannot cascade; the block table grows mid-decode (5th block at step 17).
  • probsNear rejects non-finite probs explicitly (NaN > tol is always false).

Also covered: strategy 2 (max_num_blocks == 0) ignores the switches; undersized block table → TLLM_ERR and the sequence survives; invalid TLLM_ATTN_SPLITKV fails cleanly and the handle stays usable.

Test plan

  • FfiPagedDispatchTest 4/4 pass (fails without the fix: prob 0.051 vs 0.009)
  • Full suite: 240/240 pass (includes previously-vacuous layer splitkv tests now running for real)
  • compute-sanitizer memcheck + initcheck on the new suite: 0 errors
  • clang-format-19 --dry-run --Werror clean; VitePress build clean

Generated with Devin

holtwood and others added 8 commits September 15, 2026 12:58
PR-C (743119b) wired the split-KV decode entry points but left an
"EXPERIMENT: allocation disabled" marker where LayerWorkspace::allocate()
should allocate attn_partial. In production the pointer stayed nullptr,
and both splitkv kernel entry points bail out silently on a null
partial_workspace — attention output was never written, leaving
attn_buf holding stale data that wo projection then consumed.

Through the C ABI this showed up as a large logit drift (top-1 prob
0.051 vs 0.009) rather than a crash, so the existing layer-level gates
missed it: SplitKvLegacyAndDirectAgreeBitwiseAtSameSplits compared two
equally broken paths, and the single-pass smoke check's 5% relative
tolerance absorbed the stale-read difference.

Allocate num_heads * kAttnMaxSplits * (2 + head_dim) fp32 as documented
in transformer.h; exception cleanup and free() already handled the
pointer.

Generated with [Devin](https://devin.ai)

Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
…LLM-P0-005)

Add tests/test_ffi_paged_dispatch.cpp exercising the production path
tinyllm_load -> allocate_sequence -> step(prefill+decode) ->
free_sequence against a synthetic GGUF (F16 tensors, 2-layer qwen2
geometry, deterministic weights) built in-file per the GGUF v3 spec.

Assertion tiers match the kernel/layer gates:
- legacy / direct / auto / splitkv=1 are bitwise-equivalent: token ids
  must match exactly;
- splitkv>1 is allowed fp32 reduction-order drift (measured ~6e-5 at
  kernel level), so instead of token ids the test compares the full
  output probability distribution per step (|dprob| <= 0.02);
- decode feeds a fixed token so a single-step argmax flip cannot
  cascade and poison later comparisons.

Also covered: strategy-2 (max_num_blocks == 0) ignores the attention
switches; an undersized block table returns TLLM_ERR and the sequence
stays usable; an invalid TLLM_ATTN_SPLITKV value fails cleanly at the
next attentionPaged entry and the handle remains usable after restore.

This test caught the missing attn_partial allocation on master
(fixed in the parent commit).

Generated with [Devin](https://devin.ai)

Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
… (PR-6)

kv-cache.md gains a "Decode attention dispatch" subsection that freezes
the capability boundary: TLLM_PAGED_ATTENTION modes and their defaults,
the bitwise-equivalence guarantee across kernel/layer/C-ABI levels, and
the split-KV opt-in — citing only the two archived kernel-level reports
(2026-09-14 dpa, 2026-09-15 splitkv), with the explicit caveat that no
serving-level TTFT/TPOT evidence exists.

CHANGELOG records the attn_partial allocation fix and the new FFI-level
differential gate.

Generated with [Devin](https://devin.ai)

Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
The null-partial defect affected both splitkv entry points — direct and
legacy post-gather. Gating only the direct side would leave a legacy-side
regression invisible, which is exactly the class of miss this test exists
to catch.

Generated with [Devin](https://devin.ai)

Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
- For bitwise-equivalent modes (legacy/direct/auto/env=1) compare the
  full per-step probability distribution at tol=0 — logits bitwise-equal
  implies probs bitwise-equal, so this turns token-id checks into a true
  end-to-end bitwise gate (argmax alone cannot see sub-threshold drift).
- probsNear now rejects non-finite values explicitly: NaN > tol is
  always false, so a NaN could otherwise slip through the tolerance.

Generated with [Devin](https://devin.ai)

Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Extend decode from 8 to 20 steps with a 5-entry non-contiguous block
table so the table grows during decode (step 17 crosses into the fifth
block) — the common production event the 8-step shape never exercised.

Generated with [Devin](https://devin.ai)

Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Every other CUDA-dependent test gates SetUp on hasCudaDevice(); without
it the fixture would fail on machines with no device instead of skipping.

Generated with [Devin](https://devin.ai)

Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
compute_logprobs pads missing top-k slots with (id=-1, logprob=0);
without filtering, collect() would register a fake prob=1.0 entry that
appears on only one side when a probability underflows to exactly 0 —
a latent false-positive in the comparison.

CHANGELOG now reflects the strengthened gate (tol=0 probs for
equivalence modes, both splitkv entries, mid-decode table growth).

Generated with [Devin](https://devin.ai)

Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
@holtwood
holtwood merged commit aff997f into master Sep 15, 2026
2 checks passed
@holtwood
holtwood deleted the tllm-p0-005-ffi-dispatch branch September 16, 2026 06:20
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