Sliding Window Attention (SWA)

Description

Mistral 7B uses grouped-query attention (GQA) and sliding window attention (SWA) to improve the inference speed and reduce the graphics memory.

Currently, the FA operator uses all KVs for inference. SWA can be used to improve the inference speed. The longer the distance, the less important the current location. Two tokens whose distance exceeds the window size are not involved in the attention score calculation.

When the mask is designed to implement n > m or m-n > windowLen (n is kvId and m is qId), the attention score calculation is skipped. For example, in Figure 1, if windowSize is set to 3, the attention score between "the" and "the" will not be calculated.

Figure 1 SWA mask function

The decoder allows only the latest KVs whose length is windowLen to participate in the calculation.

KV cache optimization: For example, if the window size is 4, when the fifth token needs to be cached, the first token is directly replaced. In this way, the KV cache has a maximum value (the window size) and does not increase infinitely.

How to Enable

windowSize > 0
maskType must be MASK_TYPE_SLIDING_WINDOW_NORM or MASK_TYPE_SLIDING_WINDOW_COMPRESS.

Constraints

After the feature is enabled, cacheType can be CACHE_TYPE_NORM or CACHE_TYPE_SWA. If the feature is disabled, cacheType can only be CACHE_TYPE_NORM.
The SWA feature does not support dynamic batch, high precision, clamp scaling, QKV full quantization, mla, logN scaling, and BNSD data layout.
In the SWA feature, when calcType is set to DECODER, maskType cannot be set to MASK_TYPE_SLIDING_WINDOW_COMPRESS, and attentionMask cannot be passed.
To enable this feature, the following conditions must be met: If only one condition is met, an error is reported.
- windowSize > 0.
- maskType must be MASK_TYPE_SLIDING_WINDOW_NORM or MASK_TYPE_SLIDING_WINDOW_COMPRESS.

SWA Mask Generation Example

When windowSize is greater than or equal to seqlen, the mask is the same as that when SWA is disabled, that is, the upper triangle. Otherwise, the mask is generated as follows:

maskType=MASK_TYPE_SLIDING_WINDOW_NORM, dtype is float16.

For example, max(seqLen) is 5 and windowSize is 3.

Figure 2 Mask example when maskType is MASK_TYPE_SLIDING_WINDOW_NORM and dtype is float16

Python code example:

swa_mask = np.ones(shape=[maxseq,maxseq]) * -65536.0
triu_mask = np.triu(swa_mask, 1)
tril_mask = np.tril(swa_mask, -window_size)
swa_mask = triu_mask + tril_mask

C++ code example:

static constexpr uint32_t  SWA_MASK_SIZE = maxSeq;
std::vector<float> create_attention_mask(uint32_t windowSize, uint32_t embeddim) {
    std::vector<float> attention_mask(SWA_MASK_SIZE * SWA_MASK_SIZE, -65536.0);

    for (uint32_t i = 0; i < SWA_MASK_SIZE; ++i) {
        uint32_t offset = i >= windowSize ? (i - windowSize + 1) : 0;
        for (uint32_t j = offset; j < i + 1; ++j) {
            attention_mask[i * SWA_COMPRESS_MASK_SIZE + j] = 0.0;
        }
    }
    return attention_mask;
}

maskType=MASK_TYPE_SLIDING_WINDOW_NORM, dtype is bf16.

For example, max(seqLen) is 5 and windowSize is 3.

Figure 3 Mask example when maskType=MASK_TYPE_SLIDING_WINDOW_NORM and dtype is bf16

Python code example:

swa_mask = np.ones(shape=[maxseq,maxseq]) * 1.0
triu_mask = np.triu(swa_mask, 1)
tril_mask = np.tril(swa_mask, -window_size)
swa_mask = triu_mask + tril_mask

C++ code example:

