Traced forward pass — tiny-shakespeare

Live Computation Graph

Prompt: ROMEO:
Input tokens: 6
Forward pass: 319.36 ms
Device: cuda
Modules traced: 28
Traced: 2026-09-14T10:29:52

#	Module	Type	Input shape		Output shape	Params
1	token_embedding	Embedding	1×6	→	1×6×384	24,960
2	position_embedding	Embedding	6	→	6×384	98,304
3	block0.ln1	LayerNorm	1×6×384	→	1×6×384	768
4	block0.attn	MultiHeadAttention	1×6×384	→	1×6×384	590,208
5	block0.ln2	LayerNorm	1×6×384	→	1×6×384	768
6	block0.ffwd	FeedForward	1×6×384	→	1×6×384	1,181,568
7	block1.ln1	LayerNorm	1×6×384	→	1×6×384	768
8	block1.attn	MultiHeadAttention	1×6×384	→	1×6×384	590,208
9	block1.ln2	LayerNorm	1×6×384	→	1×6×384	768
10	block1.ffwd	FeedForward	1×6×384	→	1×6×384	1,181,568
11	block2.ln1	LayerNorm	1×6×384	→	1×6×384	768
12	block2.attn	MultiHeadAttention	1×6×384	→	1×6×384	590,208
13	block2.ln2	LayerNorm	1×6×384	→	1×6×384	768
14	block2.ffwd	FeedForward	1×6×384	→	1×6×384	1,181,568
15	block3.ln1	LayerNorm	1×6×384	→	1×6×384	768
16	block3.attn	MultiHeadAttention	1×6×384	→	1×6×384	590,208
17	block3.ln2	LayerNorm	1×6×384	→	1×6×384	768
18	block3.ffwd	FeedForward	1×6×384	→	1×6×384	1,181,568
19	block4.ln1	LayerNorm	1×6×384	→	1×6×384	768
20	block4.attn	MultiHeadAttention	1×6×384	→	1×6×384	590,208
21	block4.ln2	LayerNorm	1×6×384	→	1×6×384	768
22	block4.ffwd	FeedForward	1×6×384	→	1×6×384	1,181,568
23	block5.ln1	LayerNorm	1×6×384	→	1×6×384	768
24	block5.attn	MultiHeadAttention	1×6×384	→	1×6×384	590,208
25	block5.ln2	LayerNorm	1×6×384	→	1×6×384	768
26	block5.ffwd	FeedForward	1×6×384	→	1×6×384	1,181,568
27	ln_f	LayerNorm	1×6×384	→	1×6×384	768
28	lm_head	Linear	1×6×384	→	1×6×65	25,025

Every row is a real nn.Module that fired during this exact forward pass, in the order it fired, with the actual tensor shapes it saw — captured via register_forward_hook, not diagrammed by hand. Shaded bands mark one transformer block each.