--performance crashes with an internal KeyError when a value produced inside
an inner loop is consumed one scope up, and the producer is a call whose own
body has a loop.
File ".../tilefoundry/analysis/timeline.py", line 279, in build
ready = max(
File ".../tilefoundry/analysis/timeline.py", line 280, in <genexpr>
(occurrences[key].end_ns for key in occurrence.predecessors),
~~~~~~~~~~~^^^^^
KeyError: 131376854249216
--compute-cost --memory --roofline on the same program are fine; only
performance fails, and it fails before any report is written.
Where it comes from
occurrences is the scope-local dict build() is filling, but
occurrence.predecessors holds ids representative() resolved, which can name
an expression that is not in that dict:
def representative(producer: int, scope: _Chain) -> int | None:
host = hosts[producer]
if host == scope:
return producer
if len(host) > len(scope) and host[: len(scope)] == scope:
return host[len(scope)] # <-- an id from a deeper chain
return NoneThen:
ready = max(
(occurrences[key].end_ns for key in occurrence.predecessors),
default=origin_ns,
)occurrences[key] assumes every predecessor is a direct occurrence of the
scope being built. A .get(key) with the missing case skipped would not crash,
but the interesting question is which chain the id belongs to -- the value here
is two loops deeper than its consumer, so host[len(scope)] names a loop that
this scope does not hold directly.
Reproducer
Complete program, 40 lines. NB=1 and GL=2 are enough; the crash needs the
producer to be a call whose body has a loop -- the same shape written with a
plain h = h + h inside the inner loop is accepted.
from __future__ import annotations
from tilefoundry import func, module
from tilefoundry.dsl import DimVar, Mesh, Tensor, tf
from tilefoundry.dsl.tf import * # noqa: F401, F403
from tilefoundry.ir.types.shard import Topology
from tilefoundry.target import CudaTarget
GRID, HP, NHS, RTG, NB, GL = 128, 128, 16, 256, 1, 2
S = DimVar("seq_len", 0, 8193)
_H200, _CTA = CudaTarget("nvidia.h200_sxm"), Topology("cta", GRID)
@module(entry="walk", target=_H200, topologies=(_CTA,))
class M:
@func
def stage(
x: Tensor[(NHS, S, HP), "bf16"],
out: Tensor[(NHS, S, HP), "bf16"],
) -> Tensor[(NHS, S, HP), "bf16"]:
with Mesh(("cta",), layout=(NHS, 8), names=("strip", "tile")) as m:
acc = out
for t in tile(S, RTG):
base = t + 0
xv = tf.reshard(x[:, base:base + RTG, :],
(NHS @ m.strip, RTG @ m.tile, HP), "smem")
acc = tf.insert_slice(acc, xv + xv, (0, base, 0))
return acc
@func
def walk(
x: Tensor[(NHS, S, HP), "bf16"],
b1: Tensor[(NHS, S, HP), "bf16"],
b2: Tensor[(NHS, S, HP), "bf16"],
) -> Tensor[(NHS, S, HP), "bf16"]:
with Mesh(("cta",), layout=(GRID,), names=("u",)) as mm:
h = x
for b in range(NB):
for g in range(GL):
h = stage(h, b1)
h = stage(h, b2) # <-- consumed one scope up: KeyError
return htilefoundry analyze case.py:M /tmp/r.txt \
--compute-cost --memory --roofline --performance --dim seq_len=8192
Removing the marked line, or moving it inside the g loop, makes it pass.
Why it matters here
This is the shape every layered model has: a block of N layers in an inner
loop, then something at the block level that reads what the block produced.
Ours is ten blocks of (three linear-attention layers, one full-attention
layer), and the attention layer sits exactly where the marked line is -- so
--performance is unavailable for the whole model while --roofline reports
it fine.
tilefoundry 0.0.2.dev43+g45032b967 (45032b9), target nvidia.h200_sxm.
--performancecrashes with an internalKeyErrorwhen a value produced insidean inner loop is consumed one scope up, and the producer is a call whose own
body has a loop.
--compute-cost --memory --rooflineon the same program are fine; onlyperformancefails, and it fails before any report is written.Where it comes from
occurrencesis the scope-local dictbuild()is filling, butoccurrence.predecessorsholds idsrepresentative()resolved, which can namean expression that is not in that dict:
Then:
occurrences[key]assumes every predecessor is a direct occurrence of thescope being built. A
.get(key)with the missing case skipped would not crash,but the interesting question is which chain the id belongs to -- the value here
is two loops deeper than its consumer, so
host[len(scope)]names a loop thatthis scope does not hold directly.
Reproducer
Complete program, 40 lines.
NB=1andGL=2are enough; the crash needs theproducer to be a call whose body has a loop -- the same shape written with a
plain
h = h + hinside the inner loop is accepted.Removing the marked line, or moving it inside the
gloop, makes it pass.Why it matters here
This is the shape every layered model has: a block of N layers in an inner
loop, then something at the block level that reads what the block produced.
Ours is ten blocks of (three linear-attention layers, one full-attention
layer), and the attention layer sits exactly where the marked line is -- so
--performanceis unavailable for the whole model while--rooflinereportsit fine.
tilefoundry 0.0.2.dev43+g45032b967(45032b9), targetnvidia.h200_sxm.