@@ -1,13 +1,25 @@
""" Comparative figures for milestone-gqa-decode-long-ctx-4cases.
Reads sweep.json (emitted by ``kernbench run --bench
milestone-gqa-decode-long-ctx-4cases``) and writes four PNGs into
``docs/report/1H-codesign-paper/figures/`` :
Reads sweep_decode .json (emitted by the milestone-1h-gqa bench) and
writes four PNGs into the same bench-output dir
(src/kernbench/benches/1H_milestone_output/gqa/long_ctx/) :
gqa_decode_long_ctx_4 cases_latency.png end-to-end latency per case
gqa_decode_long_ctx_4 cases_traffic.png ipcq/dma op-count breakdown
gqa_decode_long_ctx_4 cases_memory.png per-PE KV bytes per case
gqa_decode_long_ctx_4 cases_parallelism.png per-PE S_local (compute work)
gqa_decode_long_ctx_6 cases_latency.png end-to-end latency per case
gqa_decode_long_ctx_6 cases_traffic.png ipcq/dma op-count breakdown
gqa_decode_long_ctx_6 cases_memory.png per-PE KV bytes per case
gqa_decode_long_ctx_6 cases_parallelism.png per-PE S_local (compute work)
Filename still says " 4cases " for backwards compat, but the script now
covers all SIX kv-sharding strategies from the analytical chart
(`gqa_4cases_summary.png`) — the original 4 plus the two new
d_head-TP variants:
Case 1 Cube-Repl × ' t shard KV)
Case 2 Cube-SP ×
Case 3 Cube-Repl ×
Case 4 Cube-SP ×
Case 5 Cube-TP-d_head ×
Case 6 ★ Cube-SP ×
Run (after the bench):
GQA_DECODE_LONG_CTX_4CASES_RUN=1 python -m kernbench.cli.main run \\
@@ -32,16 +44,27 @@ _FIG_DIR = (
)
_SWEEP_JSON = _FIG_DIR / " sweep_decode.json "
# Panel name → (short label, case ordinal for left-to-right plot order).
# Panel name → (short label, case ordinal, accent flag) using the
# analytical chart's memory-descending ordering. PE-TP doesn't shard
# KV memory, so the cube_repl_pe_tp panel maps to Case 1 (no
# sharding, KV-wise) and cube_sp_pe_tp panel maps to Case 2.
_NORMAL , _OVERFLOW , _PARETO = " normal " , " overflow " , " pareto "
_CASE_INFO = {
" single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp " : (
" Case 1 \n Cube-SP × TP " , 1 ) ,
" single_kv_group_decode_long_ctx_gqa_cube_repl_pe_tp " : (
" Case 2 \n Cube-Repl × T P " , 2 ) ,
" single_kv_group_decode_long_ctx_gqa_cube_repl_pe_sp " : (
" Case 3 \n Cube-Repl × " , 3 ) ,
" single_kv_group_decode_long_ctx_gqa_cube_sp_pe_sp " : (
" Case 4 ★ \n Cube-SP × " , 4 ) ,
# panel name label ord flag
" single_kv_group_decode_long_ctx_gqa_cube_repl_pe_tp " : ( " Case 1 \n Cube-Repl × repl " , 1 , _OVERFLOW ) ,
" single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp " : ( " Case 2 \n Cube-SP × " , 2 , _OVERFLOW ) ,
" single_kv_group_decode_long_ctx_gqa_cube_repl_pe_sp " : ( " Case 3 \n Cube-Repl × S P " , 3 , _OVERFLOW ) ,
" single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp_dhead " : ( " Case 4 \n Cube-SP × " , 4 , _NORMAL ) ,
" single_kv_group_decode_long_ctx_gqa_cube_tp_dhead_pe_sp " : ( " Case 5 \n Cube-TP-d_head × " , 5 , _NORMAL ) ,
" single_kv_group_decode_long_ctx_gqa_cube_sp_pe_sp " : ( " Case 6 ★ \n Cube-SP × " , 6 , _PARETO ) ,
}
# Bar fill colour per flag (used by every panel).
_FLAG_COLOR = {
_NORMAL : " #888888 " , # neutral grey
_OVERFLOW : " #c0504d " , # red — fails the per-PE HBM budget
_PARETO : " #3b6ea5 " , # blue — Pareto-best
}
@@ -53,23 +76,26 @@ def _sorted_by_case(rows: list[dict]) -> list[dict]:
return sorted ( rows , key = lambda r : _CASE_INFO [ r [ " panel " ] ] [ 1 ] )
def _bar_colors ( rows : list [ dict ] ) - > list [ str ] :
return [ _FLAG_COLOR [ _CASE_INFO [ r [ " panel " ] ] [ 2 ] ] for r in rows ]
def _plot_latency ( rows : list [ dict ] ) - > Path :
rows = _sorted_by_case ( rows )
labels = [ _CASE_INFO [ r [ " panel " ] ] [ 0 ] for r in rows ]
lat_us = [ r [ " latency_ns " ] / 1e3 for r in rows ]
colors = [ " #888 " , " #888 " , " #888 " , " #3b6ea5 " ] # Case 4 highlighted
fig , ax = plt . subplots ( figsize = ( 8.0 , 4.5 ) )
bars = ax . bar ( labels , lat_us , color = colors , width = 0.6 )
fig , ax = plt . subplots ( figsize = ( 12.0 , 4.8 ) )
bars = ax . bar ( labels , lat_us , color = _bar_colors ( rows ) , width = 0.6 )
ax . set_ylabel ( " end-to-end latency (µs) " )
ax . set_title (
" Long-context decode 4 -cases — end-to-end latency per case \n "
" Long-context decode 6 -cases — end-to-end latency per case \n "
" LLaMA-3.1-70B single-KV-head group (8 cubes × "
)
ax . bar_label ( bars , fmt = " %.1f " , padding = 3 , fontsize = 9 )
ax . grid ( axis = " y " , ls = " : " , alpha = 0.5 )
ax . set_ylim ( 0 , max ( lat_us ) * 1.15 )
fig . tight_layout ( )
out = _FIG_DIR / " gqa_decode_long_ctx_4cases_latency.png "
out = _FIG_DIR / " gqa_decode_long_ctx_6cases_latency.png "
fig . savefig ( out , dpi = 150 )
plt . close ( fig )
return out
@@ -83,18 +109,18 @@ def _plot_traffic(rows: list[dict]) -> Path:
disp = [ " IPCQ copy " , " DMA read " , " DMA write " ]
colors = [ " #c0504d " , " #9bbb59 " , " #8064a2 " ]
w = 0.25
fig , ax = plt . subplots ( figsize = ( 9.0 , 4.5 ) )
fig , ax = plt . subplots ( figsize = ( 11.0 , 4.5 ) )
for i , ( k , d , c ) in enumerate ( zip ( keys , disp , colors ) ) :
vals = [ r [ " op_log_summary " ] [ k ] for r in rows ]
ax . bar ( [ xi + ( i - 1 ) * w for xi in x ] , vals , width = w , label = d , color = c )
ax . set_xticks ( list ( x ) )
ax . set_xticklabels ( labels , fontsize = 9 )
ax . set_ylabel ( " op count " )
ax . set_title ( " Long-context decode 4 -cases — op-count breakdown per case " )
ax . set_title ( " Long-context decode 6 -cases — op-count breakdown per case " )
ax . legend ( fontsize = 9 )
ax . grid ( axis = " y " , ls = " : " , alpha = 0.5 )
fig . tight_layout ( )
out = _FIG_DIR / " gqa_decode_long_ctx_4cases_traffic.png "
out = _FIG_DIR / " gqa_decode_long_ctx_6cases_traffic.png "
fig . savefig ( out , dpi = 150 )
plt . close ( fig )
return out
@@ -103,27 +129,41 @@ def _plot_traffic(rows: list[dict]) -> Path:
def _s_local_per_pe ( panel : str , * , S_kv : int , C : int , P : int ) - > int :
""" S_local (token count) each PE attends over locally.
Encodes the cube/pe sharding axes from the panel name:
cube_sp_pe_tp (Case 1 ): S_kv / C (pe=replicate within cub e)
cube_repl_pe_tp (Case 2 ): S_kv (pe=replicate; only 1 PE works)
cube_repl_pe_sp (Case 3 ): S_kv / P (pe=row_wise within cube )
cube_s p_pe_sp (Case 4 ): S_kv / (C·P) (★ 64-way split )
cube_repl_pe_tp (Case 1): S_kv (no sharding, KV-wise)
cube_sp_pe_tp (Case 2 ): S_kv / C (cube splits S_kv, PEs replicat e)
cube_repl_pe_sp (Case 3 ): S_kv / P
cube_sp_pe_tp_dhead (Case 4 ): S_kv / C (cube splits S_kv, PE splits d_head )
cube_t p_dhead_ pe_sp (Case 5 ): S_kv / P (cube splits d_head, PE splits S_kv )
cube_sp_pe_sp (Case 6 ★): S_kv / (C·P)
"""
S_per_cube = S_kv if " cube_repl " in panel else S_kv / / C
return S_per_cube / / P if " pe_sp " in panel else S_per_cube
cube_splits_s = " cube_sp " in panel
pe_splits_s = " pe_sp " in panel
S_per_cube = S_kv / / C if cube_splits_s else S_kv
return S_per_cube / / P if pe_splits_s else S_per_cube
def _d_head_per_pe ( panel : str , * , d_head : int , C : int , P : int ) - > int :
""" d_head dims each PE owns (Cases 4 and 5 shard d_head). """
if " cube_tp_dhead " in panel : # Case 5: cube shards d_head
return d_head / / C
if " pe_tp_dhead " in panel : # Case 4: PE shards d_head
return d_head / / P
return d_head # Cases 1, 2, 3, 6: full d_head per PE
def _active_pe_count ( panel : str , * , C : int , P : int ) - > int :
""" Number of PEs doing non-idle attention work.
cube_sp _pe_tp (Case 1): C (PE 0 of each cube; 7 PEs idle per cube )
cube_repl _pe_tp (Case 2): 1 (only PE 0 of CUBE 0 )
cube_repl_pe_sp (Case 3): C·P (all PEs busy, but cubes ar e redundant)
cube_sp_pe_sp (Case 4): C·P (all 64 PEs doing unique work )
cube_repl _pe_tp (Case 1): 1 (PE-TP idle for B=1; only one PE works )
cube_sp _pe_tp (Case 2): C (PE 0 of each cube; 7 PEs idle per cube )
cube_repl_pe_sp (Case 3): C·P (all PEs busy, cube-sid e redundant)
cube_sp_pe_tp_dhead (Case 4): C·P (PE shards d_head — all 64 active )
cube_tp_dhead_pe_sp (Case 5): C·P (PE shards S_kv — all active)
cube_sp_pe_sp (Case 6 ★): C·P (all 64 PEs doing unique work)
"""
if " cube_repl " in panel and " pe_tp " in panel :
if " cube_repl " in panel and " pe_tp " in panel and " dhead " not in panel :
return 1
if " cube_sp " in panel and " pe_tp " in panel :
if " cube_sp " in panel and " pe_tp " in panel and " dhead " not in panel :
return C
return C * P
@@ -132,11 +172,12 @@ def _kv_bytes_per_pe(panel: str, *, S_kv: int, h_kv: int,
d_head : int , C : int , P : int ) - > int :
""" KV bytes a single PE references (K + V, f16, 2 B/elem). """
s_local = _s_local_per_pe ( panel , S_kv = S_kv , C = C , P = P )
return 2 * s_local * h_kv * d_head * 2
d_local = _d_head_per_pe ( panel , d_head = d_head , C = C , P = P )
return 2 * s_local * h_kv * d_local * 2
def _plot_memory ( rows : list [ dict ] ) - > Path :
""" Per-PE KV bytes — Case 4 wins (64-way split). """
""" Per-PE KV bytes — Case 6 ★ wins (64-way split). """
rows = _sorted_by_case ( rows )
labels = [ _CASE_INFO [ r [ " panel " ] ] [ 0 ] for r in rows ]
mib_per_pe = [
@@ -146,26 +187,25 @@ def _plot_memory(rows: list[dict]) -> Path:
) / ( 1024 * 1024 )
for r in rows
]
colors = [ " #888 " , " #c0504d " , " #888 " , " #3b6ea5 " ] # 4 highlighted, 2 marked red
fig , ax = plt . subplots ( figsize = ( 8.0 , 4.5 ) )
bars = ax . bar ( labels , mib_per_pe , color = colors , width = 0.6 )
fig , ax = plt . subplots ( figsize = ( 12.0 , 4.8 ) )
bars = ax . bar ( labels , mib_per_pe , color = _bar_colors ( rows ) , width = 0.6 )
ax . set_ylabel ( " KV bytes per PE (MiB, K + V, f16) " )
ax . set_title (
" Long-context decode 4 -cases — KV memory per PE \n "
" Long-context decode 6 -cases — KV memory per PE \n "
" (one KV-head group; per-layer, per-token state) "
)
ax . bar_label ( bars , fmt = " %.3f " , padding = 3 , fontsize = 9 )
ax . grid ( axis = " y " , ls = " : " , alpha = 0.5 )
ax . set_ylim ( 0 , max ( mib_per_pe ) * 1.15 )
fig . tight_layout ( )
out = _FIG_DIR / " gqa_decode_long_ctx_4cases_memory.png "
out = _FIG_DIR / " gqa_decode_long_ctx_6cases_memory.png "
fig . savefig ( out , dpi = 150 )
plt . close ( fig )
return out
def _plot_parallelism ( rows : list [ dict ] ) - > Path :
""" Total active PE-token compute load — exposes Case 3 ' s redundancy ."""
""" Total active PE-token compute load — exposes redundant-work cases . """
rows = _sorted_by_case ( rows )
labels = [ _CASE_INFO [ r [ " panel " ] ] [ 0 ] for r in rows ]
total_work = [
@@ -173,19 +213,19 @@ def _plot_parallelism(rows: list[dict]) -> Path:
* _s_local_per_pe ( r [ " panel " ] , S_kv = r [ " S_kv " ] , C = r [ " C " ] , P = r [ " P " ] )
for r in rows
]
colors = [ " #888 " , " #888 " , " #c0504d " , " #3b6ea5 " ] # 4 highlighted, 3 marked red
fig , ax = plt . subplots ( figsize = ( 8.0 , 4.5 ) )
bars = ax . bar ( labels , total_work , color = colors , width = 0.6 )
fig , ax = plt . subplots ( figsize = ( 12.0 , 4.8 ) )
bars = ax . bar ( labels , total_work , color = _bar_colors ( rows ) , width = 0.6 )
ax . set_ylabel ( " active-PE × " )
ax . set_title (
" Long-context decode 4 -cases — total compute load across active PEs \n "
" (Case 3 replicates the full K/ V across 8 cubes ⇒ 8× ) "
" Long-context decode 6 -cases — total compute load across active PEs \n "
" (Case 3 replicates K V across 8 cubes → 8× ; "
" Case 6 ★ is fully parallel without replication) "
)
ax . bar_label ( bars , fmt = " %d " , padding = 3 , fontsize = 9 )
ax . grid ( axis = " y " , ls = " : " , alpha = 0.5 )
ax . set_ylim ( 0 , max ( total_work ) * 1.15 )
fig . tight_layout ( )
out = _FIG_DIR / " gqa_decode_long_ctx_4cases_parallelism.png "
out = _FIG_DIR / " gqa_decode_long_ctx_6cases_parallelism.png "
fig . savefig ( out , dpi = 150 )
plt . close ( fig )
return out
