PracHub
QuestionsCoachesLearningGuidesInterview Prep

Quick Overview

This question evaluates implementation-level mastery of transformer architectures, including multi-head masked self-attention, position-wise feed-forward networks, residual connections, and precise tensor-shape reasoning for a decoder-only language model, within the Coding & Algorithms domain of deep learning and neural network engineering.

  • medium
  • Amazon
  • Coding & Algorithms
  • Machine Learning Engineer

Implement decoder-only GPT-style transformer

Company: Amazon

Role: Machine Learning Engineer

Category: Coding & Algorithms

Difficulty: medium

Interview Round: Onsite

### Goal Implement a simplified **decoder-only Transformer language model** (similar in spirit to GPT) for next-token prediction. The implementation should be modular, using four main classes: 1. `MultiHeadAttention` 2. `FeedForward` 3. `DecoderLayer` 4. `GPT` (the full model) You may assume a deep learning framework (e.g., PyTorch or TensorFlow), but you should clearly specify tensor shapes and operations. --- ### Model details and requirements Assume: - Batch size: `B` - Sequence length: `T` - Embedding dimension: `d_model` - Number of attention heads: `num_heads` (assume `d_model` is divisible by `num_heads`) - Vocabulary size: `V` #### 1. Multi-head self-attention (`MultiHeadAttention`) Implement a class `MultiHeadAttention` that performs **masked self-attention**: - Inputs: - `x`: Tensor of shape `(B, T, d_model)` (token embeddings or layer outputs). - Internally learnable parameters: - Linear projections to compute queries `Q`, keys `K`, and values `V` for all heads. - Operations: 1. Project `x` into `Q`, `K`, and `V` with shapes `(B, T, d_model)`. 2. Split them into `num_heads` heads, each of dimension `d_head = d_model / num_heads`. 3. Compute scaled dot-product attention for each head: \[ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_{head}}} + M\right)V, \] where `M` is a **causal mask** that prevents a position from attending to future positions (positions `> t`). 4. Concatenate the outputs of all heads back to shape `(B, T, d_model)`. 5. Apply a final linear projection to return to `(B, T, d_model)`. - Output: - Tensor of shape `(B, T, d_model)`. Ensure you correctly implement the **causal (autoregressive) mask** so that position `t` can only attend to positions `0..t`. #### 2. Position-wise feed-forward network (`FeedForward`) Implement a class `FeedForward` for the position-wise MLP: - Inputs: - `x`: Tensor of shape `(B, T, d_model)`. - Internal structure (typical choice): - Linear layer from `d_model` to `d_ff` (e.g., `4 * d_model`). - Non-linear activation (e.g., GELU or ReLU). - Linear layer from `d_ff` back to `d_model`. - All operations are applied **independently** to each position in the sequence. - Output: - Tensor of shape `(B, T, d_model)`. #### 3. Decoder layer (`DecoderLayer`) Implement a single Transformer **decoder block** that combines attention, feed-forward, residual connections, and layer normalization: - Inputs: - `x`: Tensor of shape `(B, T, d_model)`. - Components: 1. LayerNorm `ln1` before self-attention. 2. `MultiHeadAttention` block (masked self-attention). 3. Residual connection: `x = x + attention_output`. 4. LayerNorm `ln2` before feed-forward. 5. `FeedForward` block. 6. Residual connection: `x = x + ff_output`. - Output: - Tensor of shape `(B, T, d_model)`. Your `DecoderLayer` should be reusable so that multiple layers can be stacked. #### 4. GPT model (`GPT`) Implement a `GPT` class representing the full decoder-only Transformer: - Inputs: - `input_ids`: Integer tensor of shape `(B, T)` representing token indices. - Components: 1. **Token embedding** layer mapping token IDs to vectors of size `d_model`. 2. **Positional encodings** (learned or fixed) of shape `(T, d_model)` added to the token embeddings. 3. A stack of `N` identical `DecoderLayer` blocks. 4. A final linear layer mapping from `d_model` to vocabulary size `V`. - Forward pass: 1. Embed tokens, add positional encodings to obtain `(B, T, d_model)`. 2. Pass through the `N` decoder layers sequentially. 3. Apply the final linear layer to obtain logits of shape `(B, T, V)`. - Output: - Logits for next-token prediction: `(B, T, V)`. #### Additional notes - You should define the `forward` method for each class with correct tensor transformations. - Be careful about tensor shapes when splitting/combining heads. - Ensure the causal mask is correctly broadcast and applied in attention. - You do **not** need to implement training loops or optimization; focus on correct and clean model implementation.

Quick Answer: This question evaluates implementation-level mastery of transformer architectures, including multi-head masked self-attention, position-wise feed-forward networks, residual connections, and precise tensor-shape reasoning for a decoder-only language model, within the Coding & Algorithms domain of deep learning and neural network engineering.

Part 1: Causal Multi-Head Self-Attention Forward Pass

Implement the forward pass of simplified GPT-style multi-head self-attention using nested Python lists instead of a deep learning framework. Given an input tensor x of shape B x T x d_model, projection matrices Wq, Wk, Wv, and Wo, compute causal scaled dot-product self-attention. Heads use contiguous chunks of the projected dimension. Position t may only attend to positions 0 through t. Use a numerically stable softmax and return every numeric value rounded to 6 decimal places.

