Feature request: d128 attention for 16 KV heads (32q/16kv GQA, ratio 2) on XDNA2

Hi FastFlowLM team — thanks for FLM, it's a genuinely excellent runtime. On Strix Halo / XDNA2 under Linux it gives us correct, fast decode on your shipped models with no elevated privileges.

We're trying to run a Llama-2-style decoder (internal name "Laneformer-2B") on the NPU through your d128 attention path. We've done the weight-side work we can validate from public artifacts and black-box runs, and now need your guidance on whether the remaining shape differences are covered by existing builds or need FastFlowLM-side artifacts. Please treat the notes below as outside observations — correct us anywhere we're inferring too much.

Why this matters to us (kog — a local-first product on Ryzen AI)

We're building kog, a local-first / on-device assistant whose model is Laneformer-2B and whose target runtime on Linux + Ryzen AI is FastFlowLM. We picked a Ryzen AI MAX+ 395 machine specifically to run this on the NPU; today we fall back to the iGPU, which rather defeats the purpose. KV16 attention support — plus confirmation about non-zoo vocab/MLP dimensions — is the missing piece that would let a real, shipping local-first product do its decode on your NPU stack, and we'd be glad to showcase FLM as the runtime that made it possible.

Environment

• AMD Ryzen AI MAX+ 395 (Strix Halo), XDNA2 NPU, /dev/accel/accel0
• Linux kernel 7.0.0-27-generic, amdxdna in-kernel driver
• FLM v0.9.43, runs without sudo

What already works here (your runtime is healthy on our box)

• Qwen3-1.7B-NPU2 — d128, GQA ratio 2 (16q/8kv), the closest shipped d128 ratio-2 comparator we see — decodes correctly ("…Paris").
• Llama-3.1-8B-NPU2 — d128 (ratio 4) — also decodes correctly here.
• We emit weights in your .q4nx / I8 on-disk layout and have matched them against your shipped models, so weight conversion is not the issue.

What we want to run

Laneformer-2B is a standard Llama-2-style decoder:

Its attention is GQA ratio 2 at head_dim 128 (after padding head_dim 96→128). The closest shipped comparator we see is Qwen3-1.7B's d128 ratio-2 config (16q/8kv); our model doubles both Q and KV head counts (32q/16kv).

Two things block us

1. 32q/16kv prefill doesn't complete for us. A 32q/16kv config deadlocks prefill on the shipped d128 xclbin even with a 1-token prompt (so it doesn't look like a context/KV-size overflow). We have no successful 16-KV run to point to — the only 16-KV observation is this deadlock — which makes us suspect the shipped attention tiles are built for a fixed KV-head count (8). The question that decides how big this is: is num_key_value_heads a build-time constant of the d128 attention kernel, or a runtime parameter? If it's a build constant, our guess (and it is only a guess, since we can't see past the deadlock) is that a d128 ratio-2 attention xclbin recompiled for 16 KV heads might drop straight in with little or no host-side change.

2. Decode also hangs when vocab_size or intermediate_size differ from the shipped dimensions. Using Llama-3.1-8B-NPU2 and changing one dimension at a time on our box: truncating vocab_size to 32000 hangs decode, while full 128256 decodes fine — and as a control, the same conversion/run pipeline round-trips cleanly at 128256, so the hang tracks the requested vocab_size rather than a broken setup. Lowering intermediate_size below the shipped value also hangs. We can pad intermediate_size inertly (no quality cost), but padding vocab_size 32000→128256 isn't viable: the added zero rows can win the argmax (→ garbage), and we don't see a way to mask them while keeping config.vocab_size equal to the weight rows the I8 layout requires.

The ask

1. Is num_key_value_heads runtime-parametric in the d128 attention path, or fixed at build? If fixed, could you publish a d128 attention xclbin built for 32q/16kv (ratio 2) — with native head_dim-96 32q/16kv as a secondary nice-to-have so we can drop the padding?
2. Are vocab_size / intermediate_size runtime-parametric for the decode path (lm_head / MLP), or would a non-zoo model like ours (vocab 32000, intermediate 12288) need builds for those dimensions?

We're happy to share the full config.json, a minimal flm run repro of the 1-token prefill deadlock, and the .q4nx weights privately — or to publish the model to a public Hugging Face repo if that fits your zoo-onboarding flow better.

Thanks for the great work on this runtime.

_{Prepared as a Mixture-of-Experts collaboration: drafted and cross-reviewed by Claude Opus 4.8 and Codex 5.5-xhigh.}

Start by reproducing the reported 1-token prefill deadlock with the 32q/16kv configuration and the decode hangs after changing vocab_size or intermediate_size. Inspect the shipped d128 xclbin and the runtime configuration path to determine which dimensions are build-time versus runtime parameters. Done means documenting the supported shapes and, if required, producing the requested 32q/16kv attention build or defining the needed artifacts.