wehub-resource-sync
212 lines
12 KiB
Markdown
212 lines
12 KiB
Markdown
# bertQKVToContextPlugin [DEPRECATED]
|
|
**This plugin has been deprecated since TensorRT 10.15 and will be removed in a future release. Use `IAttention` (via `INetworkDefinition::addAttention` / `addAttentionV2`) as the out-of-the-box replacement. See [Migration to IAttention](#migration-to-iattention) for the steps to adapt.**
|
|
|
|
**Table Of Contents**
|
|
- [Description](#description)
|
|
* [Structure](#structure)
|
|
- [Migration to IAttention](#migration-to-iattention)
|
|
- [Parameters](#parameters)
|
|
- [Additional resources](#additional-resources)
|
|
- [License](#license)
|
|
- [Changelog](#changelog)
|
|
- [Known issues](#known-issues)
|
|
|
|
|
|
## Description
|
|
|
|
Takes query, key and value tensors and computes scaled multi-head attention - computes scaled dot product attention scores `softmax(K'Q/sqrt(HeadSize))` and returns values weighted by these attention scores.
|
|
|
|
|
|
|
|
### Structure
|
|
|
|
The `bertQKVToContextPlugin` takes two inputs; `input`, and optionally `input_mask`.
|
|
|
|
`input`
|
|
input is a tensor with shape `[S, B, 3 * E, 1, 1]` where `B` is the batch size and `E` is the hidden size. The input has two trailing dimensions in order to generate an output with the same trailing dimensions for a FC optimiazation further down the network.
|
|
This plugin makes strong assumptions about its input:
|
|
- The input tensor contains all 3 matrices Q, K, V
|
|
- This input tensor is computed by multiplying a tensor of size `[S, B, E]` with the weights `W_qkv` of size `[E, 3 * E]`
|
|
- The weight matrix W_qkv is NOT just the vertical concatenation of individual matrices `W_tmp = [W_q', W_k', W_v']'`, but to start with `W_tmp`, reshaping it into `[E, 3, N, H]` (where `N * H = E` and `N` is number of heads, `H` is head size) transposing it into `[E, N, 3, H]` and reshaping it back to `[E, 3 * E]`. The interpretation is to layout the k-th heads of Q, K and V next to each other, instead of first all N heads of Q, then all N heads of K, then all heads of V
|
|
|
|
`input_mask`
|
|
The input mask is encoded in the `maskIdx` format described in `embLayerNormPlugin`, and contains the number of valid elements from the start of the sequence.
|
|
If provided, the attention scores, i.e. the softmax distribution, are only computed over the elements designated as valid by the input mask
|
|
|
|
|
|
The `bertQKVToContextPlugin` generates the following output:
|
|
|
|
`output`
|
|
output is a tensor with shape `[S, B, E, 1, 1]` where `B` is the batch size.
|
|
|
|
|
|
## Migration to [IAttention](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/operators/Attention.html)
|
|
|
|
`IAttention` is the out-of-the-box attention layer in TensorRT and is functionally equivalent to this plugin for the multi-head attention computation, but it is not a drop-in replacement. Adapter steps:
|
|
|
|
1. **Split the fused QKV input.** This plugin consumes one `[S, B, 3*E, 1, 1]` tensor with the K-th heads of Q, K, V interleaved. `IAttention` takes three separate tensors with shape `[B, N, S, H]`. Reshape the fused QKV to `[S, B, N, 3, H]`, slice along the size-3 axis, then transpose each slice to `[B, N, S, H]`. If the upstream FC weight `W_qkv` was packed using the head-interleaved layout described in the Structure section, no weight repacking is needed. Otherwise, repack `W_qkv` to produce three contiguous `[B, N, S, H]` tensors directly.
|
|
|
|
```cpp
|
|
// [S, B, 3*E, 1, 1] -> [B, N, S, 3, H], then slice axis 3.
|
|
auto* shuf = network->addShuffle(*fusedQKV);
|
|
shuf->setReshapeDimensions(Dims{5, {S, B, N, 3, H}});
|
|
shuf->setSecondTranspose(Permutation{{1, 2, 0, 3, 4}});
|
|
ITensor* t = shuf->getOutput(0);
|
|
|
|
auto pick = [&](int32_t i) -> ITensor* {
|
|
auto* sl = network->addSlice(*t, Dims{5, {0,0,0,i,0}},
|
|
Dims{5, {B,N,S,1,H}}, Dims{5, {1,1,1,1,1}});
|
|
auto* sq = network->addShuffle(*sl->getOutput(0));
|
|
sq->setReshapeDimensions(Dims{4, {B, N, S, H}});
|
|
return sq->getOutput(0);
|
|
};
|
|
ITensor *q = pick(0), *k = pick(1), *v = pick(2);
|
|
```
|
|
|
|
2. **Convert the mask.** This plugin uses the `maskIdx` format (count of valid tokens from the start of each sequence). `IAttention` takes a 4D boolean or additive mask of shape `[B, N, S_q, S_kv]` (broadcastable over `B` and `N`). For non-causal BERT-style self-attention, materialize a boolean mask from the valid-token counts. For causal generation, use `IAttention::setCausalKind(CausalMaskKind)` instead of a mask.
|
|
|
|
```cpp
|
|
// Boolean mask from maskIdx ([B] int32). Broadcasts to [B, NQ, S_q, S_kv].
|
|
auto* idx = network->addFill(Dims{1, {S}}, FillOperation::kLINSPACE, DataType::kINT32);
|
|
idx->setAlpha(0.0); idx->setBeta(1.0);
|
|
auto* less = network->addElementWise(*idx->getOutput(0), *maskIdx,
|
|
ElementWiseOperation::kLESS);
|
|
attn->setMask(*less->getOutput(0));
|
|
|
|
// For causal autoregressive decode, skip the mask and use:
|
|
// attn->setCausalKind(CausalMaskKind::kLOWER_RIGHT);
|
|
```
|
|
|
|
3. **Reshape the output.** This plugin produces `[S, B, E, 1, 1]`. `IAttention` produces `[B, N, S_q, H]`. Transpose and reshape as needed for downstream layers.
|
|
|
|
```cpp
|
|
// [B, N, S_q, H] -> [S, B, E, 1, 1] for downstream layers expecting the plugin layout.
|
|
auto* out = network->addShuffle(*attn->getOutput(0));
|
|
out->setFirstTranspose(Permutation{{2, 0, 1, 3}}); // -> [S, B, N, H]
|
|
out->setReshapeDimensions(Dims{5, {S, B, E, 1, 1}});
|
|
```
|
|
|
|
4. **INT8/FP8 quantization.** The plugin folded the entire INT8 attention pipeline into a single op, exposing only `type_id==2` (implicit per-tensor scales) or the explicit-INT8 attribute set. With `IAttention`, you express the QDQ pattern in the network so the builder can fuse it into the FP8/INT8 attention kernel. See also [Explicit Quantization](https://docs.nvidia.com/deeplearning/tensorrt/latest/inference-library/work-quantized-types.html#explicit-quantization).
|
|
|
|
```cpp
|
|
// QDQ Q, K, V (input_qkv_scale).
|
|
auto qdq = [&](ITensor* x) -> ITensor* {
|
|
auto* qn = network->addQuantize(*x, *qkvScale);
|
|
qn->setOutputType(0, DataType::kINT8);
|
|
auto* dq = network->addDequantize(*qn->getOutput(0), *qkvScale);
|
|
return dq->getOutput(0);
|
|
};
|
|
ITensor *qq = qdq(q), *kq = qdq(k), *vq = qdq(v);
|
|
|
|
auto* attn = network->addAttentionV2(*qq, *kq, *vq,
|
|
AttentionNormalizationOp::kSOFTMAX,
|
|
CausalMaskKind::kNONE);
|
|
// dq_probs -> NormalizationQuantizeScale.
|
|
attn->setNormalizationQuantizeScale(*dqProbsScale);
|
|
|
|
// QDQ attention output (output_ctx_scale).
|
|
auto* qOut = network->addQuantize(*attn->getOutput(0), *outScale);
|
|
qOut->setOutputType(0, DataType::kINT8);
|
|
auto* dqOut = network->addDequantize(*qOut->getOutput(0), *outScale);
|
|
```
|
|
|
|
5. **Variable sequence length.** This plugin accepts a `cu_seqlen` packed layout. `IAttention` does not natively consume packed cu-seqlen input; users relying on this path must either pad to a fixed sequence length or implement their own packing/unpacking around `IAttention`. See below for an example.
|
|
|
|
```cpp
|
|
// `scatterIdx`, `gatherIdx` are application-specific and are typically built from `cu_seqlen`
|
|
|
|
// Packed input: [total_tokens, E]; cu_seqlen: [B+1] int32 prefix sums.
|
|
// Scatter packed -> padded [B, S, E] (pad slots are filled with zeros at a sentinel index).
|
|
auto* padded = network->addGather(*packedInput, *scatterIdx, /*axis=*/0); // [B*S, E]
|
|
auto* padBSE = network->addShuffle(*padded->getOutput(0));
|
|
padBSE->setReshapeDimensions(Dims{3, {B, S, E}});
|
|
// ...split into Q/K/V (step 1), build mask from validLen = cu_seqlen[b+1]-cu_seqlen[b] (step 2),
|
|
// call addAttentionV2 (step 4 if quantized)...
|
|
|
|
// Gather attention output [B, N, S_q, H] -> [B, S, E] -> packed [total_tokens, E].
|
|
auto* outBSE = network->addShuffle(*attn->getOutput(0));
|
|
outBSE->setFirstTranspose(Permutation{{0, 2, 1, 3}});
|
|
outBSE->setReshapeDimensions(Dims{2, {B * S, E}});
|
|
auto* repacked = network->addGather(*outBSE->getOutput(0), *gatherIdx, /*axis=*/0);
|
|
```
|
|
|
|
|
|
## Parameters
|
|
|
|
`bertQKVToContextPlugin` has plugin creator class `QKVToContextPluginDynamicCreator` and plugin class `CustomQKVToContextPluginDynamic`.
|
|
|
|
The parameters are defined below and consists of the following attributes:
|
|
|
|
| Type | Parameter | Version | Description
|
|
|----------|-----------------------------------------|-----------------------------------|-------------------------------------------------------------------
|
|
|`int` |`type_id` | 1, 2, 4, 5 |Integer encoding the DataType (0: FP32, 1: FP16, 2: INT8)
|
|
|`int` |`hidden_size` | 1, 2, 3, 4, 5, 6 |The hidden size, denoted by `E` above.
|
|
|`int` |`num_heads` | 1, 2, 3, 4, 5, 6 |The number of self-attention heads.
|
|
|`bool` |`has_mask` | 1, 2, 4, 5 |Whether to use the input_mask input.
|
|
|`float` |`dq_probs` | 1, 2, 3, 4, 5, 6 |inner layer scale factor when run in int8 precision, default 1.f/127.f.
|
|
|`int` |`var_seqlen` | 2, 5 |Whether to use variable sequence length (0: disable, 1: enable), default 0.
|
|
|`int` |`use_int8_scale_max` | 2, 3, 5, 6 |Whether to use INT8 scale factors to optimize softmax MAX reduction. Only active when `type_id==2`. (0: disable, 1: enable), default 1.
|
|
|`int` |`use_explicit_int8` | 3, 6 |Whether to use explicit INT8, (0: disable, 1: enable), default 0.
|
|
|`float` |`input_qkv_scale` | 3, 6 |The int8 scale for the input qkv tensor when explicit precision is used, default 1.f.
|
|
|`float` |`output_ctx_scale` | 3, 6 |The int8 scale for the output context tensor when explicit precision is used, default 1.f.
|
|
|
|
|
|
## Additional resources
|
|
|
|
**Networks:**
|
|
- [Transformer](https://arxiv.org/abs/1706.03762)
|
|
|
|
|
|
## License
|
|
|
|
For terms and conditions for use, reproduction, and distribution, see the [TensorRT Software License Agreement](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sla/index.html)
|
|
documentation.
|
|
|
|
|
|
## Changelog
|
|
April 2026
|
|
Provide `IAttention` as the out-of-the-box replacement.
|
|
|
|
Oct 2025
|
|
Deprecated this plugin. No alternatives are planned to be provided.
|
|
|
|
Jan 2025
|
|
Added precompiled kernel cubins for gb100 (compute capability 10.0) and gb20x (compute capability 12.0) platforms.
|
|
|
|
Aug 2024
|
|
Ported v3 variant (IPluginV2 design) to v6 (IpluginV3 design), but preserves identical attributes and I/O to v3.
|
|
Ported v2 variant (IPluginV2 design) to v5 (IpluginV3 design), but preserves identical attributes and I/O to v2.
|
|
Ported v1 variant (IPluginV2 design) to v4 (IpluginV3 design), but preserves identical attributes and I/O to v1.
|
|
|
|
Feb 2024
|
|
The issue of the V2 plugin not supporting head sizes of 32 or less and variable sequences of 64, 96, and 384 has been resolved.
|
|
|
|
Oct 2023
|
|
Support explicit int8.
|
|
|
|
April 2023
|
|
Optimize the GPU memory usage by using the cublas handle from attachToContext.
|
|
|
|
October 2022
|
|
Add IGMMA/HGMMA sm90 fmha_v2 kernels
|
|
|
|
August 2022
|
|
Expose `use_int8_scale_max` plugin attribute (for versions 2 and 3 of the plugin). This allows users to enable/disable the usage of INT8 scale factors to optimize softmax MAX reduction.
|
|
This optimization is applied always when `type_id==2` by default.
|
|
|
|
September 2021
|
|
Add sequence length 512 support in v2 plugin
|
|
Add head size 32 support when sequence length is 128, 256 or 512 in v2 plugin
|
|
|
|
October 2020
|
|
Add v2 plugin that supports variable sequence length.
|
|
Add v3 plugin that supports int8 interleaved variable sequence length.
|
|
|
|
November 2019
|
|
This is the first release of this `README.md` file.
|
|
|
|
|
|
## Known issues
|
|
|
|
This plugin only supports GPUs with compute capability >= 7.0. For more information see the [CUDA GPU Compute Capability Support Matrix](https://developer.nvidia.com/cuda-gpus#compute)
|