Step 2 and 3 are wrong.

Let two input sequences $X \in \mathbb{R}^{T\times C}$ and $X' \in \mathbb{R}^{T'\times C}$, where $X$ consists $T$ tokens, each with $C$ dimensions, and $X'$ has $T'$ number of $C$-dimensions tokens.

The attention mechanism maps the query $Q \in \mathbb{R}^{T\times D_k}$ against the keys $K \in \mathbb{R}^{T\times D_k}$, and retrive the weighted values $V \in \mathbb{R}^{T'\times D_{out}}$, that is, for $T$ number of query, we also got the $T$ number of values, with dimensions changing from $D_k$ to $D_{out}$.

For h-th attention head: $$ \begin{align} \text{head}_h=\text{Attention}_h(XW_h^Q, X'W_h^K, X'W_h^V) \\ =\text{Attention}_h(Q_h, K_h, V_h) \\ =\text{softmax}(\frac{Q_h K_h^T}{\sqrt{d_k}})V_h \label{eq1}\tag{1}\\ =A_h V_h \label{eq2}\tag{2} \end{align} $$ $A_h$ is h-th attention matrix, and $$ Q_h=XW_h^Q,\;\; W_h^Q \in \mathbb{R}^{C\times D_{k}},\;\; Q_h\in \mathbb{R}^{T\times D_{k}}\\ K_h=X'W_h^K,\;\; W_h^K \in \mathbb{R}^{C\times D_{k}},\;\; K_h\in \mathbb{R}^{T'\times D_{k}}\\ V_h=X'W_h^V,\;\; W_h^V \in \mathbb{R}^{C\times D_{out}},\;\; V_h\in \mathbb{T}^{T'\times D_{out}}\\ $$

In softmax function in equation-\ref{eq1}, the attention matrix $Q_h K_h^T$ is computed with dimensions $(T, D_{k})\times(D_{k}, T')=(T, T')$. Thus, the attention matrix $A_h$(referred to as tensor A in your post) should be (B, T, T') instead of (B, T, T).

After creating the h-th attention matrix, in equation-\ref{eq2}. The output dimensions of the $A_h V_h$ are $(T, T')\times(T', D_{out})=(T, D_{out})$. This represents the dimensions of the h-th attention head $\text{head}_h$.

For multi-head setting, we concatenate all heads and linearly project the dimension down to the final output. $$ \text{MultiHead}(Q, K, V)=\underset{h \in [H]}{\text{Concat}}\left(head_h\right)W^o $$

Step by step, the concatenation is performed along the feature $D_{out}$ dimension. For a total of $H$ heads, yields dimensions $(T, H\cdot D_{out})$. Following this, the linear projection $W^o \in \mathbb{R}^{H\cdot D_{out} \times D_{final}}$ then gives a final output with dimensions $(T, D_{final})$.

During inference

The fact that keys and values in the decoder have different dimensions from queries is irrelevant. By design, the input sequences $X$ and $X'$ can have different token lengths, $T$ and $T'$. In the decoder's cross-attention, the length $T'$ of the keys and values cancels out in equation-\ref{eq2}.