Constraints

  • 0 <= B <= 10
  • 0 <= T <= 50
  • 1 <= d_model <= 64 when T > 0
  • 1 <= num_heads <= d_model
  • d_model is divisible by num_heads
  • All projection matrices are d_model x d_model
  • Input values are finite numbers

Examples

Input: ([], 1, [], [], [], [])

Expected Output: []

Explanation: Empty batch produces an empty output.

Input: ([[[0], [1]]], 1, [[1]], [[1]], [[1]], [[1]])

Expected Output: [[[0], [0.731059]]]

Explanation: The second token attends to values [0, 1] with softmax scores [0, 1].

Input: ([[[1, 3], [2, 5], [4, 7]]], 1, [[0, 0], [0, 0]], [[0, 0], [0, 0]], [[1, 0], [0, 1]], [[1, 0], [0, 1]])

Expected Output: [[[1, 3], [1.5, 4], [2.333333, 5]]]

Explanation: Zero query/key projections make the causal attention uniform over all allowed previous positions.

Input: ([[[1, 2, 10, 20], [3, 4, 30, 40]]], 2, [[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]], [[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]], [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]], [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])

Expected Output: [[[1, 2, 10, 20], [2, 3, 20, 30]]]

Explanation: Two heads are split across dimensions [0,1] and [2,3]; both use the same causal uniform averaging here.

Hints

  1. After projection, head h uses dimensions h * d_head through (h + 1) * d_head - 1.
  2. For each query position t, compute scores only against key positions j <= t; future positions should not appear in the softmax denominator.

Part 2: Position-Wise Feed-Forward Network

Implement the forward pass of the Transformer position-wise feed-forward network using nested lists. For every token vector independently, compute hidden = ReLU(x * W1 + b1), then output = hidden * W2 + b2. The same feed-forward network is applied to each batch item and sequence position.

Constraints

  • 0 <= B <= 100
  • 0 <= T <= 100
  • 1 <= d_model <= 128 when T > 0
  • 1 <= d_ff <= 512
  • All numeric values are finite
  • Use ReLU: ReLU(z) = max(0, z)

Examples

Input: ([[[1, 2]]], [[1, -1, 0], [0, 2, 1]], [0, 1, -3], [[1, 0], [0, 1], [1, 1]], [0, 0])

Expected Output: [[[1, 4]]]

Explanation: The hidden vector is ReLU([1, 4, -1]) = [1, 4, 0].

Input: ([[[2], [-1]]], [[2, -1]], [0, 3], [[1], [-2]], [5])

Expected Output: [[[7], [-3]]]

Explanation: Both sequence positions use the same MLP independently.

Input: ([[]], [[1]], [0], [[1]], [0])

Expected Output: [[]]

Explanation: A batch with an empty sequence returns an empty sequence.

Input: ([[[1, -1]], [[0, 3]]], [[1, 1], [-1, 1]], [0, -2], [[2, 1], [1, -1]], [1, 0])

Expected Output: [[[5, 2]], [[2, -1]]]

Explanation: Different batch elements are processed independently with shared weights.

Hints

  1. There is no mixing across sequence positions; process each vector x[b][t] independently.
  2. Treat each vector as a row vector: output[j] = sum_i hidden[i] * W2[i][j] + b2[j].

Part 3: Pre-Norm Transformer Decoder Layer

Implement one simplified GPT decoder layer. The layer uses pre-layer-normalization, causal multi-head self-attention, residual addition, another layer normalization, a ReLU feed-forward network, and a second residual addition. This problem is framework-free: tensors are nested lists, linear layers have no bias except the feed-forward biases, and LayerNorm is computed over the last dimension. For LayerNorm, use variance without epsilon; if variance is zero, use normalized value 0 for every feature before applying gamma and beta.

Constraints

  • 0 <= B <= 10
  • 0 <= T <= 50
  • 1 <= d_model <= 64 when T > 0
  • 1 <= d_ff <= 256
  • d_model is divisible by num_heads
  • LayerNorm gamma and beta each have length d_model
  • Use causal masking inside self-attention

Examples

Input: ([[[1, 3]]], 1, [1, 1], [0, 0], [[0, 0], [0, 0]], [[0, 0], [0, 0]], [[1, 0], [0, 1]], [[1, 0], [0, 1]], [1, 1], [0, 0], [[1, 0], [0, 1]], [0, 0], [[1, 0], [0, 1]], [0, 0])

Expected Output: [[[0, 5]]]

Explanation: LN([1,3]) = [-1,1], attention adds [-1,1], then the feed-forward residual adds [0,1].

Input: ([[[1, 3], [3, 1]]], 1, [1, 1], [0, 0], [[0, 0], [0, 0]], [[0, 0], [0, 0]], [[1, 0], [0, 1]], [[1, 0], [0, 1]], [1, 1], [0, 0], [[0, 0], [0, 0]], [0, 0], [[0, 0], [0, 0]], [0, 0])

