Repository navigation
Add explainer: Support dynamic language models on NPUs via WebNN - #969
wangw-1991 wants to merge 1 commit into
Conversation
This PR adds an explainer for running dynamic language models on NPUs through WebNN. It is based on our proof of concept on Intel NPU. The explainer describes the current limitations of WebNN in supporting dynamic language models, along with the proposals our proof-of-concept relies on to address them. Specifically, we introduce two new proposals: the stateful model and build-time metadata. Open questions and considered alternatives are called out explicitly in the explainer and feedback on this explainer would be very welcome.
|
@reillyeon @fdwr @huningxin @matthargett @sushraja-msft Please refer to this explainer for more detailed information. Feedback is welcome, thanks. |
mklimenko-nv
left a comment
mklimenko-nv
left a comment
There was a problem hiding this comment.
Thank you for the explainer, it is a valuable account of what Intel NPUs need to run LLMs through WebNN.
That said, much of the justification follows one implementation, OpenVINO's NPUW plugin. The metadata fields mirror NPUW_LLM_MAX_PROMPT_LEN and NPUW_LLM_MIN_RESPONSE_LEN, and the rejected alternatives are argued from how NPUW manages state internally. Even within OpenVINO, the GPU plugin runs the same stateful models without that decomposition or those properties.
Before this hardens into API shape, I would like it balanced against other stacks: GPU runtimes such as TensorRT-RTX, which exposes a fixed-capacity KV cache update with aliased input and output and keeps no state in the engine, and other accelerator stacks such as Core ML on the ANE (stateful models with fixed-shape state buffers), QNN for Snapdragon X's NPU, or even TensorRT on DLA, which is fully static-shape.
| ### Growing decode KV cache | ||
| After prefill, decode emits one token per step. Each step computes keys and values for a single position and attends over every position seen so far, so the cache grows by one entry per step, per layer. For a stateless model, that cache is an input and an output of every dispatch. With `MLTensor`, the page can keep the cache in device memory across dispatches, passing a handle rather than bytes, so the cost is not a transfer between the page and the device. It is that the graph still produces the whole cache on every step. Three costs follow: | ||
|
|
||
| - **Copying.** The compute of a single decode step is tiny, but the graph reads the whole old cache and writes the whole updated cache into another tensor, because a tensor cannot be both an input and an output of one dispatch. The traffic is proportional to the whole conversation, so the step time grows linearly with context length for reasons that have nothing to do with the model. |
There was a problem hiding this comment.
"because a tensor cannot be both an input and an output of one dispatch"
This is a WebNN rule, not a property of the backends, that could potentially be revised if the working group agrees. This would allow to pre-allocate the KV-cache tensor in advance an modify it in-place, like it's done in TensorRT's IKVCacheUpdateLayer.
It'd be also good to consider whether the cache has to leave the graph interface for the backends that have no state of their own.
|
|
||
| - **Copying.** The compute of a single decode step is tiny, but the graph reads the whole old cache and writes the whole updated cache into another tensor, because a tensor cannot be both an input and an output of one dispatch. The traffic is proportional to the whole conversation, so the step time grows linearly with context length for reasons that have nothing to do with the model. | ||
| - **Allocation and concatenation.** `newCache = concat(oldCache, newEntry)` needs a fresh buffer every step and copies the old contents into it — again linear in the context, and it doubles peak memory at the moment of the copy. Reusing a pair of tensors avoids the allocation but not the copy (see [Ping-pong MLTensors between dispatches](#ping-pong-mltensors-between-dispatches)). | ||
| - **Shape churn.** The cache length changes on every step, so an underlying implementation that specializes per shape has to re-specialize each time. |
There was a problem hiding this comment.
This is stronger on TensorRT-RTX than the text suggests. Kernels are specialized per runtime shape at first use, as well as CUDA graph capture is per shape. With the fixed-capacity form every decode step has the same shape, so it is specialized once and can run from a captured CUDA graph. With the concat form past_len changes every step, so each step is a new shape, resulting in an inference latency increase.
| for (let i = 0; i < numLayers; i++) { | ||
| // Ordinary inputs: the cache as of the previous dispatch. | ||
| const pastK = builder.input(`past_key_${i}`, kvDesc); | ||
| const pastV = builder.input(`past_value_${i}`, kvDesc); |
There was a problem hiding this comment.
The example and the validation rule take present = concat(past, new) with past_len and total_len as the normal form. TensorRT-RTX executes a different form: a cache of static length s_max ([b, kvHeads, s_max, headDim], same shape on input and output), an update of [b, kvHeads, seq_len, headDim], and a writeIndices tensor of shape [b], with attention reading the valid prefix through per-batch key/value lengths
The growing-concat form is the ONNX export convention, not what any of these backends run, and I'm not convinced that this is the best option memory- and performance-wise.
|
|
||
| No interface returns a cache's contents or shape, so the KV cache is not readable from script. This is a privacy property as much as a performance one; see [Privacy & Security Considerations](#privacy--security-considerations). | ||
|
|
||
| ### Build-time metadata |
There was a problem hiding this comment.
For TensorRT-RTX, we'd need range of values for the dynamic shapes. I agree with @reillyeon's point in #965 that this is better represented as per-dimension bounds.
|
|
||
| Both can fail on a long context: with metadata the failure is a resource limit of the underlying implementation, as described in [Metadata design](#metadata-design), and with a bound it is part of the contract. We did not choose bounds because they promise more than we want to require, and because they carry less than the page knows. | ||
|
|
||
| - **A bound is a guarantee in both directions.** Every underlying implementation must accept every shape within it and reject every shape beyond it, and one that cannot support the bound must reject the build up front. That makes a page's stated size a requirement on every implementation, and it forbids an implementation that has spare capacity from accepting a longer context. `maxPromptLength` only asks for at least that much room: an implementation may size from it, from its own defaults, or grow on demand, and may ignore it. |
There was a problem hiding this comment.
"it forbids an implementation that has spare capacity from accepting a longer context" is not a cost on TensorRT-RTX, Core ML RangeDim or OpenVINO, since all of them already reject shapes beyond the compiled maximum.
This bullet can be dropped or restated as "a bound is validated at dispatch".
|
|
||
| 1. **WebNN's backend adapts the pairs.** Each underlying implementation has its own requirements, so WebNN's backend for it checks the graph against those requirements as thoroughly as it can and rewrites the pairs into the form that underlying implementation expects. A failed backend check does not fail the build: the backend builds the graph stateless and reports `isStateful` as false. The more the backend checks up front, the fewer stateful builds reach the underlying implementation only to fail there. | ||
|
|
||
| 2. **The underlying implementation converts the pairs.** WebNN's backend passes the adapted graph and its pairs down, and only the underlying implementation turns each pair into internal state, the `ReadValue` and `AssignValue` of the figure in [Stateful model](#stateful-model). If the conversion or compilation still fails in the underlying implementation, the failure can be handled like any other compilation failure: [`MLGraphBuilder.build()`](https://www.w3.org/TR/webnn/#api-mlgraphbuilder-build) rejects with an `OperationError`. |
There was a problem hiding this comment.
On a GPU runtime there is no "underlying implementation converts the pairs" step. TensorRT-RTX gives the backend an in-place update layer and nothing else. The buffers, the counter, resetKVCache(), zero-fill and the capacity check are all WebNN-backend code, which is rejected in this explained "to keep WebNN thin".
I'd propose to rewrite the mechanism to allow WebNN implementation to own this and not give everything to the backend.
| #### ONNX Runtime | ||
| The ONNX format has no notion of state, so whether a pair can become state depends on the execution provider (EP). The EPs below are examples, not a complete list: | ||
|
|
||
| - **OpenVINO EP.** With the `enable_causallm` provider option, the EP converts the model to a stateful one with OpenVINO's `MakeStateful` transformation, so the cache becomes `ReadValue`/`Assign` state and, on the NPU, is managed by [NPUW](#npuw) as described in [KV cache management](#kv-cache-management). The EP does not receive the pairs explicitly. It finds them by naming convention, matching an output named like `present_key_0` and pairing it with the corresponding `past_key_0` input. WebNN's ONNX Runtime backend therefore has to make sure that, when `useStateful` is true, the model it generates meets this naming requirement; otherwise [`MLGraphBuilder.build()`](https://www.w3.org/TR/webnn/#api-mlgraphbuilder-build) fails. On the NPU, the metadata members map onto [NPUW](#npuw)'s options: `maxPromptLength` onto `NPUW_LLM_MAX_PROMPT_LEN` and `minResponseLength` onto `NPUW_LLM_MIN_RESPONSE_LEN`. |
There was a problem hiding this comment.
How is it handled for CPU and GPU devices on OpenVINO EP? Could you please list it here for the sake of completeness?
|
|
||
| - **OpenVINO EP.** With the `enable_causallm` provider option, the EP converts the model to a stateful one with OpenVINO's `MakeStateful` transformation, so the cache becomes `ReadValue`/`Assign` state and, on the NPU, is managed by [NPUW](#npuw) as described in [KV cache management](#kv-cache-management). The EP does not receive the pairs explicitly. It finds them by naming convention, matching an output named like `present_key_0` and pairing it with the corresponding `past_key_0` input. WebNN's ONNX Runtime backend therefore has to make sure that, when `useStateful` is true, the model it generates meets this naming requirement; otherwise [`MLGraphBuilder.build()`](https://www.w3.org/TR/webnn/#api-mlgraphbuilder-build) fails. On the NPU, the metadata members map onto [NPUW](#npuw)'s options: `maxPromptLength` onto `NPUW_LLM_MAX_PROMPT_LEN` and `minResponseLength` onto `NPUW_LLM_MIN_RESPONSE_LEN`. | ||
| - **WebGPU EP and CPU EP.** These EPs have no state support, so WebNN's ONNX Runtime backend builds the graph stateless, `isStateful` is false, and the app carries the cache. | ||
|
|
There was a problem hiding this comment.
I'd add TensorRT-RTX EP to give the more balanced picture of what's going on and how it's implemented:
| **TensorRT-RTX (NvTensorRTRTX EP, or direct)**. TensorRT-RTX has no model state, buffers are owned by the caller. Since 1.3 it has `IKVCacheUpdateLayer`: a cache of fixed capacity `[b, kvHeads, s_max, headDim]` is a network input and a network output bound to the same device address, and each run writes `[b, kvHeads, seq_len, headDim]` at a per-batch `writeIndices` position. WebNN's backend owns the buffers, the position counter and the reset, rewrites each pair into this form, and sizes `s_max` and the optimization profile from the page's bounds. Without them it has to choose both at build time. Dynamic dimensions need a `MIN`/`OPT`/`MAX` profile, and kernels are specialized per shape at first use, so a fixed-capacity decode step has one shape and can be captured in a CUDA graph. |
|
|
||
| - **The KV cache must be initialized, never recycled.** Cache memory is zero-initialized on allocation and on `resetKVCache()`. A new allocation must never expose a previous allocation's contents — the ordinary uninitialized-memory requirement, applied to a buffer that may hold the previous turn of a conversation. | ||
|
|
||
| - **The KV cache is not readable.** No API returns a cache's contents or shape, so a compromised renderer gains no new read primitive. |
There was a problem hiding this comment.
The privacy benefit is small for the KV cache: the renderer already holds the conversation the cache encodes, and the cache is derived from inputs the page supplied. It should not be the argument that decides between hidden state and page-visible aliased tensors, since the latter keeps save/restore and multi-conversation use (lines 574 to 581) without a new analysis.
| ## Future Considerations | ||
| The proposed API covers the two use cases with the smallest surface we could find. The extensions below are not needed for them, but the design should leave room for each: | ||
|
|
||
| - **Several states per graph.** A page may want several independent conversations on one model, such as one per open document, without building the graph once per conversation. A separate state object, described in [A separate state object instead of graph-bound state](#a-separate-state-object-instead-of-graph-bound-state), would allow it once underlying implementations can hold several states for one compiled model. |
There was a problem hiding this comment.
"once underlying implementations can hold several states for one compiled model" already holds for TensorRT-RTX: one engine serves many execution contexts and any number of caller-owned buffers, and writeIndices is per batch element.
The restriction is NPUW's and LiteRT's.
This PR adds an explainer for running dynamic language models on NPUs through WebNN. It is based on our proof of concept on Intel NPU.
The explainer describes the current limitations of WebNN in supporting dynamic language models, along with the proposals our proof-of-concept relies on to address them. Specifically, we introduce two new proposals: the stateful model and build-time metadata.
Open questions and considered alternatives are called out explicitly in the explainer and feedback on this explainer would be very welcome.