Encoder-Decoder Greedy Decode

Implement seq2seq_greedy_decode for one source sequence. Encode the source once, then choose one target token at a time using the largest vocabulary score at the last decoder position.

The source context stays fixed throughout this call. At each generation step, run the decoder over the entire target prefix built so far. This exercise does not use a decoder key/value cache.

Inputs

• src_ids: (S,), integer-valued source IDs in [0, V_src), delivered as floats.
• src_emb: (V_src, d), source token embeddings.
• tgt_emb: (V_tgt, d), target token embeddings.
• enc_pos_embed: (S, d), source position embeddings.
• dec_pos_embed: (P, d), target position embeddings.
• enc_blocks: (E, 6, d, d), encoder weights.
• dec_blocks: (L, 12, d, d), decoder weights.
• w_head: (d, V_tgt), target vocabulary projection.
• num_heads: a positive integer dividing d.
• start_token_id and eos_token_id: integers in [0, V_tgt).
• max_new_tokens: an integer limiting how many tokens to append after the start token.

S, d and both vocabulary sizes are positive. E and L may be zero. Let B = max(0, max_new_tokens); the target position table needs at least B rows. All parameter tensors are finite and share a floating-point dtype and device; src_ids is on that device. Shapes agree. Leave inputs unchanged.

Generation and stopping rules

Begin with seq = [start_token_id]. Repeat at most B times:

1. Decode the current seq against the fixed encoder output, using position rows 0 through len(seq)-1.
2. Project the final decoder position with w_head and choose its largest-scoring token. Break exact ties by choosing the smallest token ID.
3. Append that ID. If it equals eos_token_id, stop immediately.

Return a one-dimensional integer tensor on the input device. Include the start token and any generated EOS token. The output has between 1 and 1+B elements. With a nonpositive budget, return only the start token. An initial start token that equals EOS does not stop the loop: the EOS check applies after appending a new token. Do not pad the result, sample randomly or apply a vocabulary softmax.

The last appended token is returned without another decoder pass, so B position rows suffice even when the result has 1+B tokens. For example, if the choices are [4, 2] and 2 is EOS, start token 1 produces [1, 4, 2].

Encoder and decoder architecture

Use the following post-normalization model. Cast source IDs to integers and start with enc = src_emb[src_ids] + enc_pos_embed. Each encoder block holds [w_q, w_k, w_v, w_o, w_mlp1, w_mlp2]. Apply blocks in order:

q, k, v = enc @ w_q, enc @ w_k, enc @ w_v
enc = LN(enc + attention(q, k, v) @ w_o)
enc = LN(enc + GELU(enc @ w_mlp1) @ w_mlp2)

Encoder attention can read every source position. Keep its final output fixed for the generation loop. With E=0, use source token-plus-position embeddings.

For each target prefix, start afresh with dec = tgt_emb[seq] + dec_pos_embed[:len(seq)]. Decoder slots 0 through 3 hold self-attention [w_q, w_k, w_v, w_o], slots 4 through 7 hold cross-attention [w_q_cross, w_k_cross, w_v_cross, w_o_cross], and slots 8 and 9 hold [w_mlp1, w_mlp2]. Ignore slots 10 and 11, even when nonzero. Each block applies:

q, k, v = dec @ w_q, dec @ w_k, dec @ w_v
dec = LN(dec + causal_attention(q, k, v) @ w_o)
q, k, v = dec @ w_q_cross, enc @ w_k_cross, enc @ w_v_cross
dec = LN(dec + attention(q, k, v) @ w_o_cross)
dec = LN(dec + GELU(dec @ w_mlp1) @ w_mlp2)

Target self-attention at position t reads positions 0 through t, including itself. Exclude later positions with negative infinity before softmax. Cross-attention reads every source position and has no causal mask. With L=0, project target token-plus-position embeddings directly with w_head.

Shared operations

Split projected features into consecutive heads of width d/num_heads. Divide q @ k.T by sqrt(d/num_heads), apply stable softmax over key positions, combine value vectors and concatenate heads in their original order before the output projection. No token ID, including zero, is treated as padding.

LN uses the feature mean and population variance of each token: (token - mean) / sqrt(variance + 1e-5), with no learned scale or bias. GELU(u) = 0.5u(1 + tanh(sqrt(2/pi)(u + 0.044715u**3))). The feed-forward width is d. There are no biases, dropout, extra embedding scaling or additional final normalization.