LLaMA Components › Implementation: A Complete LLaMA Block

class LLaMABlock:
    """
    Complete LLaMA transformer block.

    Combines RMSNorm, RoPE attention, and SwiGLU FFN with
    pre-norm configuration and residual connections.
    """

    def __init__(self, d_model, n_heads, d_ff, rope_base=10000, eps=1e-6):
        """
        Args:
            d_model: Model dimension
            n_heads: Number of attention heads
            d_ff: FFN hidden dimension (typically ~2.67 * d_model for SwiGLU)
            rope_base: Base for RoPE frequency computation
            eps: Epsilon for RMSNorm
        """
        self.d_model = d_model
        self.n_heads = n_heads

        # Normalization layers
        self.norm1 = RMSNorm(d_model, eps)
        self.norm2 = RMSNorm(d_model, eps)

        # Attention with RoPE
        self.attention = RoPEAttention(d_model, n_heads, rope_base)

        # SwiGLU FFN
        self.ffn = SwiGLU(d_model, d_ff)

    def __call__(self, x, mask=None):
        """
        Forward pass through the block.

        Args:
            x: Input tensor of shape (seq_len, d_model)
            mask: Optional causal mask of shape (seq_len, seq_len)

        Returns:
            Output tensor of shape (seq_len, d_model)
        """
        # Pre-norm attention with residual
        h = x + self.attention(self.norm1(x), mask)

        # Pre-norm FFN with residual
        y = h + self.ffn(self.norm2(h))

        return y

# Configuration matching LLaMA 7B proportions (scaled down)
d_model = 128
n_heads = 4
d_ff = int(d_model * 2.67)  # SwiGLU hidden dimension
seq_len = 32

# Create a block
block = LLaMABlock(d_model=d_model, n_heads=n_heads, d_ff=d_ff)

# Create sample input
x = np.random.randn(seq_len, d_model).astype(np.float32)

# Create causal mask (lower triangular)
mask = np.tril(np.ones((seq_len, seq_len), dtype=bool))

# Forward pass
y = block(x, mask=mask)