% \resizebox{1.0\linewidth}{!}{
\begin{tikzpicture}

\tikzstyle{box} = [rectangle, minimum height=0.5cm, rounded corners=.08cm, inner sep=2pt, text=black, font=\small]
\tikzstyle{lm} = [box, minimum width=0.7cm, fill=black!20]
\tikzstyle{kv} = [box, minimum width=0.7cm, fill=dg!60]
\tikzstyle{arrow} = [-stealth, line width=0.25 mm]
\tikzstyle{txt} = [font=\small]
\tikzstyle{label} = [font=\scriptsize]


% Profiling KV importance
\node[kv] (kv) {$\kvc$};
\node[lm, right=0.3cm of kv] (lm) {$\lm$};
\node[label, above=0.05cm of kv.north] {$n_c$};
\draw[{Bar[width=3pt]}-{Bar[width=3pt]}] ([yshift=0.1cm]kv.north west) -- ([yshift=0.1cm]kv.north east);

% Arrows
\draw[arrow] (kv) -- (lm);

\node[label, below=0.54cm of lm, anchor=center] (q) {Repeat prompt {\scriptsize\rplus} {\vspace{0.4em} \textcolor{dr}{$i$-th} context chunk}};
\draw[arrow] ($(q.north) + (0,-0.04)$) -- (lm);


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

\node[right=0.8cm of lm] (arr1) {\Large $\Rightarrow$};

\tikzstyle{cell} = [minimum width=0.85cm, minimum height=0.3cm, draw=black]
\tikzstyle{key} = [fill=orange!60]
\tikzstyle{key_off} = [fill=orange!20]
\tikzstyle{query} = [minimum width=0.3cm, minimum height=0.85cm, draw=black, fill=red!50]


\node[query, right=0.9cm of arr1, yshift=-0.35cm] (query) {};
\node[font=\tiny, left=0.25cm of query.west, anchor=center, rotate=90] {$n_\text{in}$};
\draw[{Bar[width=3pt]}-{Bar[width=3pt]}] ([xshift=-0.1cm]query.south west) -- ([xshift=-0.1cm]query.north west);
\node[label, left=0.7cm of query.center, anchor=center, rotate=90] {Query};

\node[label, above=0.45cm of query, fill=gray!15, xshift=0.3cm] {Each attention head};
\matrix (keys) [right=0.1cm of query,
             anchor=south west,
             matrix of nodes,
             nodes in empty cells,
             column sep=0.0cm,
             row sep=0.0cm] {
  \node[cell, key_off] (k1) {}; & \node[cell, key_off] (k2) {}; & \node[cell, key] (k3) {}; & \node[cell, key_off] (k4) {}; & \node[cell, key] (k5) {};\\
  };
\node[label, above=0.16cm of keys.north] {Key};



\node[label, above=0.05cm of k1.north] {$m$};
\node[label, above=0.05cm of k2.north] {$m$};
\node[label, above=0.05cm of k3.north] {$m$};
\node[label, above=0.05cm of k4.north] {$m$};
\node[label, above=0.02cm of k5.north] {$n_\text{in}$};

\draw[|-|] ([yshift=0.1cm]k1.north west) -- ([yshift=0.1cm]k1.north east);
\draw[-|] ([yshift=0.1cm]k2.north west) -- ([yshift=0.1cm]k2.north east);
\draw[-|] ([yshift=0.1cm]k3.north west) -- ([yshift=0.1cm]k3.north east);
\draw[-|] ([yshift=0.1cm]k4.north west) -- ([yshift=0.1cm]k4.north east);
\draw[-|] ([yshift=0.1cm]k5.north west) -- ([yshift=0.1cm]k5.north east);


\node[label, below=-0.06cm of k3.south] {\textcolor{dr}{$i$-th} chunk};
\node[label, below=-0.07cm of k5.south] {input};

\draw[|-|] ([yshift=-0.4cm]k1.south west) -- node[label, below=-0.05cm] {$n_c$} ([yshift=-0.4cm]k4.south east);


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Attention
\tikzstyle{cell} = [minimum size=0.2cm, draw=black]
\tikzstyle{query} = [fill=blue!40]
\tikzstyle{pad} = [fill=gray!60]

\node[right=5.9cm of arr1] (arr2) {\Large $\Rightarrow$};

% Matrix of nodes
\matrix (m) [right=0.6cm of arr2,
             matrix of nodes,
             nodes in empty cells,
             column sep=0.0cm,
             row sep=0.0cm] {
  \node[cell, query] (m11) {}; & \node[cell, query] {}; & \node[cell, query] {}; & \node[cell, query] {}; & \node[cell, query] {}; & \node[cell, pad] {}; & \node[cell, pad] {}; & \node[cell, pad] (m3) {}; \\
  \node[cell, query] {}; & \node[cell, query] {}; & \node[cell, query] {}; & \node[cell, query] {}; & \node[cell, query] {}; & \node[cell, query] {}; & \node[cell, pad] {}; & \node[cell, pad] {}; \\
  \node[cell, query] {}; & \node[cell, query] {}; & \node[cell, query] {}; & \node[cell, query] {}; & \node[cell, query] {}; & \node[cell, query] {}; & \node[cell, query] {}; & \node[cell, pad] {}; \\  
  \node[cell, query] (m41) {}; & \node[cell, query] {}; & \node[cell, query] {}; & \node[cell, query] (m42) {}; & \node[cell, query] {}; & \node[cell, query] {}; & \node[cell, query] {}; & \node[cell, query] (m44) {}; \\  
};

% Red dashed rectangle for the attention window
\node[draw, red, very thick, dashed, fit=(m11)(m42), inner xsep=2pt, inner ysep=2pt] {};

% Labels
\node[txt, above=-0.04cm of m.north] {Attention matrix};

\node[font=\tiny, left=0.25cm of m.west, anchor=center, rotate=90] {$n_\text{in}$};
\draw[{Bar[width=3pt]}-{Bar[width=3pt]}] ([xshift=-0.2cm]m41.south west) -- ([xshift=-0.2cm]m11.north west);
\node[label, left=0.51cm of m.west, anchor=center, rotate=90] {Query};

\node[label, below=0.16cm of m.south] {Key};
\node[label, below=0.04cm of m.south] {\hspace{-0.85cm}$m$};
\draw[{Bar[width=3pt]}-{Bar[width=3pt]}] ([yshift=-0.2cm]m41.south west) -- ([yshift=-0.2cm]m42.south east);
\node[label, below=0.04cm of m.south] {\hspace{1cm}$n_\text{in}$};
\draw[-{Bar[width=3pt]}] ([yshift=-0.2cm]m42.south east) -- ([yshift=-0.2cm]m44.south east);

\draw[<->, dr, line width=0.2 mm] ([xshift=0.2cm]m3.north east) -- ([xshift=0.2cm]m44.south east);
\node[label, right=1.16cm of m.center] {max};

\end{tikzpicture}
% }