Skip to content

npu: LFM2's gated short convolution as host arithmetic, with known answers — first slice of serving the lfm2 packs (Refs #1786) - #2161

Merged
bong-water-water-bong merged 1 commit into
mainfrom
cl-flm/lfm2-shortconv-reference
Nov 11, 2026
Merged

bong-water-water-bong merged 1 commit into
mainfrom
cl-flm/lfm2-shortconv-reference

Conversation

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator

Refs #1786 (backlog). First, device-free slice of serving the lfm2 family; the design note for the whole thing is on the issue. Nothing in the forward changes: #2158's refusal stays the behaviour until the rest lands.

What it adds

npu/forward/include/lfm2_shortconv.h: the arithmetic between an LFM2 conv layer's two projections, and the state it carries.

bcx = in_proj(norm(x))      H -> 3H: B | C | X          (device GEMM, not here)
bx  = B * X
conv[h] = sum_j taps[h][j] * seq[t + j][h]              seq = the last L-1 values of bx, then this token's
y   = C * conv              -> out_proj                 (device GEMM, not here)
  • ShortConvState: (L-1) * H floats per conv layer, zeros at position 0.
  • shortconv_step(): one token. shortconv_rows(): M tokens, defined as M steps in order, so a batched prefill cannot disagree with decode.

The definition is the graph our own llama.cpp fork builds for this family (src/models/lfm2.cpp, build_shortconv_block): chunk order B, C, X; the present token meets the last tap; the gate multiplies the convolution's output. No vendor code was read for it.

Why a header before the forward

Three places will need this (token-by-token decode, batched prefill, and the state snapshot the prefix cache takes), and the failure it guards against is the quiet one: a prefill that advances the state differently from decode is wrong at every later position and looks like a weak model. Writing it once, with the equalities asserted, is cheaper than finding that on the device.

Tests

tests/lfm2_shortconv_test.cpp (ctest lfm2_shortconv, a build dependency of onebit so CI builds it), no model and no device, for kernel lengths 1 to 4, H 37, 23 tokens with zeros and both signs:

  • token by token equals the causal depthwise convolution written out over the whole sequence, bit for bit;
  • one batch of all tokens gives the same bits and leaves the same state;
  • uneven chunks (5, 1, 9) followed by single steps give the same bits and state;
  • a copy of the state at position 11 resumes to the same bits;
  • the state is the last L-1 values of B*X, oldest first;
  • a three-tap, one-channel case by hand: the present token meets the last tap, and C gates the result;
  • an unsized state, null buffers and a negative row count are refused.

lfm2_shortconv: ok. Two mutants from a copy of the header: the present token on the first tap fails 15 checks; the gate dropped fails 18.

tools/copyright.py --check rc 0; ci_targets_check.py --check baseline 9, none new.

What this does not show

That the definition matches the weights in FastFlowLM's packs. It matches our fork's graph; whether the container's shortconv.conv.weight is [H][L] with the last tap on the present token is settled only by running the model and comparing its text with FastFlowLM's runtime on the same file, which is step (c) of the note on the issue.

🤖 Generated with Claude Code

…swers (Refs #1786)

npu/forward/include/lfm2_shortconv.h: the block between an LFM2 conv
layer's two projections (B*X, a causal depthwise convolution over the
last L-1 values, gated by C), with its per-layer state, defined once for
decode, batched prefill and state snapshots. Not wired into the forward
yet; this is the reference the wiring will be checked against.

Co-Authored-By: Claude Fable 5.2 <noreply@anthropic.com>
@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

ready for review at 7eb633c

@github-actions

Copy link
Copy Markdown

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🎫 Ticket compliance analysis 🔶

1786 - Partially compliant

Compliant requirements:

  • it imports with zero gaps and its self-check passes against a named reference
  • the row names the kernel set and the engine commit that produced it

Non-compliant requirements:

  • it serves on --device npu
  • it answers a known-answer prompt with the model's own chat template, 3 identical requests returning identical text
  • decode tok/s and time-to-first-token measured under npu-hold

Requires further human verification:

  • it serves on --device npu
  • it answers a known-answer prompt with the model's own chat template, 3 identical requests returning identical text
  • decode tok/s and time-to-first-token measured under npu-hold
⏱️ Estimated effort to review: 3 🔵🔵🔵⚪⚪
🧪 PR contains tests
🔒 No security concerns identified
⚡ Recommended focus areas for review