Expected Output: [[[0, 4], [3, 1]]]

Explanation: The second position causally averages normalized values from positions 0 and 1; the feed-forward block is zero.

Input: ([[[5]]], 1, [1], [0], [[0]], [[0]], [[0]], [[0]], [1], [0], [[10]], [2], [[3]], [1])

Expected Output: [[[12]]]

Explanation: For a single feature, LayerNorm variance is zero, so the normalized vector is [0].

Input: ([], 1, [1], [0], [[1]], [[1]], [[1]], [[1]], [1], [0], [[1]], [0], [[1]], [0])

Expected Output: []

Explanation: Empty input produces empty output.

Hints

  1. This is a pre-norm block: normalize before attention and before the feed-forward network, not after the residual.
  2. Keep the original x for the first residual: x1 = x + attention(LN1(x)), then output = x1 + feedforward(LN2(x1)).

Part 4: Full GPT-Style Decoder-Only Forward Pass

Implement a simplified full GPT forward pass using nested lists. First add token embeddings and positional embeddings. Then apply N pre-norm decoder layers, each using causal multi-head self-attention and a ReLU feed-forward network. Finally project each hidden vector to vocabulary logits. Layer parameters are provided as parallel arrays indexed by layer number. Return logits rounded to 6 decimals.

Constraints

  • 0 <= B <= 10
  • 0 <= T <= 50
  • 1 <= vocabulary size V <= 500
  • 1 <= d_model <= 64 when T > 0
  • 0 <= N <= 12
  • d_model is divisible by num_heads
  • Every input id is in the range [0, V - 1]
  • All numeric parameters are finite

Examples

Input: ([[]], [[1, 2], [3, 4]], [], 1, [], [], [], [], [], [], [], [], [], [], [], [], [[1, 0], [0, 1]], [0, 0])

Expected Output: [[]]

Explanation: An empty sequence has no logits.

Input: ([[1]], [[1, 0], [0, 2]], [[10, 10]], 1, [], [], [], [], [], [], [], [], [], [], [], [], [[1, 0, 1], [0, 1, -1]], [0, 1, 0])

Expected Output: [[[10, 13, -2]]]

Explanation: With no decoder layers, logits are computed directly from token plus positional embedding.

Input: ([[0, 1], [1, 0]], [[1, 0], [0, 1]], [[0, 0], [10, 10]], 1, [], [], [], [], [], [], [], [], [], [], [], [], [[1, 0], [0, 1]], [0, 0])

Expected Output: [[[1, 0], [10, 11]], [[0, 1], [11, 10]]]

Explanation: The same positional embedding for index 1 is added in both batch elements.

Input: ([[0]], [[1, 3]], [[0, 0]], 1, [[1, 1]], [[0, 0]], [[[0, 0], [0, 0]]], [[[0, 0], [0, 0]]], [[[1, 0], [0, 1]]], [[[1, 0], [0, 1]]], [[1, 1]], [[0, 0]], [[[1, 0], [0, 1]]], [[0, 0]], [[[1, 0], [0, 1]]], [[0, 0]], [[1, 0], [0, 1]], [0, 0])

Expected Output: [[[0, 5]]]

Explanation: One decoder layer transforms the single hidden vector before the final projection.

Input: ([[0, 1]], [[1, 3], [3, 1]], [[0, 0], [0, 0]], 1, [[1, 1]], [[0, 0]], [[[0, 0], [0, 0]]], [[[0, 0], [0, 0]]], [[[1, 0], [0, 1]]], [[[1, 0], [0, 1]]], [[1, 1]], [[0, 0]], [[[0, 0], [0, 0]]], [[0, 0]], [[[0, 0], [0, 0]]], [[0, 0]], [[1, 0], [0, 1]], [0, 0])

Expected Output: [[[0, 4], [3, 1]]]

Explanation: The decoder layer uses causal attention, so position 0 cannot depend on position 1.

Hints

  1. The initial hidden state is token_embedding[input_ids[b][t]] + position_embedding[t].
  2. Apply decoder layers sequentially; the output of layer i becomes the input to layer i + 1.
Last updated: Jul 9, 2026

Loading coding console...

PracHub

Master your tech interviews with 8,500+ real questions from top companies.

Product

  • Questions
  • Learning Tracks
  • Interview Guides
  • Resources
  • Premium
  • For Universities
  • Student Access

Browse

  • By Company
  • By Role
  • By Category
  • Topic Hubs
  • SQL Questions
  • AI Coding Questions
  • Compare Platforms
  • Discord Community

Support

  • support@prachub.com
  • (916) 541-4762

Legal

  • Privacy Policy
  • Terms of Service
  • About Us

© 2026 PracHub. All rights reserved.

Related Coding Questions

  • Find Two-Word Compound Words - Amazon (hard)
  • Minimum Path Length Through a Grid With One Allowed Cell Conversion - Amazon (medium)
  • Circular Drone Hub Delivery Route - Amazon (hard)
  • Leaf Domain Cumulative Scores - Amazon (medium)
  • Kth Largest Perfect Binary Subtree - Amazon (medium)