Linux

1. chmod

   - chmod 400 file - Read by owner
   - chmod 040 file - Read by group
   - chmod 004 file - Read by world
   
   - chmod 200 file - Write by owner
   - chmod 020 file - Write by group
   - chmod 002 file - Write by world
   
   - chmod 100 file - execute by owner
   - chmod 010 file - execute by group
   - chmod 001 file - execute by world
   
   - chmod 444 file - Allow read permission to owner and group and world
   - chmod 777 file - Allow everyone to read, write, and execute file
   

ARCH Linux Setup

1. Setup WIFI

   iwctl
   

   station wlan0 get-networks
   station wlan0 connect <Network name>
   

2. Arch Install

   • GUI

     archinstall
     

   or

   • CLI

     Create partition

     cfdisk /dev/nvme0n1
     
     800 M for EFI System
     > 20 GB for Linux filesystem
     ... for Linux swap
     

     Format

     mkfs.fat -F32 /dev/<EFI System>
     mkfs.ext4 /dev/<Linux filesystem>
     mkswap /dev/<swap>
     

     Mount

     #root
     mount /dev/<linux filesystem /mnt
     mkdir /mnt/boot
     mount /dev/<EFI system> /mnt/boot
     swapon /dev/<swap>
     

     Install

     pacstrap -i /mnt base base-devel linux-zen linux-firmware git sudo neofetch htop intel-ucode nano vim bluez bluez-utils networkmanager 
     

     genfstab -U /mnt >> /mnt/etc/fstab
     cat /mnt/etc/fstab
     

     Enter the system

     arch-chroot /mnt
     
     # change root password
     passwd
     
     # create user
     useradd -m -g users -G wheel,storage,power,video,audio -s /bin/bash <username>
     passwd <username>
     
     EDITOR=vim visudo
     # uncomment line %wheel ALL=(ALL:ALL) ALL
     

     Timezone

     ln -sf /usr/share/zoneinfo/... /etc/localtime
     hwclock --systohc
     vim /etc/locale.gen #uncomment en_US ...
     locale-gen
     vim /etc/locale.conf # add "LANG=en_US.UTF-8
     

     Hostname

     vim /etc/hostname # add hostname
     
     vim /etc/hosts
     # add this line:
     127.0.0.1   localhost
     ::1         localhost
     127.0.1.1    <hostname>.localdomain   <hostname>
     

     Bootloader

     pacman -S grub efibootmgr dosfstools mtools
     
     grub-install --traget=x86_64-efi --efi-directory=/boot --bootloader-id=GRUB
     grub-mkconfig -o /boot/grub/grub.cfg
     

     Finish

     systemctl enable bluetooth
     systemctl enable NetworkManager
     exit
     umount -lR /mnt
     

Unplug the USB drive and boot to the system

3. Setup Enable radio wifi

   nmcli dev status
   nmcli radio wifi on
   nmcli dev wifi list
   sudo nmcli dev wifi connect <name> password "<password>"
   
   # update
   sudo pacman -Syu
   

   Install Desktop GUI

   sudo pacman -S xorg sddm plasma-meta plasma-workspace kde-applications
   
   sudo systemctl enable sddm
   sudo systemctl start sddm
   

4. Fix Backend Fix Discover App

   sudo pacman -Sy flatpak
   

   Install Nvidia Driver

   lspci | grep -E "NVIDIA"
   
   sudo pacman -Sy nvidia
   

5. Edit boot loader

   sudo pacman -Sy os-prober
   sudo vim /etc/default/grub
   # change following line
   # GRUB_TIMEOUT=20
   # uncomment GRUB_DISABLE_OS_PROBER=false
   
   sudo grub-mkconfig -o /boot/grub/grub.cfg
   