static constexpr uint32_t  SWA_MASK_SIZE = maxSeq;
std::vector<float> create_attention_mask(uint32_t windowSize, uint32_t embeddim) {
    std::vector<float> attention_mask(SWA_MASK_SIZE * SWA_MASK_SIZE, 1.0);

    for (uint32_t i = 0; i < SWA_MASK_SIZE; ++i) {
        uint32_t offset = i >= windowSize ? (i - windowSize + 1) : 0;
        for (uint32_t j = offset; j < i + 1; ++j) {
            attention_mask[i * SWA_COMPRESS_MASK_SIZE + j] = 0.0;
        }
    }
    return attention_mask;
}

maskType=MASK_TYPE_SLIDING_WINDOW_COMPRESS, dtype is float16.

Python code example:

swa_mask = np.ones(shape=(1, 512, 512)) * -65536.0
pp_n = 128 if head_size < = 128 else 64 # head_size is a size of an embedding vector of each attention head.
if window_size <= pp_n * 3:
    true_size = window_size
else:
    if window_size % pp_n == 0:
        true_size = pp_n * 3
    else:
        true_size = pp_n * 2 + window_size % pp_n
triu_mask = np.triu(swa_mask, 1)
tril_mask = np.tril(swa_mask, -true_size)
swa_mask = triu_mask + tril_mask

C++ code example:

static constexpr uint32_t  SWA_COMPRESS_MASK_SIZE = 512;
std::vector<float> create_attention_mask(uint32_t windowSize, uint32_t embeddim) {
    std::vector<float> attention_mask(SWA_COMPRESS_MASK_SIZE * SWA_COMPRESS_MASK_SIZE, -65536.0);
    uint32_t blockSize = embeddim > 128 ? (16384 / embeddim / 16 * 16) : 128;
    uint32_t compressWindow = windowSize > 3 * blockSize ? (2 * blockSize + windowSize % blockSize) : windowSize;
    for (uint32_t i = 0; i < SWA_COMPRESS_MASK_SIZE; ++i) {
        uint32_t offset = i >= compressWindow ? (i - compressWindow + 1) : 0;
        for (uint32_t j = offset; j < i + 1; ++j) {
            attention_mask[i * SWA_COMPRESS_MASK_SIZE + j] = 0.0;
        }
    }
    return attention_mask;
}

maskType=MASK_TYPE_SLIDING_WINDOW_COMPRESS, dtype is bf16.

Python code example:

swa_mask = np.ones(shape=(1, 512, 512)) * 1
pp_n = 128 if head_size < = 128 else 64 # head_size is a size of an embedding vector of each attention head.
if window_size <= pp_n * 3:
    true_size = window_size
else:
    if window_size % pp_n == 0:
        true_size = pp_n * 3
    else:
        true_size = pp_n * 2 + window_size % pp_n
triu_mask = np.triu(swa_mask, 1)
tril_mask = np.tril(swa_mask, -true_size)
swa_mask = triu_mask + tril_mask

C++ code example:

static constexpr uint32_t  SWA_COMPRESS_MASK_SIZE = 512;
std::vector<float> create_attention_mask(uint32_t windowSize, uint32_t embeddim) {
    std::vector<float> attention_mask(SWA_COMPRESS_MASK_SIZE * SWA_COMPRESS_MASK_SIZE, 1.0);
    uint32_t blockSize = embeddim > 128 ? (16384 / embeddim / 16 * 16) : 128;
    uint32_t compressWindow = windowSize > 3 * blockSize ? (2 * blockSize + windowSize % blockSize) : windowSize;
    for (uint32_t i = 0; i < SWA_COMPRESS_MASK_SIZE; ++i) {
        uint32_t offset = i >= compressWindow ? (i - compressWindow + 1) : 0;
        for (uint32_t j = offset; j < i + 1; ++j) {
            attention_mask[i * SWA_COMPRESS_MASK_SIZE + j] = 0.0;
        }
    }
    return attention_mask;
}

Parent topic: Functions