[object Object][object Object][object Object]undefined[object Object]

• Description: The DenseLightningIndexerGradKlLoss operator is the backward operator of LightningIndexer and integrates the loss calculation function. The LightningIndexer operator selects the top K tokens with the highest intrinsic relationship between the query token and key token, reducing the amount of attention computation in long-sequence scenarios and accelerating the inference and training performance of long-sequence networks. In the dense scenario, the inputs query, key, query_index, and key_index of LightningIndexerGrad do not need to be sparsified.

• Formulas:

  1. The formula for calculating the top-k value is as follows:

     $$I_{t,:}=W_{t,:}@ReLU(\tilde{q}_{t,:}@\tilde{K}_{:t,:}^\top)$$
     • $W_{t,:}$ is the $weights$ corresponding to the $t$th token.
     • $\tilde{q}_{t,:}$ is the result of combining $G$ query heads of the $t$th token in the $\tilde{q}$ matrix.
     • $\tilde{K}_{:t,:}$ is the $t$th row of the $\tilde{K}$ matrix.

  2. The forward Softmax formula is as follows:

     $$p_{t,:} = \text{Softmax}(q_{t,:} @ K_{:t,:}^\top/\sqrt{d})$$
     • $p_{t,:}$ is the Softmax result corresponding to the $t$th token.
     • $q_{t,:}$ is the result of combining $G$ query heads of the $t$th token in the $q$ matrix.
     • ${K}_{:t,:}$ is the $t$th row of the $K$ matrix.

  3. npu_lightning_indexer is trained independently. The corresponding loss function is as follows:

     $$Loss{=}\sum_tD_{KL}(p_{t,:}||Softmax(I_{t,:}))$$

     $p_{t,:}$ is the target distribution, which is obtained by summing up all heads of the main attention score and then performing L1 regularization on the sum result along the context direction. $D_{KL}$ is the KL divergence, and its expression is as follows:

     $$D_{KL}(a||b){=}\sum_ia_i\mathrm{log}{\left(\frac{a_i}{b_i}\right)}$$

  4. The gradient expression of the loss can be obtained by derivation:

  $$dI\mathop{{}}\nolimits_{{t,:}}=Softmax \left( I\mathop{{}}\nolimits_{{t,:}} \left) -p\mathop{{}}\nolimits_{{t,:}}\right. \right.$$

  The gradients of the weights, query, and key matrices can be calculated using the chain rule.

  $$dW\mathop{{}}\nolimits_{{t,:}}=dI\mathop{{}}\nolimits_{{t,:}}\text{@} \left( ReLU \left( S\mathop{{}}\nolimits_{{t,:}} \left) \left) \mathop{{}}\nolimits^{\top}\right. \right. \right. \right.$$ $$d\mathop{{\tilde{q}}}\nolimits_{{t,:}}=dS\mathop{{}}\nolimits_{{t,:}}@\tilde{K}\mathop{{}}\nolimits_{{:t,:}}$$ $$d\tilde{K}\mathop{{}}\nolimits_{{:t,:}}=\left(dS\mathop{{}}\nolimits_{{t,:}} \left) \mathop{{}}\nolimits^{\top}@\tilde{q}\mathop{{}}\nolimits_{{:t, :}}\right. \right.$$

  $S$ is the result of the matrix multiplication of $\tilde{q}$ and $K$.

[object Object][object Object]

Each operator has calls. First, [object Object] is called to obtain the input parameters and compute the required workspace size based on the process. Then, [object Object] is called to perform computation.

[object Object]

• Parameters:

  [object Object]

• Returns

  [object Object]: status code. For details, see .

  The first-phase API implements input parameter validation. The following error codes may be returned.

  [object Object]

[object Object]

• Parameters:

  [object Object]

• Returns

  [object Object]: status code. For details, see .

[object Object]

• The data types of the query, key, queryIndex, and keyIndex parameters must be the same.

• If the weights parameter is not of type float32, the data types of the query, key, queryIndex, keyIndex, and weights parameters must be the same.

• Common constraints

  • Deterministic computation: [object Object] defaults to non-deterministic implementation. You can call [object Object] to enable deterministic computation.
  • Processing when the input parameter is empty:
    • If the query, key, query_index, key_index, or weight is an empty tensor, the current version does not support this operation and an error will be reported.
  [object Object]

• Constraints

  [object Object]

• Typ.

  [object Object]

[object Object]

The following example is for reference only. For details, see .