Transformer Block Assembly › Implementation: Production-Ready Block

class ConfigurableTransformerBlock:
    """
    Production-style transformer block with configurable components.

    Supports:
        - RMSNorm or LayerNorm
        - Pre-norm or post-norm placement
        - Standard or SwiGLU FFN
        - Optional causal masking
    """

    def __init__(self, config):
        """
        Args:
            config: dict with keys:
                - d_model: int, model dimension
                - n_heads: int, number of attention heads
                - d_ff: int, FFN hidden dimension
                - norm_type: 'rmsnorm' or 'layernorm'
                - norm_placement: 'pre' or 'post'
                - ffn_type: 'standard' or 'swiglu'
                - causal: bool, whether to apply causal masking
                - eps: float, norm epsilon (default 1e-6)
        """
        self.config = config
        self.d_model = config["d_model"]
        self.causal = config.get("causal", False)

        # Initialize normalization layers
        if config["norm_type"] == "rmsnorm":
            self.norm1 = RMSNorm(config["d_model"], config.get("eps", 1e-6))
            self.norm2 = RMSNorm(config["d_model"], config.get("eps", 1e-6))
        else:
            # Simplified LayerNorm for demo
            self.norm1 = self._create_layernorm(config["d_model"])
            self.norm2 = self._create_layernorm(config["d_model"])

        # Attention
        self.attention = MultiHeadSelfAttention(
            config["d_model"], config["n_heads"]
        )

        # FFN
        if config["ffn_type"] == "swiglu":
            self.ffn = SwiGLUFFN(config["d_model"], config["d_ff"])
        else:
            self.ffn = StandardFFN(config["d_model"], config["d_ff"])

        self.norm_placement = config["norm_placement"]

    def _create_layernorm(self, dim):
        """Create a simple LayerNorm."""

        class LayerNorm:
            def __init__(self, dim, eps=1e-6):
                self.weight = np.ones(dim)
                self.bias = np.zeros(dim)
                self.eps = eps

            def __call__(self, x):
                mean = x.mean(axis=-1, keepdims=True)
                var = x.var(axis=-1, keepdims=True)
                return (
                    self.weight * (x - mean) / np.sqrt(var + self.eps)
                    + self.bias
                )

        return LayerNorm(dim)

    def __call__(self, x):
        """Forward pass through the block."""
        seq_len = x.shape[0]
        mask = create_causal_mask(seq_len) if self.causal else None

        if self.norm_placement == "pre":
            # Pre-norm: Norm -> Sublayer -> Residual
            h = x + self.attention(self.norm1(x), mask)
            y = h + self.ffn(self.norm2(h))
        else:
            # Post-norm: Sublayer -> Residual -> Norm
            h = self.norm1(x + self.attention(x, mask))
            y = self.norm2(h + self.ffn(h))

        return y


# Create LLaMA-style block
llama_config = {
    "d_model": 256,
    "n_heads": 4,
    "d_ff": 688,
    "norm_type": "rmsnorm",
    "norm_placement": "pre",
    "ffn_type": "swiglu",
    "causal": True,
}

llama_block = ConfigurableTransformerBlock(llama_config)

# Test
x_config_test = np.random.randn(8, 256) * 0.02
y_config_test = llama_block(x_config_test)