% \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!66]
\tikzstyle{kvp} = [box, minimum width=0.7cm, fill=dg!28]
\tikzstyle{arrow} = [-stealth, line width=0.25 mm]
\tikzstyle{txt} = [font=\small]
\tikzstyle{label} = [font=\scriptsize]
\tikzstyle{bar} = [{Bar[width=3pt]}-{Bar[width=3pt]}]


% Prefill
\node[lm] (lm) {$\lm$};
\node[txt, below=0.16cm of lm, anchor=north, text=dr] (ctx) {Context};
\draw[arrow] ($(ctx.north) + (0,-0.08)$) -- (lm);


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Profile
\node[right=0.3cm of lm, yshift=-0.1cm] (arr) {\Large $\Rightarrow$};
\node[label, above=0.0cm of arr] {\textbf{Prefill}};

\tikzstyle{cell} = [minimum size=0.15cm, draw=black, fill=dg!60, inner sep=0pt, outer sep=0pt, line width=0.3pt]
\matrix (kv) [right=0.5cm of arr,
             yshift=0.15cm,
             matrix of nodes,
             nodes in empty cells,
             column sep=0.0cm,
             row sep=0.0cm] {
  \node[cell] (m11) {}; & \node[cell] {}; & \node[cell] {}; & \node[cell] {}; & \node[cell] {}; & \node[cell] (m3) {}; \\
  \node[cell] {}; & \node[cell] {}; & \node[cell] {}; & \node[cell] {}; & \node[cell] {}; & \node[cell] {}; \\
  \node[cell] {}; & \node[cell] {}; & \node[cell] {}; & \node[cell] {}; & \node[cell] {}; & \node[cell] {}; \\  
  \node[cell] (m41) {}; & \node[cell] {}; & \node[cell] (m42) {}; & \node[cell] {}; & \node[cell] {}; & \node[cell] (m44) {}; \\  
};
\node[label, above=-0.16cm of kv] {$\kvc$};
\node[lm, right=0.3cm of kv, yshift=-0.05cm] (lm) {$\lm$};
\draw[arrow] ([xshift=-0.1cm, yshift=-0.05cm]kv.east) -- (lm);

\node[font=\tiny, left=0.15cm of kv.west, anchor=center, rotate=90] {$LH$};
\draw[bar] ([xshift=-0.12cm]m41.south west) -- ([xshift=-0.12cm]m11.north west);

\node[font=\tiny, below=0.1cm of kv.south, anchor=center] {$n_c$};
\draw[bar] ([yshift=-0.11cm]m41.south west) -- ([yshift=-0.11cm]m44.south east);

\node[box, below=0.33cm of lm, anchor=north, fill=orange!80!blue!9, inner sep=0pt, minimum width=3.5cm, minimum height=0.38cm, xshift=-0.9em] {};
\node[txt, below=0.28cm of lm, anchor=north] (q) {\hspace{-2em}Repeat prompt {\scriptsize\rplus} \textcolor{dr}{Context}};
\draw[arrow] ([yshift=-0.05cm]q.north) -- (lm);


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% KV importance
\node[right=0.6cm of lm, yshift=-0.1cm] (arr1) {\Large $\Rightarrow$};
\node[label, above=0.7em of arr1] {\hspace{-0.5em}\textbf{Measure} max};
\node[label, above=0.0cm of arr1] {\hspace{-0.1em}cross-attention};

\tikzstyle{cell} = [minimum size=0.2cm, draw=black]
\tikzstyle{high} = [fill=blue!40]
\tikzstyle{low0} = [fill=blue!20]
\tikzstyle{low} = [fill=blue!16]
\tikzstyle{low2} = [fill=blue!12]

\matrix (m) [right=0.95cm of arr1, yshift=0.1cm,
             matrix of nodes,
             nodes in empty cells,
             column sep=0.0cm,
             row sep=0.0cm] {
  \node[cell, low0] (m11) {}; & \node[cell, high] {}; & \node[cell, low] {}; & \node[cell, low0] {}; & \node[cell, high] {}; & \node[cell, low] (m3) {}; \\
  \node[cell, high] {}; & \node[cell, low] {}; & \node[cell, high] {}; & \node[cell, low0] {}; & \node[cell, low2] {}; & \node[cell, low] {}; \\
  \node[cell, high] {}; & \node[cell, low0] {}; & \node[cell, low] {}; & \node[cell, low2] {}; & \node[cell, high] {}; & \node[cell, high] {}; \\  
  \node[cell, low] (m41) {}; & \node[cell, high] {}; & \node[cell, high] {}; & \node[cell, high] (m42) {}; & \node[cell, low2] {}; & \node[cell, low] (m44) {}; \\  
};

