Inspiration

Agentic ReAct loops spend most of their token budget on low-entropy, strictly-formatted output — JSON braces, keys, function signatures — not on the free-form reasoning that actually needs a large model's judgment. Standard speculative decoding is supposed to help here, but its draft model still guesses against the full vocabulary distribution. On structured tool-call output it constantly proposes a plausible-looking but structurally wrong token, gets rejected by the target model, and the speedup evaporates. We kept seeing acceptance rates in the 30-40% range on exactly the kind of boilerplate that should be the easiest thing for a draft model to get right, and that gap — "this is the most predictable part of the output, why is it the least accepted part?" — is what we set out to close, specifically on Arm64 cloud hardware where every wasted rejection costs real SVE2 cycles.

What it does

Kleidi-Agent is a C++ inference engine that applies a single, identical grammar bitmask to both the draft and target model at every decoding step, not just the draft — the "Shared-Mask" technique. Each tool's JSON schema is compiled once into a finite state machine; that FSM produces the token bitmask; the same mask constrains both models simultaneously, which is what keeps rejection sampling mathematically sound instead of just biasing the draft and hoping. The mask itself is evaluated on custom Arm SVE2 intrinsics (svand, svsel), checking 128 tokens' legality per cycle, and every matmul is routed through Arm's own KleidiAI SVE2 micro-kernels. Because the draft model can no longer propose a structurally invalid token, the two models are forced into agreement on all boilerplate syntax — which is what drives boilerplate acceptance to 100% and blended acceptance to 56.6%, measured on the full 625-case BFCL multi-turn JSON split running natively on Arm64.

How we built it

The engine is split into four layers with a deliberate "novelty boundary": engine/grammar/ compiles JSON schemas to FSMs, engine/kernels/ holds our actual IP — the SVE2 masking intrinsics — and engine/matmul/ is a thin, clearly-labeled integration against Arm's KleidiAI, not something we claim credit for. engine/spec_decode/ ties it together into the shared draft/verify/accept loop. On top of that, bench/ runs the BFCL harness against a live engine instance and bench/performix_report.py turns the raw telemetry into the instruction-density/TTFT and throughput plots. CI runs on GitHub's ubuntu-24.04-arm hosted runner — real Neoverse silicon, not QEMU — gated by a verify-arm step that prints uname -m and the lscpu SVE2 flags into the log before anything is allowed to run, so hardware provenance is part of the transcript rather than a claim in a README. The whole pipeline collapses to one command, make bench: verify hardware, run the 625-case benchmark, generate the report.

Challenges we ran into

The hardest part of this project wasn't the SVE2 kernel — it was trusting our own numbers. While validating the benchmark pipeline we found and fixed four real measurement bugs, one of which had inflated reported throughput by 742x before we caught it. Rather than quietly patch that and move on, we kept a corrections log in BENCHMARKS.md documenting exactly what was wrong and how we found it, because a benchmark you can't independently verify isn't a benchmark. Separately, we hit a much more mundane but equally real bug: a bare Makefile line in .gitignore, meant to ignore the CMake-generated build/Makefile, was also silently matching and excluding our real hand-written root Makefile from every commit — so make bench, the exact command our own README told people to run, didn't exist in a fresh clone. We caught it, scoped the ignore rule to build/Makefile only, and made sure CI itself exercises the same Makefile path a judge would run, instead of a parallel hand-inlined workaround.

Accomplishments that we're proud of

100% boilerplate token acceptance across all 625 cases, driving a 56.6% blended acceptance rate against a typical 30-40% unconstrained baseline — measured on real Arm64 hardware, not projected. We're equally proud of the discipline around that number: every reported figure is traceable to a specific run, hardware is verified in-transcript rather than claimed, and every bug we found in our own pipeline is logged rather than silently fixed. We also went further than most speculative-decoding submissions in separating what's genuinely our contribution (the shared-mask SVE2 kernel and FSM compiler) from what we correctly built on top of (KleidiAI's matmul kernels) — that novelty boundary is explicit in the docs, not left for a judge to untangle.

What we learned

That the gap between "projected" and "measured" is where most of the real engineering lives. Our original target was 65-72% blended acceptance and a 2.6x speedup; what we actually measured is 56.6% and a speedup we haven't yet computed against a proper unconstrained baseline. Publishing the honest number, with the projection kept visible for context rather than quietly dropped, taught us more about where the remaining gap actually is — value-token acceptance at 47.9% is close to the unconstrained baseline and is clearly the next lever, not boilerplate, which is already saturated at 100%. We also learned that a benchmarking pipeline needs the same rigor as the kernel it's measuring: a stale telemetry file or a silently-ignored build file can undermine months of kernel work just as easily as a wrong intrinsic can.

What's next for Kleidi-Agent

Run a proper unconstrained-baseline benchmark on the same 625 cases so the end-to-end speedup claim is measured rather than left as "not yet computed." Close the value-token acceptance gap (47.9% today) toward the unconstrained ceiling, likely through tighter FSM state pruning on numeric/free-text fields. Extend grammar coverage beyond JSON-schema-first toward general context-free grammars, and add an SME2 execution path once cloud Arm instances actually expose it — the masking kernel is architected to make that a natural addition rather than a rewrite. Longer term: multi-schema batching, so a single engine instance can serve concurrent agents constrained against different tool schemas without recompiling the FSM per request.

Built With

  • arm-kleidiai
  • arm-sve2-intrinsics
  • aws-graviton
  • azure
  • bfcl
  • c++
  • cobalt
  • docker
  • fsm-based-grammar-constraints
  • github-actions
  • json-schema
  • performix
  • python
  • speculative-decoding
Share this project:

Updates