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 import func, module
from tilefoundry.analysis import analyze as run_analysis
from tilefoundry.dsl import ConstTensor, DimVar, DimVarRangePat, Mesh, Tensor, tf
from tilefoundry.dsl.tf import * # noqa: F401, F403 - bare tile() in the fused body
from tilefoundry.inspection.analysis_report import render_analysis, render_text
from tilefoundry.ir.types.shard import Topology
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 + 1)
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", "matmul(v33, 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
# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline
# compute-cost flops=bf16:328896@328896,f32:200448@200448 service=special:1024@1024
# traffic traffic=gmem:r2225620/w806592@r2225620/w806592
# peak-footprint=gmem:1675788
# roofline ideal-ns=632 bound-by=memory
v0 = matmul(hidden, w_q, a_layout="MK", b_layout="KN") # Tensor[(1, 1, 256), "bf16"]; compute-cost flops=bf16:131072@131072; traffic traffic=gmem:r131584/w512@r131584/w512 operands=0:r512/w0,1:r131072/w0,result:r0/w512; roofline ideal-ns=28 bound-by=memory
v11 = cache_update(k_cache, cur_pos, write_len, v10) # Tensor[(1, 128, 2, 32), "bf16"]; compute-cost; traffic traffic=gmem:r136/w128@r136/w128 operands=0:r0/w0,1:r4/w0,2:r4/w0,3:r128/w0,result:r0/w128; roofline ideal-ns=1 bound-by=memory
v34 = matmul(v33, w_o, a_layout="MK", b_layout="KN") # Tensor[(1, 1, 256), "bf16"]; compute-cost flops=bf16:131072@131072; traffic traffic=gmem:r131584/w512@r131584/w512 operands=0:r512/w0,1:r131072/w0,result:r0/w512; roofline ideal-ns=28 bound-by=memory
The first number before @ is global work or traffic. The number after @ is the
per-CTA projection. With no authored split they are equal. The annotated lines printed by
the previous cell 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 "))
traffic = next(line for line in lines if line.startswith("# traffic "))
peak = next(line for line in lines if line.startswith("# peak-footprint="))
roofline = next(line for line in lines if line.startswith("# roofline "))
f32 = re.search(r"f32:([^ ]+)", compute).group(1)
traffic_value = traffic.removeprefix("# traffic traffic=")
gmem_peak = re.search(r"gmem:([^,]+)", peak).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 `global@CTA` | traffic `global@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 global@CTA |
traffic global@CTA |
peak gmem bytes | ideal ns | bound |
|---|---|---|---|---|---|
| 128 | 200448@200448 |
gmem:r2225620/w806592@r2225620/w806592 |
1675788 | 632 | memory |
| 512 | 799488@799488 |
gmem:r4744660/w3202752@r4744660/w3202752 |
2572812 | 1656 | memory |
| 1024 | 1598208@1598208 |
gmem:r8103380/w6397632@r8103380/w6397632 |
3768844 | 3022 | memory |
| 2048 | 3195648@3195648 |
gmem:r14820820/w12787392@r14820820/w12787392 |
6160908 | 5752 | memory |
| 4096 | 6390528@6390528 |
gmem:r28255700/w25566912@r28255700/w25566912 |
10945036 | 11214 | memory |
| 8192 | 12780288@12780288 |
gmem:r55125460/w51125952@r55125460/w51125952 |
20513292 | 22136 | 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 capacity boundary¶
The full-cache sharded program in Stage2_Sharded places one local query head's K and V cache in smem. The boundary is derived 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
Stage1_Specialized expresses the dispatch as two half-open DimVarRangePat variants: [1, 1816) and [1816, 8193). 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(DimVarRangePat("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(DimVarRangePat("ctx_len", SPECIALIZE_T, ROPE_CONTEXT + 1))
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 valid boundary report. The following Python cell loads the file and prints its report header:
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
# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline
# compute-cost flops=bf16:2629376@328672,f32:2833728@354216 service=special:14528@1816
# traffic traffic=gmem:r5563796/w3721856@r3936212/w3721408,smem:r9597248/w9480000@r1199656/w1185000
# peak-footprint=gmem:3701260,smem:472160
# roofline ideal-ns=1935 bound-by=memory
A larger context crosses the stated capacity. The Bash cell preserves the non-zero CLI refusal in a file; the following Python cell loads and prints the actual error output:
set -euo pipefail
mkdir -p tutorial-reports
set +e
tilefoundry analyze attn_layer.py:Stage2_Sharded \
tutorial-reports/stage2-1820.txt \
--compute-cost --memory --roofline --dim ctx_len=1820 \
2> tutorial-reports/stage2-1820.err
status=$?
set -e
if [ "$status" -eq 0 ]; then
echo "expected Stage2_Sharded to refuse ctx_len=1820" >&2
exit 1
fi
from pathlib import Path
error = Path("tutorial-reports/stage2-1820.err").read_text(encoding="utf-8")
print(error.rstrip())
tilefoundry: error: function 'gqa_decode': value 'v16:240' needs 232960 B in smem, which exceeds the 232448 B the target states for that level
The formula chooses the dispatch boundary. It is not a benchmark-tuned magic number.
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(tf.cast(k_sh, dtype="f32"), perm=(0, 1, 3, 2))
values = tf.transpose(tf.cast(v_sh, dtype="f32"), perm=(0, 1, 2, 3))
scaled_q = queries * tf.full_like(queries, value=SCALE)
scores = tf.matmul(scaled_q, keys)
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(weights, values)
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 ", "# traffic ", "# peak-footprint=", "# roofline ")):
print(line)
# compute-cost flops=bf16:328896@328896,f32:200448@200448 service=special:1024@1024
# traffic traffic=gmem:r2225620/w806592@r2225620/w806592
# peak-footprint=gmem:1675788
# roofline ideal-ns=632 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 ", "# traffic ", "# peak-footprint=", "# roofline ")):
print(line)
# compute-cost flops=bf16:2629376@328672,f32:200448@25056 service=special:1024@128
# traffic traffic=gmem:r1674644/w264832@r1559508/w264384,smem:r684608/w675392@r85576/w84424
# peak-footprint=gmem:1540620,smem:33280
# roofline ideal-ns=405 bound-by=memory
The f32 work and special work divide by the eight head CTAs. The small bf16 difference is
placement and gather overhead 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.
@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"]:
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 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(tf.cast(kb, dtype="f32"), perm=(0, 2, 3, 1))
values = tf.transpose(tf.cast(vb, dtype="f32"), perm=(0, 2, 1, 3))
scores = tf.matmul(scaled, keys)
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.matmul(weights, values)
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
# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline
# compute-cost flops=bf16:4392896@328672,f32:6418944@200592 service=integer:288@9,special:33152@1036
# traffic traffic=gmem:r3476884/w4196992@r2558932/w4196544,rmem:r656/w72@r656/w72,smem:r5839296/w5662784@r682464/w672676
# peak-footprint=gmem:5047436,rmem:16,smem:33408
# roofline ideal-ns=1599 bound-by=memory
v6 = cache_update(k_cache, cur_pos, write_len, v5) # Tensor[(1, 4096, 2, 32), "bf16"]; compute-cost; traffic traffic=gmem:r136/w128@r136/w128 operands=0:r0/w0,1:r4/w0,2:r4/w0,3:r128/w0,result:r0/w128; 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.matmul(hidden_local, wq_local)
k_projected = tf.matmul(hidden_local, wk_local)
v_projected = tf.matmul(hidden_local, wv_local)
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.matmul(
attended_local, w_o_local
)
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
# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline
# compute-cost flops=bf16:337408@42176,f32:51124224@6390528 service=special:262144@32768
# traffic traffic=gmem:r28254676/w25566912@r27967956/w25565792,smem:r331008/w329984@r43168/w42144
# peak-footprint=gmem:10945036,smem:16960
# roofline ideal-ns=11213 bound-by=memory
v1 = reshard(w_q, layout=ShardLayout(
layout=Layout((1, 256, 8, 32), None),
attrs=(S(2),),
mesh=cta,
), storage=smem) # Tensor[(1, 256, 256), "bf16", ((1, 256, 8 @ cta.head, 32), (0, 32, 0, 1)), "smem"]; compute-cost; traffic traffic=gmem:r131072/w0@r16384/w0,smem:r0/w131072@r0/w16384 operands=0:r131072/w0,result:r0/w131072; roofline ideal-ns=28 bound-by=memory
v42 = reshard(w_o, layout=ShardLayout(
layout=Layout((1, 256, 8, 32), None),
attrs=(S(2),),
mesh=cta,
), storage=smem) # Tensor[(1, 256, 256), "bf16", ((1, 256, 8 @ cta.head, 32), (0, 32, 0, 1)), "smem"]; compute-cost; traffic traffic=gmem:r131072/w0@r16384/w0,smem:r0/w131072@r0/w16384 operands=0:r131072/w0,result:r0/w131072; roofline ideal-ns=28 bound-by=memory
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.
@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"]:
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 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(tf.cast(kb, dtype="f32"), perm=(0, 2, 3, 1))
values = tf.transpose(tf.cast(vb, dtype="f32"), perm=(0, 2, 1, 3))
scores = tf.matmul(scaled, keys)
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.matmul(weights, values)
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(
tf.cast(k_current, dtype="f32"), perm=(0, 2, 3, 1)
)
current_values = tf.transpose(
tf.cast(v_current, dtype="f32"), perm=(0, 2, 1, 3)
)
current_scores = tf.matmul(scaled, current_keys)
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.matmul(
current_weights, current_values
)
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
# selection requested=compute-cost,memory,roofline executed=compute-cost,memory,roofline
# compute-cost flops=bf16:1246400@328672,f32:6410280@801285 service=integer:256@32,special:33040@4130
# traffic traffic=gmem:r7672476/w4198272@r4001116/w4197824,rmem:r2560/w0@r2560/w0,smem:r21788704/w21486432@r2723588/w2685804
# peak-footprint=gmem:2950412,rmem:0,smem:33536
# roofline ideal-ns=2474 bound-by=memory
v21 = slice(k_cache, (0, v20, 0, 0), sizes=(1, 128, 2, 32), strides=(1, 1, 1, 1)) # Tensor[(1, 128, 2, 32), "bf16"]; compute-cost; traffic traffic=rmem:r32/w0@r32/w0 operands=0:r0/w0,1:r32/w0,result:r0/w0; roofline
v6 = cache_update(k_cache, cur_pos, write_len, v5) # Tensor[(1, 4096, 2, 32), "bf16"]; compute-cost; traffic traffic=gmem:r136/w128@r136/w128 operands=0:r0/w0,1:r4/w0,2:r4/w0,3:r128/w0,result:r0/w128; 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(DimVarRangePat) |
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 traffic
--memory residency and peak footprint
--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["function"],
"topology": report["topology"],
"requested": report["requested"],
"executed": report["executed"],
"totals": report["totals"],
}
print(json.dumps(summary, indent=2, sort_keys=True))
{
"executed": [
"compute-cost",
"memory",
"roofline"
],
"function": "gqa_decode",
"module": "Stage0_Naive",
"requested": [
"compute-cost",
"memory",
"roofline"
],
"target": "nvidia.h200_sxm",
"topology": "cta",
"totals": {
"flops": {
"bf16": 328896,
"f32": 200448
},
"traffic": {
"gmem": {
"read": 2225620,
"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.