Skip to content

Npu support mm bert - #306

Open
zhaixuejun1993 wants to merge 14 commits into
ravi9:dev_backend_openvinofrom
zhaixuejun1993:npu-support-mmBERT
Open

Npu support mm bert#306
zhaixuejun1993 wants to merge 14 commits into
ravi9:dev_backend_openvinofrom
zhaixuejun1993:npu-support-mmBERT

Conversation

@zhaixuejun1993

Copy link
Copy Markdown
Collaborator

openvino: support cacheless encoder models on NPU

  1. Packed QKV views used by mmBERT were rejected by the ROPE support check. This split Q/K RoPE onto CPU, prevented cacheless attention detection, and sent fragmented encoder graphs through the decoder-oriented NPUW path.
  2. Accept packed QKV RoPE views, detect cacheless attention from its mask, and run these models as a single full-sequence prefill without NPUW or a decode graph. Also provide static mask, output index, and mean-pooling shapes and inputs.

openvino: optimize norm and RoPE translation

  1. Replace the decomposed mean/variance normalization graph with an opset6 MVN operation. This preserves the GGML epsilon placement while allowing OpenVINO plugins to compile normalization as one operation with fewer intermediate tensors.
  2. Cache RoPE sine and cosine outputs in the graph-wide tensor map. Build the cache key from all RoPE parameters and the optional frequency-factor input so compatible Q/K and layer nodes share one subgraph without mixing different RoPE configurations.

mostafafaheem and others added 12 commits September 3, 2026 05:21
…state

The stateful path seeds its KV state from ggml's cache when the decode position
is ahead of what the state holds. That only works when ggml's cache is a plain
prefix, where cell i holds position i. A sliding-window layer keeps just the last
n_swa positions and drops the rest, so past the window cell i no longer holds
position i and the seeded state is wrong.

Slicing the state to the decode position also had no bounds check, so a position
past the end surfaced as a bare ov::Exception from the ROI constructor
(llama_decode ret = -3, with no reason given at default verbosity).

Refuse both cases with a clear message instead, and refuse on the compile path
too, where a new model starts with an empty state and so can only serve a
sequence from its beginning. Reproducible with llama-bench -d, which restores a
saved sequence state rather than recomputing the depth prefill.

Assisted-by: Claude Opus 5
The stateful path reinterprets ggml's KV buffer [1, 1, seq, n_heads_kv * head_size]
as [1, seq, n_heads_kv, head_size]. The head size is already taken from the
tensor's own combined dim, because gemma-4 varies it per layer type, but the head
count still came from a model-level scalar that compute_llm_params() overwrites
per attention node, so it ended up holding whatever the last layer said.

gemma-4 varies the head count per layer too: 12B has 8 x 256 sliding layers and
1 x 512 full layers, 31B has 16 x 256 and 4 x 512. So 40 of 12B's 48 layers were
split as 1 x 2048 instead of 8 x 256, and attention read the state with the wrong
head split - both models decoded garbage on CPU and GPU. E2B is unaffected, its
head count is 1 everywhere.

Record the count per layer instead and look it up by the cache_k_l<N> leaf name.
Key it by layer, not by layer type: the sliding/full classification comes from
cache extents, which tie at a small -c, while the head count does not.

The stateful state trim now derives its sequence axis per state for the same
reason, since pass::KVStateSeqAxis matches per state on the head count.

Assisted-by: Claude Opus 5
pass::KVStateSeqAxis was limited to states with a single KV head, where moving
the sequence axis from dim 1 to dim 2 is a pure metadata change. The limit was
also based on a measurement showing no gain for a multi-head model, but that was
taken at depth 0, which is the one depth where this change does nothing.

With several heads the pass does more than move metadata: it drops the reader
side transpose of the whole accumulated state, which the graph otherwise redoes
every token at a cost that grows with the context length, and replaces it with a
transpose of the single new row. Measured on GPU, tg128, alternating arms:
gemma-4-12B 6.27 -> 9.11 t/s at depth 8192 (stateless is 7.69, so stateful now
wins at depth instead of losing), Llama-3.2-1B 47.8 -> 59.6 t/s. Both are within
noise at depth 0, which is why the earlier check saw nothing.

The state refill needs the rows copied rather than reinterpreted now: ggml stores
[seq][n_heads_kv * head_size], and a relayout state with several heads is a
different element order. Without that, a refill would seed wrong data - it is
reachable today through llama-bench -d.

Assisted-by: Claude Opus 5
    Packed QKV views used by mmBERT were rejected by the ROPE support check. This split Q/K RoPE onto CPU, prevented cacheless attention detection, and sent fragmented encoder graphs through the decoder-oriented NPUW path.

    Accept packed QKV RoPE views, detect cacheless attention from its mask, and run these models as a single full-sequence prefill without NPUW or a decode graph. Also provide static mask, output index, and mean-pooling shapes and inputs.
    Replace the decomposed mean/variance normalization graph with an opset6 MVN operation. This preserves the GGML epsilon placement while allowing OpenVINO plugins to compile normalization as one operation with fewer intermediate tensors.

    Cache RoPE sine and cosine outputs in the graph-wide tensor map. Build the cache key from all RoPE parameters and the optional frequency-factor input so compatible Q/K and layer nodes share one subgraph without mixing different RoPE configurations.

    Expose NodeContext::put_shared() to publish translator-created outputs for graph-level reuse.
@ravi9
ravi9 force-pushed the dev_backend_openvino branch from fd9bc04 to 33237ab Compare September 3, 2026 21:10
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.

3 participants