% Labels
\node[label, left=1cm of m.center, anchor=center, rotate=90] {Heads {\tiny($LH$)}};
\node[label, below=-0.15cm of m.south] {Sequence ($n_c$)};
\node[txt, above=-0.04cm of m.north, xshift=-0.1em] {KV importance};


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% evict
\node[right=0.87cm of m.east, yshift=-0.1cm] (arr2) {\Large $\Rightarrow$};
\node[label, above=1.3em of arr2, xshift=-0.1cm] {\hspace{-1.8em}\textbf{Evict} KV};
\node[label, above=0.6em of arr2, xshift=-0.1cm] {\hspace{0.7em}with low scores};
\node[label, above=-0.2em of arr2, xshift=-0.1cm] {\hspace{1.2em}(pair-/head-level)};

\tikzstyle{cell} = [minimum size=0.15cm, draw=black, fill=dg!60, inner sep=0pt, outer sep=0pt, line width=0.3pt]
\tikzstyle{empty} = [minimum size=0.15cm, inner sep=0pt, outer sep=0pt]
\tikzstyle{dottededge} = [
  draw=black,
  line width=0.3pt,
  dotted,
  dash pattern=on 0.2pt off 1.2pt
]
\matrix (kvp) [right=0.78cm of arr2, yshift=0.1cm,
             matrix of nodes,
             nodes in empty cells,
             column sep=0.0cm,
             row sep=0.0cm] {
  \node[empty] (m11) {}; & \node[cell] {}; & \node[empty] (m13) {}; & \node[empty] (m14) {}; & \node[cell] {}; & \node[empty] (m16) {}; \\
  \node[cell] {}; & \node[empty] (m22) {}; & \node[cell] {}; & \node[empty] (m24) {}; & \node[empty] (m25) {}; & \node[empty] (m26) {}; \\
  \node[cell] {}; & \node[empty] (m32) {}; & \node[empty] (m33) {}; & \node[empty] (m34) {}; & \node[cell] {}; & \node[cell] {}; \\  
  \node[empty] (m41) {}; & \node[cell] {}; & \node[cell] (m42) {}; & \node[cell] {}; & \node[empty] (m45) {}; & \node[empty] (m46) {}; \\  
};
\foreach \nodename in {m11,m34} {
    \draw[dottededge] (\nodename.north west)--(\nodename.north east); % north edge
    \draw[dottededge] (\nodename.north west)--(\nodename.south west); % west edge
}
\foreach \nodename in {m14} {
    \draw[dottededge] (\nodename.north west)--(\nodename.north east); % north edge
}
\foreach \nodename in {m13,m16,m24,m26,m32,m46} {
    \draw[dottededge] (\nodename.north west)--(\nodename.north east); % north edge
    \draw[dottededge] (\nodename.north east)--(\nodename.south east); % east edge
}
\foreach \nodename in {m25} {
    \draw[dottededge] (\nodename.north east)--(\nodename.south east); % east edge
}
\foreach \nodename in {m41,m46} {
    \draw[dottededge] (\nodename.south west)--(\nodename.south east); % south edge
    \draw[dottededge] (\nodename.north west)--(\nodename.south west); % west edge
}
\foreach \nodename in {m45} {
    \draw[dottededge] (\nodename.south west)--(\nodename.south east); % south edge
}


% \matrix (kvp) [right=1.0cm of arr2,
%              matrix of nodes,
%              nodes in empty cells,
%              column sep=0.0cm,
%              row sep=0.0cm] {
%   \node[cell] {}; & \node[cell] {}; \\
%   \node[cell] {}; & \node[cell] {}; \\
%   \node[cell] {}; & \node[cell] {}; & \node[cell] {};\\  
%   \node[cell] {}; & \node[cell] {}; & \node[cell] {};\\  
% };
\node[label, above=-0.18cm of kvp] {$\kvp$};

% \node[kvp, right=0.7cm of arr2] (kvp) {\raisebox{-0.7em}{$\kvp$}};
\node[lm, right=0.35cm of kvp] (lm2) {$\lm$};
\draw[arrow] ([xshift=-0.1cm]kvp.east) -- (lm2);

\node[txt, below=0.16cm of lm2, anchor=north] (q) {Queries};
\draw[arrow] ($(q.north) + (0,-0.12)$) -- (lm2);
\node[txt, above=0.47cm of lm2, anchor=center] (ctx2) {Responses};
\draw[arrow] (lm2) -- node[label, anchor=west] {decode} ($(ctx2.south)+(0, 0.08)$);


\end{tikzpicture}
% }