From a98e93110b873815b4bff498972004261f957aeb Mon Sep 17 00:00:00 2001 From: eusonice <98355549+eusonice@users.noreply.github.com> Date: Tue, 23 Jun 2026 16:32:17 -0700 Subject: [PATCH] short context length composite command --- docs/sweeps/short_context_decode_sweep.csv | 17 ++ docs/sweeps/short_context_prefill_sweep.csv | 17 ++ .../_gqa_attention_decode_short_composite.py | 191 ++++++++++++++++ .../_gqa_attention_prefill_short_composite.py | 216 ++++++++++++++++++ 4 files changed, 441 insertions(+) create mode 100644 docs/sweeps/short_context_decode_sweep.csv create mode 100644 docs/sweeps/short_context_prefill_sweep.csv create mode 100644 src/kernbench/benches/_gqa_attention_decode_short_composite.py create mode 100644 src/kernbench/benches/_gqa_attention_prefill_short_composite.py diff --git a/docs/sweeps/short_context_decode_sweep.csv b/docs/sweeps/short_context_decode_sweep.csv new file mode 100644 index 0000000..d67f7f8 --- /dev/null +++ b/docs/sweeps/short_context_decode_sweep.csv @@ -0,0 +1,17 @@ +mode,kv_per_cube,C,S_kv,wall_us,n_pe,gemm_util,math_count,hbm_bw_util,hbm_read_mb,hbm_write_kb,ipcq_kb,kv_cache_per_cube_mb +A1,1,8,8192,149.2685345000007,64,0.001756190618994587,656,0.013780533230866485,16.0625,8.0,57.75,2.0 +A1,1,8,16384,282.2252665000025,64,0.0018576933472394743,1360,0.014545118695104373,32.0625,8.0,57.75,4.0 +A1,1,8,32768,548.138730500006,64,0.0019129755692392889,2768,0.014961540835691614,64.0625,8.0,57.75,8.0 +A1,1,8,65536,1080.073058500367,64,0.0019416760593127277,5584,0.01517767698303756,128.0625,8.0,57.75,16.0 +A2,2,4,8192,142.16883275000066,32,0.0036877843748065785,656,0.028881154333033524,16.03125,8.0,24.75,4.0 +A2,2,4,16384,277.0102967500018,32,0.0037853322143696724,1360,0.029609007665885423,32.03125,8.0,24.75,8.0 +A2,2,4,32768,546.8006247504073,32,0.0038353138330044275,2768,0.02998167752182655,64.03125,8.0,24.75,16.0 +A2,2,4,65536,1087.0962007518285,32,0.003858263874988177,5584,0.03015188534127058,128.03125,8.0,24.75,32.0 +A4,4,2,8192,142.21417550000066,16,0.007373217165681769,656,0.05768763888097753,16.015625,8.0,8.25,8.0 +A4,4,2,16384,280.8788035001758,16,0.007466394665122732,1360,0.058373931374247456,32.015625,8.0,8.25,16.0 +A4,4,2,32768,558.5655195010398,16,0.007509063580845671,2768,0.058686042828569145,64.015625,8.0,8.25,32.0 +A4,4,2,65536,1114.319951501857,16,0.007528006645388841,5584,0.05882332081702009,128.015625,8.0,8.25,64.0 +B,8,1,8192,147.93485600006025,8,0.014176185766516818,656,0.11085960701508589,16.0078125,8.0,0.0,16.0 +B,8,1,16384,294.3171420004266,8,0.014250967413897551,1360,0.1113900460475132,32.0078125,8.0,0.0,32.0 +B,8,1,32768,587.2722140010417,8,0.01428401991446495,2768,0.11162115018757507,64.0078125,8.0,0.0,64.0 +B,8,1,65536,1173.1823580017838,8,0.014300603725891686,5584,0.11173710472707321,128.0078125,8.0,0.0,128.0 diff --git a/docs/sweeps/short_context_prefill_sweep.csv b/docs/sweeps/short_context_prefill_sweep.csv new file mode 100644 index 0000000..bcef085 --- /dev/null +++ b/docs/sweeps/short_context_prefill_sweep.csv @@ -0,0 +1,17 @@ +mode,kv_per_cube,C,S_kv,wall_us,n_pe,gemm_util,math_count,hbm_bw_util,hbm_read_mb,hbm_write_kb,ipcq_kb,kv_cache_per_cube_mb +A1,1,8,8192,1103.359903999566,64,0.0019006962210579498,5248,0.001870649814732448,16.0625,64.0,114688.0,2.0 +A1,1,8,16384,2182.457823968016,64,0.0019218259129421755,10880,0.0018841143021603984,32.0625,64.0,229376.0,4.0 +A1,1,8,32768,4343.2822238495155,64,0.0019313983221991736,22144,0.0018898150239762977,64.0625,64.0,458752.0,8.0 +A1,1,8,65536,8666.440783650422,64,0.001935883071136998,44672,0.0018923570136126979,128.0625,64.0,917504.0,16.0 +A2,2,4,8192,1097.2535885006619,32,0.003822547535007557,2624,0.0037621203004135906,16.0625,64.0,49152.0,4.0 +A2,2,4,16384,2179.78425743804,32,0.003848366172648072,5440,0.003772850442394648,32.0625,64.0,98304.0,8.0 +A2,2,4,32768,4344.834905351148,32,0.003861416225357031,11072,0.003778279349528763,64.0625,64.0,196608.0,16.0 +A2,2,4,65536,8674.946611121853,32,0.003867969856656271,22336,0.003781003096658628,128.0625,64.0,393216.0,32.0 +A4,4,2,8192,1103.675383502932,16,0.007600611670231015,1312,0.007480460399321815,16.0625,64.0,16384.0,8.0 +A4,4,2,16384,2197.0616234262543,16,0.00763620638634706,2720,0.007486362614786297,32.0625,64.0,32768.0,16.0 +A4,4,2,32768,4383.834103368879,16,0.007654129058897454,5536,0.007489334501679554,64.0625,64.0,65536.0,32.0 +A4,4,2,65536,8757.379062883949,16,0.007663121981838932,11168,0.007490825682997995,128.0625,64.0,131072.0,64.0 +B,8,1,8192,1125.9457520018113,8,0.014900554462921552,656,0.01466500492643045,16.0625,64.0,0.0,16.0 +B,8,1,16384,2242.543191943168,8,0.014962669223294344,1360,0.014669059716747552,32.0625,64.0,0.0,32.0 +B,8,1,32768,4475.650071825993,8,0.014994216018466134,2768,0.014671388277951352,64.0625,64.0,0.0,64.0 +B,8,1,65536,8942.79183160376,8,0.015008481750139946,5584,0.014671033662702533,128.0625,64.0,0.0,128.0 diff --git a/src/kernbench/benches/_gqa_attention_decode_short_composite.py b/src/kernbench/benches/_gqa_attention_decode_short_composite.py new file mode 100644 index 0000000..a9d8875 --- /dev/null +++ b/src/kernbench/benches/_gqa_attention_decode_short_composite.py @@ -0,0 +1,191 @@ +"""GQA decode kernel — composite (second-level) variant. + +Identical mapping to ``_gqa_attention_decode_short.py`` (sequence-shard ++ chain reduce + FA2 head fusion, unified A1/A2/A4/B). Per-tile fusion +follows the ADR-0065 two-composite pattern (cf. decode_opt2): + + tile 0 primitives establish (m, ℓ, O) + tile k>0 #1 Q·Kᵀ composite → scores (pinned primary-out) + #2 softmax_merge prologue → P (pinned, auto-binds to GEMM a) + + P·V composite, ``out=O_local`` + + add epilogue folding the result into O_local + +The softmax_merge recipe (tl_recipes.py) folds max/exp/sum into the +composite, eliminating the GEMM/MATH engine bubble that a flat +primitives chain incurs. The intra-group chain reduce after the tile +loop is unchanged. + +Shard addressing, layouts, and caller contract are identical to the +first-level decode kernel (ADR-0011 D-VA1). +""" +from __future__ import annotations + + +TILE_S_KV = 1024 + + +def _validate_config(*, kv_per_cube: int, T_q: int, P: int, C: int, + h_q: int, h_kv: int, S_kv: int) -> None: + """Validate composite-decode config — caller-side, sim-cost 0. + + Mirrors first-level decode ``_validate_config``. + """ + if kv_per_cube not in (1, 2, 4, 8): + raise ValueError(f"kv_per_cube must be in {{1,2,4,8}}; got {kv_per_cube}") + if T_q != 1: + raise ValueError(f"decode requires T_q == 1; got {T_q}") + if P != 8: + raise ValueError(f"2x4 mesh chain-reduce geometry requires P == 8; got {P}") + if h_q % h_kv != 0: + raise ValueError(f"GQA group G = h_q/h_kv must be integer; got h_q={h_q}, h_kv={h_kv}") + if C != h_kv // kv_per_cube: + raise ValueError( + f"C must equal h_kv/kv_per_cube = {h_kv // kv_per_cube}; " + f"got C={C} (h_kv={h_kv}, kv_per_cube={kv_per_cube})" + ) + group_size = P // kv_per_cube + if S_kv % (group_size * TILE_S_KV) != 0: + raise ValueError( + f"S_kv ({S_kv}) must be a multiple of group_size·TILE_S_KV " + f"({group_size * TILE_S_KV}); each PE's sequence shard must " + f"be a whole number of tiles" + ) + + +def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl): + """Used only by the intra-group chain reduce below (recipe handles + per-tile fold internally).""" + m_new = tl.maximum(m_local, m_other) + scale_old = tl.exp(m_local - m_new) + scale_new = tl.exp(m_other - m_new) + l_new = l_local * scale_old + l_other * scale_new + O_new = O_local * scale_old + O_other * scale_new + return m_new, l_new, O_new + + +def gqa_attention_decode_short_composite_kernel( + q_ptr: int, + k_ptr: int, + v_ptr: int, + o_ptr: int, + T_q: int, + S_kv: int, + h_q: int, + h_kv: int, + d_head: int, + C: int, + P: int, + kv_per_cube: int, + *, + tl, +) -> None: + """Composite-GEMM decode — same mapping as first-level + tl.composite. + + Caller must invoke ``_validate_config(...)`` first. + """ + group_size = P // kv_per_cube + S_local = S_kv // group_size + n_tiles_per_pe = S_local // TILE_S_KV + + G = h_q // h_kv + pe_id = tl.program_id(axis=0) + cube_id = tl.program_id(axis=1) + pe_in_group = pe_id % group_size + group_id_in_cube = pe_id // group_size + + KV_ROW_BYTES = d_head * 2 + K_TILE_BYTES = d_head * TILE_S_KV * 2 + Q_ROW_BYTES = G * d_head * 2 + K_HEAD_BYTES = S_kv * d_head * 2 + V_HEAD_BYTES = S_kv * KV_ROW_BYTES + + # Global VA per ADR-0011 D-VA1: kernel computes its own shard base. + q_base = q_ptr + cube_id * kv_per_cube * Q_ROW_BYTES + group_id_in_cube * Q_ROW_BYTES + k_shard_base = (k_ptr + + cube_id * kv_per_cube * K_HEAD_BYTES + + group_id_in_cube * K_HEAD_BYTES + + pe_in_group * n_tiles_per_pe * K_TILE_BYTES) + v_shard_base = (v_ptr + + cube_id * kv_per_cube * V_HEAD_BYTES + + group_id_in_cube * V_HEAD_BYTES + + pe_in_group * n_tiles_per_pe * TILE_S_KV * KV_ROW_BYTES) + + Q = tl.load(q_base, shape=(G, d_head), dtype="f16") + + # ── Tile 0: establish running (m, ℓ, O) with primitives ────────── + # (Reference: decode_opt2 — running state is set up with tl.dot/MATH + # primitives, not composite. Recipe-driven composite enters in tile 1+.) + K_T = tl.load(k_shard_base, shape=(d_head, TILE_S_KV), dtype="f16") + V = tl.load(v_shard_base, shape=(TILE_S_KV, d_head), dtype="f16") + scores = tl.dot(Q, K_T) + m_local = tl.max(scores, axis=-1) + centered = scores - m_local + exp_scores = tl.exp(centered) + l_local = tl.sum(exp_scores, axis=-1) + O_local = tl.dot(exp_scores, V) + + # ── Tiles 1..n_tiles_per_pe-1: two composites per tile ────────── + # #1 Q·Kᵀ composite → scores (pinned primary-out, fed into #2). + # #2 softmax_merge prologue + P·V GEMM + add epilogue, all in one + # composite: updates (m, ℓ) in place, computes P, runs P·V with + # pinned auto-bind, and folds the result into O_local. + for tile_idx in range(1, n_tiles_per_pe): + with tl.scratch_scope(): + K_T_t = tl.load(k_shard_base + tile_idx * K_TILE_BYTES, + shape=(d_head, TILE_S_KV), dtype="f16") + V_t = tl.load(v_shard_base + tile_idx * TILE_S_KV * KV_ROW_BYTES, + shape=(TILE_S_KV, d_head), dtype="f16") + scores_t = tl.composite(op="gemm", a=Q, b=K_T_t) + tl.composite( + prologue=[{"op": "softmax_merge", "s": scores_t, + "m": m_local, "l": l_local, "O": O_local}], + op="gemm", b=V_t, out=O_local, + epilogue=[{"op": "add", "other": O_local}], + ) + + # ── Chain reduce ── + group_cols = min(4, group_size) + group_rows = (group_size + group_cols - 1) // group_cols + pe_col_in_group = pe_in_group % group_cols + pe_row_in_group = pe_in_group // group_cols + + if group_cols > 1: + if pe_col_in_group < group_cols - 1: + with tl.scratch_scope(): + m_other = tl.recv(dir="intra_E", shape=m_local.shape, dtype="f16") + l_other = tl.recv(dir="intra_E", shape=l_local.shape, dtype="f16") + O_other = tl.recv(dir="intra_E", shape=O_local.shape, dtype="f16") + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + if pe_col_in_group > 0: + tl.send(dir="intra_W", src=m_local) + tl.send(dir="intra_W", src=l_local) + tl.send(dir="intra_W", src=O_local) + + if pe_col_in_group == 0 and group_rows > 1: + if pe_row_in_group < group_rows - 1: + with tl.scratch_scope(): + m_other = tl.recv(dir="intra_S", shape=m_local.shape, dtype="f16") + l_other = tl.recv(dir="intra_S", shape=l_local.shape, dtype="f16") + O_other = tl.recv(dir="intra_S", shape=O_local.shape, dtype="f16") + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_other, l_other, O_other, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + if pe_row_in_group > 0: + tl.send(dir="intra_N", src=m_local) + tl.send(dir="intra_N", src=l_local) + tl.send(dir="intra_N", src=O_local) + + if pe_in_group == 0: + O_final = O_local / l_local + o_base = (o_ptr + + cube_id * kv_per_cube * Q_ROW_BYTES + + group_id_in_cube * Q_ROW_BYTES) + tl.store(o_base, O_final) diff --git a/src/kernbench/benches/_gqa_attention_prefill_short_composite.py b/src/kernbench/benches/_gqa_attention_prefill_short_composite.py new file mode 100644 index 0000000..385b1a3 --- /dev/null +++ b/src/kernbench/benches/_gqa_attention_prefill_short_composite.py @@ -0,0 +1,216 @@ +"""GQA prefill kernel — composite (second-level) variant. + +Identical mapping to ``_gqa_attention_prefill_short.py`` (Q-tile split + +FA2 head fusion + IPCQ KV broadcast, unified A1/A2/A4/B). + +Composite hybrid (option c1): Q·Kᵀ uses ``tl.composite(op="gemm")`` so +GEMM and any future prologue/epilogue can overlap; P·V stays a plain +``tl.dot``. The reason: prefill broadcasts V over IPCQ, so on non-root +PEs V is a recv'd TCM slot (unpinned per ADR-0065 D4). Feeding it to a +composite would trigger the scheduler's "stream from HBM" branch and +PhysAddr-decode the slot address (PE scratch, bit-61 set) → PageFault. +``tl.dot`` reads TCM in place, preserving the broadcast model. + +This means prefill cannot host the ``softmax_merge`` prologue (which +attaches to the P·V composite). Full second-level fusion would require +either: (i) pinning recv handles (sim/API change, ADR-0065 supplement), +or (ii) skipping the broadcast model — both out of scope here. + +Shard addressing, layouts, and the caller contract are identical to +the first-level kernel (see its docstring for the full Q/K/V/O layout +and ADR-0011 D-VA1 offset computation). +""" +from __future__ import annotations + + +TILE_S_KV = 1024 + + +def _validate_config(*, kv_per_cube: int, T_q: int, P: int, C: int, + h_q: int, h_kv: int, S_kv: int) -> None: + """Validate composite-prefill config — caller-side, sim-cost 0. + + Mirrors first-level prefill ``_validate_config``. + """ + if kv_per_cube not in (1, 2, 4, 8): + raise ValueError(f"kv_per_cube must be in {{1,2,4,8}}; got {kv_per_cube}") + if P != 8: + raise ValueError(f"2x4 mesh broadcast geometry requires P == 8; got {P}") + if h_q % h_kv != 0: + raise ValueError(f"GQA group G = h_q/h_kv must be integer; got h_q={h_q}, h_kv={h_kv}") + if C != h_kv // kv_per_cube: + raise ValueError( + f"C must equal h_kv/kv_per_cube = {h_kv // kv_per_cube}; " + f"got C={C} (h_kv={h_kv}, kv_per_cube={kv_per_cube})" + ) + group_size = P // kv_per_cube + if T_q < group_size: + raise ValueError( + f"T_q ({T_q}) must be >= group_size ({group_size}) so every PE " + f"has q-tile work for the broadcast chain" + ) + if S_kv % TILE_S_KV != 0: + raise ValueError( + f"S_kv ({S_kv}) must be a multiple of TILE_S_KV ({TILE_S_KV})" + ) + + +def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl): + m_new = tl.maximum(m_local, m_other) + scale_old = tl.exp(m_local - m_new) + scale_new = tl.exp(m_other - m_new) + l_new = l_local * scale_old + l_other * scale_new + O_new = O_local * scale_old + O_other * scale_new + return m_new, l_new, O_new + + +def gqa_attention_prefill_short_composite_kernel( + q_ptr: int, + k_ptr: int, + v_ptr: int, + o_ptr: int, + T_q: int, + S_kv: int, + h_q: int, + h_kv: int, + d_head: int, + C: int, + P: int, + kv_per_cube: int, + *, + tl, +) -> None: + """Composite-GEMM prefill — same mapping as first-level + tl.composite. + + Caller must invoke ``_validate_config(...)`` first. + """ + group_size = P // kv_per_cube + + G = h_q // h_kv + pe_id = tl.program_id(axis=0) + cube_id = tl.program_id(axis=1) + pe_in_group = pe_id % group_size + group_id_in_cube = pe_id // group_size + + q_start = (T_q * pe_in_group) // group_size + q_end = (T_q * (pe_in_group + 1)) // group_size + T_q_pe = q_end - q_start + + n_tiles = S_kv // TILE_S_KV + + Q_ROW_BYTES = G * d_head * 2 + KV_ROW_BYTES = d_head * 2 + K_TILE_BYTES = d_head * TILE_S_KV * 2 + K_HEAD_BYTES = n_tiles * K_TILE_BYTES + V_HEAD_BYTES = S_kv * KV_ROW_BYTES + + # Global VA per ADR-0011 D-VA1: kernel computes its own shard base. + q_base = (q_ptr + + cube_id * kv_per_cube * T_q * Q_ROW_BYTES + + (group_id_in_cube * T_q + q_start) * Q_ROW_BYTES) + k_head_shard_base = (k_ptr + + cube_id * kv_per_cube * K_HEAD_BYTES + + group_id_in_cube * K_HEAD_BYTES) + v_head_shard_base = (v_ptr + + cube_id * kv_per_cube * V_HEAD_BYTES + + group_id_in_cube * V_HEAD_BYTES) + + Q = tl.load(q_base, shape=(T_q_pe * G, d_head), dtype="f16") + + group_cols = min(4, group_size) + group_rows = (group_size + group_cols - 1) // group_cols + pe_col_in_group = pe_in_group % group_cols + pe_row_in_group = pe_in_group // group_cols + + # ── Tile 0 — broadcast + persistent (m, ℓ, O) ── + if pe_in_group == 0: + K_T = tl.load(k_head_shard_base, + shape=(d_head, TILE_S_KV), dtype="f16") + V = tl.load(v_head_shard_base, + shape=(TILE_S_KV, d_head), dtype="f16") + if group_cols > 1: + tl.send(dir="intra_E", src=K_T) + tl.send(dir="intra_E", src=V) + if group_rows > 1: + tl.send(dir="intra_S", src=K_T) + tl.send(dir="intra_S", src=V) + elif pe_col_in_group == 0 and pe_row_in_group > 0: + K_T = tl.recv(dir="intra_N", + shape=(d_head, TILE_S_KV), dtype="f16") + V = tl.recv(dir="intra_N", + shape=(TILE_S_KV, d_head), dtype="f16") + if group_cols > 1: + tl.send(dir="intra_E", src=K_T) + tl.send(dir="intra_E", src=V) + else: + K_T = tl.recv(dir="intra_W", + shape=(d_head, TILE_S_KV), dtype="f16") + V = tl.recv(dir="intra_W", + shape=(TILE_S_KV, d_head), dtype="f16") + if pe_col_in_group < group_cols - 1: + tl.send(dir="intra_E", src=K_T) + tl.send(dir="intra_E", src=V) + + # Q·Kᵀ is a composite GEMM (pinned load operands). P·V stays a + # primitive tl.dot: in the broadcast topology, V on non-root PEs is + # a recv'd TCM slot (unpinned), which the composite scheduler would + # try to stream from HBM. tl.dot reads TCM in-place, preserving + # the broadcast model. softmax_merge prologue requires a composite + # P·V to host it — so prefill keeps Q·Kᵀ-only composite (no fused + # softmax epilogue); see kernel docstring (c1). + scores = tl.composite(op="gemm", a=Q, b=K_T) + m_local = tl.max(scores, axis=-1) + centered = scores - m_local + exp_scores = tl.exp(centered) + l_local = tl.sum(exp_scores, axis=-1) + O_local = tl.dot(exp_scores, V) + + # ── Tiles 1..n_tiles-1 ── + for tile_idx in range(1, n_tiles): + with tl.scratch_scope(): + if pe_in_group == 0: + K_T_t = tl.load(k_head_shard_base + tile_idx * K_TILE_BYTES, + shape=(d_head, TILE_S_KV), dtype="f16") + V_t = tl.load(v_head_shard_base + tile_idx * TILE_S_KV * KV_ROW_BYTES, + shape=(TILE_S_KV, d_head), dtype="f16") + if group_cols > 1: + tl.send(dir="intra_E", src=K_T_t) + tl.send(dir="intra_E", src=V_t) + if group_rows > 1: + tl.send(dir="intra_S", src=K_T_t) + tl.send(dir="intra_S", src=V_t) + elif pe_col_in_group == 0 and pe_row_in_group > 0: + K_T_t = tl.recv(dir="intra_N", + shape=(d_head, TILE_S_KV), dtype="f16") + V_t = tl.recv(dir="intra_N", + shape=(TILE_S_KV, d_head), dtype="f16") + if group_cols > 1: + tl.send(dir="intra_E", src=K_T_t) + tl.send(dir="intra_E", src=V_t) + else: + K_T_t = tl.recv(dir="intra_W", + shape=(d_head, TILE_S_KV), dtype="f16") + V_t = tl.recv(dir="intra_W", + shape=(TILE_S_KV, d_head), dtype="f16") + if pe_col_in_group < group_cols - 1: + tl.send(dir="intra_E", src=K_T_t) + tl.send(dir="intra_E", src=V_t) + # Q·Kᵀ composite; P·V stays primitive (see comment above). + scores_t = tl.composite(op="gemm", a=Q, b=K_T_t) + m_tile = tl.max(scores_t, axis=-1) + centered_t = scores_t - m_tile + exp_scores_t = tl.exp(centered_t) + l_tile = tl.sum(exp_scores_t, axis=-1) + O_tile = tl.dot(exp_scores_t, V_t) + m_new, l_new, O_new = _merge_running( + m_local, l_local, O_local, m_tile, l_tile, O_tile, tl=tl, + ) + tl.copy_to(m_local, m_new) + tl.copy_to(l_local, l_new) + tl.copy_to(O_local, O_new) + + O_final = O_local / l_local + o_base = (o_ptr + + cube_id * kv_per_cube * T_q * Q_ROW_BYTES + + (group_id_in_cube * T_q + q_start) * Q_ROW_BYTES) + tl.store(o_base, O_final)