6. Chinese Character and Keyboard

   sudo pacman -S noto-fonts noto-fonts-cjk noto-fonts-extra noto-fonts-emoji ttf-dejavu ttf-liberation
   sudo pacman -S fcitx5-im fcitx5-rime
   
   cd ~/.local/share/fcitx5/rime
   git clone https://github.com/iDvel/rime-ice.git
   cp -r ./rime-ice/* .
   

SSH Configuration

1. RSA

   • RSA keys have been the default for many years and are supported by almost all SSH clients and servers. They are well-understood and trusted in various computing environments. Many systems default to RSA key lengths of 2048 or 3072 bits, though some users prefer 4096 bits for enhanced security.

     ssh-keygen -t rsa -b 4096 -C "your_email@example.com"
     

2. Ed25519

   • Ed25519 is increasingly popular due to its strong security features and efficiency. It uses elliptic curve cryptography to provide excellent security with shorter keys, resulting in faster performance and less data usage during authentication. Many modern systems and security guidelines now recommend Ed25519 as the preferred choice for new key generation.

     ssh-keygen -t ed25519 -C "your_email@example.com"
     

3. ECDSA

   • ECDSA is another commonly used type, particularly because it also offers good security with shorter key lengths compared to RSA. It’s often used where there’s a need for a balance between compatibility and modern cryptographic practices. ECDSA keys using the NIST P-256 curve (nistp256) are particularly common.

     ssh-keygen -t ecdsa -b 256 -C "your_email@example.com"
     

• note: RSA and Ed25519 are generally the most recommended, with Ed25519 often preferred for new deployments due to its robustness and efficiency. RSA remains widely used due to its long history and broad support across older and legacy systems. For new systems or updates, transitioning to Ed25519 from RSA or ECDSA is a common recommendation for enhanced security and performance.

4. Server Config

   • copy and paste the public keys to the authorized_keys file on the server.

     echo "paste-your-public-key-here" >> ~/.ssh/authorized_keys
     chmod 600 ~/.ssh/authorized_keys
     

5. Local Config

   • create a config file in .ssh folder

     Host "custom name"
     HostName "hostname -after @"
     User "username"
     IdentityFile "private key location"
     

   • after configuration use the following command to connect to the server

     ssh "custom name"
     

Install Ansible

1. Initializing
   • Create a folder inventory with a hosts file in it

     [server] # group name
     {ip address}
     {server name}
     

   • Try to ping the servers with password

     ansible -i ./inventory/hosts server -m ping --user sysadmin --ask-pass
     

   • Create a folder playboos and has a yaml file apt.yml in it

     - hosts: "*"
       become:
       tasks:
         - name: apt
           apt:
             update_cache: yes
             upgrade: 'yes'
     

   • Run the playbook

     ansible-playbook ./playbooks/apt.yml --user serveradmin --ask-pass --ask-become-pass -i ./inventory/hosts
     

   • Create a file qemu-get-agent.yml under playbooks to install a module

     - name: install latest qemu-guest-agent
       hosts: "*"
       tasks:
         - name: install qemu-guest-agent
           apt:
             name: qemu-guest-agent
             state: present
             update_cache: true
           become: true
     

   • Add the mattermost playbook

     ---
     - name: Install Mattermost Server
       hosts: all
       become: yes
       vars:
         mattermost_version: 5.31.0
         mattermost_db_name: mattermost
         mattermost_db_user: mmuser
         mattermost_db_password: mmuser_password
     
       tasks:
         - name: Install necessary packages
           apt:
             name: "{{ item }}"
             state: present
           with_items:
             - git
             - nginx
             - postgresql
             - postgresql-contrib
     
         - name: Create Mattermost user
           user:
             name: mattermost
             state: present
     
         - name: Clone Mattermost server
           git:
             repo: 'https://github.com/mattermost/mattermost-server.git'
             dest: "/opt/mattermost-server"
             version: "v{{ mattermost_version }}"
           become: yes
           become_user: mattermost
     
         - name: Configure PostgreSQL
           block:
             - name: Create Mattermost database
               postgresql_db:
                 name: "{{ mattermost_db_name }}"
                 login_user: postgres
     
             - name: Create Mattermost database user
               postgresql_user:
                 db: "{{ mattermost_db_name }}"
                 name: "{{ mattermost_db_user }}"
                 password: "{{ mattermost_db_password }}"
                 priv: ALL
                 login_user: postgres
     
         - name: Set up Mattermost configuration
           template:
             src: mattermost_config.json.j2
             dest: "/opt/mattermost-server/config/config.json"
             owner: mattermost
             mode: '0644'
     
         - name: Start Mattermost service
           systemd:
             name: mattermost
             state: started
             enabled: yes
     

Virtual Machine with Vagrant

1. Download and Install Tools

   • Download and install VirtualBox from the Official VirtualBox website.
   • Download and install Vagrant from the Official Vagrant website.

2. Get the Linux Box from Vagrant Cloud

   • Visit Vagrant Cloud to find a suitable Linux box. Alternatively, you can add a Linux box directly using the command line:

     vagrant box add [box_name]
     

     Replace [box_name] with the name of the Linux box you want to use.

3. Initialize Vagrant Environment

   • Initialize the VM with the following command:

     vagrant init [box_name]
     

     Again, replace [box_name] with the name of your chosen box.

4. Start the Virtual Machine

   • Start the VM with:

     vagrant up
     

5. Check Installed Linux Box Version

   • To check the installed Linux version and other boxes, use:

     vagrant box list
     

6. Connect to VM

   • Connect to your VM via SSH using:

     vagrant ssh
     

7. Disconnect to VM

   • suspend the VM

     vagrant suspend
     

   • resume from suspend

     vagrant resume
     

   • shutdown the VM

     vagrant halt
     

Add on features for Linux app

1. NeoVim Setup

   Requirements:

   • Install Nerd font first

     wget https://github.com/ryanoasis/nerd-fonts/releases/download/v3.2.1/Hack.zip
     unzip Hack.zip 
     mkdir -p ~/.local/share/fonts
     sudo cp Hack/*.ttf ~/.local/share/fonts/
     fc-cache -fv
     

   • Install npm

     sudo apt install npm
     

   git clone https://github.com/Henryfzh/documentation.git
   

   `gcc` - Toggles the current line using linewise comment
   `gbc` - Toggles the current line using blockwise comment
   `[count]gcc` - Toggles the number of line given as a prefix-count using linewise
   `[count]gbc` - Toggles the number of line given as a prefix-count using blockwise
   `gc[count]{motion}` - (Op-pending) Toggles the region using linewise comment
   `gb[count]{motion}` - (Op-pending) Toggles the region using blockwise comment
   

2. Theme

   Blur the windows:

   mutter-rounded
   mutter-rounded setting
   

3. mdBook

   Requirements:

   • Install Rust

     cargo install mdbook
     

4. tmux

   Install TPM:

   • Clone:

     git clone https://github.com/tmux-plugins/tpm ~/.tmux/plugins/tpm
     

   • Create ~/.tmux.conf, and add following to it:

     # List of plugins
     set -g @plugin 'tmux-plugins/tpm'
     set -g @plugin 'tmux-plugins/tmux-sensible'
     
     set -g @plugin 'catppuccin/tmux'
     set -g @catppuccin_flavour 'mocha'
     
     run '~/.tmux/plugins/tpm/tpm'
     
     set -g default-terminal 'tmux-256color'
     
     set -g mouse on
     
     # 2. Copy‑mode feels like Vim
     setw -g mode-keys vi                # h‑j‑k‑l in copy‑mode
     bind -T copy-mode-vi v  send -X begin-selection
     bind -T copy-mode-vi y  send -X copy-selection-and-cancel
     
     # 3. Quick pane movement with Ctrl‑h/j/k/l (same as Vim splits)
     bind -n C-h select-pane -L
     bind -n C-j select-pane -D
     bind -n C-k select-pane -U
     bind -n C-l select-pane -R
     
     # 4. Resize panes with Alt + arrows
     bind -n M-Left  resize-pane -L 5
     bind -n M-Right resize-pane -R 5
     bind -n M-Up    resize-pane -U 2
     bind -n M-Down  resize-pane -D 2
     
     # 5. Fast config reload:  <prefix> r
     bind r source-file ~/.tmux.conf \; display-message "✔ tmux.conf reloaded"
     
     # 6. Shorter escape delay (makes Vim feel snappier inside tmux)
     set -s escape-time 0
     

   • Install

     ctrl + B, I
     

   • Reload:

     tmux source ~/.tmux.conf
     

5. zsh fuzzy finder

   fzf
   

   .zshrc plugins

   powerlevel10k
   

   copypath copyfile copybuffer
   

   flatpak zsh on VSCode, add the lines to settings.json on VSCode

   "terminal.integrated.defaultProfile.linux": "bash",
    "terminal.integrated.profiles.linux": {
      "bash": {
        "path": "/usr/bin/flatpak-spawn",
        "overrideName": true,
        "args": ["--host", "--env=TERM=xterm-256color", "zsh"]
      }
    },
   

Docker Basics

1. Build

   sudo docker -t <target-name> -f Dockerfile .
   

2. Run

   • With bash

     sudo docker run --rm -it --device /dev/kfd --device /dev/dri --security-opt seccomp=unconfined <target-name> /bin/bash
     

   • Without bash

     sudo docker run --rm --device /dev/kfd --device /dev/dri --security-opt seccomp=unconfined rochpl.6.0  mpirun_rochpl -P 1 -Q 1 -N 45312
     

   • Mount a Directory

     sudo docker run --rm -it --device /dev/kfd --device /dev/dri --security-opt seccomp=unconfined --network host --name rochpl_node -v <directory>:/opt rochpl /usr/sbin/sshd -D
     

3. Update Docker

   • commit the changes

     docker ps # to get the container id
     

     docker commit <containerID> <imageid>
     

4. Clean

   # remove all images
   sudo docker system prune -a 
   sudo docker container prune
   sudo docker buildx prune -f
   

High Performance Linpack

1. Install OpenMPI
   • Download and unzip

     wget https://download.open-mpi.org/release/open-mpi/v5.0/openmpi-5.0.6.tar.gz
     tar xzf openmpi-5.0.6.tar.gz
     cd openmpi-5.0.6/
     

   • Compile
     Preferably installed at “/usr/local/” but require sudo access

     ./configure --prefix=<OPENMPI_INSTALL_DIRECTORY>
     make
     make install
     

2. Install OpenBLAS
   • Download and unzip

     wget https://github.com/OpenMathLib/OpenBLAS/releases/download/v0.3.28/OpenBLAS-0.3.28.tar.gz
     tar xzf OpenBLAS-0.3.28.tar.gz
     cd OpenBLAS-0.3.28/
     

   • Compile
     Preferably installed at “/usr/local/” but require sudo access

     make
     make PREFIX=<OPEN_BLAS_INSTALL_DIRECTORY> install
     

3. Update Path
   • Update path to OpenMPI and OpenBLAS in .bashrc or .zshrc

     export PATH=<OPENMPI_INSTALL_DIRECTORY>/bin:$PATH
     export LD_LIBRARY_PATH=<OPENMPI_INSTALL_DIRECTORY>/lib:$LD_LIBRARY_PATH
     export LD_LIBRARY_PATH=<OPEN_BLAS_INSTALL_DIRECTORY>/lib:$LD_LIBRARY_PATH
     

     source ~/.bashrc
     

     or

     source ~/.zshrc
     

4. Download HPL
   • Using wget or curl download from official website:

     wget http://www.netlib.org/benchmark/hpl/hpl-2.3.tar.gz
     

     or

     curl -O http://www.netlib.org/benchmark/hpl/hpl-2.3.tar.gz
     

   • Unzip

     tar -xf hpl-2.3.tar.gz
     

   1. Compile with CPU

   2. Compile with AMD GPU

High Performance Linpack with CPU

1. Compile HPL
   • Copy the template Makefile:

     cp setup/Make.Linux_Intel64 Make.Linux_Intel64
     

   • Edit the make file and change following lines:

     TOPdir = <hpl-2.3 top folder directory>
     
     MPdir =  <openmpi file directory>
     MPinc = -I$(MPdir)/include
     MPlib = -L$(MPdir)/lib -lmpi
     
     LAdir = <openblas file directory>
     LAinc = -I$(LAdir)/include
     LAlib = $(LAdir)/lib/libopenblas.a
     
     CC = mpicc
     CCNOOPT = $(HPL_DEFS)
     CCFLAGS = $(HPL_DEFS) -O3 -w -z noexecstack -z relro -z now -Wall # modify this according to the cpu
     
     LINKFLAGS = $(CCFLAGS) $(OMP_DEFS)
     

   • Compile

     make arch=Linux_Intel64
     

   • If you want to clean:

     make clean arch=Linux_Intel64
     

2. Run HPL

   • Edit the file bin/Linux_Intel64/HPL.dat inside the top folder.
     Here is an example with 8GB RAM and 4 Cores CPU:

     HPLinpack benchmark input file
     Innovative Computing Laboratory, University of Tennessee
     HPL.out      output file name (if any) 
     6            device out (6=stdout,7=stderr,file)
     1            # of problems sizes (N)
     29184         Ns
     1            # of NBs
     192           NBs
     0            PMAP process mapping (0=Row-,1=Column-major)
     1            # of process grids (P x Q)
     2            Ps
     2            Qs
     16.0         threshold
     1            # of panel fact
     2            PFACTs (0=left, 1=Crout, 2=Right)
     1            # of recursive stopping criterium
     4            NBMINs (>= 1)
     1            # of panels in recursion
     2            NDIVs
     1            # of recursive panel fact.
     1            RFACTs (0=left, 1=Crout, 2=Right)
     1            # of broadcast
     1            BCASTs (0=1rg,1=1rM,2=2rg,3=2rM,4=Lng,5=LnM)
     1            # of lookahead depth
     1            DEPTHs (>=0)
     2            SWAP (0=bin-exch,1=long,2=mix)
     64           swapping threshold
     0            L1 in (0=transposed,1=no-transposed) form
     0            U  in (0=transposed,1=no-transposed) form
     1            Equilibration (0=no,1=yes)
     8            memory alignment in double (> 0)
     ##### This line (no. 32) is ignored (it serves as a separator). ######
     0                               Number of additional problem sizes for PTRANS
     1200 10000 30000                values of N
     0                               number of additional blocking sizes for PTRANS
     40 9 8 13 13 20 16 32 64        values of NB
     

   • To tune the parameters, can reference the website here. It is not guaranteed to be the optimized setup. Try to tune the parameter by yourself.

     The following parameters are probably you need to tune:

     • Ps * Qs: the number of cores
     • Ns: the problem size
     • NBs: the block size

   • Run benchmark

     mpirun -np <number of cores> ./xhpl
     

High Performance Linpack with AMD GPU

Prepare (Download the Dockerfile)

Dockerfile for AMD GPU

1. Build Dockerfile

sudo docker build -t rochpl -f Dockerfile .

2. Setup Docker image

• Node A

  docker save -o rochpl_image.tar rochpl
  scp rochpl_image.tar user@10.0.0.12:~
  

• Node B

  docker load -i ~/rochpl_image.tar
  

• Both nodes

  sudo docker run --rm -it \
  --device /dev/kfd \
  --device /dev/dri \
  --security-opt seccomp=unconfined \
  --network=host \
  --name=rochpl_node \
  rochpl /bin/bash
  

• Setup SSH keys

  # Both Nodes
  ssh-keygen -t rsa -f ~/.ssh/id_rsa -q -N ""
  

  # Both Nodes
  vim /etc/ssh/sshd_config 
  # change the line --- PasswordAuthentication yes
  # add this line --- PermitRootLogin yes
  

  # Node A
  ssh-copy-id -p 2222 root@10.0.0.12
  

  # Node B
  ssh-copy-id -p 2222 root@10.0.0.14
  

• Add following to both nodes

  vim ~/.ssh/config
  

  Host 10.0.0.14
      Port 2222
      User root
  
  Host 10.0.0.12
      Port 2222
      User root
  

  • Test if it works

    ssh 10.0.0.14 hostname
    

3. Run HPL

• Add the rochpl_hostfile on both node

  10.0.0.14 slots=4
  10.0.0.12 slots=4
  

• Run HPL using this command (modify the arguments to suit your environment)

  export OMPI_MCA_pmix=pmix
  

  mpirun --hostfile rochpl_hostfile -np 8 --bind-to none   -x HIP_VISIBLE_DEVICES=0,1,2,3   --mca pml ucx --mca btl ^vader,tcp,openib,uct   ./run_rochpl -P 2 -Q 4 -N 256000 --NB 512
  

Dockerfile(AMD)

ARG UBUNTU_VERSION="jammy"

FROM ubuntu:${UBUNTU_VERSION}

ARG ROCM_URL="https://repo.radeon.com/amdgpu-install/6.1.1/ubuntu/jammy/amdgpu-install_6.1.60101-1_all.deb"
ARG UCX_BRANCH="v1.16.0"
ARG UCC_BRANCH="v1.3.0"
ARG OMPI_BRANCH="v5.0.3"
ARG APT_GET_APPS=""
ARG GPU_TARGET="gfx908,gfx90a,gfx942"

# Update and Install basic Linux development tools
RUN apt-get update \
    && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
        ca-certificates \
        git \
        ssh \
	openssh-client \
	openssh-server \
        make \
        vim \
        nano \
        libtinfo-dev\
        initramfs-tools \
        libelf-dev \
        numactl \
        curl \
        wget \
        tmux \
        build-essential \
        autoconf \
        automake \
        libtool \
        pkg-config \
        libnuma-dev \
        gfortran \
        flex \
        hwloc \
        libstdc++-12-dev \
        libxml2-dev \
        python3-dev \
        python3-pip \
        python3-distutils \
        unzip ${APT_GET_APPS}\
    && apt-get clean

RUN wget -qO- https://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor | tee /etc/apt/trusted.gpg.d/rocm.gpg \
    && wget -O rocm.deb ${ROCM_URL} \
    && apt install -y ./rocm.deb \
    && amdgpu-install --usecase=rocm,hiplibsdk --no-dkms -y

RUN bash -c """IFS=',' read -r -a ARCH <<<${GPU_TARGET} \
        &&  for gpu_arch in \${ARCH[@]}; do \
            echo \$gpu_arch  >> /opt/rocm/bin/target.lst; \
        done""" \
    && chmod a+r /opt/rocm/bin/target.lst 

# # Requires cmake > 3.22 
RUN mkdir -p /opt/cmake  \
  && wget --no-check-certificate --quiet -O - https://cmake.org/files/v3.27/cmake-3.27.7-linux-x86_64.tar.gz | tar --strip-components=1 -xz -C /opt/cmake

ENV ROCM_PATH=/opt/rocm \
    UCX_PATH=/opt/ucx \
    UCC_PATH=/opt/ucc \
    OMPI_PATH=/opt/ompi \
    GPU_TARGET=${GPU_TARGET}

# Adding rocm/cmake to the Environment 
ENV PATH=$ROCM_PATH/bin:/opt/cmake/bin:$PATH \
    LD_LIBRARY_PATH=$ROCM_PATH/lib:$ROCM_PATH/lib64:$ROCM_PATH/llvm/lib:$LD_LIBRARY_PATH \
    LIBRARY_PATH=$ROCM_PATH/lib:$ROCM_PATH/lib64:$LIBRARY_PATH \
    C_INCLUDE_PATH=$ROCM_PATH/include:$C_INCLUDE_PATH \
    CPLUS_INCLUDE_PATH=$ROCM_PATH/include:$CPLUS_INCLUDE_PATH \
    CMAKE_PREFIX_PATH=$ROCM_PATH/lib/cmake:$CMAKE_PREFIX_PATH

# Create the necessary directory for SSH
RUN mkdir /var/run/sshd

# Set root password for login
RUN echo 'root:redhat' | chpasswd

# Allow root login and password authentication
RUN sed -i 's/#PasswordAuthentication no/PasswordAuthentication yes/' /etc/ssh/sshd_config && \
    sed -i 's/PermitRootLogin prohibit-password/PermitRootLogin yes/' /etc/ssh/sshd_config && \
    echo "StrictModes no" >> /etc/ssh/sshd_config

# Change the SSH port to 2222
RUN sed -i 's/#Port 22/Port 2222/' /etc/ssh/sshd_config

# Expose the new SSH port
EXPOSE 2222

# Start the SSH service and keep the container running
ENTRYPOINT service ssh restart && bash

WORKDIR /tmp


# Install UCX
RUN git clone https://github.com/openucx/ucx.git -b ${UCX_BRANCH} \
    && cd ucx \
    && ./autogen.sh \
    && mkdir build \
    && cd build \
    && ../contrib/configure-release --prefix=$UCX_PATH \
        --with-rocm=$ROCM_PATH \
        --without-knem \
        --without-xpmem  \
        --without-cuda \
        --enable-optimizations  \
        --disable-logging \
        --disable-debug \
        --disable-examples \
    && make -j $(nproc)  \
    && make install

# Install UCC
RUN git clone -b ${UCC_BRANCH} https://github.com/openucx/ucc \
    && cd ucc \
    && ./autogen.sh \
    && sed -i 's/memoryType/type/g' ./src/components/mc/rocm/mc_rocm.c \
    # offload-arch=native builds the local architecutre, which may not be present at build time for a container. 
    && sed -i 's/--offload-arch=native//g' ./cuda_lt.sh \
    && mkdir build \
    && cd build \
    && ../configure --prefix=${UCC_PATH} --with-rocm=${ROCM_PATH} --with-ucx=${UCX_PATH} --with-rccl=no  \
    && make -j $(nproc) \
    && make install


# Install OpenMPI
RUN git clone --recursive https://github.com/open-mpi/ompi.git -b ${OMPI_BRANCH} \
    && cd ompi \
    && ./autogen.pl \
    && mkdir build \
    && cd build \
    && ../configure --prefix=$OMPI_PATH --with-ucx=$UCX_PATH \
        --with-ucc=${UCC_PATH} \
        --enable-mca-no-build=btl-uct  \
        --without-verbs \
        --with-pmix=internal \
        --enable-mpi \
        --enable-mpi-fortran=yes \
        --disable-man-pages \
        --disable-debug \
    && make -j $(nproc) \
    && make install

# Adding OpenMPI, UCX, and UCC to Environment
ENV PATH=$OMPI_PATH/bin:$UCX_PATH/bin:$UCC_PATH/bin:$PATH \
    LD_LIBRARY_PATH=$OMPI_PATH/lib:$UCX_PATH/lib:$UCC_PATH/lib:$LD_LIBRARY_PATH \
    LIBRARY_PATH=$OMPI_PATH/lib:$UCX_PATH/lib:$UCC_PATH/lib:$LIBRARY_PATH \
    C_INCLUDE_PATH=$OMPI_PATH/include:$UCX_PATH/include:$UCC_PATH/include:$C_INCLUDE_PATH \
    CPLUS_INCLUDE_PATH=$OMPI_PATH/include:$UCX_PATH/include:$UCC_PATH/include:$CPLUS_INCLUDE_PATH \
    PKG_CONFIG_PATH=$OMPI_PATH/lib/pkgconfig:$UCX_PATH/lib/pkgconfig/:$PKG_CONFIG_PATH  \
    OMPI_ALLOW_RUN_AS_ROOT=1 \
    OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1 \
    UCX_WARN_UNUSED_ENV_VARS=n

# Install Additional Apps Below
ARG HPL_BRANCH="main" 
 
WORKDIR /opt 
 
# Installing rocHPL 
RUN git clone -b ${HPL_BRANCH} https://github.com/ROCmSoftwarePlatform/rocHPL.git \ 
    && cd rocHPL \ 
    && ./install.sh \ 
        --prefix=/opt/rochpl \ 
        --with-rocm=/opt/rocm/ \ 
        --with-mpi=/opt/ompi \ 
    && rm -rf /tmp/rocHPL 
 
ENV PATH=$PATH:/opt/rochpl:/opt/rochpl/bin
ENV HIP_VISIBLE_DEVICES=0,1,2,3


#CMD ["/usr/sbin/sshd", "-D"]
CMD ["/bin/bash"]

Machine Learning

Classic machine-learning material. The large-language-model and LLM-application material has moved to its own section: LLM & AI Applications.

• Computer Vision

Computer Vision

1. VGG16

   • VGG16 has a total of 138 million parameters. The important point to note here is that all the conv kernels are of size 3x3 and maxpool kernels are of size 2x2 with a stride of two.

2. ResNet

   • Resnet18 has around 11 million trainable parameters. It consists of CONV layers with filters of size 3x3 (just like VGGNet). Only two pooling layers are used throughout the network one at the beginning and the other at the end of the network. Identity connections are between every two CONV layers. The solid arrows show identity shortcuts where the dimension of the input and output is the same, while the dotted ones present the projection connections where the dimensions differ.

3. Architecture Differences:

   • VGG16: VGG16 is a deep convolutional network with a straightforward and uniform architecture, consisting of 16 layers with very small (3x3) convolution filters. It is known for its simplicity and has been a popular choice for image classification tasks.
   • ResNet: ResNet, particularly ResNet-50, uses residual connections that help mitigate the vanishing gradient problem, allowing for the training of much deeper networks. ResNet architectures are typically deeper and more complex than VGG16, which generally results in better feature extraction and higher accuracy in many tasks.

4. Performance:

   • Accuracy: ResNet models, due to their depth and residual connections, generally outperform VGG16 in many image recognition tasks, including object detection. They are able to learn more complex features and provide better accuracy.
   • Computation and Memory: ResNet models are usually more computationally expensive and require more memory compared to VGG16. This can be a consideration if you have limited computational resources.

5. Application in Object Detection:

   • Object detection frameworks such as Faster R-CNN, SSD, and YOLO have utilized both VGG and ResNet as backbone feature extractors. In many cases, ResNet-based models have shown better performance in terms of both precision and recall.
   • For instance, Faster R-CNN with a ResNet-50 or ResNet-101 backbone generally performs better than the same framework with a VGG16 backbone.

Practical Considerations:

• ResNet Advantages:

  • Better accuracy and feature representation due to deeper network architecture.
  • Residual connections help in training deeper networks, resulting in improved performance.

• VGG16 Advantages:

  • Simpler architecture which can be easier to implement and train.
  • Less computationally intensive compared to ResNet.

Conclusion:

In general, ResNet models tend to be better than VGG16 for object detection tasks due to their superior feature extraction capabilities and higher accuracy. However, this comes at the cost of increased computational requirements.

If computational resources are not a constraint, it is recommended to use ResNet (e.g., ResNet-50 or ResNet-101) for better performance in object detection. However, if you need a simpler and less resource-intensive model, VGG16 is still a viable option and can achieve good results.

LLM & AI Applications

A working textbook on large language models and the applications built on top of them — from why self-attention works, through how a model is aligned with reinforcement learning, to how you assemble a retrieval pipeline or an agent and prove that it works.

What this section is for

Most interview material stops at the definition. That is not enough. An interviewer who knows the field will immediately ask why — and the answer to “why” is almost always a constraint: a memory limit, a gradient problem, a failure that showed up in production. This section is written so that every topic lands on:

1. Why it exists — the problem it solves, and what breaks without it.
2. How it works — the mechanism, with the math.
3. How you would build it — runnable code, and the knobs that matter.
4. What it costs — memory, compute, latency, quality, with numbers.
5. How it fails — the pitfalls, and the symptom of each.
6. When not to use it — the decision framing that matters more than the definition.

How to use it

If you have an interview in a few days. Read the TL;DR block at the top of each chapter, then the Interview drill at the bottom. Then work the Question Index, which collects all 255 drill questions in one place, so you can see at a glance which ones you cannot answer.

If you are preparing properly. Work the chapters in order — later material assumes earlier material. Part 03 (alignment and RL) in particular is not readable without SFT, and it builds the policy/value vocabulary it needs in its own first three chapters rather than assuming it from Part 01.

If you are building something. Jump straight to the part you need — 05 for agents, 06 for RAG, 07 for applications — and read the Using it in a project section of each chapter first.

A note on numbers

Every concrete figure in this section is either derived (a calculation you can check), cited (from a paper or official doc, with a link) or labelled (example). Nothing is written as if it were a measured result from a project that does not exist. When you have your own numbers, replace the examples — that is the point of labelling them.

Contents

Part 01 — Foundations

How the model is built.

• The Transformer
• Attention
• Positional Encoding
• Modern Architecture Deltas
• BERT vs GPT vs Encoder-Decoder

Part 02 — Pretraining and Fine-Tuning

How it is trained.

• Pretraining Pipeline
• Supervised Fine-Tuning (SFT)
• LoRA and QLoRA
• Quantization

Part 03 — Alignment and Reinforcement Learning

How it is made to behave.

• RLHF: The Three Stages
• Reward Modeling
• PPO for LLMs
• GRPO
• DPO
• Monte Carlo, TD, and Search
• RL Frameworks

Part 04 — Inference and Serving

How it is run.

• KV Cache
• Serving: vLLM and PagedAttention
• Ray

Part 05 — Agents

How it is used to do things.

• What an Agent Actually Is
• ReAct
• Tool / Function Calling
• LangChain
• LangGraph
• Multi-Agent Patterns
• Memory and Context Management
• MCP (Model Context Protocol)
• Guardrails

Part 06 — Retrieval-Augmented Generation

How it is grounded in facts.

• The RAG Pipeline
• Chunking
• Embedding Models
• Retrieval: Hybrid Search and Reranking
• Query Transformation
• RAG Evaluation

Part 07 — Applications

What you actually ship.

• Intent Recognition and Routing
• NL2SQL (Text-to-SQL)
• AI-Assisted Coding

Part 08 — Handwritten Code

What you write on a whiteboard.

• Multi-Head Attention from Scratch

Part 09 — Interview

How to deploy all of the above.

• The Interview Playbook
• Question Index

Status. Complete. All 39 leaf chapters across Parts 01–09 are written, each to the same eight-section template, and every relative link resolves. Research notes for the material live in docs/research/; the writing rules are in docs/AUTHORING.md.

Part 01 — Foundations

How the model is built.

The goal of this part is narrow and mechanical: by the end you should be able to draw the inside of a transformer block from memory and justify every piece of it. Nothing here is trivia — every later chapter is a consequence of these mechanisms. LoRA exists because of how the weight matrices are shaped. The KV cache exists because of the attention equation. Every context-extension trick exists because of how RoPE rotates positions.

Read the chapters in order; each assumes the previous one.

1. The Transformer — the end-to-end block: attention, feed-forward, residual stream, layer norm, and why the residual stream is the right mental model for everything that follows.
2. Attention — the core mechanism, derived rather than asserted. This is the reference chapter for the whole book: it shows the depth every other chapter is written to.
3. Positional Encoding — attention is permutation-invariant, so it cannot tell “dog bites man” from “man bites dog”. Sinusoidal, learned, ALiBi and RoPE, and why RoPE won.
4. Modern Architecture Deltas — the five substitutions between the 2017 paper and a 2026 model: pre-norm, RMSNorm, SwiGLU, GQA, and FlashAttention. Each gets the problem it solved at scale and the cost it added. (MoE appears in the interview drill, not as a full delta.)
5. BERT vs GPT vs Encoder-Decoder — the three architectures, their training objectives, and why decoder-only won for general-purpose models while encoders still own retrieval and classification.

Depends on: nothing; this is where the book starts.

Where this leads: Part 02 takes these mechanisms as given and asks how the model is trained. Nothing later is readable without the residual stream and the attention shape from here.

The Transformer

TL;DR — A Transformer is a stack of identical blocks, each of which is just two things: a mixing step (attention, which moves information between positions) and a computing step (a position-wise MLP, which processes each position independently). Around every one of those sits a residual connection, so the block’s input is added to its output and the stream flowing up the stack is never overwritten — only read from and added to. That residual stream is the mental model that makes the rest of the architecture obvious: norm layers keep its scale sane, the pre-norm placement keeps its gradient path clean, and depth is what lets the model compose operations. Attention is the only part that is $O(n^2)$; the MLP holds most of the parameters and most of the FLOPs at realistic context windows.

1. Why this exists

The 2017 paper’s own framing is the right one: this was not a better layer, it was a better graph.

A recurrent network computes $h_t=f(h_{t-1},x_t)$. Training it on a 1000-token sequence means 1000 dependent steps; the GPU’s parallel width is useless along the time axis, and the gradient from token 1000 to token 1 must survive 999 multiplies. A convolutional stack removes the sequential dependency but pays in path length: to relate two positions far apart you need $O(n/k)$ layers of contiguous kernels, or $O({\log }_{k}n)$ with dilations. Both designs force you to choose between parallelism and reach.

The Transformer’s claim in Table 1 of the paper is narrower and more useful than “it’s better”: self-attention is $O(n^2\cdot d)$ per layer with $O(1)$ sequential operations and $O(1)$ maximum path length between any two positions. You pay a quadratic bill in exchange for removing the architectural barrier to parallelism and long-range mixing. Everything else — the hardware that got built, the scaling laws that fit, the fact that GPT-4-class models are decoder-only Transformer stacks — follows from that trade.

Here is the concrete failure the original design fixes. Take a task where the answer depends on two tokens 800 apart, e.g. “the error code in section 3 is which one, given the fix described in section 40?” An LSTM must route that dependency through 800 cell-state updates; in practice the signal is gone and the model attends to recency instead. Symptom in the product: the model answers using the last thing it read rather than the right thing, and retrieval quality is blamed when the architecture is the culprit.

What this chapter is and is not. Attention derives the mixing operation itself. This chapter is the block that wraps it, the stream that carries it, and the reason a 32-deep stack trains at all. If you can draw the block from memory and justify every arrow, this chapter did its job.

2. How it works

2.1 The one idea: the residual stream

Write ${\mathbf{x}}_i^{(\ell )}\in {\mathbb{R}}^d$ for the vector at position $i$ after layer $\ell$. Every modern block has the shape

$${\mathbf{x}}^{(\ell +1)}={\mathbf{x}}^{(\ell )}+\mathrm{Sublayer}(\mathrm{Norm}({\mathbf{x}}^{(\ell )}))$$

and the whole point is the leading ${\mathbf{x}}^{(\ell )}$.

• ${\mathbf{x}}^{(\ell )}$ — the residual stream. It is a $d$-dimensional vector per position that is created once by the embedding and then never replaced. Every sublayer reads from it and writes an increment back into it. Without this identity path, each block would be a function that destroys its input: the only way information from layer 1 could reach layer 32 would be through 31 nonlinear transforms, and the gradient would have to survive the same journey backwards.
• $\mathrm{Sublayer}(\cdot )$ — the increment: either the attention output or the MLP output. It is not the state; it is a correction to the state. Because the residual is additive, sublayers compose by summation, which is what makes the “write to a shared bus” interpretation (Elhage et al., A Mathematical Framework for Transformer Circuits, 2021, https://transformer-circuits.pub/2021/framework/index.html) legitimate rather than merely poetic.
• The absence of a scale on the residual — no $\alpha {\mathbf{x}}^{(\ell )}$ with $\alpha <1$ — is deliberate. Any per-layer decay would bound the effective path length between layer 1 and layer $L$, reintroducing exactly the vanishing-gradient problem the design removed.

If you delete the residual connection, a 12-layer Transformer still trains but noticeably worse and becomes far more sensitive to learning rate; at 32 layers it typically does not converge. It is the cheapest ablation in the field.

2.2 The block, end to end

flowchart TB
    X["x (n × d)"] --> R1["+"]
    X --> N1["Norm"]
    N1 --> ATT["Self-attention<br/>(mixes across positions)"]
    ATT --> R1
    R1 --> R2["+"]
    R1 --> N2["Norm"]
    N2 --> MLP["Feed-forward<br/>(each position alone)"]
    MLP --> R2
    R2 --> OUT["x (n × d) → next block"]

    style R1 fill:#2E7A5A,stroke:#1E6A4A,color:#fff
    style R2 fill:#2E7A5A,stroke:#1E6A4A,color:#fff

Formally, with $n$ the sequence length and $d$ the model width:

$$Z=X+\mathrm{Attention}(\mathrm{Norm}(X)),X^{\prime }=Z+\mathrm{MLP}(\mathrm{Norm}(Z))$$

• $X\in {\mathbb{R}}^{n\times d}$ — the current block’s input: one row per position, all rows carried on the same residual stream.
• $\mathrm{Attention}(\cdot )$ — the only operation in the block that lets row $i$ see row $j$. Its internals are the subject of Attention; here it matters only that it is position-aware (via positional encoding) and sequence-mixing.
• $\mathrm{MLP}(\cdot )$ — applied to each row independently and with the same weights, so it is a per-position transform. It cannot move information between positions; it can only reshape what is already there.
• Two norms, not one, and one per sublayer. The norm belongs to the sublayer it feeds, not to the block.

The sequencing is the load-bearing part: the block alternates mix → think. Attention lets positions average each other; the MLP then applies a learned nonlinearity to the result. Neither is sufficient alone — attention with no MLP can only re-weight existing embeddings, and an MLP with no attention is a per-token function, no better than a bag of words for anything order-dependent.

2.3 The feed-forward network is not “just an MLP”

$$\mathrm{MLP}(x)=W_2\sigma (W_{1}x+b_1)+b_2$$

• $W_1\in {\mathbb{R}}^{d_{ff}\times d}$ — the up-projection. It lifts the residual stream into a wider space ($d_{ff}=2048$ against $d=512$ in the 2017 base model; $4d$ is the classic ratio). The expansion is what creates room for the activation to act per-feature; without it, the activation is applied in the same basis as the residual reads and the layer has far less representational room.
• $\sigma$ — the elementwise nonlinearity. The 2017 paper uses $\max (0,\cdot )$ (ReLU); modern models largely use gated variants (SwiGLU, Shazeer 2020, https://arxiv.org/abs/2002.05202). Without some nonlinearity the two linears collapse into one and the block becomes linear — attention included, since the whole thing is then a matrix product.
• $W_2\in {\mathbb{R}}^{d\times d_{ff}}$ — the down-projection back to the residual width, so the increment can be added to the stream.
• $b_1,b_2$ — biases. The 2017 paper has them; most modern LLMs drop them from the MLP and attention projections (LLaMA, PaLM, Mistral, Qwen). The reason is partly that a norm sits on the sublayer’s input, so a constant shift can be partly absorbed by $\beta$ of the preceding norm, and partly that removing them saves parameters and a kernel launch for no measured loss.

Two numbers make the MLP matter. At $d_{ff}=4d$ it holds $8d^2$ parameters while the attention projections hold $4d^2$ — so $8d^2$ of the block’s $12d^2$ total, or two-thirds of every block’s weights, live in the MLP. That is why “the model’s knowledge is in the MLP” is a reasonable first guess when you are deciding what to freeze, quantize, or ablate.

2.4 Pre-norm vs post-norm: where the norm goes

There are two placements, and the difference is not cosmetic.

Post-norm (the 2017 paper). $$x^{\prime }=\mathrm{Norm}(x+\mathrm{Sublayer}(x))$$

Here the norm is on the residual path: the identity ${\mathbf{x}}^{(\ell )}$ you just added gets renormalised before it continues. Xiong et al. 2020 (On Layer Normalization in the Transformer Architecture, ICML 2020, https://arxiv.org/abs/2002.04745) prove with mean-field analysis that at initialization the expected gradients of parameters near the output layer of a post-norm stack are large, and that this is precisely why a learning-rate warmup is mandatory. The practical consequences:

• You need a warmup schedule or the run diverges.
• The required warmup grows with depth, so scaling the stack means retuning optimization, not just the model.

Pre-norm (most modern LLMs). $$x^{\prime }=x+\mathrm{Sublayer}(\mathrm{Norm}(x))$$

Now the norm is inside the increment, and the residual path is a pure identity. Xiong et al. show the gradients are well-behaved at initialization and the warmup can be dropped. That is why the pre-2020 encoder-decoder recipes treated warmup as mandatory and why modern recipes treat it as an optimization detail: warmup is still almost always used — it rarely hurts — but a pre-norm run without it no longer diverges.

The cost of pre-norm: the residual stream’s norm grows monotonically with depth (each block adds an increment), so the final norm before the output head is doing real work — it is not decoration. A side effect worth knowing: because the stream grows, per-layer increments become relatively smaller as you go deeper, which is one motivation for scaled residual-init schemes.

2.5 Layer norm: what each symbol does

$$\mathrm{LN}(x)=\gamma \odot \frac{x-\mu }{\sqrt{{\sigma }^2+ϵ}}+\beta$$

• $x\in {\mathbb{R}}^d$ — one position’s vector, on the residual stream. The norm is applied per position, normalizing across the $d$ features; it never mixes across the sequence dimension.
• $\mu =\frac{1}{d}{\sum }_j{x}_j$ — the mean across features. Subtracting it re-centers the vector at the origin. Removing this term gives RMSNorm (below).
• ${\sigma }^2=\frac{1}{d}{\sum }_j(x_j-\mu {)}^2$ — the variance across features. Dividing by $\sqrt{{\sigma }^2+ϵ}$ rescales the vector to unit scale. This is the term that actually keeps training stable: it makes the sublayer’s input statistics independent of how large the residual stream has grown.
• $ϵ$ — a small constant (commonly $10^{-5}$ to $10^{-6}$) that keeps the denominator away from zero when the vector is (near-)constant. Without it, a saturated or collapsed activation produces a division by ~0 and NaNs.
• $\gamma \in {\mathbb{R}}^d$ — the learned scale. It restores the model’s ability to represent features of different magnitudes after normalization; without it, every sublayer would be forced to read unit-variance inputs.
• $\beta \in {\mathbb{R}}^d$ — the learned shift. It restores a per-feature offset. This is the parameter RMSNorm drops, on the empirical finding (Zhang & Sennrich 2019, https://arxiv.org/abs/1910.07467) that the re-centering is not what is doing the work — the re-scaling is. RMSNorm is therefore cheaper and has fewer parameters, and it is what LLaMA, PaLM, Gemma, Mistral, Qwen and DeepSeek use.

The interviewer’s version of this question. “Why layer norm and not batch norm?” Strong answer: batch norm computes statistics across the batch, so at inference you must freeze running estimates, and the statistics depend on sequence padding and on batch composition — fatal for variable-length text and for the autoregressive decode loop, where batch size is 1. Layer norm normalizes each position using only that position’s own features, so train and inference behave identically and batching does not change the result.

2.6 Encoder, decoder, and cross-attention

The original architecture is a stack of $N=6$ identical encoder layers and a stack of $N=6$ identical decoder layers, plus one extra sublayer in the decoder. The three attention layouts are the whole difference:

• Cross-attention’s $Q$ has $n_{dec}$ rows while $K,V$ have $n_{enc}$ rows. The score matrix is therefore $n_{dec}\times n_{enc}$ — rectangular, which is fine, and is the reason a seq2seq model can read a source of a different length than its target.
• There is no causal mask on the encoder axis because the whole source is available at once; masking it would hide input the decoder is entitled to.
• Encoder self-attention is bidirectional: masking it costs you exactly the right-context signal that makes BERT embeddings good for retrieval.

Decoder-only models (GPT and everything since) keep only the causal self-attention column. BERT vs GPT vs Encoder-Decoder does the objective-by-objective comparison; Modern Architecture Deltas covers what changed in the block since 2017 (pre-norm, RMSNorm, SwiGLU, GQA, FlashAttention, MoE).

2.7 Two 2017 details: the $\sqrt{d}$ scale and weight tying

The paper multiplies the token embedding by $\sqrt{d_{model}}$ before adding the positional encoding. With a norm immediately downstream, the pre-norm value does not affect forward outputs at all — the norm divides it out — and it matters only for initialization calibration: the embedding rows start with small variance, and $\sqrt{d}$ puts them in the regime the norm’s $\gamma$ (initialized to 1) expects. Delete it and the model still trains; the symptom is a slightly worse early loss curve, not a break.

Weight tying is the second: the input embedding, the output embedding, and the pre-softmax projection are the same matrix. That is $d\times V$ parameters saved (16.4M of the base model’s ~60M, 27%), and it encodes the prior that a token’s input representation and its output logit direction should be related.

3. Implementation

Runnable block and stack, torch==2.4.0 (also correct on 2.1+), CPU-only:

import math
import torch
import torch.nn as nn
import torch.nn.functional as F


class TransformerBlock(nn.Module):
    """Pre-norm block: x = x + attn(norm(x)); x = x + mlp(norm(x))."""

    def __init__(self, d_model: int, n_heads: int, d_ff: int, dropout: float = 0.0):
        super().__init__()
        assert d_model % n_heads == 0
        self.n_heads, self.d_k = n_heads, d_model // n_heads

        # QKV in one matmul: 3x fewer kernel launches than three Linears.
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.out_proj = nn.Linear(d_model, d_model, bias=False)
        self.norm1 = nn.LayerNorm(d_model)   # owns the attention sublayer
        self.norm2 = nn.LayerNorm(d_model)   # owns the MLP sublayer
        self.mlp = nn.Sequential(
            nn.Linear(d_model, d_ff, bias=False),
            nn.GELU(),
            nn.Linear(d_ff, d_model, bias=False),
        )
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, causal: bool = True):
        B, T, _ = x.shape

        # --- sublayer 1: attention, operating on the NORMED copy ---
        h = self.norm1(x)
        q, k, v = self.qkv(h).view(B, T, 3, self.n_heads, self.d_k) \
                          .permute(2, 0, 3, 1, 4)          # each (B, H, T, d_k)
        attn = F.scaled_dot_product_attention(q, k, v, is_causal=causal)
        attn = attn.transpose(1, 2).reshape(B, T, -1)
        x = x + self.dropout(self.out_proj(attn))      # <-- the residual

        # --- sublayer 2: MLP, position-wise, again on the NORMED copy ---
        x = x + self.dropout(self.mlp(self.norm2(x)))  # <-- the residual
        return x


class TinyTransformer(nn.Module):
    """Decoder-only stack, pre-norm, final norm before the head."""

    def __init__(self, vocab: int, d_model: int, n_heads: int, d_ff: int,
                 n_layers: int, max_len: int = 512):
        super().__init__()
        self.tok_emb = nn.Embedding(vocab, d_model)
        self.pos_emb = nn.Embedding(max_len, d_model)   # see positional-encoding.md
        self.norm_f = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, vocab, bias=False)
        # Tie AFTER the head exists: same tensor, and it must be set before the
        # optimizer is constructed.
        self.head.weight = self.tok_emb.weight
        self.blocks = nn.ModuleList(
            TransformerBlock(d_model, n_heads, d_ff) for _ in range(n_layers)
        )
        self.d_model = d_model

    def forward(self, idx):
        B, T = idx.shape
        pos = torch.arange(T, device=idx.device)
        x = self.tok_emb(idx) * math.sqrt(self.d_model) + self.pos_emb(pos)
        for block in self.blocks:
            x = block(x)
        return self.head(self.norm_f(x))                 # (B, T, vocab) logits


if __name__ == "__main__":
    torch.manual_seed(0)
    m = TinyTransformer(vocab=1000, d_model=512, n_heads=8, d_ff=2048, n_layers=6)
    print(m(torch.randint(0, 1000, (2, 64))).shape)  # torch.Size([2, 64, 1000])
    print(sum(p.numel() for p in m.parameters()))    # 19_661_824

That last number is checkable by hand: tied tok_emb ($1000\times 512$) 512,000, plus pos_emb ($512\times 512$) 262,144, plus 6 blocks at $12d^2+4d=3,147,776$ each = 18,886,656, plus norm_f 1,024. Sum: 19,661,824. If your count is off by ~512,000 the head is not actually tied; if it is off by a factor of two, you changed d_ff.

The four knobs that matter, with the ranges you actually see:

Notes on the code that interviewers probe:

• out_proj is applied inside the residual add, not outside it. Adding before projecting changes the block’s function, not just its efficiency.
• self.head.weight = self.tok_emb.weight rebinds both to one tensor and must run before the optimizer is built. Two parameter groups with different LRs will then apply conflicting updates to the same weights.
• norm_f before the head is not optional in a pre-norm stack: the residual stream’s norm grows with depth, so without it the logit scale — and hence the effective softmax temperature — varies with depth.
• is_causal=True must match the training-time masking exactly. A model trained bidirectionally and served causally produces garbage from token one.

4. Numbers and trade-offs

Parameter count, derived

Per block, with $d_{ff}=4d$ and no biases:

Six encoder layers: $6\times 3.15\text{M}=18.9\text{M}$. Word embedding, tied to the output projection, at a 32k-token vocabulary: $32000\times 512=16.4\text{M}$. The encoder alone is therefore $\approx 35.3$ M. The decoder adds $6\times (4d_{\text{self}}^2+4d_{\text{cross}}^2+8d_{\text{mlp}}^2+6d)=6\times 4.20\text{M}=25.2\text{M}$, giving a full encoder–decoder base model of roughly $60$ M counting weights only. The familiar headline is ~65M, and the bulk of that gap is vocabulary, not biases: the paper’s BPE vocabulary is ~37k tokens rather than the 32k assumed here, and the extra $5000\times 512=2.56$M of embedding takes the derived total to ~63M. The biases this derivation omits and the fixed sinusoidal encoding (if you materialise it as a buffer) add only ~0.33M. Two observations worth carrying: the MLP is two-thirds of every block, and the embedding is nearly half of a small encoder. That share falls as $d$ grows, but not to zero: Llama-3-8B uses a 128,256-token vocabulary and untied input and output embeddings (tie_word_embeddings=False, as in Llama-2), so its single embedding matrix is $128256\times 4096=525$ M parameters — about 6.5% of an 8B model — and the input+output pair is $2Vd\approx 1.05$ B, ~13% of 8.03B. Those two are by far its largest single matrices.

Compute per token, derived

Count multiply-accumulates as 2 FLOPs and ignore the softmax (it is $O(n^2)$ with a tiny constant). Per layer, per token:

• Attention: QKV projections $3\cdot 2d^2$, output projection $2d^2$, plus the score and value matmuls at $2\cdot 2nd=4nd$. Total $8d^2+4nd$.
• MLP: $2\cdot 2\cdot d\cdot 4d=16d^2$.

So the per-token budget is $24d^2+4nd$, and the ratio of attention to MLP work is $\frac{4nd}{16d^2}=\frac{n}{4d}$. The crossover is $n=4d$. For the 2017 base model ($d=512$) that is 2048 tokens — at or above the length of most WMT sentence pairs — so 2017 was an MLP-dominated computation. For Llama-3-8B ($d=4096$) the crossover is 16384 tokens: at 8k context, attention is

$$\frac{4nd}{24d^2+4nd}{|}_{d=4096,n=8192}=\frac{134,217,728}{402,653,184+134,217,728}=0.25$$

of the FLOPs per token, so the MLP still dominates. This is the honest answer to “why is long context expensive?” — not because attention is inherently the compute cost, but because it is the only term that grows with $n$, and it grows quadratically.

Memory at inference: the KV cache, in context

The cache size is the formula derived in Attention §4: $2\times L\times n_{kv}\times d_{head}\times n\times \text{bytes}$, where $d_{head}=d_k$ is the per-head width and $n_{kv}$ is the number of KV heads (equal to the head count $h$ without GQA). Plug in the 2017 base model — 6 encoder plus 6 decoder layers, 8 heads, $d_k=64$, fp32 — and it is $2\times 12\times 8\times 64\times 4=49,152$ bytes/token $=48$ KiB/token (binary units: 1 KiB = 1024 B), about 192 MiB (201 MB) at 4096 tokens: trivial next to a modern model’s weights. The regimes where that stops being true, and the serving consequences, are in Attention §4 and vLLM and PagedAttention.

Which axis you are optimizing

5. Pitfalls

Post-norm with a pre-norm learning rate. Symptom: the loss spikes in the first few hundred steps and either recovers only with a warmup restart or diverges outright. The gradient at initialization near the output layer is simply too large for the LR you chose — Xiong et al. 2020 derives exactly this.

Forgetting the final norm in a pre-norm stack. Symptom: training is fine but the logits’ scale drifts with depth; you see the loss plateau at a value that looks like a calibration problem rather than a capacity problem, and sampling temperature has to be retuned per checkpoint.

Applying the residual add to the wrong tensor. x = x + mlp(norm2(x)) is correct; x = norm2(x + mlp(x)) is post-norm; x = x + mlp(x) (no norm) usually still trains at small scale and hides the bug. Symptom: no crash, a loss curve that is merely worse, and a model that is unusually sensitive to initialization scale.

Assuming the MLP is cheap. Symptom: profile shows the “attention” budget blown at 32k context and you optimize attention, while 80%+ of the FLOPs are in $W_1$/$W_2$ GEMMs. At $n<4d$ the MLP is the thing to optimize — quantization and MoE (see Modern Architecture Deltas) attack it directly.

Tying weights with two optimizer groups. Symptom: the loss increases for a few hundred steps and then recovers, or embedding norms diverge from head norms. The two parameter groups are applying conflicting updates to one tensor.

Assuming a causal model can be evaluated bidirectionally. Symptom: an elegant offline evaluation using masked-token prediction scores far above the model’s real generation quality, because the model has never seen a right context during training.

Treating $d_{ff}$ as a rounding detail. Symptom: you port a checkpoint, change $d_{ff}$ to a “nicer” number, and get NaN-free garbage — the shapes match and the weights do not.

6. Using it in a project

You almost never build a Transformer. What you actually control:

The checkpoint’s shape, not the architecture. Vocabulary size, $d$, $L$, $d_{ff}$, and head count are fixed when you download a model. Your only levers are which checkpoint, how many bits, and how it is served.

When NOT to use a Transformer block. Reach for something else when:

• $n$ is small and the structure is tabular. Gradient-boosted trees beat a Transformer on a 10k-row, 50-column tabular problem, they need no tokenization, and they train in seconds. The Transformer’s inductive bias buys nothing when there is no meaningful long-range order to model.
• Latency budget is single-digit milliseconds. A Transformer’s cost is dominated by $O(L)$ sequential block evaluations; a distilled CNN or a n-gram/feature-hashing model will fit in a budget the Transformer cannot.
• The task is a per-token classification with abundant labels and short inputs. A frozen encoder plus a linear head is usually the right call; a decoder-only LLM is a much more expensive way to get the same label.
• You need a hard guarantee of determinism and auditability. Transformer generations are stochastic functions of weights and sampling; if the domain demands a rule engine, no amount of prompting fixes it.
• The data is a few hundred labeled examples. Fine-tuning the block structure will overfit; few-shot prompting plus retrieval, or a linear probe on a frozen encoder, is the better first move.

When a Transformer is the right tool but the block is not the problem. If your RAG system gives bad answers, the block is nearly never at fault — retrieval recall, prompt construction, and output parsing are. See RAG. Optimizing the block before measuring those is the classic premature optimization.

Two block-level decisions you do get to make. Pre-norm vs post-norm when you train from scratch (choose pre-norm), and whether to use the fused attention kernel (choose it; F.scaled_dot_product_attention is a drop-in with the same math). Everything else is already decided by whoever trained the checkpoint.

7. How you would prove it worked

A Transformer has no product metric of its own. What you prove is that a component earns its place. Three experiments, in increasing cost:

The residual ablation. Train two identical small models, one with $\mathbf{x}+\mathrm{Sublayer}(\cdot )$ and one with $\mathrm{Sublayer}(\cdot )$ only. (example) Expect the no-residual run to plateau at a substantially higher loss and to be far more LR-sensitive. This is the one experiment in this chapter that is nearly always worth running, because it makes the central claim of the architecture measurable in an afternoon.

Pre-norm vs post-norm with warmup removed. Train both, both without warmup. (example) Post-norm shows an early loss spike and often diverges; pre-norm descends smoothly. Then add warmup back to post-norm and show it recovers — which is the demonstration that the norm placement was the variable, not the optimizer.

The mixing/thinking split. Ablate attention only, then the MLP only, at matched parameter count. (example) Removing attention hurts order-sensitive and retrieval-sensitive tasks much more than it hurts token-level tasks; removing the MLP hurts nearly everything. If your two curves look identical, you are evaluating on a task that does not depend on sequence order, and the ablation proves nothing. The cheap version of this experiment uses the permutation test from Positional Encoding §7 as the order-sensitive probe.

Reporting template (example — replace with your own numbers). “We compared pre-norm and post-norm 12-layer blocks on the same data at the same LR. Post-norm without warmup diverged at step ~400; with a 2000-step warmup it reached 2.61 validation loss. Pre-norm without warmup reached 2.58 and was stable, removing one hyperparameter and 2000 steps of schedule.”

8. Interview drill

“Explain a Transformer block in 60 seconds.” Strong answer: two sublayers — attention, which mixes information across positions, and a position-wise MLP, which transforms each position independently — each wrapped in a residual connection with a layer norm on the sublayer’s input. The residual stream is the state; the sublayers only add to it. Pre-norm placement is what makes deep stacks trainable without warmup. Common wrong answer: reciting the 2017 figure as “attention then layer norm then feed-forward then layer norm” without saying what the residual path is for. That is a diagram, not an explanation.

“Why is the residual connection load-bearing?” Strong answer: it makes the identity path between the input embedding and the output head unbroken, so gradients reach early layers unimpeded and information is added rather than overwritten — the “residual stream” reading. Removing it reintroduces depth-dependent gradient decay. Follow-up: “why does the norm then have to move inside the residual?” → because a norm on the residual path renormalises the identity path, which is exactly the post-norm instability Xiong et al. 2020 analyses.

“Why did the Transformer replace the RNN? Be specific.” Strong answer: self-attention has $O(1)$ sequential operations per layer and $O(1)$ maximum path length between positions, versus $O(n)$ and $O(n)$ for a recurrence. That removes the parallelism barrier on the time axis and the long-range gradient-decay problem at once, at the cost of $O(n^2)$ compute. Common wrong answer: “it has better accuracy.” Accuracy is an outcome; the architectural claim is about parallelism and path length, and both are exactly what made training on web-scale data feasible.

“What does the feed-forward network do that attention can’t?” Strong answer: it is the only nonlinear, per-position computation in the block. Attention only re-weights and averages existing vectors — it is a convex combination, so with softmax weights the output is bounded by the inputs’ convex hull. The MLP creates new features (an expanded hidden space, a nonlinearity, a projection back), and it holds roughly two-thirds of the block’s parameters. Common wrong answer: “it’s just a classifier” or “it adds depth.” It is a key-value memory: the up-projection’s rows act as learned feature detectors and the down-projection turns each on into a residual-stream direction, which is why the MLP — not attention — is the target of MoE and the layer that quantization damages most.

“Which is more expensive, attention or the MLP?” Strong answer: it depends on $n$ relative to $d$. They cross at $n=4d$ for $d_{ff}=4d$: below it the MLP dominates, above it attention does. For a 2017 base model at 512 tokens attention is ~14% of the FLOPs ($4nd/(24d^2+4nd)$); at 32k tokens it is ~91%. “Attention is the expensive part” is only true at long context. Common wrong answer: “attention, because it’s $O(n^2)$” — dropping the constant is fine for asymptotics and wrong for the regime you are actually in.

“Why layer norm and not batch norm?” Strong answer: layer norm computes statistics per position across features, so its behaviour is identical at training and inference and independent of batch size and padding. Batch norm couples positions in a batch, which is degenerate in autoregressive decode where the effective batch is 1, and it makes sequence padding statistically meaningful. Follow-up: “why does RMSNorm work?” → the re-centering term matters less than the re-scaling term empirically (Zhang & Sennrich 2019); RMSNorm keeps the rescaling and drops the mean and bias, saving parameters and a reduction.

“The 2017 paper scaled embeddings by $\sqrt{d}$ — why?” Strong answer: a scale calibration for the embedding’s initial statistics relative to the norm’s learned $\gamma$; with a norm immediately after, the forward output is invariant to it, so removing it costs a slightly worse early loss curve and nothing structural. Common wrong answer: “to prevent the positional encoding from dominating,” or “to keep gradients from vanishing.” Both are plausible-sounding and neither is what the paper’s own code does.

“What would you change in the block for a 2026 model?” Strong answer: pre-norm with RMSNorm, SwiGLU MLP with $d_{ff}\approx 8d/3$ (rounded for hardware), no biases, RoPE instead of additive positional encoding, GQA to shrink the KV cache, and a fused FlashAttention kernel. Each of those is a response to a specific cost — see Modern Architecture Deltas.

Next: Attention — the mixing half of the block, derived rather than asserted: why the query–key dot product, why $\sqrt{d_k}$, and why multiple heads.

Attention

TL;DR — Attention lets every position in a sequence look at every other position and pull in whatever is relevant, in a single matrix multiply. Each token emits a query, a key, and a value; the query–key dot product scores relevance, softmax turns scores into weights, and the weights take a weighted average of the values. Dividing by $\sqrt{d_k}$ keeps those scores at unit variance so softmax does not saturate. Multi-head attention runs several of these in parallel subspaces, so one layer can track several kinds of relationship at once. The cost is $O(n^2)$ time and memory in sequence length — which is why long context is expensive.

1. Why this exists

Before attention, sequences were handled by recurrent networks, which read a sequence one step at a time and carry a hidden state forward:

$$h_t=f(h_{t-1},x_t)$$

That design has three problems, and attention is essentially a response to all three.

It cannot be parallelised across time. Computing $h_t$ requires $h_{t-1}$, which requires $h_{t-2}$, and so on. A 1000-token sequence is 1000 sequential steps, and no amount of GPU width helps.

Long-range dependencies decay. The signal linking token 1 to token 1000 has to survive 999 multiplications. Gradients flowing back through that chain vanish or explode, and in practice the start of a long sequence is simply forgotten.

Everything is squeezed through a fixed-size state. The hidden state has the same width at step 5 and step 5000, so all history must be compressed into one vector of fixed dimension.

Attention attacks the first two directly. Between any two positions $i$ and $j$, the path length is one — they interact directly, no matter how far apart they are. And because attention is expressed as matrix products over the whole sequence at once, there is no sequential dependency at all within a layer.

Where this matters for you. If someone asks “why did transformers replace RNNs?”, the expected answer is not “better accuracy”. It is: parallelism over the sequence dimension (which is what made training on web-scale data feasible) and constant path length between distant tokens.

2. How it works

2.1 The lookup analogy

Attention is a soft dictionary lookup. Imagine a dictionary where each entry has a key and a value. A normal lookup takes a query, finds the entry whose key matches, and returns its value. Attention does the same thing, except instead of returning one exact match it returns a weighted blend of every value, weighted by how well each key matches the query.

• Query ($Q$) — what this token is looking for.
• Key ($K$) — what each token advertises about itself.
• Value ($V$) — the content each token actually contributes.

All three are learned linear projections of the same input, which is why this is called self-attention:

$$Q=X{W}_Q,K=X{W}_K,V=X{W}_V$$

where $X\in {\mathbb{R}}^{n\times d_{model}}$ is the sequence of $n$ token vectors, and $W_Q,W_K,W_V\in {\mathbb{R}}^{d_{model}\times d_k}$ are the learned projections.

2.2 Scaled dot-product attention

$$\mathrm{Attention}(Q,K,V)=\mathrm{softmax}\left(\frac{Q{K}^{\top }}{\sqrt{d_k}}\right)V$$

Read it in three steps:

1. Score. $Q{K}^{\top }\in {\mathbb{R}}^{n\times n}$ — entry $(i,j)$ is the dot product of query $i$ with key $j$. A large positive value means “token $j$ is relevant to token $i$”.
2. Normalise. Softmax over each row turns raw scores into a probability distribution: how much attention token $i$ pays to each token $j$. Rows sum to 1.
3. Mix. Multiplying by $V$ produces, for each row, a weighted average of all value vectors — the token’s new representation.

flowchart LR
    X["X (n × d_model)"] --> WQ["× W_Q"]
    X --> WK["× W_K"]
    X --> WV["× W_V"]
    WQ --> Q["Q"]
    WK --> K["K"]
    WV --> V["V"]
    Q --> S["QKᵀ / √d_k"]
    K --> S
    S --> SM["softmax (row-wise)"]
    SM --> A["attention weights (n × n)"]
    V --> O["× V"]
    A --> O
    O --> OUT["output (n × d_k)"]

2.3 Why divide by $\sqrt{d_k}$

This is the single most-asked question about attention, and the answer is a variance argument — not a heuristic.

Assume the entries of $q$ and $k$ are independent, zero-mean, unit-variance. Their dot product is

$$q\cdot k=\sum _{i=1}^{d_k}{q}_i{k}_i$$

Each term $q_i{k}_i$ has mean $0$ and variance $\mathbb{E}[q_i^2]\mathbb{E}[k_i^2]-(\mathbb{E}[q_i]\mathbb{E}[k_i]{)}^2=1$. The $d_k$ terms are independent, so the sum has variance $d_k$ and standard deviation $\sqrt{d_k}$.

So the magnitude of the scores grows with the head dimension. Feed scores with standard deviation $\sqrt{d_k}$ into a softmax and, for large $d_k$, the distribution becomes nearly one-hot: one weight near 1, the rest near 0. In that regime the softmax is saturated, its local gradient is close to zero, and the attention weights stop moving. Dividing by $\sqrt{d_k}$ rescales the scores back to unit variance, keeping the softmax in the region where it has a useful gradient.

This is also why the scaling is not $\sqrt{d_{model}}$. It is the size of the per-head projection $d_k$ that governs the variance of the dot product, so that is what the score is divided by.

2.4 Masking

A decoder must not look at future tokens — otherwise, predicting token $t$ would be trivial, since the answer is sitting in the input. The fix is to add $-\infty$ to the scores in the forbidden positions before the softmax:

$$S_{ij}\leftarrow \{\begin{array}{ll}{S}_{ij}& j\le i\\ -\infty & j>i\end{array}$$

Since $e^{-\infty }=0$, those positions get exactly zero attention weight and the rows still sum to 1. The result is a causal (lower-triangular) attention pattern:

$$\left(\begin{array}{cccc}■& & & \\ ■& ■& & \\ ■& ■& ■& \\ ■& ■& ■& ■\end{array}\right)$$

Padding masks are applied the same way, but to columns: padding tokens must not be attended to.

2.5 Multi-head attention

One attention operation computes a single weighted average. That is limiting: a token often needs to relate to several different things at once — the subject of the sentence, the most recent noun, a matching bracket. Multi-head attention runs $h$ attention operations in parallel on different learned projections, then concatenates and projects the results:

$$\mathrm{MultiHead}(X)=\mathrm{Concat}({\mathrm{head}}_1,\dots ,{\mathrm{head}}_h)W_O$$ $$\text{where }{\mathrm{head}}_i=\mathrm{Attention}(X{W}_Q^{(i)},X{W}_K^{(i)},X{W}_V^{(i)})$$

The key design choice: each head works in a subspace of width $d_k=d_{model}/h$, so the total compute is roughly the same as a single head of full width. Multi-head is not a wider model; it is the same width partitioned into $h$ independent relationship-trackers.

flowchart TB
    X["X"] --> H1["head 1: Attention(d_k = d_model/h)"]
    X --> H2["head 2: Attention(d_k)"]
    X --> HD["⋯ head h"]
    H1 --> C["Concat (n × d_model)"]
    H2 --> C
    HD --> C
    C --> WO["× W_O"]
    WO --> OUT["output (n × d_model)"]

2.6 What it costs

Let $n$ be sequence length and $d$ the model width. For the standard implementation:

The quadratic memory term is the important one: it is what makes naive attention unusable at long context, and it is precisely what FlashAttention (see Modern Architecture Deltas) eliminates by never materialising the full matrix.

3. Implementation

A correct, readable from-scratch implementation:

import math
import torch
import torch.nn as nn
import torch.nn.functional as F


class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int, dropout: float = 0.0):
        super().__init__()
        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        self.W_o = nn.Linear(d_model, d_model, bias=False)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        B, T, _ = x.shape

        # (B, T, d_model) -> (B, n_heads, T, d_k)
        # .view splits d_model into n_heads x d_k; .transpose moves heads
        # next to the batch dimension so the matmuls below are batched.
        Q = self.W_q(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
        K = self.W_k(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
        V = self.W_v(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)

        # (B, n_heads, T, T)
        scores = Q @ K.transpose(-2, -1) / math.sqrt(self.d_k)

        if mask is not None:
            # mask: True where attention is ALLOWED.
            # Use the dtype's most-negative finite value, not -inf: a fully
            # masked row would be softmax over all -inf, i.e. 0/0 -> NaN.
            scores = scores.masked_fill(~mask, torch.finfo(scores.dtype).min)

        attn = self.dropout(F.softmax(scores, dim=-1))
        out = attn @ V                                   # (B, n_heads, T, d_k)

        # merge heads: (B, n_heads, T, d_k) -> (B, T, d_model)
        out = out.transpose(1, 2).contiguous().view(B, T, self.d_model)
        return self.W_o(out)


def causal_mask(T, device):
    """Lower-triangular boolean mask: position i may attend to j <= i."""
    return torch.tril(torch.ones(T, T, dtype=torch.bool, device=device))

Four details in that code are exactly what interviewers probe:

• .transpose(1, 2) moves the head dimension next to the batch dimension so that Q @ K.transpose(-2, -1) is a batched matrix multiply over (B, n_heads) — the heads are computed in parallel rather than in a Python loop.
• / math.sqrt(self.d_k) is applied to the scores, not to $Q$ or $K$ individually. Mathematically the two are equivalent; numerically, scaling the scores is what the derivation describes.
• Masking applies after scaling and before the softmax — applying it after the softmax would give a distribution that no longer sums to 1. Note the code uses torch.finfo(dtype).min rather than -inf; see the pitfalls below for why.
• .contiguous() before .view() — after transpose, the tensor is not contiguous in memory, and view requires it to be. Omitting this raises a runtime error (or silently produces a wrong reshape with older APIs).

The version you should actually ship

Since PyTorch 2.0, F.scaled_dot_product_attention dispatches to a fused kernel (FlashAttention / memory-efficient attention) and never materialises the $n\times n$ matrix:

out = F.scaled_dot_product_attention(Q, K, V, is_causal=True)

It is faster and uses less memory than the explicit version above, at the same result. Write the explicit version on a whiteboard; ship this one.

4. Numbers and trade-offs

The KV cache is the real memory cost at inference

During autoregressive decoding you generate one token at a time. Naively you would recompute keys and values for the entire prefix on every step. Instead you cache them — the KV cache — and the cache size is:

$$\text{bytes}=2\times n_{layers}\times n_{heads}\times d_{head}\times \text{seq\_len}\times \text{bytes\_per\_element}$$

The factor of 2 is for keys and values. For LLaMA-2 7B (32 layers, 32 heads, $d_{head}=128$) in fp16, that is 512 KiB per token (binary units throughout: 1 KiB = 1024 B):

$$2\times 32\times 32\times 128\times 2\text{ B}=512\text{ KiB/token}$$

At 4096 tokens that is 2 GiB per sequence — frequently more than the model weights once you batch. This is why long context is expensive, and why PagedAttention and continuous batching (see Serving) exist.

Head dimension

Per-head width $d_k$ is usually kept in the 64–128 range even as models grow, with head count increasing instead. Keeping $d_k$ fixed means the per-head compute stays constant and the $\sqrt{d_k}$ scaling stays in a well-behaved range, while the number of independent relationship-trackers grows with capacity.

Which axis you are optimising

Attention is a three-way trade-off and it is worth naming which axis a given decision moves:

5. Pitfalls

Forgetting the $\sqrt{d_k}$ scale. The model still trains — just worse, and the damage is invisible until you compare against a properly scaled baseline. Symptom: attention weights saturate to one-hot, gradients through the softmax collapse.

Masking with 0 instead of -inf. masked_fill(mask, 0) leaves masked positions at score 0, which softmax interprets as “moderately relevant” — not “forbidden”. The model silently leaks future information and reports a suspiciously good training loss that does not transfer to generation.

Using float('-inf'), or a value fp16 cannot represent. A fully-masked row (every position masked) becomes softmax over all $-\infty$, which is 0/0 → NaN, and NaN propagates through the whole model. The tempting fix is a large finite negative such as -1e9 — and in bf16 that works, because bf16 has fp32’s exponent range. In fp16 it does not: the largest finite magnitude is 65,504, so -1e9 is not representable and masked_fill rounds it to -inf, producing exactly the NaN you were trying to avoid. Use a representable value such as torch.finfo(torch.float16).min ($=-65504$) — which is what the code above does — or fall back to fp32 for the mask fill.

Forgetting .contiguous() before .view(). After transpose the tensor’s memory layout no longer matches its logical shape; view will raise, or in older code paths produce a scrambled reshape.

Confusing self-attention with cross-attention. In self-attention, $Q$, $K$, and $V$ all come from the same sequence. In cross-attention, $Q$ comes from the decoder and $K$, $V$ come from the encoder — that is how a seq2seq model reads its input. Getting this backwards is a common slip when describing an encoder–decoder stack.

Assuming nn.MultiheadAttention is a drop-in for the above. PyTorch’s built-in expects inputs in (seq_len, batch, d_model) by default (batch_first=False) and handles masking with a different convention. It is not a drop-in replacement, and mixing the two conventions is a classic source of silently wrong results.

6. Using it in a project

You almost never implement attention. In practice you control four things:

Context length. This is your dominant cost lever, because both attention time and KV-cache memory scale with it. Before increasing it, ask what actually needs to be in context — often the answer is “retrieved chunks”, which is the RAG pattern rather than a longer window.

Attention implementation. Prefer the fused kernel (F.scaled_dot_product_attention, or the equivalent in your serving stack). If your serving framework supports GQA, using a GQA checkpoint can cut KV-cache memory several-fold for a small quality cost — the single highest-leverage serving decision.

Head configuration, when fine-tuning or training small. If you are training a model from scratch, $d_k$ of 64–128 is the well-trodden range. If you are fine-tuning, head structure is fixed by the checkpoint; do not fight it.

Precision. Attention score computation is where fp16 overflows most often. If you see NaN losses early in training, this is the first place to look.

When not to touch any of this. For a 7B model at short context, attention is not your bottleneck — retrieval quality, prompt design, or the output parser usually is. Optimising attention before measuring is a classic premature optimisation.

7. How you would prove it worked

Attention has no product metric of its own; what you prove is that a design choice about attention matters. Three concrete experiments:

The scaling ablation. Train two identical small transformers on the same data, one with the $\sqrt{d_k}$ divisor and one without. Plot the loss curves. (example) The unscaled run typically shows a flat or barely-decreasing loss and much higher attention entropy variance, because the softmax has saturated. This is the experiment that turns “because the paper says so” into “because I watched it break”.

Attention entropy as a diagnostic. Log the entropy of the attention distribution per layer during training. A healthy model starts fairly uniform and sharpens; a model whose entropy collapses to near-zero at layer 0 early in training is usually saturated, and often indicates a scaling or learning-rate problem.

Long-context degradation. Evaluate the same task at increasing sequence lengths and plot quality against $n$. (example) Quality is typically flat up to the training length and then falls off sharply — the “lost in the middle” effect is partly an attention-pattern phenomenon. This is the measurement that justifies investing in retrieval rather than a bigger window.

Reporting template (example — replace with your own numbers). “We ablated the attention kernel: switching from the eager implementation to FlashAttention at 8k context cut peak memory by 40% (example) at identical output quality, and raised throughput 2.1× (example). We kept the eager path as a fallback for older GPUs.”

8. Interview drill

“Why divide by $\sqrt{d_k}$?” Strong answer: the dot products of independent unit-variance vectors have variance $d_k$; unnormalised, the softmax saturates and its gradient vanishes; dividing by $\sqrt{d_k}$ restores unit variance. Common wrong answer: “to keep the numbers small so it doesn’t overflow.” That is a consequence, not the reason — and it misses the gradient argument entirely.

“Why multiple heads instead of one big one?” Strong answer: the same total width split into subspaces lets the layer track several independent relationships simultaneously; it is not more compute. Follow-up: “how do the heads get combined?” → concatenated along the feature dimension, then projected by $W_O$. Heads are not averaged.

“Why is attention $O(n^2)$?” Because the score matrix is $n\times n$ — every token against every token. Follow-up: “what do you do about it?” → FlashAttention to remove the memory term; GQA to shrink the KV cache; sliding window or sparse attention to reduce the work; or retrieval, to avoid needing a long context in the first place.

“How does the causal mask work?” Strong answer: add $-\infty$ to scores above the diagonal before the softmax, so future positions receive exactly zero weight. Trap: saying “multiply by a triangular matrix afterwards” — that is a different operation and does not produce a normalised distribution.

“Self-attention vs cross-attention?” Strong answer: in self-attention $Q$, $K$, $V$ all come from the same sequence; in cross-attention $Q$ comes from the decoder and $K$, $V$ from the encoder. Follow-up: “which does a decoder-only model like GPT use?” → only self-attention; there is no separate encoder to cross-attend to.

“What is the difference between your implementation and nn.MultiheadAttention?” Strong answer: tensor layout conventions differ (batch_first), and the built-in handles masking differently; they are not interchangeable, and the most common bug is mixing the two conventions.

Next: Positional Encoding — attention as described here is permutation-invariant, so it cannot tell “dog bites man” from “man bites dog” without help.

Positional Encoding

TL;DR — Attention is permutation-equivariant: shuffle the tokens and the output rows shuffle with them, so “dog bites man” and “man bites dog” are the same input. Position must therefore be injected explicitly. Three families exist: add a vector to the embedding (sinusoidal, fixed; learned, a table), add a bias to the attention logits (ALiBi), or rotate the query and key vectors by an angle proportional to position (RoPE). RoPE won because $R_m^{\top }{R}_n=R_{n-m}$ — absolute positions cancel and attention scores become a function of the relative offset, with zero extra parameters, zero extra activation memory, and no change to attention’s form. Extending context is then a matter of rescaling the frequencies or the position indices.

1. Why this exists

Write out what self-attention does to a sequence, ignoring the projections:

$$y_i=\sum _j\frac{\exp (q_i\cdot k_j)}{{\sum }_{j^{\prime }}\exp (q_i\cdot k_{j^{\prime }})}{v}_j$$

Nothing in that expression refers to $i$ or $j$ except as indices into the sum. Apply a permutation $\pi$ to the input rows and every output row is permuted the same way: $\mathrm{Attn}(PX)=P\mathrm{Attn}(X)$. The model is permutation-equivariant — it is, to a first approximation, a very expensive bag of words.

That is a hard failure, not a stylistic one. “The server rejected the request because the token expired” and “the token rejected the server because the request expired” produce identical attention outputs. Any task whose answer depends on order — basically every task worth doing — is unreachable without position information. The failure is also silent: the model trains, the loss goes down, and it learns to answer order-insensitive questions well enough to look like it is working. Symptom: your model does fine on topic classification and fails on any question whose answer is a sequence of instructions, and fine-tuning on more data does not fix it.

The demo that makes this concrete takes two minutes and is worth running before you read the rest of this chapter:

# transformers==4.44.0, torch==2.4.0
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

tok = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2").eval()

def logprob_of(text: str, continuation: str) -> float:
    ids = tok(text, return_tensors="pt").input_ids
    cont = tok(continuation, add_special_tokens=False, return_tensors="pt").input_ids
    full = torch.cat([ids, cont], dim=1)
    with torch.no_grad():
        logits = model(full).logits
    lp = torch.log_softmax(logits[0, len(ids[0]) - 1:-1], dim=-1)
    return lp.gather(-1, cont[0].unsqueeze(-1)).sum().item()

print(logprob_of("The dog bites the man.", " The reverse is"))
# (example) clearly higher after a correct forward pass than:
print(logprob_of("The man bites the dog.", " The reverse is"))

Now imagine the same model with no positional encoding at all: both prompts would produce identical logits, and the difference above would be exactly zero. That is what §1 formalizes.

2. How it works

2.1 The three places you can put position

There are only three places a positional signal can enter a Transformer, and every published scheme is one of them:

1. Added to the token embedding, before the first block: $x_i\leftarrow x_i+p_i$. (Sinusoidal, learned absolute, and every “position embedding” variant.)
2. Added to the attention logits, per head, at score time: $s_{ij}\leftarrow s_{ij}+b(i,j)$. (ALiBi, T5 relative bias, Shaw et al. 2018.)
3. Mixed into $q$ and $k$ themselves, as a function of position. (RoPE — a rotation, not an addition.)

The distinction between (1) and (3) is the one that took the field five years to absorb. In (1), position enters through the content: the vector the model sees at position $i$ is a different vector, and everything downstream — $W_Q$, $W_K$, $W_V$, the MLP — has to learn to disentangle “this is about the token” from “this is about the position.” In (3), position enters only through the $q\cdot k$ product, which is exactly where the model needs it and nowhere else.

2.2 Sinusoidal encodings: what each symbol is for

$${\mathrm{PE}}_{(pos,2i)}=\sin \left(\frac{pos}{10000^{2i/d}}\right),{\mathrm{PE}}_{(pos,2i+1)}=\cos \left(\frac{pos}{10000^{2i/d}}\right)$$

• $pos\in \{0,1,\dots ,n-1\}$ — the position index. It is integer and unbounded: the function is defined for any $pos$, which is precisely why it was chosen over a lookup table.
• $i\in \{0,\dots ,d/2-1\}$ — the frequency index. Each value of $i$ selects a different sinusoid, and the pair $(2i,2i+1)$ shares one frequency as its sine and cosine. Note that $i$ indexes frequency pairs, so there are $d/2$ frequencies across $d$ dimensions.
• $2i/d$ — the frequency exponent. It maps the frequency index into $[0,1)$, which makes the divisor $10000^{2i/d}$ run geometrically from $10000^0=1$ (fastest) to $10000^{1-2/d}\approx 10000$ (slowest). The paper states the consequence directly: “the wavelengths form a geometric progression from $2\pi$ to $10000\cdot 2\pi$.” So the fastest dimension rotates once every ~6.3 tokens and the slowest once every ~62,832 tokens. This is the multi-scale clock idea: fine dimensions resolve local order, coarse dimensions resolve position in the document, and using a geometric spacing rather than a linear one is what means no single scale has to carry the whole range.
• $10000$ — the base, a magic constant in 2017 that became a tuned hyperparameter with RoPE (LLaMA-3 raises it to 500,000; see §2.5). Raising it stretches every wavelength by the same factor, which is how you buy context length.
• The addition $\mathrm{PE}(pos)$ to the embedding is itself load-bearing and is the part people get wrong. Because the pair is $(\sin \theta ,\cos \theta )$ for the same $\theta$, the encoding vector at position $pos+k$ is a fixed linear map (a block-diagonal rotation, one $2\times 2$ block per frequency pair) applied to the encoding at $pos$. In other words $\mathrm{PE}(pos+k)=R_k\mathrm{PE}(pos)$ for a $k$-dependent rotation, independent of $pos$. That property is the entire justification the paper gives — “we hypothesized it would allow the model to easily learn to attend by relative positions.” A learned table has no such structure: there is no guarantee that $p_{m+k}$ relates to $p_m$ at all, so the model must discover relative position from scratch from data.

2.3 Why RoPE: the derivation in one line

RoPE starts from a requirement rather than a construction. The requirement is

$$⟨f_q(x_m,m),f_k(x_n,n)⟩=g(x_m,x_n,m-n)$$

That is: the attention score between a query at absolute position $m$ and a key at absolute position $n$ must depend on their positions only through the difference $m-n$. Su et al. 2021 (RoFormer, https://arxiv.org/abs/2104.09864, Eq. 11) derive that in 2D the only solution of this form is a rotation.

Concretely: pair up the dimensions of each head vector, $(x_{2i},x_{2i+1})$, and treat each pair as a complex number. Multiply by $e^{im{\theta }_i}$ — i.e. rotate by angle $m{\theta }_i$:

$$R_{m,i}=\left(\begin{array}{cc}\cos (m{\theta }_i)& -\sin (m{\theta }_i)\\ \sin (m{\theta }_i)& \cos (m{\theta }_i)\end{array}\right),{\theta }_i={\Theta }^{-2i/d_k}$$

• $m$ — the absolute position of the query (or $n$ for the key). It appears only inside the rotation angle.
• ${\theta }_i$ — the rotation frequency of dimension pair $i$. It is the exact analogue of $10000^{-2i/d}$ in the sinusoidal formula: the same geometric progression, with fast pairs resolving local order and slow pairs resolving global order.
• $\Theta$ — the base, 10,000 in the paper (and in LLaMA-1/2). Because ${\theta }_i={\Theta }^{-2i/d_k}$, raising $\Theta$ decreases every ${\theta }_i$ and therefore stretches every wavelength — the knob long-context extension turns.
• $d_k$ — the per-head dimension. Note it is $d_k$, not $d_{model}$: RoPE is applied per head, to the projected $q$ and $k$, never to the residual stream. Getting this wrong (rotating the $d_{model}$-wide hidden state) changes which dimensions pair with which.

The identity that makes it work:

$$(R_{m}q{)}^{\top }(R_{n}k)=q^{\top }{R}_m^{\top }{R}_{n}k=q^{\top }{R}_{n-m}k$$

because rotations compose additively: $R_m^{\top }{R}_n=R_{n-m}$. The absolute positions $m$ and $n$ have cancelled; what survives is the relative offset. This is why RoPE is a relative position scheme implemented with absolute angles — the property that took the field a while to appreciate.

Three consequences follow:

• No parameters. ${\theta }_i$ are constants. Compare a learned relative position table $b(i,j)$ (T5) or per-head relative vectors (Shaw et al. 2018), which add parameters and, worse, are undefined for offsets outside the trained range.
• No added activation memory. The rotation is applied in place to $q$ and $k$ at each layer. An additive encoding materialises an $n\times d$ matrix first.
• The fast/slow split gives natural decay. The paper (§3.3, “long-term decay”) shows that with geometrically spaced ${\theta }_i$ the expected inner product between rotated $q$ and $k$ decays as $|m-n|$ grows — a free prior that nearby tokens are more likely to be related.

2.4 One detail that trips up every implementation: which pairs

Mathematically the pairs are $(x_0,x_1),(x_2,x_3),\dots$ with ${\theta }_i={\Theta }^{-2i/d_k}$. The efficient implementation (Su et al. §3.4.2) uses the half-split form instead:

$$q\leftarrow q\odot \cos (m\theta )+rotate\_half(q)\odot \sin (m\theta )$$

where $rotate\_half([q_0,\dots ,q_{d/2-1},q_{d/2},\dots ,q_{d-1}])=[-q_{d/2},\dots ,-q_{d-1},q_0,\dots ,q_{d/2-1}]$ and $\theta =({\theta }_0,\dots ,{\theta }_{d/2-1})$ is repeated twice to fill $d$ dimensions. The two forms are related by a fixed permutation of dimensions, so they compute the same function up to a basis change — but a checkpoint trained with one convention and loaded with a kernel that assumes the other produces silently wrong logits, not an error. This is the single most common RoPE bug, and §5 lists its symptom.

2.5 Long context is a RoPE-frequency problem

Every long-context extension is an intervention on ${\theta }_i$ or on the position index $m$:

The pattern: PI scales all frequencies equally and damages local resolution; NTK-aware scaling exploits the fact that you only need the slow dimensions to cover the new range, leaving the fast ones alone. The honest summary [as of 2026-02] is that frontier models train natively at longer context and use YaRN/NTK only for further stretching; the RoPE base is raised, not the positions rescaled.

flowchart TB
    SUB["q, k after projection (per head)"] --> PAIR["Treat as d_k/2 complex pairs"]
    PAIR --> ROT["Rotate pair i by angle m·θ_i<br/>θ_i = Θ^(-2i/d_k)"]
    ROT --> QR["q_m, k_n rotated"]
    QR --> DOT["q_mᵀ k_n = f(x_m, x_n, m−n)"]
    DOT --> SM["softmax → attention weights"]
    subgraph CONTEXT ["Long-context knobs (all touch θ_i or m)"]
        PI["PI: m → m·(L_train/L_new)"]
        NTK["NTK: Θ → larger"]
        YARN["YaRN: NTK-by-parts + temperature"]
    end
    CONTEXT -.-> ROT

3. Implementation

Two implementations: the dense, readable one (so you can see the math), then the version you actually ship.

# torch==2.4.0
import torch


def sinusoidal_pe(n: int, d: int, base: float = 10000.0) -> torch.Tensor:
    """(n, d) sinusoidal encoding. Note it is ADDED to the embedding."""
    assert d % 2 == 0
    i = torch.arange(0, d // 2, dtype=torch.float32)        # frequency index
    theta = base ** (-2.0 * i / d)                          # (d/2,) frequencies
    pos = torch.arange(n, dtype=torch.float32)              # (n,)
    angles = pos[:, None] * theta[None, :]                  # (n, d/2)
    # interleave sin/cos so dimension pair (2i, 2i+1) shares one frequency
    pe = torch.stack([angles.sin(), angles.cos()], dim=-1)  # (n, d/2, 2)
    return pe.reshape(n, d)


def rope_tables(seq_len: int, d_k: int, base: float = 10000.0):
    """cos/sin tables for RoPE, in the half-split (HF) convention."""
    assert d_k % 2 == 0
    inv_freq = base ** (-torch.arange(0, d_k, 2, dtype=torch.float32) / d_k)  # (d_k/2,)
    t = torch.arange(seq_len, dtype=torch.float32)
    freqs = torch.outer(t, inv_freq)              # (seq_len, d_k/2)
    emb = torch.cat([freqs, freqs], dim=-1)       # (seq_len, d_k) — repeated twice
    return emb.cos(), emb.sin()


def rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat([-x2, x1], dim=-1)


def apply_rope(q, k, cos, sin):
    """q, k: (B, H, T, d_k); cos, sin: (T, d_k). Rotates q and k only — never v."""
    cos = cos[None, None, :, :]     # broadcast over batch and head
    sin = sin[None, None, :, :]
    q_out = q * cos + rotate_half(q) * sin
    k_out = k * cos + rotate_half(k) * sin
    return q_out, k_out


if __name__ == "__main__":
    torch.manual_seed(0)
    d_k, total = 64, 32

    # Check 1: a rotation must preserve norms.
    q = torch.randn(1, 4, total, d_k)
    cos, sin = rope_tables(total, d_k)
    qr, _ = apply_rope(q, q.clone(), cos, sin)
    assert torch.allclose(qr.norm(dim=-1), q.norm(dim=-1), atol=1e-5)

    # Check 2: the score must depend on the OFFSET (m - n), not on absolute m.
    # Use the SAME q and k for every pair so only the rotation differs.
    q0 = torch.randn(1, 1, 1, d_k)
    k0 = torch.randn(1, 1, 1, d_k)
    cos, sin = rope_tables(128, d_k)

    def score(m: int, n: int) -> float:
        qm, _ = apply_rope(q0, q0, cos[m:m + 1], sin[m:m + 1])
        _, kn = apply_rope(k0, k0, cos[n:n + 1], sin[n:n + 1])
        return (qm @ kn.transpose(-1, -2)).item()

    print(score(0, 0), score(7, 7))                      # offset 0 pair: identical to fp error
    print(score(0, 7), score(100, 107), score(7, 114))   # offset +7 triple: also identical

Those prints are the whole theorem. The score depends only on the offset $n-m$: score(0, 0) and score(7, 7) share offset 0 and must be equal, and score(0, 7), score(100, 107) and score(7, 114) all share offset +7 and must equal each other. What must not match is the offset-0 value against the offset-+7 value. If two scores that share an offset differ by more than floating-point noise, you have a pairing or sign error — check that before suspecting the model.

It is worth running that check once even if you use a library, because the failure mode is silent and library conventions genuinely differ (see §5).

The version you ship. Every mainstream stack already has this, and you should never hand-roll it in production:

• HuggingFace transformers — LlamaRotaryEmbedding (per head, applied in LlamaAttention.forward), configured by rope_theta and rope_scaling in config.json. [as of 2026-02] the RoPE config was being consolidated under a rope_parameters key, and checkpoints written by one major version do not always load correctly in an older one — pin the transformers version alongside any checkpoint you ship.
• torchtune — exposes rope_base directly: llama3_8b() is built with rope_base=500_000, num_kv_heads=8, intermediate_dim=14336 (https://meta-pytorch.org/torchtune/main/_modules/torchtune/models/llama3/_model_builders.html).
• vLLM / SGLang / TensorRT-LLM — these must be given the same RoPE configuration as the checkpoint. A mismatch here produces fluent but subtly-wrong output rather than an error.

The knobs that matter, with realistic ranges:

4. Numbers and trade-offs

Parameter and activation overhead

For $d=4096$ and a 128k vocabulary:

The activation row is not trivial at long context. An additive encoding needs an $n\times d$ tensor; at $n=100,000$, $d=4096$, bf16, that is 800 MB per sequence — on top of the KV cache, and only at the input. RoPE adds nothing, which is one reason it is the default even before you get to quality.

Wavelengths, derived

The fastest pair in a head of $d_k=128$ has period $2\pi /{\theta }_0=2\pi /1\approx 6.28$ tokens. The slowest has $2\pi /{\theta }_{63}=2\pi {\Theta }^{(d_k-2)/d_k}$. With $\Theta =10^4$ that is $\approx 5.4\times 10^4$ tokens; with $\Theta =5\times 10^5$ it is $\approx 2.6\times 10^6$ tokens. (Asymptotically the slowest period is $2\pi \Theta$; the $(d_k-2)/d_k$ exponent is why the exact figure is a little smaller, and it is the exact figure that governs behaviour.)

That calculation is the whole long-context story. With $\Theta =10^4$, the slowest dimension completes one revolution in ~54k tokens — so a 128k-token context asks the model to distinguish positions using an angle that has already wrapped around more than twice, and any two positions separated by exactly one wavelength are indistinguishable by that dimension alone. Raising $\Theta$ to $5\times 10^5$ moves the slowest period to ~2.6M tokens and keeps the whole 128k range inside a single revolution (0.05 of one, in fact). The cost: the faster dimensions have longer periods too, which is why naive base scaling blurs local resolution and why NTK-aware/YaRN schemes scale the frequencies unevenly instead.

Attention score magnitude is unchanged

A rotation preserves norms: $\Vert R_{m}q\Vert =\Vert q\Vert$. So the variance argument behind the $1/\sqrt{d_k}$ scaling in Attention is unaffected — RoPE does not change the score distribution’s scale and does not interact with that constant. This is a real advantage over relative-bias schemes, which do add to the logits and therefore change their effective temperature (T5-style biases are typically initialised so this is small, but it is a thing to check).

Which axis you are optimizing

5. Pitfalls

Applying position after the first block (or not at all). Symptom: the model trains and reaches a plausible loss, but any task requiring order — following multi-step instructions, code, arithmetic — is at chance, and the loss curve looks suspiciously similar to a bag-of-words baseline. Diagnose by permuting the input tokens and checking the logits change; if they do not, your position signal isn’t reaching attention.

Rotating $v$ as well as $q$ and $k$. Symptom: the model trains but converges to a strictly worse loss, with no error and no shape mismatch. The output of attention is a weighted sum of values; rotating values by their absolute position injects a position-dependent transform into the output that no subsequent $W_O$ can undo (it varies per position). RoPE goes on $q$ and $k$ only — always.

Mismatched pairing convention between training and inference. Symptom: the model generates fluent text that is subtly wrong — right topic, wrong facts, degraded instruction-following — and perplexity is a few points up rather than catastrophically broken. This is the classic half-split-vs-interleaved bug and its signature is that almost everything works. The one-line test is the offset-invariance check in §3.

Caching unrotated keys. Symptom: generation degrades as the sequence grows, and prefix caching gives different results from no caching. The KV cache must store $R_n{k}_n$ — the already-rotated key at its own position — not the raw $k_n$. If you rotate at read time, you must rotate by the position the key occupied, not by the current one.

Adding both RoPE and an additive positional embedding. Symptom: slow convergence and a model that is oddly sensitive to sequence length; you have made position information twice as strong and partly redundant, and the optimizer spends capacity separating the two signals. Check config.json/model code — some models genuinely use both, but doing it by accident is common when porting a BERT-family training script to a LLaMA-family checkpoint.

Mismatched rope_scaling between the checkpoint and the server. Symptom: output quality depends on the serving framework. The same model served by transformers and by vLLM at 32k context gives measurably different answers, with no error in either log. Always read the RoPE config out of the checkpoint and pass it explicitly.

Assuming a large max_position_embeddings means the model works there. Symptom: needle-in-a-haystack passes but multi-hop and aggregation questions at the same length fail badly. The window fits, but the effective context is much shorter. Measuring this is the topic of §7.

Sinusoidal encodings with a learned base. Symptom: training-time loss looks normal but quality collapses at lengths beyond what you trained on, worse than a properly configured RoPE model. The frequencies are constants in the original design; treating $10000$ as a hyperparameter turns it into an untuned long-context intervention without the rest of the recipe.

6. Using it in a project

In 2026 you almost never choose a positional encoding. A checkpoint comes with one burned in, and if it is RoPE you have exactly one real decision: the context window, and whether the checkpoint’s RoPE config matches how you are serving it.

Reach for a positional-encoding decision when:

• You are fine-tuning a model to a longer context than it was trained on. This is the common case, and the answer is a RoPE intervention (YaRN or NTK scaling) plus a modest amount of long-context fine-tuning data — usually far less than continued pretraining would need.
• You are training from scratch, in which case RoPE with a base matched to your target context is the default and there is no serious competitor. [as of 2026-02]
• You are serving a checkpoint through a different stack than it was trained in, and need to verify the RoPE configuration survives the port.
• You need a hard context ceiling you can reason about: a learned absolute table gives you an exact wall at $n_{max}$; RoPE gives you graceful degradation whose shape you must measure.

When NOT to touch any of this:

• When your task fits in the model’s native window. Almost all production RAG and chat traffic fits in 8k–32k. Adding YaRN scaling to squeeze out 128k context, when 95% of your requests are under 4k and the long ones are summarization, buys quality loss on the common case for a capability the rare case does not need.
• When the real problem is retrieval. If the answer isn’t in the context, a larger context window does not help. Measure retrieval recall@k first — see RAG and Retrieval.
• When a streaming/multi-turn design would work better. Repositioning or re-summarizing a conversation is often cheaper and more reliable than opening the window further.
• When the model is not RoPE-based. Bidirectional encoders used for retrieval embeddings often use learned or sinusoidal absolute positions plus a fixed max_position_embeddings; the RoPE playbook simply does not apply, and forcing a RoPE checkpoint into that role loses the properties you were relying on.

The one transferable mental model. Position is a property of the attention score, not of the token vector. Adding to the residual stream is a workaround for not having found a way to put position in the score. That makes the RoPE derivation obvious and the failure modes predictable: anything added to the residual stream can be learned around; anything that changes $q\cdot k$ cannot.

7. How you would prove it worked

Positional encoding has one direct metric and one strong diagnostic.

Per-position loss. Compute the model’s loss bucketed by token position on a held-out set. Plot position on the x-axis. Two things to look for:

• Within the trained range, the curve should be roughly flat. A curve that rises steadily with position inside the training window indicates your position signal is decaying or your attention is diluting — not a long-context problem at all.
• Beyond the trained range, the curve tells you the extrapolation quality. With a learned absolute table it turns into a wall (the embeddings do not exist). With RoPE at a mismatched rope_scaling, it rises smoothly and steeply. With a correctly configured extension, it stays flat and then degrades.

The permutation test. (example) Take 200 evaluation examples, shuffle each input’s tokens at random, and measure the drop in task accuracy. A model with a working positional encoding should lose most or all of its accuracy on order-sensitive tasks; a model that barely notices has a broken or absent position signal. Report the drop, not the absolute number — the absolute number depends on the task.

The long-context ablation, done honestly. Needle-in-a-haystack saturates on almost every model that advertises a large window and therefore proves very little. Use an aggregation task instead: place k items at random positions and ask for a count or a sum. (example) A model can pass 128k needle retrieval and fall to near-chance on aggregation at 16k. Report the length at which that metric drops below its short-context baseline; that is the number to put in front of a stakeholder, not the advertised window.

Reporting template (example — replace with your own numbers). “At the checkpoint’s native 8k window we saw flat per-position loss. Extending to 32k with YaRN at factor 4, no fine-tuning, raised per-position loss beyond 8k by 0.9 nats (example) and dropped aggregation accuracy from 91% to 44% (example). After 400 steps of long-context SFT on 8k–32k sequences, loss was within 0.05 nats of the short-context baseline and aggregation recovered to 87% (example). We kept the native config for all traffic under 8k, which is 96% of requests.”

8. Interview drill

“Why does a Transformer need positional encoding at all?” Strong answer: self-attention is permutation-equivariant — the output is a weighted sum over positions with weights that depend only on $q\cdot k$ content, so reordering the input reorders the output identically. Without a position signal the model cannot distinguish “dog bites man” from “man bites dog”, and it will happily train to a plausible loss while failing every order-sensitive task. Common wrong answer: “because attention has no concept of distance.” Attention has no concept of anything positional, including order — saying “distance” suggests the problem is only long-range, when it is already broken at two tokens.

“Walk me through the sinusoidal encoding.” Strong answer: for each position $pos$ and each dimension pair $i$, emit $\sin$ and $\cos$ of $pos/10000^{2i/d}$. The $d/2$ frequencies form a geometric progression from $2\pi$ to $10000\cdot 2\pi$, giving a multi-scale clock. The reason for the specific pairing is that for a fixed offset $k$, $\mathrm{PE}(pos+k)$ is a linear function of $\mathrm{PE}(pos)$, so relative offsets are linearly recoverable. Follow-up: “so why did it lose?” → it is added to the residual stream, so position and content are entangled in the same vector from layer 0 onward, every weight has to learn to separate them, and it does not extrapolate well because the model has never seen long-range position values.

“Explain RoPE to me as if I have 90 seconds.” Strong answer: pair the head dimensions and rotate each pair by an angle proportional to position. Because $R_m^{\top }{R}_n=R_{n-m}$, the query–key inner product depends only on the relative offset, so attention becomes a function of $m-n$ while the implementation still uses absolute angles. Zero parameters, zero added memory, no change to attention’s form, and the base $\Theta$ is a single knob for context extension. Common wrong answer: “it adds a rotary positional vector to the embedding.” RoPE adds nothing to the embedding — it multiplies $q$ and $k$ after the projection, per head.

“Why is RoPE applied to $q$ and $k$ but not $v$?” Strong answer: the position information needs to affect the attention weights, which are the $q\cdot k$ products. The output is a weighted sum of values, and rotating a value would apply a position-dependent linear transform to the output that no downstream weight can invert — it injects noise, not signal. Follow-up: “so how does position reach the output?” → through which values get weighted, i.e. through the softmax distribution, which is exactly where you want it.

“I trained a 2k-context model and I need 32k. What do you do?” Strong answer: raise the effective context with a RoPE intervention — YaRN or NTK-aware scaling, which touch the frequencies rather than the positions — then run a short fine-tuning pass with long sequences in the data. Measure per-position loss and an aggregation task at 32k, not needle retrieval. Common wrong answer: “increase max_position_embeddings and the RoPE base.” Changing the base without matching training degrades quality within the original window too, and declaring a bigger window in config does not create the capability — you get fluent, confident, badly-positioned output.

“ALiBi vs RoPE — why did RoPE win?” Strong answer: both give relative position without per-offset parameters, and ALiBi extrapolates well. RoPE won on ecosystem gravity plus two concrete properties: it composes equally well with the KV cache and with FlashAttention (a rotation of $q$ and $k$ composes with any kernel that is already doing a matmul), and its frequency structure gives a natural, tunable decay and a clean long-context extension story. ALiBi’s bias has to be added to the logits, which some fused kernels handle less flexibly. Common wrong answer: “RoPE is more accurate.” On many benchmarks ALiBi is competitive; the decisive factors are engineering and extendability, and saying “more accurate” invites a follow-up you cannot answer.

“Does RoPE change the $1/\sqrt{d_k}$ scaling?” Strong answer: no. Rotations preserve norms, so the score variance is unchanged and the same constant applies. Common wrong answer: “yes, you have to rescale by the base.” That confusion comes from YaRN, which adds a temperature correction on top of RoPE as part of its extension recipe — that is a property of YaRN, not of RoPE.

Next: Modern Architecture Deltas — what changed between the 2017 paper and a 2026 model: pre-norm, RMSNorm, SwiGLU, GQA, FlashAttention, MoE, and the cost each delta was paying down.

Modern Architecture Deltas

TL;DR — A 2026 LLM is the 2017 transformer with five substitutions. The norm moved before the sublayer (pre-norm) so gradients reach depth without a hand-tuned warmup; LayerNorm became RMSNorm because re-centering was doing no work and re-scaling was; the FFN became a 3-matrix SwiGLU gate at $8d/3$ hidden width; multi-head attention became grouped-query attention so the KV cache shrank by the query-to-KV head ratio; and attention moved into an IO-aware fused kernel (FlashAttention) that is exact, not approximate — it removes memory traffic, not FLOPs. Each delta is a response to a specific thing that broke at scale, and each has a cost that shows up somewhere.

1. Why this exists

The 2017 transformer in The Transformer trains a 6-layer encoder–decoder on 2014-era translation data. Scale that recipe to 100+ layers and trillions of tokens and four things break, each in a different way:

Depth refuses to train. Post-LN transformers — the original design, with the norm applied after the residual add — diverge at depth unless you spend a long, fragile learning-rate warmup. Xiong et al. proved the cause: at initialization the expected gradients near the output layer grow with depth, so the early steps take far too large a step. That is a training-stability failure, and it is why essentially every modern model is pre-LN (Xiong et al. 2020, ICML).

Normalization is on the critical path of every layer. LayerNorm computes a mean and a variance, and at 100+ layers × every token × every step that reduction is measurable wall-clock. Zhang & Sennrich’s finding was that the re-centering half of LayerNorm is dispensable — only the re-scaling mattered — so half the reduction can be deleted for free (Zhang & Sennrich 2019, NeurIPS).

The FFN is the parameter budget and it is paying for a weaker function. The position-wise FFN holds roughly two-thirds of a transformer’s parameters. It spent them on act(xW_1)W_2; Shazeer showed that multiplying in a second, gate-controlled projection is strictly better at equal parameter count and equal FLOPs (Shazeer 2020).

Inference is memory-bandwidth-bound, not compute-bound. Once you serve the model, the throughput ceiling is how many bytes you must move per token. Two separate byte counts dominate: the KV cache (attention) and the attention score matrix (training and prefill). GQA attacks the first (Ainslie et al. 2023); FlashAttention attacks the second (Dao et al. 2022).

The pattern to internalise. None of these is “a better idea someone had”. Each is a named mechanism deleting a named cost. If you can recite the cost, the mechanism follows.

2. How it works

2.1 The residual stream, and why normalization position is a real choice

Everything below is easier if you hold one picture: the residual stream. The model is $L$ blocks, each of which adds to a running vector and never replaces it:

$$x_{l+1}=x_l+\mathrm{Sublayer}(\cdot )$$

Sublayers (attention, FFN) read from the stream and write back into it. Read capacity, not computation, is the right intuition: the stream is a shared workspace and each block appends a correction.

Normalization has to happen somewhere, and there are exactly two natural places. Post-norm (Vaswani 2017) normalizes the sum; pre-norm normalizes the input to the sublayer, leaving the sum alone:

$$x_{l+1}=\mathrm{LN}\left(x_l+F_l(x_l)\right)\text{vs}{x}_{l+1}=x_l+F_l\left(\mathrm{LN}(x_l)\right)$$

The difference is structural, not cosmetic. In post-norm, the identity path is broken at every layer: $x_l$ passes through an LN before it can reach $x_{l+1}$. LN divides by the standard deviation of its input, so $\partial x_{l+1}/\partial x_l$ carries a factor from that LN’s Jacobian at every one of the $L$ layers. The product of $L$ such factors is what Xiong et al. show has the wrong behaviour at initialization: the expected gradient magnitude near the output grows with depth, so the step at step 0 is scaled wrong and the optimizer needs warmup to survive it. In pre-norm, $\partial x_{l+1}/\partial x_l$ contains an exact identity term, so gradient reaches every layer without passing through a normalization Jacobian; Xiong et al. prove the consequence and remove warmup entirely for pre-LN with comparable results.

There is a second, forward-pass reason GPT-2 discovered independently. With the identity path intact, each block adds to the stream. If each block’s contribution has roughly unit variance and is roughly independent of the others, the stream’s variance after $L$ blocks is

$$\mathrm{Var}(x_L)\approx \mathrm{Var}(x_0)+\sum _{l=0}^{L-1}\mathrm{Var}(F_l)\approx L$$

so the magnitude grows like $\sqrt{L}$: a 100-layer model ends up with outputs an order of magnitude larger than a 12-layer one, and the learning rate that suited one does not suit the other. GPT-2’s fix, still in modeling_gpt2.py, shrinks each residual branch’s output projection at initialization by $1/\sqrt{2L}$ — the 2 counts the two branches (attention, FFN) per block — so the accumulated variance stays $O(1)$ regardless of depth:

# ... scale the weights of residual layers at initialization by a factor of
# 1/√N where N is the # of residual layers.   -- GPT-2
std = self.config.initializer_range / math.sqrt(2 * self.config.n_layer)
init.normal_(module.c_proj.weight, mean=0.0, std=std)

Why this is an interview-grade distinction. “Pre-norm is more stable” is a fact. “Pre-norm keeps an identity term in $\partial x_{l+1}/\partial x_l$, so no normalization Jacobian enters the depth-accumulated product; post-norm puts an LN on the identity path at every layer” is the mechanism. The follow-up is always “then why does GPT-2 still scale the init?” — pre-norm fixes the backward path and does nothing about the forward variance accumulation.

2.2 RMSNorm: deleting the mean from LayerNorm

LayerNorm, written out, is

$$\mathrm{LN}(x)=\gamma \odot \frac{x-\mu }{\sqrt{{\sigma }^2+ϵ}}+\beta ,\mu =\frac{1}{d}\sum _{i=1}^d{x}_i,{\sigma }^2=\frac{1}{d}\sum _{i=1}^d(x_i-\mu {)}^2$$

Symbol by symbol, with the justification:

• $x\in {\mathbb{R}}^d$ — one token’s activation in the residual stream. Normalization is per-token, not per-batch: a token’s normalization cannot depend on other examples in the batch, which is what makes the model usable at batch size 1 during decoding.
• $\mu$ — the mean over the $d$ features of that one token. Subtracting it re-centers the vector so its entries sum to zero.
• ${\sigma }^2$ — the variance across the same $d$ features. Dividing by $\sqrt{{\sigma }^2+ϵ}$ re-scales to unit standard deviation, making the sublayer’s input distribution independent of how large the residual stream happened to grow. This is the half that does the work.
• $ϵ$ — a small constant ($10^{-5}$ or $10^{-6}$), added inside the square root. Without it, a constant vector — or an all-zero position after masking — divides by zero and produces NaN. It costs a tiny scale bias; it buys numerical survival.
• $\gamma \in {\mathbb{R}}^d$, $\beta \in {\mathbb{R}}^d$ — learned per-feature gain and bias. Normalization discards the input’s scale and offset; $\gamma$ and $\beta$ let the network relearn which features should be large and where the zero point sits, so the layer is not locked into a rigid unit-variance representation.

RMSNorm asks: which of those terms is doing the work? Its claim is that re-centering is not, and that the re-scaling is:

$$\mathrm{RMSNorm}(x)=\gamma \odot \frac{x}{\sqrt{\frac{1}{d}{\sum }_{i=1}^d{x}_i^2+ϵ}}$$

• The denominator is the root mean square, $\sqrt{\frac{1}{d}\sum x_i^2}$, not the standard deviation $\sqrt{\frac{1}{d}\sum (x_i-\mu {)}^2}$. They differ only by the mean: ${\mathrm{RMS}}^2={\sigma }^2+{\mu }^2$. So RMSNorm is invariant to re-scaling ($\mathrm{RMSNorm}(cx)=\mathrm{RMSNorm}(x)$, $c>0$) but not to re-centering ($\mathrm{RMSNorm}(x+c1)\ne \mathrm{RMSNorm}(x)$). That is the trade: one invariance given up for a cheaper reduction.
• $\beta$ is gone. A bias followed by a rescale is partly redundant; Zhang & Sennrich found no benefit from it, and dropping it removes $d$ parameters per norm layer.
• $\gamma$ stays, for the original reason: it restores the per-feature expressivity pure normalization removes. Removing $\gamma$ too would make the norm a fixed, parameterless transformation, and the network could no longer learn feature scales.

Why is that a win rather than mere simplification? A norm layer’s dominant cost is its reduction and the extra memory pass. LayerNorm needs mean and variance; RMSNorm needs one mean-of-squares. Zhang & Sennrich report 7%–64% runtime reduction at comparable quality. At 100+ norm layers that is real wall-clock, and it is largest exactly where the norm is a big fraction of a small layer’s work.

The $ϵ$ trap. In bf16, $\frac{1}{d}\sum x_i^2$ can underflow for small activations and overflow for large ones — a sum of squares over a wide dynamic range is precisely the operation that loses precision. Production kernels accumulate it in fp32 even when the activations are bf16, and place $ϵ$ so the RMS cannot swamp it. Getting this wrong produces loss spikes that look like a data problem.

Removing RMSNorm entirely does not give a small quality regression; it gives training divergence at depth, because the residual stream’s magnitude is then unconstrained. RMSNorm is not decoration.

2.3 SwiGLU: the FFN as a learned gate

The original FFN is a two-matrix bottleneck with a pointwise nonlinearity:

$$\mathrm{FFN}(x)=W_2\mathrm{act}(W_{1}x)$$

Symbols: $x\in {\mathbb{R}}^d$ is one token’s stream vector; $W_1\in {\mathbb{R}}^{d_{ff}\times d}$ expands to a wide hidden space ($d_{ff}=4d$ in the 2017 paper); $\mathrm{act}$ is elementwise; $W_2\in {\mathbb{R}}^{d\times d_{ff}}$ projects back. The hidden width is where the parameters live — $2d{d}_{ff}=8d^2$ of them at $d_{ff}=4d$.

A gated linear unit replaces the single expansion with two, and multiplies them elementwise:

$$\mathrm{SwiGLU}(x)=W_2(\underset{\text{gate}}{\underbrace{{\mathrm{Swish}}_{\beta }(x{W}_1)}}\odot \underset{\text{value}}{\underbrace{(x{W}_3)}}),{\mathrm{Swish}}_{\beta }(z)=z\sigma (\beta z)$$

• $x{W}_1\in {\mathbb{R}}^{d_{ff}}$ — the gate pre-activation. It is the branch that decides, per hidden unit, how much of the value branch to let through.
• ${\mathrm{Swish}}_{\beta }(z)=z\sigma (\beta z)$ with $\sigma$ the logistic sigmoid and $\beta$ a learned or fixed slope (Shazeer sets $\beta =1$, which is exactly SiLU). This $\beta$ is the gate’s slope and is unrelated to the LayerNorm bias $\beta$ of §2.2; the two names coincide only because both are conventional. Why this function and not ReLU? Because it is smooth and non-monotone: it dips slightly negative for small negative $z$ and asymptotes to $z$ for large positive $z$. The non-monotonicity is what lets the gate express “partially suppress this unit” instead of only “on” or “off”, and the smoothness means it has a usable gradient everywhere, including at the origin where ReLU’s gradient is discontinuous.
• $x{W}_3\in {\mathbb{R}}^{d_{ff}}$ — the value branch with no nonlinearity. It is linear so that the product’s nonlinearity comes entirely from the gate.
• $\odot$ — elementwise product. This is the whole trick: the gate learns a multiplicative, input-dependent mask. A ReLU acts like a fixed mask on a single projection; a GLU lets the mask depend on the input through a different learned projection. That is strictly more expressive per parameter.
• $W_2\in {\mathbb{R}}^{d\times d_{ff}}$ — the down-projection back to model width. Remove it and the block writes a $d_{ff}$-wide vector into the residual stream; the shapes no longer close.

What breaks if you remove the gate. Setting ${\mathrm{Swish}}_{\beta }(x{W}_1)=1$ reduces SwiGLU to $W_2(x{W}_3)$ — a single linear layer, collapsing the FFN to one low-rank linear map. Removing the nonlinearity while keeping the product gives Shazeer’s Bilinear variant, $(xW)\odot (xV)W_2$, a quadratic form that empirically underperforms the gated variants.

Where $8d/3$ comes from. Three matrices instead of two means $3d{d}_{ff}$ parameters. Holding the parameter count of the $4d$ FFN constant: $3d{d}_{ff}=2d(4d)\Rightarrow d_{ff}=\frac{8d}{3}$. Shazeer states this explicitly — “to keep the number of parameters and the amount of computation constant, we reduce the number of hidden units $d_{ff}$ by a factor of $\frac{2}{3}$” — and implementations then round up to a hardware-friendly multiple (commonly 128, sometimes 256). That rounding is why real configs contain odd widths like 11008; at $d=4096$, $8d/3=10922.67$, which rounds to 11008.

2.4 GQA: shrinking the KV cache by sharing key/value heads

MHA gives every one of the $h$ query heads its own key and value head. At training time that is what you want; at decode time it is a liability, because the KV cache is the thing you re-read on every generated token. The cache size is

$$\text{bytes}=2\cdot L\cdot n_{kv}\cdot d_{head}\cdot n_{ctx}\cdot B\cdot b$$

• $2$ — one copy for keys, one for values. (Keys and values are both cached; queries are not, because a past query is never needed again.)
• $L$ — number of layers. Every layer has its own cache.
• $n_{kv}$ — number of key/value heads. This is the only factor GQA changes.
• $d_{head}$ — per-head width (both $d_k$ and $d_v$; modern models use $d_k=d_v$).
• $n_{ctx}$ — cached context length in tokens.
• $B$ — batch size.
• $b$ — bytes per element (2 for fp16/bf16, 1 for fp8).

Multi-Query Attention (Shazeer 2019) sets $n_{kv}=1$: all query heads share one key head and one value head. The cache shrinks by a factor of $h$. Grouped-Query Attention (Ainslie et al. 2023) interpolates: split the $h$ query heads into $g$ groups and give each group its own KV head, so $n_{kv}=g$:

$$g=h\Rightarrow \text{MHA},g=1\Rightarrow \text{MQA},1<g<h\Rightarrow \text{GQA}$$

$g$ is a hard design constant in the checkpoint — the number of KV heads is baked into the weight shapes, so you cannot change it at serving time.

Concretely, the arithmetic. For a 70B-class model with $L=80$, $h=64$ query heads, $d_{head}=128$, fp16:

(Binary units: 1 MiB $=1024^2$ B, 1 GiB $=1024^3$ B, so 10.0 GiB $=1.0737\times 10^{10}$ B, 10.74 GB.) That is an 8× reduction, and it is the difference between one long-context request fitting on a node and not fitting. The quality argument for interpolating rather than collapsing to MQA is Ainslie et al.’s empirical result: uptrained GQA reaches quality close to MHA at approximately MQA speed. They also supply the recipe most teams actually use — take an existing MHA checkpoint, mean-pool the KV heads within each group, and continue training on ~5% of the original pretraining compute, rather than pretraining a GQA model from scratch.

Why this is the highest-leverage serving knob. Decode is memory-bandwidth-bound: each generated token requires reading the weights and the whole KV cache out of HBM. Shrinking $n_{kv}$ by 8× shrinks one of the two terms that set the throughput ceiling, and unlike quantization it does not touch the weights, so it does not interact with your quantized checkpoint.

MLA, one paragraph. DeepSeek-V2’s Multi-head Latent Attention takes a different route: instead of sharing KV heads across queries, it compresses the KV state into a low-rank latent vector and caches that, decompressing per layer at use time. The paper reports a 93.3% KV-cache reduction versus DeepSeek 67B alongside a 5.76× maximum generation throughput increase (DeepSeek-V2). The trade is that the cache is no longer a plain $n_{kv}\cdot d_{head}$ layout, so kernels must understand the compression — which is why MLA support lands in serving frameworks later than GQA support.

2.5 FlashAttention: what it changes, and what it does not

Standard attention computes $S=Q{K}^{\top }/\sqrt{d_k}$, then $P=\mathrm{softmax}(S)$, then $O=PV$. The $n\times n$ matrix $S$ (and $P$) is materialised in HBM. For $n=128\text{k}$, $h=64$ heads, fp16, that is

$$128000^2\times 64\times 2\text{bytes}\approx 2.1\text{TB}$$

per attention layer. Long before the FLOPs become the problem, the memory traffic does.

FlashAttention’s insight is that the bottleneck is IO between HBM and on-chip SRAM, not arithmetic. It (1) tiles the computation — load blocks of $Q$, $K$, $V$ into SRAM and compute blockwise; (2) uses an online softmax carrying a running maximum $m$ and running normalizer $\ell$, so the softmax is computed incrementally across tiles without ever seeing a full row; and (3) recomputes the $S$ blocks in the backward pass rather than storing them, because a recompute in SRAM is cheaper than a read from HBM. The result is $O(n)$ memory instead of $O(n^2)$ and far fewer HBM accesses.

FlashAttention is exact. It is not an approximation, not a sparsity pattern, not a low-rank projection. It computes the same mathematical function as the reference implementation, up to floating-point reassociation. It does not reduce the $O(n^2)$ FLOP count. It reduces memory traffic and memory footprint.

The three things it does not fix: (1) the FLOPs — prefill at 100k context is still $O(n^2)$ compute, and only sub-quadratic attention (sliding window, sparse, linear/hybrid) changes that; (2) the KV cache — it caches exactly what MHA/GQA caches and makes that cache usable at long context, without shrinking it, which is GQA/MLA’s job; (3) arbitrary attention bias — fused kernels support a subset of masks and biases, and an exotic additive bias or per-request custom mask can force a fallback to the slower math kernel, so check support before designing around one.

flowchart LR
    A["Q,K,V"] --> B["eager: write n×n scores to HBM,<br/>read back for softmax, write P, read back"]
    A --> C["FlashAttention: tile into SRAM,<br/>online softmax (running m, ℓ),<br/>n×n never hits HBM"]
    B --> D["O(n²) memory, 5+ HBM round trips"]
    C --> E["O(n) memory, exact same result"]

3. Implementation

Reference implementations of the four deltas, plus the kernel call. All of it is runnable against PyTorch 2.14 [as of 2026-10]; no CUDA required for the first three.

import torch
import torch.nn as nn
import torch.nn.functional as F


class RMSNorm(nn.Module):
    """gamma only: no mean subtraction, no beta."""
    def __init__(self, d_model, eps=1e-6):
        super().__init__()
        self.eps, self.weight = eps, nn.Parameter(torch.ones(d_model))

    def forward(self, x):
        x32 = x.float()                      # sum-of-squares needs fp32 range
        rms = torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + self.eps)
        return (x32 * rms).to(x.dtype) * self.weight


class SwiGLU(nn.Module):
    def __init__(self, d_model, d_ff=None):
        super().__init__()
        # 8/3*d keeps parameter count equal to the 4d ReLU FFN; round to x128.
        d_ff = d_ff or 128 * ((int(8 * d_model / 3) + 127) // 128)
        self.w_gate = nn.Linear(d_model, d_ff, bias=False)  # W_1
        self.w_up = nn.Linear(d_model, d_ff, bias=False)    # W_3
        self.w_down = nn.Linear(d_ff, d_model, bias=False)  # W_2

    def forward(self, x):
        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))


class Block(nn.Module):
    """Pre-norm: norm the INPUT of each sublayer; the residual add stays bare.
    GQA: h query heads, g KV heads (g == h -> MHA, g == 1 -> MQA)."""
    def __init__(self, d_model, n_heads, n_kv_heads):
        super().__init__()
        assert n_heads % n_kv_heads == 0
        self.h, self.g, self.d_head = n_heads, n_kv_heads, d_model // n_heads
        self.n1, self.n2 = RMSNorm(d_model), RMSNorm(d_model)
        self.w_q = nn.Linear(d_model, n_heads * self.d_head, bias=False)
        self.w_k = nn.Linear(d_model, n_kv_heads * self.d_head, bias=False)
        self.w_v = nn.Linear(d_model, n_kv_heads * self.d_head, bias=False)
        self.w_o = nn.Linear(n_heads * self.d_head, d_model, bias=False)
        self.ffn = SwiGLU(d_model)

    def forward(self, x, is_causal=True):
        B, T, _ = x.shape
        y = self.n1(x)
        q = self.w_q(y).view(B, T, self.h, self.d_head).transpose(1, 2)
        k = self.w_k(y).view(B, T, self.g, self.d_head).transpose(1, 2)
        v = self.w_v(y).view(B, T, self.g, self.d_head).transpose(1, 2)
        # enable_gqa expands KV inside the kernel; doing it yourself with
        # repeat_interleave allocates the memory GQA exists to save.
        attn = F.scaled_dot_product_attention(
            q, k, v, is_causal=is_causal, enable_gqa=(self.h != self.g))
        x = x + self.w_o(attn.transpose(1, 2).contiguous().view(B, T, -1))
        return x + self.ffn(self.n2(x))

The knobs that matter, with realistic ranges

The kernel

F.scaled_dot_product_attention dispatches at runtime among three backends [as of PyTorch 2.14]: FlashAttention-2, Memory-Efficient Attention, and a plain C++ math implementation. All are enabled by default and the fastest compatible one is selected. Two behaviours that bite in production, both from the official docs:

• Boolean mask semantics are inverted relative to nn.MultiheadAttention. In SDPA, True means “participate”; in MultiheadAttention’s key_padding_mask, True means “mask out”. Migrating without inverting the mask produces silently wrong attention.
• Output can differ between backends. Fusing floating-point operations reassociates them; the docs state plainly that the output may differ depending on which kernel is chosen. This is why bit-exact reproducibility across GPU generations is not available for attention.

Force or forbid a specific backend when you need determinism or when a fused kernel’s mask support is missing:

from torch.nn.attention import sdpa_kernel, SDPBackend

with sdpa_kernel([SDPBackend.FLASH_ATTENTION]):
    out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
# A backend that cannot run raises RuntimeError ("No available kernel. Aborting
# execution.") rather than warning and silently falling back -- use this during
# bring-up, not in production loops.

4. Numbers and trade-offs

Pre-norm vs post-norm: what you actually pay

The forward cost is identical — one normalization per sublayer either way. The difference is when you can stop tuning the learning rate: post-norm needs a warmup schedule whose length you must get right, and getting it wrong at 100 layers means divergence, whereas pre-norm lets Xiong et al. remove warmup with comparable results. You trade a little representational rigidity for a large reduction in hyperparameter risk. At 6 layers on a laptop, post-norm is fine.

RMSNorm vs LayerNorm: measured, not asserted

Zhang & Sennrich’s headline is 7%–64% faster runtime at comparable quality across RNN, seq2seq translation, image-caption retrieval and CIFAR-10 CNN experiments. The range’s shape matters: the saving is largest where the norm is a big fraction of a small layer’s work, and smallest where matmuls dominate. At 7B+ transformer scale it is a few percent of wall-clock — real, but not why RMSNorm won. It won for being no worse and strictly cheaper. A useful cross-check when someone claims a big norm speedup: divide the norm’s FLOPs by the layer’s total FLOPs. If the norm is 2% of the work, no implementation of it can give you 20%.

SwiGLU: equal-parameter comparison, not equal-width

The comparison that matters is at equal parameters and equal FLOPs, which is why $8d/3$ exists. A SwiGLU at $d_{ff}=4d$ versus a ReLU FFN at $d_{ff}=4d$ is not a comparison of activations — the SwiGLU has 50% more parameters and of course wins. Shazeer’s ablations hold parameters constant; quote those.

GQA: exact, derivable memory numbers

Use the formula, not a rule of thumb: the number of layers and d_head differ per model and dominate the answer. The 70B table in §2.4 gives $0.31$ MB/token for GQA-8 and $2.5$ MB/token for MHA, exactly $64/8=8\times$. Against a 140 GB fp16 weight footprint, a single 32k-context request costs 10 GB of cache — and at batch 16 it costs 163 GB, larger than the weights. That crossover is why GQA is not optional for long-context serving.

FlashAttention: which axis moves

These compose, and they attack different terms in the same bandwidth equation, so their speedups multiply roughly rather than overlap. The paper’s own end-to-end numbers: 15% on BERT-large at seq len 512, 3× on GPT-2 at 1K, 2.4× on long-range arena at 1K–4K, plus 0.7 better perplexity on GPT-2 and 6.4 points on long-document classification from the longer context it made affordable (FlashAttention). The small BERT number is the honest one to remember: at short sequence lengths, where the $n\times n$ matrix is small, FlashAttention buys little.

Version note. FlashAttention-2 (Dao 2023, 2307.08691) improves work partitioning and reports ~2× over v1; FlashAttention-3 (Shah et al. 2024, 2407.08608) targets Hopper specifically with fp8 support. The PyPI package flash-attn is at 2.8.3.post1 and requires Python ≥ 3.9 [as of 2026-10]. The FA3/Blackwell generation moves fast enough that you should treat any specific kernel version claim as stale within months. [UNVERIFIED: exact FA4/Blackwell release status.]

5. Pitfalls

Pre-norm with an unscaled residual init at depth. Symptom: the final pre-head norm is doing a lot of work, training loss is fine but the model needs an unusually low learning rate, and loss spikes appear when you raise it. The residual stream magnitude grows like $\sqrt{L}$; scale each residual branch’s output projection by $1/\sqrt{2L}$ at initialization, as GPT-2 does.

RMSNorm’s $ϵ$ inside the square root but applied after rsqrt. Symptom: NaN or Inf losses that appear only under bf16 autocast and disappear in fp32. Compute the mean-of-squares in fp32 and place $ϵ$ inside the denominator, not as a post-hoc clamp.

Dropping $\gamma$ from RMSNorm “since it’s just a norm”. Symptom: the model trains but plateaus well below the baseline on every task, with no instability to explain it. Without a learned gain the layer cannot re-adjust per-feature scale after normalization.

SwiGLU at the ReLU-FFN hidden size. Symptom: the FFN has 50% more parameters than the baseline, so every comparison drawn from it is invalid. Use $8d/3$ when the claim is “same budget”.

SwiGLU weight-naming mismatch. Symptom: a shape error on w_down, or worse, a silent load of the gate into the value slot and a loss worse than the released model reports. Published checkpoints use gate_proj/up_proj/down_proj (w1/w3/w2); map names explicitly and assert on shapes.

Assuming n_kv_heads is tunable at load time. Symptom: RuntimeError on weight shapes, or enable_gqa doing nothing because the checkpoint is plain MHA. $g$ is baked into the checkpoint; you can uptrain MHA into GQA, but you cannot flip it at serving time.

Materialising the expanded K/V instead of using enable_gqa. Symptom: peak memory is unchanged after moving to a GQA checkpoint, because repeat_interleave allocated the $h$-head tensors you were trying to avoid. Confirm the fused kernel path is actually taken.

Treating FlashAttention as an approximation. Symptom: you avoid it out of fear and pay the memory cost for nothing, or you “A/B it for quality”, see a difference within noise, and conclude it is approximate. The differences are floating-point reassociation. A/B it for throughput and memory, never quality.

Expecting bit-exact reproducibility with a fused kernel. Symptom: the same seed produces a slightly different loss on a different GPU or PyTorch version. Documented behaviour, not a bug; SDPBackend.MATH is the escape hatch at a large throughput cost.

Forgetting SDPA’s inverted boolean-mask convention. Symptom: training loss looks excellent and generation is incoherent, because the “allowed” mask was inverted and the model attended to exactly what it should not see.

6. Using it in a project

This chapter is mostly not actionable when you call an API. It becomes actionable the moment you train from scratch or self-host.

Fine-tuning an open-weight checkpoint. None of these are decisions — pre-norm, RMSNorm, SwiGLU and the GQA head counts are all fixed by the checkpoint. You choose the checkpoint, and the two questions worth asking are: what is my KV cache per token at my target context? (compute it from the config) and does my serving stack have a fused kernel for this architecture? An architecture your framework falls back to the math kernel on will be several times slower than its benchmarks suggest.

Training a small model from scratch. Adopt all five. Pre-norm and RMSNorm are better defaults, SwiGLU at $8d/3$ is free at equal budget, and FlashAttention is a one-line change through SDPA.

Serving long context. GQA/MLA first — largest single memory lever, composes with everything else. Then prefix caching for shared system prompts, then KV quantization.

When NOT to use each of these:

• Pre-norm: when reproducing a specific 2017-era result, or on a shallow model where post-norm’s instability never manifests. Pre-norm is not a strict superset — post-norm’s optimization properties can be preferable at small depth; the case for pre-norm strengthens with depth.
• RMSNorm: when you need re-centering invariance (some retrieval and embedding setups care that the representation is mean-centred), or when a tuned LayerNorm $ϵ$ cannot be re-validated.
• SwiGLU: when the deployment target’s kernels have no gated-FFN path — some edge and NPU runtimes only optimize the 2-matrix FFN, and a gate on the slow path can be slower overall than a plain FFN on the fast path. Also at very small widths, where the third matmul’s launch overhead is not amortized.
• GQA: on short-context, high-batch workloads where the KV cache is not the bottleneck and MQA’s quality cost is not worth paying. Do not pick it if you need per-head KV differentiation for interpretability work.
• FlashAttention: when you need bit-exact reproducibility across hardware, or you use an attention variant (custom learned bias, unusual mask) the fused kernels do not support. The math backend is exact and portable; it is just slow.

7. How you would prove it worked

Every claim above is an ablation with a stated falsifier.

Ablation 1 — pre-norm vs post-norm at depth. Two identical 24-layer models, one each, trained at a fixed learning rate with no warmup. Metric: loss vs step. Baseline: the post-norm run. Prediction: post-norm diverges or spikes; pre-norm trains. Falsified if post-norm trains fine — in which case the model is too shallow or the LR too low for the effect to appear, and you have not tested the claim. A useful extra signal is the per-layer gradient norm at step 0: roughly flat for pre-norm, growing toward the output for post-norm.

Ablation 2 — RMSNorm vs LayerNorm. Swap only the norm layer; hold $d$, depth, data, schedule, seed fixed. Metrics: tokens/sec and loss at fixed step count. Baseline: LayerNorm. Prediction: loss within noise, tokens/sec higher (a few percent at transformer scale, more at small scale). Falsified if a consistent loss gap appears in either direction — then the change was not free.

Ablation 3 — SwiGLU vs ReLU FFN at equal parameters. Build three configs: ReLU at $d_{ff}=4d$; SwiGLU at $d_{ff}=4d$ (more parameters — keep it only as a sanity check that parameters help); SwiGLU at $d_{ff}=8d/3$. Metric: loss at matched parameter count and matched tokens. Falsified as evidence for gating if the $8d/3$ SwiGLU does not beat the ReLU baseline — and if you only ran the $4d$ SwiGLU, you have proven that more parameters help, which nobody doubted.

Ablation 4 — GQA-8 vs MHA on a 7B-class model. Metrics: measured KV bytes/token vs the formula; decode tokens/sec at batch 1, 8, 32; eval accuracy. Baseline: MHA. Derivable prediction: the cache is exactly $h/g=8\times$ smaller, and decode throughput gains most at small batch (bandwidth-bound) and least at large batch (compute-bound). Falsified if accuracy drops consistently beyond your eval’s noise floor across tasks — publish the drop if you see it, since MQA’s known weakness is exactly this.

Ablation 5 — FlashAttention on/off. Metrics: peak allocated memory and tokens/sec at $n\in \{1\text{k},4\text{k},16\text{k}\}$. Baseline: the eager/math path. Prediction: the memory gap widens with $n$; the throughput gap is small at 1k and large at 16k. Falsified if quality differs beyond numerical noise on a fixed seed — which should not happen, and if it does, check the mask convention or scale argument, not the kernel.

Reporting template (example — replace with your own numbers). “We moved a 7B-class decoder from MHA to GQA-8 by uptraining on 5% of the pretraining tokens. Measured KV memory fell from 0.63 MB/token (example) to 0.08 MB/token (example) — 8.0×, matching the formula — and decode throughput at batch 4 rose 1.7× (example). MMLU moved within noise (−0.4, n=5 seeds, example). FlashAttention was on throughout; it affects neither number.”

8. Interview drill

“Why does everyone use pre-norm instead of post-norm?” Strong: post-norm puts LayerNorm on the residual identity path, so $\partial x_{l+1}/\partial x_l$ accumulates $L$ normalization Jacobians; Xiong et al. show the expected gradient near the output grows with depth at initialization, which is why post-norm needs tuned warmup. Pre-norm normalizes only the sublayer’s input, leaving an exact identity term in the Jacobian. Wrong: “pre-norm converges faster” — the controlled result is about stability and removing warmup, not a better final loss at matched budget. Mentioning GPT-2’s $1/\sqrt{2L}$ residual init completes the answer.

“What does RMSNorm drop relative to LayerNorm, and why is that allowed?” Strong: mean subtraction and the bias $\beta$; the RMS rescale and the gain $\gamma$ stay. Zhang & Sennrich hypothesise — with ablation support — that re-centering is dispensable and the rescaling is what stabilizes training. The gain is one reduction instead of two: 7%–64% faster on their workloads. Wrong: “it’s faster and just as good” — true, but says nothing about what was deleted or why that is safe. Also wrong: that RMSNorm removes the need for a gain.

“Derive the $8d/3$ hidden size.” Strong: a GLU FFN has three matrices, so $3d{d}_{ff}$ parameters; the ReLU FFN has two at hidden $4d$, equal to $8d^2$. Setting $3d{d}_{ff}=8d^2$ gives $d_{ff}=8d/3$, rounded to a multiple of 128. Wrong: “it’s 2/3 of 4d because there are three matrices instead of two” — the right ratio as numerology. What matters is that parameter count is held constant, because that makes the comparison fair.

“Why does adding a gate improve the FFN?” Strong: the gate makes the hidden-unit mask input-dependent. A ReLU is a fixed nonlinearity on one projection; a GLU multiplies $\mathrm{Swish}(x{W}_1)$ with $(x{W}_3)$, so how much of each value-branch unit passes is decided by a different learned projection of the same input — more expressive per parameter, and visible in Shazeer’s equal-parameter ablations. Wrong: “it adds a nonlinearity” — the FFN already had one. The gate adds multiplicative, input-conditioned control.

“How does GQA reduce memory, exactly?” Strong: the cache is $2\cdot L\cdot n_{kv}\cdot d_{head}\cdot n_{ctx}\cdot B\cdot b$, and GQA reduces only $n_{kv}$, from $h$ to $g$ — so the saving is exactly $h/g$. Weights and FLOPs are untouched. Follow-up: “why not MQA ($g=1$)?” → the largest saving, but Ainslie et al. report quality degradation; GQA interpolates, and their uptraining recipe (~5% of pretraining compute, mean-pool the KV heads) recovers MHA-level quality. Wrong: “GQA makes attention cheaper” — it changes the memory read during decode, not the FLOPs.

“Is FlashAttention approximate?” Strong: no — exact, up to floating-point reassociation. It removes the materialized $n\times n$ matrix from HBM via tiling and online softmax, and recomputes blocks in the backward pass. It reduces neither the $O(n^2)$ FLOPs nor the KV cache. Wrong: “it trades a bit of quality for speed” — that describes sparse or linear attention. Second wrong answer: “it reduces FLOPs” — it reduces IO.

“Then what limits long-context inference?” Strong: two separate things. Prefill compute is $O(n^2)$ FLOPs, made affordable — not smaller — by FlashAttention; decode is bandwidth-bound and dominated by re-reading the KV cache, addressed by GQA/MLA and KV quantization. Quality at long context is a third, separate problem: fitting the tokens does not mean using them well. Wrong: naming one bottleneck for both phases.

“What is MoE and when is it the wrong choice?” Strong: a Mixture-of-Experts replaces the dense FFN with $N$ expert FFNs and a router sending each token to the top-$k$ (usually 1–2), decoupling total parameters from FLOPs per token. It is wrong when you are memory-bound at small batch: every expert must be resident, so a 671B-parameter model with 37B active parameters still needs memory for 671B, and a single-request workload never amortizes it. Wrong: treating active-parameter count as the memory footprint.

Next: BERT vs GPT vs Encoder-Decoder — the deltas above describe how a modern block is wired; the next chapter is about what objective you train it with, and why one of the three attention patterns took the general-purpose crown.

BERT vs GPT vs Encoder-Decoder

TL;DR — Three attention patterns, three objectives. BERT is bidirectional with masked language modeling: predict the tokens you hid, seeing both sides. GPT is causal with next-token prediction: predict each token from the prefix only, which makes a single loss factorize the whole joint distribution and gives you generation for free. T5/BART stack both and add cross-attention, so the encoder reads the source once and the decoder attends to it while generating. Decoder-only won the general-purpose crown because one objective covers every task expressed as text and scales cleanly; encoder-only still owns embedding, reranking and classification, because bidirectional context produces a better fixed representation per token — and nothing you can generate with.

1. Why this exists

Pretraining is expensive, and the objective decides what the pretrained weights are worth. Three concrete failures motivate the split.

A model trained to generate cannot be asked to read a sentence whole. If you train with a causal mask, position $t$ sees only $\mathrm{1..}t-1$. Take the sentence “The trophy would not fit in the brown suitcase because it was too ___” and ask for the referent of “it” — an encoder that sees the whole sentence can use the following word “suitcase”; a causal model at the word “it” cannot. This is not a quality argument but a masking argument: the architecture forbids the information, so no amount of scale recovers it at that position.

A model trained to fill in blanks cannot be sampled from coherently. Masked language modeling corrupts the input by replacing tokens with a sentinel and predicts the originals independently — each masked position is predicted from the same corrupted context, with no dependency between the predictions. You cannot write down $p(x_1,\dots ,x_n)$ from that, so you cannot autoregressively generate. Bidirectional encoders are very good at scoring and very bad at writing.

A model trained on next-token prediction has no place to put the source. When the input and output are genuinely different sequences drawn from different distributions — source German to target English; a document to its summary — a decoder-only model handles it by concatenating them into one stream and trusting the causal mask to keep them apart. That works, but it spends the decoder’s capacity re-encoding the source. An encoder that reads the source once and cross-attends to it is the structurally cheaper arrangement, and it is why T5 and BART held the seq2seq crown long after GPT-style models took over general-purpose text.

What you are choosing. Not “which model is smarter” but which conditional distribution the loss is a maximum-likelihood fit of, which in turn fixes the set of operations the model can perform well. Generation requires a factorized joint. Scoring requires bidirectional context. Seq2seq wants both, separated.

2. How it works

2.1 Three attention patterns

The architectures differ in exactly one mechanical thing: which positions the attention mask allows, and where the queries come from.

flowchart TB
    subgraph ENC["Encoder-only (BERT)"]
        E1["tokens"] --> E2["bidirectional self-attention × L"]
        E2 --> E3["contextual embeddings"]
    end
    subgraph DEC["Decoder-only (GPT)"]
        D1["prefix tokens"] --> D2["causal self-attention × L"]
        D2 --> D3["next-token distribution"]
    end
    subgraph ED["Encoder-decoder (T5/BART)"]
        C1["source tokens"] --> C2["bidirectional encoder × L"]
        C2 --> C3["encoder states (K, V)"]
        C4["output so far"] --> C5["causal decoder × L"]
        C3 --> C5
        C5 --> C6["next-token distribution"]
    end

Three mask shapes, drawn as $n\times n$ attention patterns (■ = allowed):

$$\text{bidirectional}=\left(\begin{array}{cccc}■& ■& ■& ■\\ ■& ■& ■& ■\\ ■& ■& ■& ■\\ ■& ■& ■& ■\end{array}\right)\text{causal}=\left(\begin{array}{cccc}■& & & \\ ■& ■& & \\ ■& ■& ■& \\ ■& ■& ■& ■\end{array}\right)$$

The third shape is cross-attention: the queries come from the decoder and the keys and values come from the encoder, so the rectangle is $n_{dec}\times n_{enc}$ and it is not triangular — every decoder position may look at the entire source. The mask mechanics are unchanged from Attention; only which axis is masked differs.

2.2 The masked language modeling loss

BERT corrupts a fraction of positions and predicts the originals. Let $M$ be the set of masked positions and $\tilde{x}$ the corrupted sequence:

$${\mathcal{L}}_{\mathrm{MLM}}=-\sum _{i\in M}\log p_{\theta }\left(x_i\mid \tilde{x}\right)$$

• $x=(x_1,\dots ,x_n)$ — the original token sequence, already subword-tokenized.
• $\tilde{x}$ — the corrupted input. BERT’s recipe: choose 15% of positions uniformly; of those, replace 80% with [MASK], 10% with a uniformly random vocabulary token, and leave 10% unchanged. The random and unchanged branches exist because [MASK] never appears in downstream text — a model trained to expect it would face a train/test mismatch. The cost is that some targets are deliberately unlearnable noise, which is accepted.
• $M$ — the set of masked positions, $|M|\approx 0.15n$. The sum runs only over masked positions. This is the load-bearing detail for every efficiency claim about MLM: 85% of the positions contribute no gradient at all. ELECTRA names this directly as the reason MLM is compute-inefficient — its replaced-token-detection task is defined over all input tokens rather than “just the small subset that was masked out”.
• $p_{\theta }(x_i\mid \tilde{x})$ — the model’s distribution over the vocabulary at position $i$, conditioned on the entire corrupted sequence, left and right. This is where bidirectionality enters: nothing in the mask forbids position $i-1$ or $i+1$, so both contribute to the representation used to predict $x_i$.

What this loss actually is. It is a sparse reconstruction objective with an independence assumption baked in: the model maximizes ${\sum }_{i\in M}\log p(x_i\mid \tilde{x})$, which is not a factorization of $\log p(x)$ — no chain rule appears anywhere. Removing the masking uses the true context on both sides, but you cannot sample from what you trained, because you never modelled a joint. Getting this distinction right is the usual first filter in interviews.

Why 15%. Prediction and representation are in tension: mask more and the corrupted context yields worse representations, mask less and too few positions produce gradient. 15% is BERT’s empirical setting from an ablation, not a derivation.

2.3 The causal language modeling loss

GPT factorizes the joint distribution with the chain rule and fits each factor:

$${\mathcal{L}}_{\mathrm{CLM}}=-\sum _{t=1}^n\log p_{\theta }\left(x_t\mid x_1,\dots ,x_{t-1}\right)$$

• $x_t$ — the target token at position $t$.
• $x_{<t}=x_1,\dots ,x_{t-1}$ — the prefix. The causal mask makes this the entire conditioning set; position $t$’s query cannot attend to $t$ or later.
• $p_{\theta }(x_t\mid x_{<t})$ — the output distribution after the final projection to vocabulary size, softmaxed.

Because ${\prod }_{t=1}^{n}p(x_t\mid x_{<t})=p(x_1,\dots ,x_n)$ exactly, this one loss is a proper sequence likelihood. That buys three things MLM cannot have: generation (sample $x_t$, append it, repeat), arbitrary conditioning (prefix anything and the model continues it), and a signal on every position.

The price is that the model never sees the future while representing a position, so an encoder’s representation of token $i$ is strictly more informed. Encoder models pay the opposite price: better representations, no joint.

2.4 Encoder-decoder: span corruption and cross-attention

T5 unifies everything as text-to-text and replaces MLM with span corruption: contiguous spans of the source are replaced by unique sentinel tokens, and the decoder must emit the removed spans. Concretely, with spans (1,2) and (3) removed:

encoder input : Thank you <X> me to your party <Y> week .
decoder target: <X> for inviting <Y> last <Z>

• <X>, <Y>, <Z> — sentinel tokens, one per corrupted span. They are content-addressable markers: the decoder knows where in the original sequence to place what it generates, which is what lets one pass emit several non-adjacent spans.
• The decoder loss is ordinary causal language modeling over the target, with the encoder states injected through cross-attention at every layer.
• Why spans rather than independent positions? Predicting a contiguous span requires modelling within-span dependencies, so the objective is closer to a real sequence likelihood than MLM’s independent-blank-filling, and the target sequence is shorter than the source (the encoder sees the whole input, the decoder emits only the removed spans), which shifts compute to the encoder — the right place, since the encoder is the only fully parallel half.

BART takes a different route to the same architecture — a denoising autoencoder with several noise functions (token deletion, infilling, sentence permutation, document rotation) and a standard next-token decoder loss — but the shape is the same: bidirectional encoder, causal decoder, cross-attention between them.

2.5 Why decoder-only won for general-purpose models

Five reasons, in rough order of weight.

Universal task interface. Every task becomes “continue the text.” Classification is a label in the continuation; extraction is JSON in the continuation; translation is the translation. No task-specific head, no new architecture per task — BERT needs a fresh classification head fine-tuned per task, which means a separate training run and a separate artifact per task.

In-context learning emerged with scale, and only on generative models. GPT-3 demonstrated strong few-shot performance without any gradient updates, with tasks specified purely by text in the prompt. The mechanism is not available to an encoder: an MLM model cannot produce the demonstration completions that the prompt’s structure relies on.