State Management in shortconv_step

The shortconv_step function modifies the ShortConvState history by sliding the window. The implementation uses std::memmove for the sliding operation when L > 2, but for L == 2, it directly assigns the new value to the last element. This could lead to inconsistent behavior or potential memory issues if not carefully handled, especially in a performance-critical environment like an NPU inference engine.

if (L > 1) {
    if (L > 2) std::memmove(st.history.data(), st.history.data() + H, size_t(L - 2) * size_t(H) * sizeof(float));
    float* last = st.history.data() + size_t(L - 2) * size_t(H);
    for (int h = 0; h < H; ++h) last[h] = B[h] * X[h];
}
return true;
Test Coverage for Edge Cases

The test suite covers various kernel lengths and hidden sizes, but does not explicitly test edge cases such as very large kernel lengths or very small hidden sizes. These edge cases could reveal issues in memory handling or performance that are not caught by the current tests.

{   // the present token meets the LAST tap, and the gate C multiplies the result, not the input
    ShortConvState st;
    st.reset(1, 3);
    const float taps[3] = {100.0f, 10.0f, 1.0f};
    float y = 0.0f;
    const float t0[3] = {2.0f, 5.0f, 3.0f};   // B, C, X -> bx 6
    check(shortconv_step(st, t0, taps, &y) && y == 5.0f * (1.0f * 6.0f), "at position 0 only the last tap sees anything");
    const float t1[3] = {1.0f, 2.0f, 4.0f};   // bx 4
    check(shortconv_step(st, t1, taps, &y) && y == 2.0f * (10.0f * 6.0f + 1.0f * 4.0f), "the token before meets the middle tap");
    const float t2[3] = {1.0f, 1.0f, 1.0f};   // bx 1
    check(shortconv_step(st, t2, taps, &y) && y == 1.0f * (100.0f * 6.0f + 10.0f * 4.0f + 1.0f * 1.0f),
          "and the one before that the first tap");
}
{   // refusals: an unsized state, null buffers, a negative row count
    ShortConvState st;
    float b[3] = {1, 1, 1}, k[1] = {1}, y = 7.0f;
    check(!shortconv_step(st, b, k, &y) && y == 7.0f, "an unsized state is refused and nothing is written");
    st.reset(1, 1);
    check(!shortconv_step(st, nullptr, k, &y) && !shortconv_step(st, b, nullptr, &y) && !shortconv_step(st, b, k, nullptr),
          "null buffers are refused");
    check(!shortconv_rows(st, b, -1, k, &y) && shortconv_rows(st, b, 0, k, &y), "a negative row count is refused; zero rows is a no-op");
}

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VERDICT at 7eb633c: READY

Reviewed head 7eb633c (refs #1786). This is the first verdict.

The definition matches the fork's graph. I read src/models/lfm2.cpp's build_shortconv_block at the engine's pin (204493d40) next to lfm2_shortconv.h:

  • bcx is viewed as B | C | X (chunk offsets 0, 1, 2), and the header reads B = bcx, C = bcx + H, X = bcx + 2H.
  • bx = B ⊙ X.
  • In the causal case the state is concatenated before bx, ggml_ssm_conv runs over [state | bx] (so the present token meets the last tap, k[L-1] * bx), and the new state is the last d_conv columns. The header's oldest-first history with the newest at L-2 is the same thing.
  • y = C ⊙ conv, before out_proj.
  • The per-channel kernel taps[h·L + j] matches ggml's conv weight (L taps contiguous per channel).

The fork also has a non-causal branch (symmetric padding) that the header does not implement. That is fine for LFM2's causal text models, and worth one line in the header so nobody reuses it for a bidirectional encoder.

Checked

  • git merge-tree onto main: no conflicts. lfm2_shortconv_test, built from this head: "lfm2_shortconv: ok".
  • It is a build dependency of onebit (CMakeLists.txt:1382), and CI runs it: #183: lfm2_shortconv … Passed, 231/231. CI build, hrx_build and CodeQL green.
  • Mutants of the header, beyond the body's two, each rebuilt and run:
mutant result
B and C swapped rc 1, 21 failing checks
the newest B*X written to the oldest history slot rc 1, 12 failing checks
the state never advanced rc 1, 17 failing checks

I restored the file (git status clean).

Nothing in the forward changes, so #2158's load refusal stays the behaviour, as the body says.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant