Visualizing LLM tensors: from tokens to batch matrices

0Strings in, matrices out

What actually sits inside a transformer at step zero is not words. It is a rectangular slab of floats with a fixed shape. This note traces how individual text prompts get tokenized, transformed into dense matrices, padded for static batching, and routed through attention, ending with the exact logits row that hides the next word prediction.¹

1Single prompt: 1D array to 3D tensor

Trace one prompt: "The cat sat".

Step A: tokenization (1D array)

The raw string becomes a 1D list of integer token IDs using the tokenizer's vocabulary map:

Tokens: [ 128, 567, 2432 ]
Shape:  (3,)

Step B: the batch dimension (2D matrix)

Deep learning libraries (PyTorch, JAX) require a batch dimension at the very front of the tensor. Even with a single prompt, we wrap it in a nested array:²

Batch Tensor: [
  [ 128, 567, 2432 ]
]
Shape: (1, 3)  ->  (Batch Size = 1, Sequence Length = 3)

Step C: embedding lookup (3D tensor)

An embedding layer (a lookup table) maps each integer token ID to a dense vector of size D (hidden dimension, here D = 8):

Token ID       Embedding Vector (Size D = 8)
-------------------------------------------------------------------------------
[ 128  ] --->  [  0.12, -0.45,  0.89,  0.03,  0.77, -0.19,  0.05, -0.66 ]  (Vector 1)
[ 567  ] --->  [ -0.98,  0.22, -0.11,  0.74,  0.33,  0.50, -0.88,  0.12 ]  (Vector 2)
[ 2432 ] --->  [  0.55,  0.01,  0.67, -0.32, -0.09,  0.82,  0.41,  0.95 ]  (Vector 3)

Stack the vectors together and they form a 2D matrix of size 3 × 8:

         Col 1   Col 2   Col 3   Col 4   Col 5   Col 6   Col 7   Col 8
        +-------+-------+-------+-------+-------+-------+-------+-------+
Token 1 |  0.12 | -0.45 |  0.89 |  0.03 |  0.77 | -0.19 |  0.05 | -0.66 |
        +-------+-------+-------+-------+-------+-------+-------+-------+
Token 2 | -0.98 |  0.22 | -0.11 |  0.74 |  0.33 |  0.50 | -0.88 |  0.12 |
        +-------+-------+-------+-------+-------+-------+-------+-------+
Token 3 |  0.55 |  0.01 |  0.67 | -0.32 | -0.09 |  0.82 |  0.41 |  0.95 |
        +-------+-------+-------+-------+-------+-------+-------+-------+

Including the batch size of 1, the actual shape inside the model is (1, 3, 8). Every transformer layer performs high-speed matrix multiplications directly against this 3 × 8 token matrix.

2Static batching: multiple inputs

Suppose two users submit prompts at the same time:

• User 1: "The cat" → [128, 567] (length 2)
• User 2: "A dog barked loudly" → [45, 902, 1821, 3445] (length 4)

A GPU requires a perfect rectangular grid of inputs to process items in parallel, so we must pad the shorter sequence.

Finding max length and padding

The longest sequence has length 4. Pad User 1 with a placeholder <pad> token (represented as 0) to match:

User 1 (Padded): [ 128,  567,    0,    0 ]
User 2:          [  45,  902, 1821, 3445 ]

Creating the 2D batch matrix

Stack the sequences vertically to form a single matrix of shape (Batch Size, Max Sequence Length) → (2, 4):

                Token 1    Token 2    Token 3    Token 4
              +----------+----------+----------+----------+
Row 1 (User 1)|   128    |   567    |    0     |    0     |  <-- padded tokens
              +----------+----------+----------+----------+
Row 2 (User 2)|    45    |   902    |   1821   |   3445   |
              +----------+----------+----------+----------+

Embedding to 3D tensor

When this (2, 4) token matrix passes through the embedding layer, each slot is replaced by an 8-dimensional vector, giving shape (2, 4, 8):

                |--- Token 1 ---| |--- Token 2 ---| |--- Token 3 ---| |--- Token 4 ---|
Row 1 (User 1)  [ e_128 vector ] [ e_567 vector ] [  e_0 vector  ] [  e_0 vector  ]
Row 2 (User 2)  [ e_45 vector  ] [ e_902 vector ] [ e_1821 vector] [ e_3445 vector]

3The attention mask

During self-attention, the Query (Q) and Key (K) matrices multiply to produce attention scores. For User 1, the model must not pay attention to the padded zero columns. We apply an attention mask of shape (2, 4):

Scores where the mask is 0 are set to negative infinity (−∞), which makes them evaluate to exactly 0 after the softmax activation. The padding participates in the tensor shapes but contributes nothing to the result.

4Extracting the next token prediction

After the forward pass, the model yields output logits of shape (Batch Size, Max Sequence Length, Vocab Size). Suppose the vocabulary size is 5 for this visualization:

[ USER 1 LOGITS ]            vocab = ["The", "cat", "sat", "dog", "barked"]
Token 1 (128)  |  0.1 | 0.8 | 0.0 | 0.1 | 0.0 |
Token 2 (567)  |  0.0 | 0.1 | 0.7 | 0.1 | 0.1 |  <-- EXTRACT (next word after "cat")
Token 3 (pad)  |  0.2 | 0.2 | 0.2 | 0.2 | 0.2 |     ignored
Token 4 (pad)  |  0.2 | 0.2 | 0.2 | 0.2 | 0.2 |     ignored

[ USER 2 LOGITS ]
Token 1 (45)   |  0.0 | 0.1 | 0.0 | 0.9 | 0.0 |
Token 2 (902)  |  0.1 | 0.1 | 0.1 | 0.1 | 0.6 |
Token 3 (1821) |  0.0 | 0.0 | 0.0 | 0.2 | 0.8 |
Token 4 (3445) |  0.7 | 0.1 | 0.1 | 0.0 | 0.1 |  <-- EXTRACT (next word after "loudly")

Finding the predictions

• For User 1, the last actual token sits at index 1 ("cat"). Extract row index 1: [0.0, 0.1, 0.7, 0.1, 0.1]. The highest value is 0.7 under "sat".
• For User 2, the last actual token sits at index 3 ("loudly"). Extract row index 3: [0.7, 0.1, 0.1, 0.0, 0.1]. The highest value is 0.7 under "The".

Lesson

Padding buys rectangular shapes; masking buys correct math. Every padded slot costs real compute, which is why serving stacks fight so hard over sequence packing and why the extraction index (last real token) matters as much as the logits themselves.

5References

1. Attention Is All You Need (Vaswani et al., 2017): scaled dot-product attention and masking.
2. Transformers tokenizer documentation: padding, attention masks, and batch encoding.
3. torch.nn.Embedding: the lookup table used in step C.
4. Local source: bin/blogs/token_to_matrix_visualization.md, the original markdown note this page was rewritten from.