Putting it all to work on one decode kernel¶
This page is generated by executing showcase.ipynb.
The notebook owns the Markdown cells, Python cells, and analysis outputs; rerun the
renderer after changing a cell to refresh the installed Markdown page.
write a complete shape
|
v
analyze one size and sweep ctx_len
|
v
change one placement decision
|
v
analyze again -> keep the evidence -> choose the next decision
The six programs are Python cells in this notebook:
| stage | source | decision |
|---|---|---|
| 0 | Stage0_Naive |
one unsplit CTA, open ctx_len |
| 1 | Stage1_Specialized |
dispatch on a measured range |
| 2 | Stage2_Sharded |
split query heads across CTAs |
| 3 | Stage3_Fused |
split the KV scan and combine online state |
| 4 | Stage4_WeightPrepared |
stage projection weights by output slice |
| 5 | Stage5_CachePrepared |
stream cache blocks through smem |
Every displayed analysis output belongs to the visible Python cell immediately above it.
The preceding %%bash cell runs the same tilefoundry analyze CLI a reader can run after
extracting attn_layer.py; the Python cell loads the report or JSON file and prints the
displayed result. The reports are static analysis results; no CUDA kernel is launched.
Executable cells¶
This is an ipynb-style rendering of the executable tutorial source. Notebook cell boundaries separate the executable Python blocks from the prose. Each Python cell is embedded as one fenced code block at the point where the tutorial uses it. The page does not repeat the complete source in a second heredoc.
To run this installed page, extract its Python cells into one executable file:
set -euo pipefail
awk '
/^<!-- tilefoundry-source: attn_layer.py -->$/ { source_block=1; next }
source_block && /^```python$/ { in_python=1; next }
in_python && /^```$/ { in_python=0; source_block=0; next }
in_python { print }
' showcase.md > attn_layer.py
chmod +x attn_layer.py
Setup cell¶
#!/usr/bin/env python3
"""A small GQA decode attention ladder for the showcase tutorial.
The six Modules keep one public shape and change one placement decision at a
time. The dimensions are intentionally small, but the query/KV ratio matches
the GQA shape used by the published Qwen attention model.
"""
from __future__ import annotations
import ast
import math
import re
import sys
from pathlib import Path
from tilefoundry.dsl import *
from tilefoundry.analysis import analyze as run_analysis
from tilefoundry.inspection.analysis_report import render_analysis, render_text
from tilefoundry.target import CudaTarget
HIDDEN = 256
QUERY_HEADS = 8
KV_HEADS = 2
HEAD_DIM = 32
KV_DIM = KV_HEADS * HEAD_DIM
GQA_GROUP = QUERY_HEADS // KV_HEADS
ROPE_CONTEXT = 8192
CTX = DimVar("ctx_len", 1, ROPE_CONTEXT)
SCALE = 1.0 / math.sqrt(HEAD_DIM)
WORKERS = 4
BLOCK = 128
_H200 = CudaTarget("nvidia.h200_sxm")
_CTA = Topology("cta", 132)
0. Start with the complete shape¶
The small teaching shape keeps the published GQA ratio:
| value | extent |
|---|---|
| hidden | 256 |
| query heads | 8 |
| KV heads | 2 |
| head dimension | 32 |
| prior cache | ctx_len in [1, 8193) |
The entry includes q/k/v projection, RoPE, one cache append, an attention scan,
and the output projection. The weights are ConstTensor parameters. The starting
point has no explicit mesh and no storage transition.
@module(entry="gqa_decode", target=_H200, topologies=(_CTA,))
class Stage0_Naive:
"""One unsplit CTA reads the complete attention sublayer."""
@func
def gqa_decode(
hidden: Tensor[(1, 1, HIDDEN), "bf16"],
w_q: ConstTensor[(1, HIDDEN, HIDDEN), "bf16"],
w_k: ConstTensor[(1, HIDDEN, KV_DIM), "bf16"],
w_v: ConstTensor[(1, HIDDEN, KV_DIM), "bf16"],
w_o: ConstTensor[(1, HIDDEN, HIDDEN), "bf16"],
k_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), "bf16"],
v_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), "bf16"],
cur_pos: Tensor[(1,), "i32"],
write_len: Tensor[(1,), "i32"],
pos_ids: Tensor[(1,), "i32"],
cos_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), "bf16"],
sin_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), "bf16"],
) -> Tensor[(1, 1, HIDDEN), "bf16"]:
q = tf.reshape(tf.matmul(hidden, w_q), new_shape=(1, 1, QUERY_HEADS, HEAD_DIM))
k = tf.reshape(tf.matmul(hidden, w_k), new_shape=(1, 1, KV_HEADS, HEAD_DIM))
v = tf.reshape(tf.matmul(hidden, w_v), new_shape=(1, 1, KV_HEADS, HEAD_DIM))
q_rope, k_rope = tf.rope(q, k, cos_cache, sin_cache, pos_ids)
k_all = tf.cache_update(k_cache, cur_pos, write_len, k_rope)
v_all = tf.cache_update(v_cache, cur_pos, write_len, v)
k_heads = tf.repeat_interleave(k_all, repeats=GQA_GROUP, axis=2)
v_heads = tf.repeat_interleave(v_all, repeats=GQA_GROUP, axis=2)
k_heads = tf.transpose(k_heads, perm=(0, 2, 1, 3))
v_heads = tf.transpose(v_heads, perm=(0, 2, 1, 3))
q_f32 = tf.cast(q_rope, dtype="f32")
k_f32 = tf.cast(k_heads, dtype="f32")
v_f32 = tf.cast(v_heads, dtype="f32")
scaled_q = q_f32 * tf.full_like(q_f32, value=SCALE)
q_e = tf.reshape(scaled_q, new_shape=(1, 1, QUERY_HEADS, 1, HEAD_DIM))
k_e = tf.reshape(k_f32, new_shape=(1, 1, QUERY_HEADS, CTX, HEAD_DIM))
v_e = tf.reshape(v_f32, new_shape=(1, 1, QUERY_HEADS, CTX, HEAD_DIM))
scores = tf.reduce(q_e * k_e, axes=(-1,), keepdim=True, kind="sum")
peak = tf.reduce(scores, axes=(-2,), keepdim=True, kind="max")
weights = tf.exp(scores - peak)
normalizer = tf.reduce(weights, axes=(-2,), keepdim=False, kind="sum")
weighted = tf.reduce(weights * v_e, axes=(-2,), keepdim=False, kind="sum")
attended = weighted / normalizer
attended_bf16 = tf.cast(attended, dtype="bf16")
return tf.matmul(tf.reshape(attended_bf16, new_shape=(1, 1, HIDDEN)), w_o)
gqa_decode = Stage0_Naive.entry_function()
Run one point like this. The Bash cell writes the report; the following Python cell loads it and prints its header plus selected annotated calls:
set -euo pipefail
mkdir -p tutorial-reports
tilefoundry analyze attn_layer.py:Stage0_Naive \
tutorial-reports/stage0-128.txt \
--compute-cost --memory --roofline --operands --dim ctx_len=128
from pathlib import Path
report = Path("tutorial-reports/stage0-128.txt").read_text(encoding="utf-8")
header, separator, annotated = report.partition("\n\n")
print(header.rstrip())
print()
for needle in ("matmul(hidden, w_q", "cache_update(k_cache", ", w_o)"):
line = next(line for line in annotated.splitlines() if needle in line)
print(line.rstrip())
# analysis target=nvidia.h200_sxm module=Stage0_Naive function=gqa_decode topology=cta wave=132/132
# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline
# compute-cost flops=bf16:328896@logical,328896@total,328896@cta;f32:200448@logical,200448@total,200448@cta other-ops=special:1024@logical,1024@total,1024@cta precision=exact
# memory traffic=gmem:r1.12MB/w787.69KB@logical,r1.12MB/w787.69KB@total,r1.12MB/w787.69KB@cta footprint=hidden:512B;w_q:128.00KB;v0:1.00KB;w_k:32.00KB;v2:256B;cos_cache:1.00KB;sin_cache:1.00KB;pos_ids:4B;rope[0]:1.00KB;rope[1]:256B;q_f32:2.00KB;v8:1.00KB;scaled_q:2.00KB;cur_pos:4B;write_len:4B;k_all:16.12KB;k_heads:128.00KB;k_heads_2:128.00KB;k_f32:256.00KB;v18:256.00KB;scores:8.00KB;peak:64B;v21:8.00KB;weights:8.00KB;w_v:32.00KB;v23:256B;v_all:16.12KB;v_heads:128.00KB;v_heads_2:128.00KB;v_f32:256.00KB;v30:256.00KB;weighted:2.00KB;normalizer:64B;attended:2.00KB;attended_bf16:1.00KB;w_o:128.00KB;v36:512B footprint-precision=upper_bound peak=gmem:1.54MB persistent=gmem:1.34MB
# roofline ideal-ns=412 bound-by=memory
v0 = tf.matmul(hidden, w_q) # Tensor[(1, 1, 256), "bf16"]; compute-cost flops=bf16:131072@logical,131072@total,131072@cta precision=exact; memory traffic=gmem:r128.50KB/w512B@logical,r128.50KB/w512B@total,r128.50KB/w512B@cta footprint=hidden:512B;w_q:128.00KB;v0:512B footprint-precision=exact operands=0:r512B/w0;1:r128.00KB/w0;result:r0/w512B; roofline ideal-ns=28 bound-by=memory
k_all = tf.cache_update(k_cache, cur_pos, write_len, k_rope) # Tensor[(1, 128, 2, 32), "bf16"]; compute-cost precision=exact; memory traffic=gmem:r136B/w128B@logical,r136B/w128B@total,r136B/w128B@cta footprint=cur_pos:4B;write_len:4B;rope[1]:128B;k_all:128B footprint-precision=upper_bound operands=0:r0/w0;1:r4B/w0;2:r4B/w0;3:r128B/w0;result:r0/w128B; roofline ideal-ns=1 bound-by=memory
v36 = tf.matmul(v35, w_o) # Tensor[(1, 1, 256), "bf16", (1, 1, 256)]; compute-cost flops=bf16:131072@logical,131072@total,131072@cta precision=exact; memory traffic=gmem:r128.50KB/w512B@logical,r128.50KB/w512B@total,r128.50KB/w512B@cta footprint=attended_bf16:512B;w_o:128.00KB;v36:512B footprint-precision=exact operands=0:r512B/w0;1:r128.00KB/w0;result:r0/w512B; roofline ideal-ns=28 bound-by=memory
@logical is the authored request before loop replication; @total is the
whole execution. After them comes one share per topology level the program names,
each under that level's own name: what one CTA does, what one thread does. With no
authored split a level's share is the total. The annotated lines come from the same report
file, so the call-level operands and roofline numbers stay tied to the command
that produced them.
w_q and w_o are each 256 * 256 * 2 = 131072 bytes. w_k and w_v are each 256 * 64 * 2 = 32768 bytes, so the four projection weights total 327680 bytes. That fixed amount is separate from the cache scan.
1. Sweep the open dimension¶
The same command, with a different report path, produces the table. The Bash cell runs all six CLI calls; the following Python cell loads each fresh report and prints the selected fields as Markdown:
set -euo pipefail
mkdir -p tutorial-reports
for ctx in 128 512 1024 2048 4096 8192; do
tilefoundry analyze attn_layer.py:Stage0_Naive \
tutorial-reports/stage0-$ctx.txt \
--compute-cost --memory --roofline --operands --dim ctx_len=$ctx
done
import re
from pathlib import Path
def metrics(ctx_len):
report = Path(f"tutorial-reports/stage0-{ctx_len}.txt").read_text(encoding="utf-8")
lines = report.splitlines()
compute = next(line for line in lines if line.startswith("# compute-cost "))
memory = next(line for line in lines if line.startswith("# memory "))
roofline = next(line for line in lines if line.startswith("# roofline "))
f32 = re.search(r"f32:([^ ]+)", compute).group(1)
traffic_value = re.search(r"traffic=(.*?)(?: footprint=| peak=)", memory).group(1)
gmem_peak = re.search(r"peak=.*?gmem:([^;, ]+)", memory).group(1)
ideal, bound = re.search(r"ideal-ns=([^ ]+) bound-by=([^ ]+)", roofline).groups()
return f32, traffic_value, gmem_peak, ideal, bound
print("| `ctx_len` | f32 flops `logical@total@CTA` | traffic `logical@total@CTA` | peak gmem bytes | ideal ns | bound |")
print("|---:|---:|---|---:|---:|---|")
for ctx_len in (128, 512, 1024, 2048, 4096, 8192):
f32, traffic, peak, ideal, bound = metrics(ctx_len)
print(f"| {ctx_len} | `{f32}` | `{traffic}` | {peak} | {ideal} | {bound} |")
ctx_len |
f32 flops logical@total@CTA |
traffic logical@total@CTA |
peak gmem bytes | ideal ns | bound |
|---|---|---|---|---|---|
| 128 | 200448@logical,200448@total,200448@cta |
gmem:r1.12MB/w787.69KB@logical,r1.12MB/w787.69KB@total,r1.12MB/w787.69KB@cta |
1.54MB | 412 | memory |
| 512 | 799488@logical,799488@total,799488@cta |
gmem:r3.53MB/w3.05MB@logical,r3.53MB/w3.05MB@total,r3.53MB/w3.05MB@cta |
2.20MB | 1431 | memory |
| 1024 | 1598208@logical,1598208@total,1598208@cta |
gmem:r6.73MB/w6.10MB@logical,r6.73MB/w6.10MB@total,r6.73MB/w6.10MB@cta |
3.09MB | 2789 | memory |
| 2048 | 3195648@logical,3195648@total,3195648@cta |
gmem:r13.14MB/w12.20MB@logical,r13.14MB/w12.20MB@total,r13.14MB/w12.20MB@cta |
4.88MB | 5506 | memory |
| 4096 | 6390528@logical,6390528@total,6390528@cta |
gmem:r25.95MB/w24.38MB@logical,r25.95MB/w24.38MB@total,r25.95MB/w24.38MB@cta |
8.44MB | 10940 | memory |
| 8192 | 12780288@logical,12780288@total,12780288@cta |
gmem:r51.57MB/w48.76MB@logical,r51.57MB/w48.76MB@total,r51.57MB/w48.76MB@cta |
15.56MB | 21808 | memory |
The table says:
per-CTA work = global work -> no authored split yet
weight bytes = fixed -> projection weights are a staging target
cache scan = grows with ctx_len -> a full-cache residency decision will fail first
2. Specialize at the cache-only capacity estimate¶
The full-cache sharded program in Stage2_Sharded places one local query head's K and V cache in smem. A cache-only estimate follows from the target capacity:
bytes per ctx per CTA = K + V
= 2 * HEAD_DIM * sizeof(bf16)
= 2 * 32 * 2
= 128 B
T = floor(232448 B / 128 B)
= 1816
This is not the complete placement peak: other simultaneously resident values also occupy smem. Stage1_Specialized still uses the estimate to express two closed RangePattern variants, [1, 1816] and [1817, 8192]; the CP placement report below decides whether either variant actually fits. The Stage1 body is deliberately the unsplit baseline, so the dispatch contract can be read independently from the later implementations.
SMEM_BUDGET = 232448
CACHE_BYTES_PER_CONTEXT_PER_CTA = 2 * HEAD_DIM * 2
SPECIALIZE_T = SMEM_BUDGET // CACHE_BYTES_PER_CONTEXT_PER_CTA
@module(entry="gqa_decode", target=_H200, topologies=(_CTA,))
class Stage1_Specialized:
"""Dispatch the same unsplit kernel at the measured context boundary."""
@func
def _decode_core(
hidden: Tensor[(1, 1, HIDDEN), "bf16"],
w_q: ConstTensor[(1, HIDDEN, HIDDEN), "bf16"],
w_k: ConstTensor[(1, HIDDEN, KV_DIM), "bf16"],
w_v: ConstTensor[(1, HIDDEN, KV_DIM), "bf16"],
w_o: ConstTensor[(1, HIDDEN, HIDDEN), "bf16"],
k_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), "bf16"],
v_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), "bf16"],
cur_pos: Tensor[(1,), "i32"],
write_len: Tensor[(1,), "i32"],
pos_ids: Tensor[(1,), "i32"],
cos_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), "bf16"],
sin_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), "bf16"],
) -> Tensor[(1, 1, HIDDEN), "bf16"]:
q = tf.reshape(tf.matmul(hidden, w_q), new_shape=(1, 1, QUERY_HEADS, HEAD_DIM))
k = tf.reshape(tf.matmul(hidden, w_k), new_shape=(1, 1, KV_HEADS, HEAD_DIM))
v = tf.reshape(tf.matmul(hidden, w_v), new_shape=(1, 1, KV_HEADS, HEAD_DIM))
q_rope, k_rope = tf.rope(q, k, cos_cache, sin_cache, pos_ids)
k_all = tf.cache_update(k_cache, cur_pos, write_len, k_rope)
v_all = tf.cache_update(v_cache, cur_pos, write_len, v)
k_heads = tf.transpose(
tf.repeat_interleave(k_all, repeats=GQA_GROUP, axis=2), perm=(0, 2, 1, 3)
)
v_heads = tf.transpose(
tf.repeat_interleave(v_all, repeats=GQA_GROUP, axis=2), perm=(0, 2, 1, 3)
)
q_f32 = tf.cast(q_rope, dtype="f32")
k_f32 = tf.cast(k_heads, dtype="f32")
v_f32 = tf.cast(v_heads, dtype="f32")
scaled_q = q_f32 * tf.full_like(q_f32, value=SCALE)
q_e = tf.reshape(scaled_q, new_shape=(1, 1, QUERY_HEADS, 1, HEAD_DIM))
k_e = tf.reshape(k_f32, new_shape=(1, 1, QUERY_HEADS, CTX, HEAD_DIM))
v_e = tf.reshape(v_f32, new_shape=(1, 1, QUERY_HEADS, CTX, HEAD_DIM))
scores = tf.reduce(q_e * k_e, axes=(-1,), keepdim=True, kind="sum")
peak = tf.reduce(scores, axes=(-2,), keepdim=True, kind="max")
weights = tf.exp(scores - peak)
normalizer = tf.reduce(weights, axes=(-2,), keepdim=False, kind="sum")
weighted = tf.reduce(weights * v_e, axes=(-2,), keepdim=False, kind="sum")
attended = tf.cast(weighted / normalizer, dtype="bf16")
return tf.matmul(tf.reshape(attended, new_shape=(1, 1, HIDDEN)), w_o)
@func
def gqa_decode(
hidden: Tensor[(1, 1, HIDDEN), "bf16"],
w_q: ConstTensor[(1, HIDDEN, HIDDEN), "bf16"],
w_k: ConstTensor[(1, HIDDEN, KV_DIM), "bf16"],
w_v: ConstTensor[(1, HIDDEN, KV_DIM), "bf16"],
w_o: ConstTensor[(1, HIDDEN, HIDDEN), "bf16"],
k_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), "bf16"],
v_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), "bf16"],
cur_pos: Tensor[(1,), "i32"],
write_len: Tensor[(1,), "i32"],
pos_ids: Tensor[(1,), "i32"],
cos_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), "bf16"],
sin_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), "bf16"],
) -> Tensor[(1, 1, HIDDEN), "bf16"]:
pass
@gqa_decode.specialize(RangePattern("ctx_len", 1, SPECIALIZE_T))
def short_context(
hidden: Tensor[(1, 1, HIDDEN), "bf16"],
w_q: ConstTensor[(1, HIDDEN, HIDDEN), "bf16"],
w_k: ConstTensor[(1, HIDDEN, KV_DIM), "bf16"],
w_v: ConstTensor[(1, HIDDEN, KV_DIM), "bf16"],
w_o: ConstTensor[(1, HIDDEN, HIDDEN), "bf16"],
k_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), "bf16"],
v_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), "bf16"],
cur_pos: Tensor[(1,), "i32"],
write_len: Tensor[(1,), "i32"],
pos_ids: Tensor[(1,), "i32"],
cos_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), "bf16"],
sin_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), "bf16"],
) -> Tensor[(1, 1, HIDDEN), "bf16"]:
return _decode_core(
hidden, w_q, w_k, w_v, w_o, k_cache, v_cache,
cur_pos, write_len, pos_ids, cos_cache, sin_cache,
)
@gqa_decode.specialize(RangePattern("ctx_len", SPECIALIZE_T + 1, ROPE_CONTEXT))
def long_context(
hidden: Tensor[(1, 1, HIDDEN), "bf16"],
w_q: ConstTensor[(1, HIDDEN, HIDDEN), "bf16"],
w_k: ConstTensor[(1, HIDDEN, KV_DIM), "bf16"],
w_v: ConstTensor[(1, HIDDEN, KV_DIM), "bf16"],
w_o: ConstTensor[(1, HIDDEN, HIDDEN), "bf16"],
k_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), "bf16"],
v_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), "bf16"],
cur_pos: Tensor[(1,), "i32"],
write_len: Tensor[(1,), "i32"],
pos_ids: Tensor[(1,), "i32"],
cos_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), "bf16"],
sin_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), "bf16"],
) -> Tensor[(1, 1, HIDDEN), "bf16"]:
return _decode_core(
hidden, w_q, w_k, w_v, w_o, k_cache, v_cache,
cur_pos, write_len, pos_ids, cos_cache, sin_cache,
)
gqa_decode_specialized = Stage1_Specialized.entry_function()
The Bash cell writes the complete report at the cache-only estimate. Analysis still returns when the solved placement exceeds capacity; the report header carries a non-fatal # error= line:
set -euo pipefail
mkdir -p tutorial-reports
tilefoundry analyze attn_layer.py:Stage2_Sharded \
tutorial-reports/stage2-1816.txt \
--compute-cost --memory --roofline --dim ctx_len=1816
from pathlib import Path
report = Path("tutorial-reports/stage2-1816.txt").read_text(encoding="utf-8")
print(report.partition("\n\n")[0].rstrip())
# analysis target=nvidia.h200_sxm module=Stage2_Sharded function=gqa_decode topology=cta wave=132/132
# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline
# compute-cost flops=bf16:2203264@logical,4503744@total,562968@cta;f32:44352@logical,44352@total,5544@cta other-ops=special:14528@logical,14528@total,1816@cta precision=exact
# memory traffic=gmem:r4.31MB/w3.55MB@logical,r22.05MB/w28.39MB@total,r2.76MB/w3.55MB@cta;rmem:r315.19KB/w201.69KB@logical,r315.19KB/w201.69KB@total,r39.40KB/w25.21KB@cta;smem:r3.55MB/w3.55MB@logical,r3.55MB/w3.55MB@total,r454.62KB/w454.62KB@cta footprint=hidden:512B;w_q:128.00KB;v0:1.00KB;w_k:32.00KB;v2:256B;cos_cache:1.00KB;sin_cache:1.00KB;pos_ids:4B;rope[0]:1.00KB;rope[1]:256B;cur_pos:4B;write_len:4B;k_all:227.12KB;v16:1.77MB;k_heads:1.77MB;w_v:32.00KB;v25:256B;v_all:227.12KB;v28:1.77MB;v_heads:1.77MB;attended_2:1.00KB;w_o:128.00KB;v39:512B footprint-precision=upper_bound peak=gmem:3.53MB;rmem:10.77KB;smem:227.12KB persistent=gmem:1.76MB
# buffer=w_q holds=7.86MB time=none space=cta.head reuse=896.00KB fits=yes precision=upper_bound
# buffer=w_o holds=7.86MB time=none space=cta.head reuse=896.00KB fits=yes precision=upper_bound
# buffer=w_k holds=7.86MB time=none space=cta.head reuse=224.00KB fits=yes precision=upper_bound
# buffer=w_v holds=7.86MB time=none space=cta.head reuse=224.00KB fits=yes precision=upper_bound
# buffer=cos_cache holds=7.86MB time=none space=cta.head reuse=7.00KB fits=yes precision=upper_bound
# buffer=sin_cache holds=7.86MB time=none space=cta.head reuse=7.00KB fits=yes precision=upper_bound
# buffer=hidden holds=7.86MB time=none space=cta.head reuse=3.50KB fits=yes precision=upper_bound
# buffer=pos_ids holds=7.86MB time=none space=cta.head reuse=28B fits=yes precision=upper_bound
# buffer=cur_pos holds=7.86MB time=none space=cta.head reuse=28B fits=yes precision=upper_bound
# buffer=write_len holds=7.86MB time=none space=cta.head reuse=28B fits=yes precision=upper_bound
# error="smem placement peak 227.12KB exceeds capacity 227.00KB"
# roofline ideal-ns=1717 bound-by=memory
A larger context increases the solved peak. The CLI still succeeds and writes the complete report, including its capacity error:
set -euo pipefail
mkdir -p tutorial-reports
tilefoundry analyze attn_layer.py:Stage2_Sharded \
tutorial-reports/stage2-1820.txt \
--compute-cost --memory --roofline --dim ctx_len=1820
from pathlib import Path
report = Path("tutorial-reports/stage2-1820.txt").read_text(encoding="utf-8")
print(report.partition("\n\n")[0].rstrip())
# analysis target=nvidia.h200_sxm module=Stage2_Sharded function=gqa_decode topology=cta wave=132/132
# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline
# compute-cost flops=bf16:2207392@logical,4507872@total,563484@cta;f32:44448@logical,44448@total,5556@cta other-ops=special:14560@logical,14560@total,1820@cta precision=exact
# memory traffic=gmem:r4.32MB/w3.56MB@logical,r22.09MB/w28.45MB@total,r2.76MB/w3.56MB@cta;rmem:r315.88KB/w202.12KB@logical,r315.88KB/w202.12KB@total,r39.48KB/w25.27KB@cta;smem:r3.56MB/w3.56MB@logical,r3.56MB/w3.56MB@total,r455.62KB/w455.62KB@cta footprint=hidden:512B;w_q:128.00KB;v0:1.00KB;w_k:32.00KB;v2:256B;cos_cache:1.00KB;sin_cache:1.00KB;pos_ids:4B;rope[0]:1.00KB;rope[1]:256B;cur_pos:4B;write_len:4B;k_all:227.62KB;v16:1.78MB;k_heads:1.78MB;w_v:32.00KB;v25:256B;v_all:227.62KB;v28:1.78MB;v_heads:1.78MB;attended_2:1.00KB;w_o:128.00KB;v39:512B footprint-precision=upper_bound peak=gmem:3.53MB;rmem:10.79KB;smem:227.62KB persistent=gmem:1.76MB
# buffer=w_q holds=7.87MB time=none space=cta.head reuse=896.00KB fits=yes precision=upper_bound
# buffer=w_o holds=7.87MB time=none space=cta.head reuse=896.00KB fits=yes precision=upper_bound
# buffer=w_k holds=7.87MB time=none space=cta.head reuse=224.00KB fits=yes precision=upper_bound
# buffer=w_v holds=7.87MB time=none space=cta.head reuse=224.00KB fits=yes precision=upper_bound
# buffer=cos_cache holds=7.87MB time=none space=cta.head reuse=7.00KB fits=yes precision=upper_bound
# buffer=sin_cache holds=7.87MB time=none space=cta.head reuse=7.00KB fits=yes precision=upper_bound
# buffer=hidden holds=7.87MB time=none space=cta.head reuse=3.50KB fits=yes precision=upper_bound
# buffer=pos_ids holds=7.87MB time=none space=cta.head reuse=28B fits=yes precision=upper_bound
# buffer=cur_pos holds=7.87MB time=none space=cta.head reuse=28B fits=yes precision=upper_bound
# buffer=write_len holds=7.87MB time=none space=cta.head reuse=28B fits=yes precision=upper_bound
# error="smem placement peak 227.62KB exceeds capacity 227.00KB"
# roofline ideal-ns=1720 bound-by=memory
The formula supplies a cache-only dispatch estimate, while the solved placement and its report error expose the complete capacity requirement.
3. Split the query heads¶
The next change is Stage2_Sharded: one cta.head owns one query head. The same-size comparison isolates placement from context growth.
@module(entry="gqa_decode", target=_H200, topologies=(_CTA,))
class Stage2_Sharded:
"""Give each CTA one query-head slice while keeping the cache whole."""
@func
def gqa_decode(
hidden: Tensor[(1, 1, HIDDEN), "bf16"],
w_q: ConstTensor[(1, HIDDEN, HIDDEN), "bf16"],
w_k: ConstTensor[(1, HIDDEN, KV_DIM), "bf16"],
w_v: ConstTensor[(1, HIDDEN, KV_DIM), "bf16"],
w_o: ConstTensor[(1, HIDDEN, HIDDEN), "bf16"],
k_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), "bf16"],
v_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), "bf16"],
cur_pos: Tensor[(1,), "i32"],
write_len: Tensor[(1,), "i32"],
pos_ids: Tensor[(1,), "i32"],
cos_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), "bf16"],
sin_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), "bf16"],
) -> Tensor[(1, 1, HIDDEN), "bf16"]:
with Mesh(("cta",), layout=(QUERY_HEADS,), names=("head",)) as cta:
q = tf.reshape(
tf.matmul(hidden, w_q), new_shape=(1, 1, QUERY_HEADS, HEAD_DIM)
)
k = tf.reshape(tf.matmul(hidden, w_k), new_shape=(1, 1, KV_HEADS, HEAD_DIM))
v = tf.reshape(tf.matmul(hidden, w_v), new_shape=(1, 1, KV_HEADS, HEAD_DIM))
q_rope, k_rope = tf.rope(q, k, cos_cache, sin_cache, pos_ids)
k_all = tf.cache_update(k_cache, cur_pos, write_len, k_rope)
v_all = tf.cache_update(v_cache, cur_pos, write_len, v)
q_sh = tf.reshard(
q_rope, (1, 1, QUERY_HEADS @ cta.head, HEAD_DIM), "smem"
)
k_heads = tf.transpose(
tf.repeat_interleave(k_all, repeats=GQA_GROUP, axis=2), perm=(0, 2, 1, 3)
)
v_heads = tf.transpose(
tf.repeat_interleave(v_all, repeats=GQA_GROUP, axis=2), perm=(0, 2, 1, 3)
)
k_sh = tf.reshard(
k_heads, (1, QUERY_HEADS @ cta.head, CTX, HEAD_DIM), "smem"
)
v_sh = tf.reshard(
v_heads, (1, QUERY_HEADS @ cta.head, CTX, HEAD_DIM), "smem"
)
queries = tf.transpose(tf.cast(q_sh, dtype="f32"), perm=(0, 2, 1, 3))
keys = tf.transpose(k_sh, perm=(0, 1, 3, 2))
values = tf.transpose(v_sh, perm=(0, 1, 2, 3))
scaled_q = queries * tf.full_like(queries, value=SCALE)
scores = tf.matmul(tf.cast(scaled_q, dtype="bf16"), keys, out_dtype="f32")
peak = tf.reduce(scores, axes=(-1,), keepdim=True, kind="max")
weights = tf.exp(scores - peak)
normalizer = tf.reduce(weights, axes=(-1,), keepdim=True, kind="sum")
weighted = tf.matmul(tf.cast(weights, dtype="bf16"), values, out_dtype="f32")
attended = tf.transpose(
tf.cast(weighted / normalizer, dtype="bf16"), perm=(0, 2, 1, 3)
)
attended = tf.reshard(attended, (1, 1, QUERY_HEADS, HEAD_DIM), "gmem")
return tf.matmul(tf.reshape(attended, new_shape=(1, 1, HIDDEN)), w_o)
gqa_decode_sharded = Stage2_Sharded.entry_function()
Baseline at ctx_len=128. The Bash cell reruns the CLI; the following Python cell loads
the output file and prints the selected report lines:
set -euo pipefail
mkdir -p tutorial-reports
tilefoundry analyze attn_layer.py:Stage0_Naive \
tutorial-reports/stage0-128-summary.txt \
--compute-cost --memory --roofline --dim ctx_len=128
from pathlib import Path
report = Path("tutorial-reports/stage0-128-summary.txt").read_text(encoding="utf-8")
for line in report.splitlines():
if line.startswith(("# compute-cost ", "# memory ", "# roofline ")):
print(line)
# compute-cost flops=bf16:328896@logical,328896@total,328896@cta;f32:200448@logical,200448@total,200448@cta other-ops=special:1024@logical,1024@total,1024@cta precision=exact
# memory traffic=gmem:r1.12MB/w787.69KB@logical,r1.12MB/w787.69KB@total,r1.12MB/w787.69KB@cta footprint=hidden:512B;w_q:128.00KB;v0:1.00KB;w_k:32.00KB;v2:256B;cos_cache:1.00KB;sin_cache:1.00KB;pos_ids:4B;rope[0]:1.00KB;rope[1]:256B;q_f32:2.00KB;v8:1.00KB;scaled_q:2.00KB;cur_pos:4B;write_len:4B;k_all:16.12KB;k_heads:128.00KB;k_heads_2:128.00KB;k_f32:256.00KB;v18:256.00KB;scores:8.00KB;peak:64B;v21:8.00KB;weights:8.00KB;w_v:32.00KB;v23:256B;v_all:16.12KB;v_heads:128.00KB;v_heads_2:128.00KB;v_f32:256.00KB;v30:256.00KB;weighted:2.00KB;normalizer:64B;attended:2.00KB;attended_bf16:1.00KB;w_o:128.00KB;v36:512B footprint-precision=upper_bound peak=gmem:1.54MB persistent=gmem:1.34MB
# roofline ideal-ns=412 bound-by=memory
Head-sharded at ctx_len=128. The Bash cell reruns the CLI; the following Python cell
loads the output file and prints the selected report lines:
set -euo pipefail
mkdir -p tutorial-reports
tilefoundry analyze attn_layer.py:Stage2_Sharded \
tutorial-reports/stage2-128-summary.txt \
--compute-cost --memory --roofline --dim ctx_len=128
from pathlib import Path
report = Path("tutorial-reports/stage2-128-summary.txt").read_text(encoding="utf-8")
for line in report.splitlines():
if line.startswith(("# compute-cost ", "# memory ", "# roofline ")):
print(line)
# compute-cost flops=bf16:461248@logical,2761728@total,345216@cta;f32:3840@logical,3840@total,480@cta other-ops=special:1024@logical,1024@total,128@cta precision=exact
# memory traffic=gmem:r613.39KB/w258.62KB@logical,r3.91MB/w2.02MB@total,r500.96KB/w258.19KB@cta;rmem:r25.06KB/w17.06KB@logical,r25.06KB/w17.06KB@total,r3.13KB/w2.13KB@cta;smem:r261.00KB/w261.00KB@logical,r261.00KB/w261.00KB@total,r32.62KB/w32.62KB@cta footprint=hidden:512B;w_q:128.00KB;v0:1.00KB;w_k:32.00KB;v2:256B;cos_cache:1.00KB;sin_cache:1.00KB;pos_ids:4B;rope[0]:1.00KB;rope[1]:256B;cur_pos:4B;write_len:4B;k_all:16.12KB;v16:128.00KB;k_heads:128.00KB;w_v:32.00KB;v25:256B;v_all:16.12KB;v28:128.00KB;v_heads:128.00KB;attended_2:1.00KB;w_o:128.00KB;v39:512B footprint-precision=upper_bound peak=gmem:1.47MB;rmem:896B;smem:16.12KB persistent=gmem:1.34MB
# roofline ideal-ns=186 bound-by=memory
The two attention contractions read bf16 operands into f32 results (out_dtype="f32"),
so their work is bf16 work: each head CTA adds 2 * 2 * 128 * 32 = 16384 matmul flops
plus 160 for the two casts that narrow their f32 inputs. Only the softmax arithmetic stays
in f32, 480 flops per CTA, and the special work divides by the eight head CTAs. The remaining bf16 work
is the projection, placement and gather work that is not head-shardable. ideal-ns also
changes because the authored storage choices change the traffic seen by the roofline
calculation.
4. Keep the scan state on chip¶
At long context, the full-cache form is the wrong residency choice. Stage3_Fused uses a two-dimensional CTA mesh. The head axis owns query heads and the worker axis owns disjoint cache blocks. Each worker keeps online (m, l, acc) state, then the worker axis is combined with an explicit log-sum-exp merge.
The single gqa_decode function prepares QKV and updates the caches inside a
single-CTA mesh, then scans them inside the head/worker mesh.
@module(entry="gqa_decode", target=_H200, topologies=(_CTA,))
class Stage3_Fused:
"""Split the cache scan across workers and combine online-softmax partials."""
@func
def gqa_decode(
hidden: Tensor[(1, 1, HIDDEN), "bf16"],
w_q: ConstTensor[(1, HIDDEN, HIDDEN), "bf16"],
w_k: ConstTensor[(1, HIDDEN, KV_DIM), "bf16"],
w_v: ConstTensor[(1, HIDDEN, KV_DIM), "bf16"],
w_o: ConstTensor[(1, HIDDEN, HIDDEN), "bf16"],
k_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), "bf16"],
v_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), "bf16"],
cur_pos: Tensor[(1,), "i32"],
write_len: Tensor[(1,), "i32"],
pos_ids: Tensor[(1,), "i32"],
cos_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), "bf16"],
sin_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), "bf16"],
) -> Tensor[(1, 1, HIDDEN), "bf16"]:
with Mesh(("cta",), layout=(1,), names=("tile",)) as _cta:
q = tf.reshape(
tf.matmul(hidden, w_q), new_shape=(1, 1, QUERY_HEADS, HEAD_DIM)
)
k = tf.reshape(tf.matmul(hidden, w_k), new_shape=(1, 1, KV_HEADS, HEAD_DIM))
v = tf.reshape(tf.matmul(hidden, w_v), new_shape=(1, 1, KV_HEADS, HEAD_DIM))
q_rope, k_rope = tf.rope(q, k, cos_cache, sin_cache, pos_ids)
k_all = tf.cache_update(k_cache, cur_pos, write_len, k_rope)
v_all = tf.cache_update(v_cache, cur_pos, write_len, v)
k_heads = tf.repeat_interleave(k_all, repeats=GQA_GROUP, axis=2)
v_heads = tf.repeat_interleave(v_all, repeats=GQA_GROUP, axis=2)
with Mesh(
("cta",), layout=(QUERY_HEADS, WORKERS), names=("head", "worker")
) as cta:
qh = tf.reshard(
q_rope, (1, 1, QUERY_HEADS @ cta.head, HEAD_DIM), "smem"
)
queries = tf.transpose(tf.cast(qh, dtype="f32"), perm=(0, 2, 1, 3))
scaled = queries * tf.full_like(queries, value=SCALE)
m_slots = tf.zeros(
Tensor[
(WORKERS @ cta.worker, QUERY_HEADS @ cta.head, 1, 1),
"f32",
"smem",
]
)
acc_slots = tf.zeros(
Tensor[
(WORKERS @ cta.worker, QUERY_HEADS @ cta.head, 1, HEAD_DIM),
"f32",
"smem",
]
)
m = tf.full_like(m_slots, value=-1e30)
l = tf.full_like(m_slots, value=0.0)
acc = tf.full_like(acc_slots, value=0.0)
for start in tf.tile(CTX, BLOCK * WORKERS):
base = start + cta.worker * BLOCK
kb = tf.reshard(
k_heads[:, base : base + BLOCK, :, :],
(1, BLOCK, QUERY_HEADS @ cta.head, HEAD_DIM),
"smem",
)
vb = tf.reshard(
v_heads[:, base : base + BLOCK, :, :],
(1, BLOCK, QUERY_HEADS @ cta.head, HEAD_DIM),
"smem",
)
keys = tf.transpose(kb, perm=(0, 2, 3, 1))
values = tf.transpose(vb, perm=(0, 2, 1, 3))
scores = tf.reshard(
tf.matmul(tf.cast(scaled, dtype="bf16"), keys, out_dtype="f32"),
(1, QUERY_HEADS @ cta.head, 1, BLOCK),
"smem",
)
block_m = tf.reduce(scores, axes=(-1,), keepdim=True, kind="max")
next_m = tf.max(m, block_m)
correction = tf.exp(m - next_m)
weights = tf.exp(scores - next_m)
l = l * correction + tf.reduce(
weights, axes=(-1,), keepdim=True, kind="sum"
)
acc = acc * correction + tf.reshard(
tf.matmul(tf.cast(weights, dtype="bf16"), values, out_dtype="f32"),
(WORKERS @ cta.worker, QUERY_HEADS @ cta.head, 1, HEAD_DIM),
"smem",
)
m = next_m
all_m = tf.reshard(
m, (WORKERS, QUERY_HEADS @ cta.head, 1, 1), "smem"
)
all_l = tf.reshard(
l, (WORKERS, QUERY_HEADS @ cta.head, 1, 1), "smem"
)
all_acc = tf.reshard(
acc, (WORKERS, QUERY_HEADS @ cta.head, 1, HEAD_DIM), "smem"
)
global_m = tf.reduce(all_m, axes=(0,), keepdim=False, kind="max")
weights = tf.exp(all_m - global_m)
global_l = tf.reduce(weights * all_l, axes=(0,), keepdim=False, kind="sum")
global_acc = tf.reduce(
weights * all_acc, axes=(0,), keepdim=False, kind="sum"
)
attended = tf.cast(global_acc / global_l, dtype="bf16")
attended = tf.reshape(
attended, new_shape=(1, 1, QUERY_HEADS, HEAD_DIM)
)
attended = tf.reshard(
attended, (1, 1, QUERY_HEADS, HEAD_DIM), "gmem"
)
return tf.matmul(tf.reshape(attended, new_shape=(1, 1, HIDDEN)), w_o)
gqa_decode_fused = Stage3_Fused.entry_function()
set -euo pipefail
mkdir -p tutorial-reports
tilefoundry analyze attn_layer.py:Stage3_Fused \
tutorial-reports/stage3-4096.txt \
--compute-cost --memory --roofline --operands --dim ctx_len=4096
from pathlib import Path
report = Path("tutorial-reports/stage3-4096.txt").read_text(encoding="utf-8")
header, separator, annotated = report.partition("\n\n")
print(header.rstrip())
print()
print(next(line.rstrip() for line in annotated.splitlines() if "cache_update(k_cache" in line))
# analysis target=nvidia.h200_sxm module=Stage3_Fused function=gqa_decode topology=cta wave=132/132
# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline
# compute-cost flops=bf16:2983360@logical,8628160@total,461024@cta;f32:94080@logical,127488@total,3984@cta other-ops=integer:9@logical,512@total,16@cta;special:33056@logical,33152@total,1036@cta precision=exact
# memory traffic=gmem:r2.32MB/w4.00MB@logical,r9.21MB/w4.02MB@total,r1.44MB/w4.00MB@cta;rmem:r64.64KB/w64.07KB@logical,r184.00KB/w164.00KB@total,r5.75KB/w5.12KB@cta;smem:r2.62MB/w2.46MB@logical,r8.90MB/w8.61MB@total,r284.71KB/w275.41KB@cta footprint=hidden:512B;w_q:128.00KB;v0:1.00KB;w_k:32.00KB;v2:256B;cos_cache:1.00KB;sin_cache:1.00KB;pos_ids:4B;rope[0]:1.00KB;rope[1]:256B;cur_pos:4B;write_len:4B;k_all:512.12KB;k_heads:2.25MB;w_v:32.00KB;v11:256B;v_all:512.12KB;v_heads:2.25MB;attended_3:1.00KB;w_o:128.00KB;v91:512B footprint-precision=upper_bound peak=gmem:6.81MB;rmem:520B;smem:16.91KB persistent=gmem:2.31MB
# buffer=w_o holds=5.82MB time=none space=cta.head,cta.worker reuse=3.88MB fits=yes precision=upper_bound
# roofline ideal-ns=2036 bound-by=memory
k_all = tf.cache_update(k_cache, cur_pos, write_len, k_rope) # Tensor[(1, 4096, 2, 32), "bf16"]; compute-cost precision=exact; memory traffic=gmem:r136B/w128B@logical,r136B/w128B@total,r136B/w128B@cta footprint=cur_pos:4B;write_len:4B;rope[1]:128B;k_all:128B footprint-precision=upper_bound operands=0:r0/w0;1:r4B/w0;2:r4B/w0;3:r128B/w0;result:r0/w128B; roofline ideal-ns=1 bound-by=memory
The embedded Stage3_Fused program is the split-K example for this page.
This program uses explicit state and does not create an authored Partial value. A split-K
algorithm and a Partial shard attribute are related ideas, not interchangeable syntax.
5. Stage projection weights¶
Stage4_WeightPrepared moves each projection weight's output slice to smem before the matmul. The q and o weight lines show the per-CTA read:
@module(entry="gqa_decode", target=_H200, topologies=(_CTA,))
class Stage4_WeightPrepared:
"""Stage projection weights by output slice before the attention scan."""
@func
def gqa_decode(
hidden: Tensor[(1, 1, HIDDEN), "bf16"],
w_q: ConstTensor[(1, HIDDEN, HIDDEN), "bf16"],
w_k: ConstTensor[(1, HIDDEN, KV_DIM), "bf16"],
w_v: ConstTensor[(1, HIDDEN, KV_DIM), "bf16"],
w_o: ConstTensor[(1, HIDDEN, HIDDEN), "bf16"],
k_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), "bf16"],
v_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), "bf16"],
cur_pos: Tensor[(1,), "i32"],
write_len: Tensor[(1,), "i32"],
pos_ids: Tensor[(1,), "i32"],
cos_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), "bf16"],
sin_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), "bf16"],
) -> Tensor[(1, 1, HIDDEN), "bf16"]:
with Mesh(("cta",), layout=(QUERY_HEADS,), names=("head",)) as cta:
wq_local = tf.reshard(
w_q, (1, HIDDEN, HIDDEN @ cta.head), "smem"
)
wk_local = tf.reshard(
w_k, (1, HIDDEN, KV_DIM @ cta.head), "smem"
)
wv_local = tf.reshard(
w_v, (1, HIDDEN, KV_DIM @ cta.head), "smem"
)
hidden_local = tf.reshard(hidden, (1, 1, HIDDEN), "smem")
q_projected = tf.cast(tf.matmul(hidden_local, wq_local, out_dtype="f32"), dtype="bf16")
k_projected = tf.cast(tf.matmul(hidden_local, wk_local, out_dtype="f32"), dtype="bf16")
v_projected = tf.cast(tf.matmul(hidden_local, wv_local, out_dtype="f32"), dtype="bf16")
q_projected = tf.reshard(q_projected, (1, 1, HIDDEN), "gmem")
k_projected = tf.reshard(k_projected, (1, 1, KV_DIM), "gmem")
v_projected = tf.reshard(v_projected, (1, 1, KV_DIM), "gmem")
q = tf.reshape(
q_projected, new_shape=(1, 1, QUERY_HEADS, HEAD_DIM)
)
k = tf.reshape(k_projected, new_shape=(1, 1, KV_HEADS, HEAD_DIM))
v = tf.reshape(v_projected, new_shape=(1, 1, KV_HEADS, HEAD_DIM))
q_rope, k_rope = tf.rope(q, k, cos_cache, sin_cache, pos_ids)
k_all = tf.cache_update(k_cache, cur_pos, write_len, k_rope)
v_all = tf.cache_update(v_cache, cur_pos, write_len, v)
k_heads = tf.transpose(
tf.repeat_interleave(k_all, repeats=GQA_GROUP, axis=2),
perm=(0, 2, 1, 3),
)
v_heads = tf.transpose(
tf.repeat_interleave(v_all, repeats=GQA_GROUP, axis=2),
perm=(0, 2, 1, 3),
)
q_f32 = tf.cast(q_rope, dtype="f32")
k_f32 = tf.cast(k_heads, dtype="f32")
v_f32 = tf.cast(v_heads, dtype="f32")
scaled_q = q_f32 * tf.full_like(q_f32, value=SCALE)
q_e = tf.reshape(
scaled_q, new_shape=(1, 1, QUERY_HEADS, 1, HEAD_DIM)
)
k_e = tf.reshape(
k_f32, new_shape=(1, 1, QUERY_HEADS, CTX, HEAD_DIM)
)
v_e = tf.reshape(
v_f32, new_shape=(1, 1, QUERY_HEADS, CTX, HEAD_DIM)
)
scores = tf.reduce(q_e * k_e, axes=(-1,), keepdim=True, kind="sum")
peak = tf.reduce(scores, axes=(-2,), keepdim=True, kind="max")
weights = tf.exp(scores - peak)
normalizer = tf.reduce(weights, axes=(-2,), keepdim=False, kind="sum")
weighted = tf.reduce(weights * v_e, axes=(-2,), keepdim=False, kind="sum")
attended = tf.cast(weighted / normalizer, dtype="bf16")
w_o_local = tf.reshard(
w_o, (1, HIDDEN, HIDDEN @ cta.head), "smem"
)
attended_local = tf.reshard(
tf.reshape(attended, new_shape=(1, 1, HIDDEN)),
(1, 1, HIDDEN),
"smem",
)
output_local = tf.cast(
tf.matmul(attended_local, w_o_local, out_dtype="f32"), dtype="bf16"
)
return tf.reshard(output_local, (1, 1, HIDDEN), "gmem")
gqa_decode_weight_prepared = Stage4_WeightPrepared.entry_function()
set -euo pipefail
mkdir -p tutorial-reports
tilefoundry analyze attn_layer.py:Stage4_WeightPrepared \
tutorial-reports/stage4-4096.txt \
--compute-cost --memory --roofline --operands --dim ctx_len=4096
from pathlib import Path
report = Path("tutorial-reports/stage4-4096.txt").read_text(encoding="utf-8")
header, separator, annotated = report.partition("\n\n")
print(header.rstrip())
lines = annotated.splitlines()
for needle in ("reshard(w_q", "reshard(w_o"):
start = next(index for index, line in enumerate(lines) if needle in line)
end = start
while end + 1 < len(lines):
end += 1
if end > start and " # " in lines[end]:
break
print()
print("\n".join(line.rstrip() for line in lines[start : end + 1]))
# analysis target=nvidia.h200_sxm module=Stage4_WeightPrepared function=gqa_decode topology=cta wave=132/132
# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline
# compute-cost flops=bf16:329536@logical,338048@total,42256@cta;f32:6390528@logical,51124224@total,6390528@cta other-ops=special:32768@logical,262144@total,32768@cta precision=exact
# memory traffic=gmem:r25.95MB/w24.38MB@logical,r205.39MB/w195.05MB@total,r25.67MB/w24.38MB@cta;rmem:r3.75KB/w3.75KB@logical,r3.75KB/w3.75KB@total,r480B/w480B@cta;smem:r322.00KB/w321.00KB@logical,r336.00KB/w328.00KB@total,r42.00KB/w41.00KB@cta footprint=hidden:512B;w_q:128.00KB;q_projected_2:1.00KB;w_k:32.00KB;k_projected_2:256B;cos_cache:1.00KB;sin_cache:1.00KB;pos_ids:4B;rope[0]:1.00KB;rope[1]:256B;q_f32:2.00KB;v15:1.00KB;scaled_q:2.00KB;cur_pos:4B;write_len:4B;k_all:512.12KB;v21:4.00MB;k_heads:4.00MB;k_f32:8.00MB;v25:8.00MB;scores:256.00KB;peak:64B;v28:256.00KB;weights:256.00KB;w_v:32.00KB;v_projected_2:256B;v_all:512.12KB;v36:4.00MB;v_heads:4.00MB;v_f32:8.00MB;v40:8.00MB;weighted:2.00KB;normalizer:64B;v43:2.00KB;attended:1.00KB;w_o:128.00KB;v50:512B footprint-precision=upper_bound peak=gmem:8.44MB;rmem:128B;smem:16.50KB persistent=gmem:2.31MB
# buffer=cos_cache holds=50.08MB time=none space=cta.head reuse=7.00KB fits=no precision=upper_bound
# buffer=sin_cache holds=50.08MB time=none space=cta.head reuse=7.00KB fits=no precision=upper_bound
# buffer=hidden holds=50.08MB time=none space=cta.head reuse=3.50KB fits=no precision=upper_bound
# buffer=pos_ids holds=50.08MB time=none space=cta.head reuse=28B fits=no precision=upper_bound
# buffer=cur_pos holds=50.08MB time=none space=cta.head reuse=28B fits=no precision=upper_bound
# buffer=write_len holds=50.08MB time=none space=cta.head reuse=28B fits=no precision=upper_bound
# advisory="l2 reuse window cta.head holds 50.08MB at a 132-unit wave, exceeding capacity 47.68MB"
# roofline ideal-ns=10940 bound-by=memory
wq_local = tf.reshard(w_q, layout=(1, 256, 8 @ mesh.head, 32), storage="smem") # Tensor[(1, 256, 256), "bf16", ((1, 256, 8 @ mesh.head, 32), (0, 32, 0, 1)), "smem"]; compute-cost precision=exact; memory traffic=gmem:r128.00KB/w0@logical,r128.00KB/w0@total,r16.00KB/w0@cta;smem:r0/w128.00KB@logical,r0/w128.00KB@total,r0/w16.00KB@cta footprint=w_q:128.00KB footprint-precision=exact operands=0:r128.00KB/w0;result:r0/w128.00KB; roofline ideal-ns=28 bound-by=memory
v2 = tf.matmul(hidden_local, wq_local, out_dtype="f32") # Tensor[(1, 1, 256), "f32", ((1, 1, 8 @ mesh.head, 32), (256, 256, 32, 1)), "rmem"]; compute-cost flops=bf16:131072@logical,131072@total,16384@cta precision=exact; memory traffic=rmem:r0/w1.00KB@logical,r0/w1.00KB@total,r0/w128B@cta;smem:r128.50KB/w0@logical,r132.00KB/w0@total,r16.50KB/w0@cta footprint-precision=exact operands=0:r512B/w0;1:r128.00KB/w0;result:r0/w1.00KB; roofline ideal-ns=1 bound-by=compute
w_o_local = tf.reshard(w_o, layout=(1, 256, 8 @ mesh.head, 32), storage="smem") # Tensor[(1, 256, 256), "bf16", ((1, 256, 8 @ mesh.head, 32), (0, 32, 0, 1)), "smem"]; compute-cost precision=exact; memory traffic=gmem:r128.00KB/w0@logical,r128.00KB/w0@total,r16.00KB/w0@cta;smem:r0/w128.00KB@logical,r0/w128.00KB@total,r0/w16.00KB@cta footprint=w_o:128.00KB footprint-precision=exact operands=0:r128.00KB/w0;result:r0/w128.00KB; roofline ideal-ns=28 bound-by=memory
v48 = tf.matmul(attended_local, w_o_local, out_dtype="f32") # Tensor[(1, 1, 256), "f32", ((1, 1, 8 @ mesh.head, 32), (256, 256, 32, 1)), "rmem"]; compute-cost flops=bf16:131072@logical,131072@total,16384@cta precision=exact; memory traffic=rmem:r0/w1.00KB@logical,r0/w1.00KB@total,r0/w128B@cta;smem:r128.50KB/w0@logical,r132.00KB/w0@total,r16.50KB/w0@cta footprint-precision=exact operands=0:r512B/w0;1:r128.00KB/w0;result:r0/w1.00KB; roofline ideal-ns=1 bound-by=compute
6. Stream the KV cache¶
Stage5_CachePrepared leaves the static projection weights in their ordinary form and changes only the cache scan. BLOCK=128 rows move through smem while (m, l, acc) stays resident. The updated cache is read at cur_pos once, so append and scan are separate traffic events.
The single gqa_decode function prepares QKV and updates the caches inside a
single-CTA mesh, then scans inside the head mesh. The historical scan uses the
original caches; the current-position read uses the updated caches.
@module(entry="gqa_decode", target=_H200, topologies=(_CTA,))
class Stage5_CachePrepared:
"""Stream cache blocks through smem while an online state stays resident."""
@func
def gqa_decode(
hidden: Tensor[(1, 1, HIDDEN), "bf16"],
w_q: ConstTensor[(1, HIDDEN, HIDDEN), "bf16"],
w_k: ConstTensor[(1, HIDDEN, KV_DIM), "bf16"],
w_v: ConstTensor[(1, HIDDEN, KV_DIM), "bf16"],
w_o: ConstTensor[(1, HIDDEN, HIDDEN), "bf16"],
k_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), "bf16"],
v_cache: Tensor[(1, CTX, KV_HEADS, HEAD_DIM), "bf16"],
cur_pos: Tensor[(1,), "i32"],
write_len: Tensor[(1,), "i32"],
pos_ids: Tensor[(1,), "i32"],
cos_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), "bf16"],
sin_cache: Tensor[(ROPE_CONTEXT, HEAD_DIM), "bf16"],
) -> Tensor[(1, 1, HIDDEN), "bf16"]:
with Mesh(("cta",), layout=(1,), names=("tile",)) as _cta:
q = tf.reshape(
tf.matmul(hidden, w_q), new_shape=(1, 1, QUERY_HEADS, HEAD_DIM)
)
k = tf.reshape(tf.matmul(hidden, w_k), new_shape=(1, 1, KV_HEADS, HEAD_DIM))
v = tf.reshape(tf.matmul(hidden, w_v), new_shape=(1, 1, KV_HEADS, HEAD_DIM))
q_rope, k_rope = tf.rope(q, k, cos_cache, sin_cache, pos_ids)
k_all = tf.cache_update(k_cache, cur_pos, write_len, k_rope)
v_all = tf.cache_update(v_cache, cur_pos, write_len, v)
with Mesh(("cta",), layout=(QUERY_HEADS,), names=("head",)) as cta:
qh = tf.reshard(
q_rope, (1, 1, QUERY_HEADS @ cta.head, HEAD_DIM), "smem"
)
queries = tf.transpose(tf.cast(qh, dtype="f32"), perm=(0, 2, 1, 3))
scaled = queries * tf.full_like(queries, value=SCALE)
template = tf.reduce(queries, axes=(-1,), keepdim=True, kind="sum")
m = tf.full_like(template, value=-1e30)
l = tf.full_like(template, value=0.0)
acc = tf.full_like(queries, value=0.0)
for start in tf.tile(CTX, BLOCK):
base = start + 0
kb = tf.reshard(
tf.repeat_interleave(
k_cache[:, base : base + BLOCK, :, :],
repeats=GQA_GROUP,
axis=2,
),
(1, BLOCK, QUERY_HEADS @ cta.head, HEAD_DIM),
"smem",
)
vb = tf.reshard(
tf.repeat_interleave(
v_cache[:, base : base + BLOCK, :, :],
repeats=GQA_GROUP,
axis=2,
),
(1, BLOCK, QUERY_HEADS @ cta.head, HEAD_DIM),
"smem",
)
keys = tf.transpose(kb, perm=(0, 2, 3, 1))
values = tf.transpose(vb, perm=(0, 2, 1, 3))
scores = tf.reshard(
tf.matmul(tf.cast(scaled, dtype="bf16"), keys, out_dtype="f32"),
(1, QUERY_HEADS @ cta.head, 1, BLOCK),
"smem",
)
block_m = tf.reduce(scores, axes=(-1,), keepdim=True, kind="max")
next_m = tf.max(m, block_m)
correction = tf.exp(m - next_m)
weights = tf.exp(scores - next_m)
l = l * correction + tf.reduce(
weights, axes=(-1,), keepdim=True, kind="sum"
)
acc = acc * correction + tf.reshard(
tf.matmul(tf.cast(weights, dtype="bf16"), values, out_dtype="f32"),
(1, QUERY_HEADS @ cta.head, 1, HEAD_DIM),
"smem",
)
m = next_m
k_current = tf.repeat_interleave(
tf.index_select(k_all, cur_pos, dim=1), repeats=GQA_GROUP, axis=2
)
v_current = tf.repeat_interleave(
tf.index_select(v_all, cur_pos, dim=1), repeats=GQA_GROUP, axis=2
)
k_current = tf.reshard(
k_current, (1, 1, QUERY_HEADS @ cta.head, HEAD_DIM), "smem"
)
v_current = tf.reshard(
v_current, (1, 1, QUERY_HEADS @ cta.head, HEAD_DIM), "smem"
)
current_keys = tf.transpose(k_current, perm=(0, 2, 3, 1))
current_values = tf.transpose(v_current, perm=(0, 2, 1, 3))
current_scores = tf.reshard(
tf.matmul(tf.cast(scaled, dtype="bf16"), current_keys, out_dtype="f32"),
(1, QUERY_HEADS @ cta.head, 1, 1),
"smem",
)
current_m = tf.max(m, current_scores)
current_correction = tf.exp(m - current_m)
current_weights = tf.exp(current_scores - current_m)
l = l * current_correction + current_weights
acc = acc * current_correction + tf.reshard(
tf.matmul(
tf.cast(current_weights, dtype="bf16"), current_values, out_dtype="f32"
),
(1, QUERY_HEADS @ cta.head, 1, HEAD_DIM),
"smem",
)
attended = tf.transpose(
tf.cast(acc / l, dtype="bf16"), perm=(0, 2, 1, 3)
)
attended = tf.reshard(
attended, (1, 1, QUERY_HEADS, HEAD_DIM), "gmem"
)
return tf.matmul(tf.reshape(attended, new_shape=(1, 1, HIDDEN)), w_o)
gqa_decode_cache_prepared = Stage5_CachePrepared.entry_function()
set -euo pipefail
mkdir -p tutorial-reports
tilefoundry analyze attn_layer.py:Stage5_CachePrepared \
tutorial-reports/stage5-4096.txt \
--compute-cost --memory --roofline --operands --dim ctx_len=4096
from pathlib import Path
report = Path("tutorial-reports/stage5-4096.txt").read_text(encoding="utf-8")
header, separator, annotated = report.partition("\n\n")
print(header.rstrip())
print()
for needle in ("slice(k_cache", "cache_update(k_cache"):
print(next(line.rstrip() for line in annotated.splitlines() if needle in line))
# analysis target=nvidia.h200_sxm module=Stage5_CachePrepared function=gqa_decode topology=cta wave=132/132
# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline
# compute-cost flops=bf16:4557512@logical,5482952@total,858241@cta;f32:117288@logical,117288@total,14661@cta other-ops=integer:32@logical,256@total,32@cta;special:33040@logical,33040@total,4130@cta precision=exact
# memory traffic=gmem:r5.32MB/w4.00MB@logical,r13.20MB/w32.02MB@total,r1.82MB/w4.00MB@cta;rmem:r163.53KB/w161.03KB@logical,r181.03KB/w161.03KB@total,r22.63KB/w20.13KB@cta;smem:r8.82MB/w8.55MB@logical,r8.85MB/w8.57MB@total,r1.11MB/w1.07MB@cta footprint=hidden:512B;w_q:128.00KB;v0:1.00KB;w_k:32.00KB;v2:256B;cos_cache:1.00KB;sin_cache:1.00KB;pos_ids:4B;rope[0]:1.00KB;rope[1]:256B;cur_pos:4B;write_len:4B;k_all:256B;w_v:32.00KB;v10:256B;v_all:256B;v71:256B;k_current:1.00KB;v84:256B;v_current:1.00KB;attended_2:1.00KB;w_o:128.00KB;v100:512B;k_cache:16.00KB;v38:128.00KB;v_cache:16.00KB;v59:128.00KB footprint-precision=upper_bound peak=gmem:3.38MB;rmem:512B;smem:16.91KB persistent=gmem:2.31MB
# buffer=w_o holds=617.76KB time=none space=cta.head reuse=896.00KB fits=yes precision=upper_bound
# buffer=k_cache holds=617.76KB time=none space=cta.head reuse=112.00KB fits=yes precision=upper_bound
# buffer=v_cache holds=617.76KB time=none space=cta.head reuse=112.00KB fits=yes precision=upper_bound
# buffer=k_all holds=617.76KB time=none space=cta.head reuse=1.75KB fits=yes precision=upper_bound
# buffer=v_all holds=617.76KB time=none space=cta.head reuse=1.75KB fits=yes precision=upper_bound
# buffer=cur_pos holds=617.76KB time=none space=cta.head reuse=28B fits=yes precision=upper_bound
# roofline ideal-ns=2037 bound-by=memory
v37 = tf.slice(k_cache, (0, base, 0, 0), sizes=(1, 128, 2, 32), strides=(1, 1, 1, 1)) # Tensor[(1, 128, 2, 32), "bf16", ((1, 128, 2, 32), (262144, 64, 32, 1))]; compute-cost precision=exact; memory traffic=rmem:r32B/w0@logical,r256B/w0@total,r32B/w0@cta footprint-precision=exact operands=0:r0/w0;1:r32B/w0;result:r0/w0
k_all = tf.cache_update(k_cache, cur_pos, write_len, k_rope) # Tensor[(1, 4096, 2, 32), "bf16"]; compute-cost precision=exact; memory traffic=gmem:r136B/w128B@logical,r136B/w128B@total,r136B/w128B@cta footprint=cur_pos:4B;write_len:4B;rope[1]:128B;k_all:128B footprint-precision=upper_bound operands=0:r0/w0;1:r4B/w0;2:r4B/w0;3:r128B/w0;result:r0/w128B; roofline ideal-ns=1 bound-by=memory
weight staging: static tensor -> one output slice -> reusable for the step
cache staging: growing context -> one block -> compute -> next block
Feature ledger¶
The page uses this embedded ladder for features orthogonal to GQA:
| feature | live program |
|---|---|
@module(entry/target/topologies) |
Stage0_Naive |
@func, Tensor, ConstTensor, DimVar |
Stage0_Naive |
pass prototype and @f.specialize(RangePattern) |
Stage1_Specialized |
single Mesh, shard sugar X @ m.axis, reshard to smem/gmem |
Stage2_Sharded |
| split-K worker mesh and online softmax state | Stage3_Fused |
| weight staging and output gather | Stage4_WeightPrepared |
cache update, block scan, matmul, rope, reduce, cast |
Stage5_CachePrepared |
nested Mesh, rmem, rank-changing reshard, multi-level Topology |
not in this embedded ladder |
| runtime weight converter | migrate |
The command surface used by this page is:
--compute-cost logical work and executed totals
--memory movement, residency, and placement peaks
--roofline ideal bound and limiting resource
--performance per-level execution projection
--operands operand split in annotated call lines
--dim bind ctx_len for one static analysis run
--json write the same report data as JSON
For example, the JSON form writes to a path just like the text form. The Bash cell runs the
CLI with --json; the following Python cell loads the JSON report and prints a stable summary.
set -euo pipefail
mkdir -p tutorial-reports
tilefoundry analyze attn_layer.py:Stage0_Naive \
tutorial-reports/stage0-128.json \
--compute-cost --memory --roofline --dim ctx_len=128 --json
import json
from pathlib import Path
report = json.loads(Path("tutorial-reports/stage0-128.json").read_text(encoding="utf-8"))
summary = {
"target": report["target"],
"module": report["module"],
"function": report["functions"][0]["function"],
"topology": report["topology"],
"requested": report["requested"],
"executed": report["executed"],
"totals": report["functions"][0]["totals"],
}
print(json.dumps(summary, indent=2, sort_keys=True))
{
"executed": [
"compute-cost",
"memory",
"roofline"
],
"function": "Stage0_Naive.gqa_decode",
"module": "Stage0_Naive",
"requested": [
"compute-cost",
"memory",
"roofline"
],
"target": "nvidia.h200_sxm",
"topology": "cta",
"totals": {
"communication": {},
"flops": {
"bf16": 328896,
"f32": 200448
},
"traffic": {
"gmem": {
"read": 1178068,
"write": 806592
}
}
}
}
- The split-K tutorial programs use explicit log-sum-exp state. There is no authored
Partialvalue in this ladder. analyzereports authored static bounds. It does not replacecheck, a GPU run, cache invalidation, or a benchmark protocol.
The source-level rules still apply: an @func body uses variants instead of Python if, a Module
entry is called through its bare binding inside that Module, and a cast that looks redundant may be
the dtype boundary required by check. The embedded ladder keeps its explicit bf16 to f32
boundary visible in the annotated reports.