Cahier

LLM Inference with vLLM

2026-05-25T00:00:00+00:00

Hi there! Biplob here. This article is about serving an LLM with vLLM on Nvidia DGX.

Check

First check nvidia gpu.

# GPU memory status
nvidia-smi

# Cuda version
nvcc --version

Install UV

Install uv: the project’s preferred package manager (if not already installed):

Run these commands:

# Install uv
curl -LsSf https://astral.sh/uv/install.sh | sh

# Reload shell to pick up uv in PATH
source $HOME/.local/bin/env

Step 1 — Create a uv Virtual Environment:

# Create project folder
mkdir -p ~/vllm-server && cd ~/vllm-server

# Create venv with Python 3.11 (stable for vLLM)
uv venv .venv --python 3.11

# Activate it
source .venv/bin/activate

# Confirm Python version
python --version

Step 2 — Install PyTorch for CUDA 13.0 / Blackwell

The GB10 is a Blackwell GPU. Use the nightly PyTorch build which has Blackwell (sm_100) support:

# Install PyTorch nightly with CUDA 12.8 wheels
# (CUDA 12.8 wheels are the latest stable; CUDA 13.0 is forward-compatible)
uv pip install --pre torch torchvision torchaudio \
  --index-url https://download.pytorch.org/whl/nightly/cu128

# Verify PyTorch sees your GB10
python -c "import torch; print('PyTorch:', torch.__version__); print('CUDA available:', torch.cuda.is_available()); print('GPU:', torch.cuda.get_device_name(0))"

Step 3 — Install vLLM

# Install latest vLLM (0.9.x+ has Blackwell/GB10 support)
uv pip install vllm

# Verify
python -c "import vllm; print('vLLM version:', vllm.__version__)"

Step 4 — Download Qwen3-14B

# Install huggingface_hub
uv pip install huggingface_hub[cli]

# Download model (AWQ quantized — fits better on single GB10)
hf download Qwen/Qwen3-14B-AWQ \
  --local-dir ~/models/Qwen3-14B-AWQ

# --- OR --- download full BF16 if your GB10 has enough VRAM
hf download Qwen/Qwen3-14B \
  --local-dir ~/models/Qwen3-14B

Step 5 — Create a starting script

#!/bin/bash
source ~/projects/vllm-server/.venv/bin/activate

nohup vllm serve ~/models/Qwen3-14B \
  --host 0.0.0.0 \
  --port 8000 \
  --dtype bfloat16 \
  --max-model-len 131072 \
  --gpu-memory-utilization 0.90 \
  --max-num-seqs 4 \
  --enable-prefix-caching \
  --served-model-name qwen3-14b \
  > ~/projects/vllm-server/vllm.log 2>&1 &

echo $! > ~/projects/vllm-server/vllm.pid
echo "vLLM started with PID: $(cat ~/projects/vllm-server/vllm.pid)"

Meaning of the Parameters:

Parameter	What it means
`--host 0.0.0.0`	Listen on all network interfaces (not just localhost). Use `127.0.0.1` to restrict to local only
`--port 8000`	Port the HTTP server listens on
`--dtype bfloat16`	Weight precision. BF16 is ideal for Blackwell/Ampere — full quality, half the memory of FP32
`--max-model-len`	Max context window in tokens (input + output combined)
`--gpu-memory-utilization`	Fraction of memory vLLM can use for KV cache (0.0–1.0). `0.90` leaves headroom for OS
`--max-num-seqs`	Max concurrent requests being processed. Higher = more throughput, more memory
`--enable-prefix-caching`	Reuses KV cache for repeated prompt prefixes (great for system prompts). Speeds up repeated calls
`--served-model-name`	The model name clients use in API calls. Can be anything — doesn’t have to match folder name

Step 6 - Start, Monitor and Check

tail -f ~/projects/vllm-server/vllm.log
# You'll know it's ready when you see:
# INFO:     Application startup complete.
# INFO:     Uvicorn running on http://0.0.0.0:8000

chmod +x ~/projects/vllm-server/start.sh
~/projects/vllm-server/start.sh

# Model loading takes 1–3 minutes. Watch the log:
tail -f ~/projects/vllm-server/vllm.log

# 1. Health check
curl http://localhost:8000/health

# 2. List loaded models
curl http://localhost:8000/v1/models

# 3. Test inference
curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "qwen3-14b",
    "messages": [
      {"role": "system", "content": "You are a helpful assistant."},
      {"role": "user", "content": "Explain unified memory in one paragraph."}
    ],
    "max_tokens": 300,
    "temperature": 0.7
  }'

Step 7 - Stop the Service

# Graceful stop using saved PID
kill $(cat ~/projects/vllm-server/vllm.pid)

# Confirm it's stopped
sleep 2 && nvidia-smi

# Clean up PID file
rm ~/projects/vllm-server/vllm.pid

# If the PID file is lost:
# Find and kill by process name
pkill -f "vllm serve"

# Check if It's Already Running
hpgrep -a -f "vllm serve"

To watch nvidia gpu usage on another terminal:

watch -n 1 -d nvidia-smi

| Flag | Meaning | |—|—| | -n 1 | Refresh every 1 second | | -d | Highlight differences between refreshes |

To exit: Press Ctrl + C.

Good practice to use .env file. Here’s how to set it up:

`.env` file

For Ollama:

LLM_API_BASE=http://localhost:11434
LLM_MODEL=ollama/qwen3:14b

For vLLM:

LLM_API_BASE=http://localhost:8000/v1
LLM_MODEL=openai/qwen3-14b

Python code (unchanged regardless of which backend)

import os
from dotenv import load_dotenv

load_dotenv()

model    = os.getenv("LLM_MODEL")
api_base = os.getenv("LLM_API_BASE")

completion = await litellm.acompletion(
    model=model,
    messages=messages,
    temperature=0.7,
    api_base=api_base,
    api_key="dummy",
)

Key difference in the URL

Backend	`LLM_API_BASE`	Why
Ollama	`http://localhost:11434`	Ollama uses its own format; LiteLLM’s `ollama/` prefix handles the `/api/chat` path internally
vLLM	`http://localhost:8000/v1`	vLLM is OpenAI-compatible; the `/v1` suffix is required for the `openai/` prefix to route correctly

So to switch backends you only ever edit the .env file — your Python code stays the same.

Training LLM with LlamaFactory on Nvidia DGX

2026-05-20T00:00:00+00:00

Hi there! Biplob here. This article is about training a small LLM on Nvidia DGX with LlamaFactory.

Installation

First clone the Llamafactory git repo.

git clone --depth 1 https://github.com/hiyouga/LlamaFactory.git
cd LlamaFactory

Install UV

If uv isn’t installed yet. Here’s the plan:

Install uv (the project’s preferred package manager)
Use uv to install LlamaFactory in an isolated environment

Run these commands:

# Install uv
curl -LsSf https://astral.sh/uv/install.sh | sh

# Reload shell to pick up uv in PATH
source $HOME/.local/bin/env

# Sync the project (creates a .venv and installs all deps)
cd ~/projects/LlamaFactory
uv sync

# Optionally install metrics extras
uv pip install -r requirements/metrics.txt

After that, you can run LlamaFactory commands with:

uv run llamafactory-cli webui
# or
# For training
uv run llamafactory-cli train examples/train_lora/qwen3_lora_sft.yaml

# For inference/chat
uv run llamafactory-cli chat examples/inference/qwen3_lora_sft.yaml

# For exporting model
uv run llamafactory-cli export examples/merge_lora/qwen3_lora_sft.yaml

To watch nvidia gpu usage on another terminal:

watch -n 0.5 nvidia-smi

-n 0.5: Refreshes the window every 0.5 seconds (twice a second).

To exit: Press Ctrl + C.

How Did I Do This?

2020-06-10T00:00:00+00:00

Hi there! In case you are wondering how to make a github page like this, it’s relatively easy! Just introduce yourself with Lanyon which is a Jekyll theme. You can also modify it according to this and that post. A little understanding of markdown and github pages should be sufficient to accomplish all these.

Hope this helps, Thanks!

GPT-style decoder-only Transformer

2020-06-08T00:00:00+00:00

Howdy! Biplob Here! Back with a fully annotated, from-scratch GPT-style decoder-only Transformer.

transformer_lm.py

Every tensor operation is commented with its shape transformation so you can follow the data through the entire forward pass.

Notation used throughout:

B = batch size
T = sequence length (number of tokens)
D = d_model (hidden / embedding dimension)
H = n_heads (number of attention heads)
Dh = D // H (per-head dimension, also called head_dim)
V = vocab_size
F = d_ffn (feed-forward inner dimension, typically 4 * D)

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


# ─────────────────────────────────────────────────────────────────────────────
# 1.  TOKEN + POSITIONAL EMBEDDINGS
# ─────────────────────────────────────────────────────────────────────────────

class Embeddings(nn.Module):
    """
    Converts a sequence of integer token IDs into a sequence of continuous
    vectors by summing a learned token embedding with a learned positional
    embedding.

    Token embedding:    maps each vocabulary index to a D-dimensional vector.
    Positional embedding: maps each sequence position (0 … T-1) to a
                          D-dimensional vector so the model knows *where*
                          each token sits in the sequence.
    """

    def __init__(self, vocab_size: int, d_model: int, seq_len: int, dropout: float):
        super().__init__()

        # Lookup table: V rows, each D-wide.
        # Indexing with a token id returns its embedding vector.
        self.token_emb = nn.Embedding(vocab_size, d_model)   # shape: (V, D)

        # Lookup table: seq_len rows, each D-wide.
        # Indexing with position index 0…T-1 returns positional embedding.
        self.pos_emb   = nn.Embedding(seq_len,    d_model)   # shape: (T_max, D)

        self.dropout   = nn.Dropout(dropout)
        self.d_model   = d_model

    def forward(self, idx: torch.Tensor) -> torch.Tensor:
        """
        idx  : (B, T)  — integer token IDs
        return (B, T, D)
        """
        B, T = idx.shape

        # ── token embeddings ─────────────────────────────────────────────────
        # Each integer in idx is replaced by its D-dimensional row from the
        # token embedding table.
        #   idx          : (B, T)
        #   tok_emb      : (B, T, D)
        tok_emb = self.token_emb(idx)                        # (B, T, D)

        # ── positional embeddings ─────────────────────────────────────────────
        # Build position indices [0, 1, 2, …, T-1] and look them up.
        #   positions    : (T,)   → broadcast to (1, T) → look up → (1, T, D)
        positions = torch.arange(T, device=idx.device)      # (T,)
        pos_emb   = self.pos_emb(positions)                  # (T, D) → broadcast (1,T,D)

        # ── sum and scale ─────────────────────────────────────────────────────
        # Adding positional to token embedding is the standard GPT approach.
        # Scaling by sqrt(D) keeps the magnitude stable (from "Attention is
        # All You Need"; sometimes omitted in GPT implementations).
        x = tok_emb + pos_emb                                # (B, T, D)  broadcast add
        x = self.dropout(x)                                  # (B, T, D)
        return x


# ─────────────────────────────────────────────────────────────────────────────
# 2.  MULTI-HEAD CAUSAL SELF-ATTENTION
# ─────────────────────────────────────────────────────────────────────────────

class CausalSelfAttention(nn.Module):
    """
    Scaled dot-product self-attention with a causal (autoregressive) mask so
    that position t can only attend to positions 0 … t.

    The QKV projections are fused into one linear layer for efficiency:

        [Q, K, V] = X @ W_qkv      (one big matmul instead of three)

    Then the output projection merges the heads back into D dimensions.

    Head splitting:
        D  total dimensions are split into H heads of Dh = D/H each.
        Each head learns to attend to a different subspace of the embedding.
    """

    def __init__(self, d_model: int, n_heads: int, seq_len: int, dropout: float):
        super().__init__()
        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"

        self.d_model  = d_model          # D
        self.n_heads  = n_heads          # H
        self.head_dim = d_model // n_heads  # Dh = D / H

        # ── fused QKV projection ──────────────────────────────────────────────
        # Input  (B, T, D)  × weight (D, 3D)  →  output (B, T, 3D)
        # The factor-of-3 gives us Q, K, and V concatenated in one go.
        self.qkv_proj = nn.Linear(d_model, 3 * d_model, bias=False)

        # ── output projection ─────────────────────────────────────────────────
        # After concatenating all heads: (B, T, D) → (B, T, D)
        self.out_proj = nn.Linear(d_model, d_model, bias=False)

        self.attn_dropout = nn.Dropout(dropout)
        self.resid_dropout = nn.Dropout(dropout)

        # ── causal mask ───────────────────────────────────────────────────────
        # Upper-triangular matrix of -inf so softmax zeroes out future tokens.
        # Registered as a buffer so it moves to the right device automatically
        # but is NOT a learned parameter.
        #   mask shape: (1, 1, T_max, T_max)  — broadcastable over (B, H)
        causal_mask = torch.triu(
            torch.full((seq_len, seq_len), float('-inf')),
            diagonal=1
        )                                                    # (T_max, T_max)
        self.register_buffer(
            'causal_mask',
            causal_mask.unsqueeze(0).unsqueeze(0)            # (1, 1, T_max, T_max)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        x      : (B, T, D)
        return   (B, T, D)
        """
        B, T, D = x.shape

        # ── step 1: project input to Q, K, V ─────────────────────────────────
        #   x        : (B, T, D)
        #   qkv_proj : weight (D, 3D)
        #   qkv      : (B, T, 3D)
        qkv = self.qkv_proj(x)                               # (B, T, 3D)

        # Split along the last axis into three equal chunks.
        #   q, k, v  : each (B, T, D)
        q, k, v = qkv.split(self.d_model, dim=-1)            # 3 × (B, T, D)

        # ── step 2: split into heads ──────────────────────────────────────────
        # Reshape from (B, T, D) to (B, T, H, Dh) then transpose to
        # (B, H, T, Dh) so the head axis is second — each head sees
        # a full sequence of Dh-dimensional vectors.
        #
        #   view(B, T, H, Dh)    : unstack D into H heads of Dh each
        #   transpose(1, 2)      : swap T and H axes
        #   → (B, H, T, Dh)
        def split_heads(t):
            return (t.view(B, T, self.n_heads, self.head_dim)  # (B, T, H, Dh)
                     .transpose(1, 2))                         # (B, H, T, Dh)

        q = split_heads(q)   # (B, H, T, Dh)
        k = split_heads(k)   # (B, H, T, Dh)
        v = split_heads(v)   # (B, H, T, Dh)

        # ── step 3: scaled dot-product attention ──────────────────────────────
        # Attention scores: how much should each position attend to each other?
        #
        #   q          : (B, H, T, Dh)
        #   k.transpose: (B, H, Dh, T)   ← swap last two dims
        #   scores     : (B, H, T, T)    ← (T queries) × (T keys)
        #
        # Dividing by sqrt(Dh) prevents the dot products from growing large
        # in magnitude (which would push softmax into near-zero gradient regions).
        scale  = math.sqrt(self.head_dim)
        scores = torch.matmul(q, k.transpose(-2, -1)) / scale  # (B, H, T, T)

        # ── step 4: apply causal mask ─────────────────────────────────────────
        # Add -inf to all future positions so softmax assigns them zero weight.
        # Slice the pre-built mask to the current sequence length T
        # (at inference time T can be shorter than the training seq_len).
        #   mask slice : (1, 1, T, T)  broadcasts over (B, H, T, T)
        scores = scores + self.causal_mask[:, :, :T, :T]      # (B, H, T, T)

        # ── step 5: softmax → attention weights ──────────────────────────────
        #   weights : (B, H, T, T)   — each row sums to 1 (over key positions)
        weights = F.softmax(scores, dim=-1)                    # (B, H, T, T)
        weights = self.attn_dropout(weights)

        # ── step 6: weighted sum of values ───────────────────────────────────
        #   weights : (B, H, T, T)
        #   v       : (B, H, T, Dh)
        #   attn_out: (B, H, T, Dh)   — each position is now a mixture of values
        attn_out = torch.matmul(weights, v)                    # (B, H, T, Dh)

        # ── step 7: merge heads back ──────────────────────────────────────────
        # Reverse the split_heads operation:
        #   transpose(1,2) : (B, T, H, Dh)
        #   contiguous()   : make memory layout contiguous before reshape
        #   view(B, T, D)  : flatten H * Dh back into D
        attn_out = (attn_out.transpose(1, 2)                   # (B, T, H, Dh)
                             .contiguous()
                             .view(B, T, D))                   # (B, T, D)

        # ── step 8: output projection ─────────────────────────────────────────
        # Linear mix of all head outputs: (B, T, D) → (B, T, D)
        out = self.out_proj(attn_out)                          # (B, T, D)
        out = self.resid_dropout(out)
        return out                                             # (B, T, D)


# ─────────────────────────────────────────────────────────────────────────────
# 3.  FEED-FORWARD NETWORK  (position-wise MLP)
# ─────────────────────────────────────────────────────────────────────────────

class FeedForward(nn.Module):
    """
    Two-layer MLP applied independently to each token position.

    Architecture:
        x  →  Linear(D → 4D)  →  GeLU  →  Linear(4D → D)  →  Dropout

    The 4× expansion gives the model extra capacity to learn non-linear
    feature interactions. The bottleneck structure means each token's
    representation is projected up, transformed, then projected back down.
    """

    def __init__(self, d_model: int, dropout: float):
        super().__init__()
        d_ffn = 4 * d_model              # inner dimension  (commonly 4 × D)

        # ── first linear: expand D → 4D ──────────────────────────────────────
        # Applied to each token independently: (B, T, D) → (B, T, 4D)
        self.fc1     = nn.Linear(d_model, d_ffn)

        # ── second linear: project 4D → D ────────────────────────────────────
        # Brings representation back to model dimension: (B, T, 4D) → (B, T, D)
        self.fc2     = nn.Linear(d_ffn, d_model)

        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        x      : (B, T, D)
        return   (B, T, D)
        """
        # fc1 + GeLU: (B, T, D) → (B, T, 4D)
        x = F.gelu(self.fc1(x))          # (B, T, 4D)

        # fc2 + dropout: (B, T, 4D) → (B, T, D)
        x = self.dropout(self.fc2(x))    # (B, T, D)
        return x


# ─────────────────────────────────────────────────────────────────────────────
# 4.  TRANSFORMER DECODER BLOCK
# ─────────────────────────────────────────────────────────────────────────────

class TransformerBlock(nn.Module):
    """
    One decoder block: pre-norm self-attention + pre-norm feed-forward.

    Pre-LayerNorm (as used in GPT-2 / GPT-3):
        x  →  LN  →  Attention  →  + residual  →  LN  →  FFN  →  + residual

    The residual connections are critical — they let gradients flow directly
    from the loss to every layer without vanishing, and they mean that
    each block only needs to learn a *correction* to the existing representation.

    Pre-LN (normalise before the sub-layer) is preferred over Post-LN
    (normalise after) because it produces more stable gradients during
    training of very deep networks.
    """

    def __init__(self, d_model: int, n_heads: int, seq_len: int, dropout: float):
        super().__init__()

        # Layer norms applied BEFORE each sub-layer (pre-LN convention)
        self.ln1  = nn.LayerNorm(d_model)   # normalises (B, T, D) → (B, T, D)
        self.ln2  = nn.LayerNorm(d_model)

        self.attn = CausalSelfAttention(d_model, n_heads, seq_len, dropout)
        self.ffn  = FeedForward(d_model, dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        x      : (B, T, D)
        return   (B, T, D)
        """
        # ── sub-layer 1: masked self-attention with residual ──────────────────
        #   ln1(x)      : (B, T, D)   normalise across D for each token
        #   attn(...)   : (B, T, D)   attend across T positions
        #   x + ...     : (B, T, D)   residual skip connection
        x = x + self.attn(self.ln1(x))   # (B, T, D)

        # ── sub-layer 2: position-wise FFN with residual ──────────────────────
        #   ln2(x)      : (B, T, D)
        #   ffn(...)    : (B, T, D)
        #   x + ...     : (B, T, D)
        x = x + self.ffn(self.ln2(x))    # (B, T, D)

        return x                         # (B, T, D)


# ─────────────────────────────────────────────────────────────────────────────
# 5.  FULL LANGUAGE MODEL
# ─────────────────────────────────────────────────────────────────────────────

class TransformerLM(nn.Module):
    """
    GPT-style decoder-only language model.

    Full forward pass pipeline:
        token IDs (B, T)
             ↓
        Embeddings: token + position  →  (B, T, D)
             ↓
        N × TransformerBlock          →  (B, T, D)   [each block: LN+Attn+res, LN+FFN+res]
             ↓
        Final LayerNorm               →  (B, T, D)
             ↓
        LM head Linear(D → V)         →  (B, T, V)   [logits over vocabulary]

    Weight tying:
        The LM head weight is shared with the token embedding matrix.
        This is standard practice in language models (Press & Wolf 2017).
        It saves V × D parameters and often improves perplexity because
        the model uses the same geometric space for both input and output.
    """

    def __init__(
        self,
        vocab_size : int   = 50257,
        d_model    : int   = 768,
        n_heads    : int   = 12,
        n_layers   : int   = 12,
        seq_len    : int   = 1024,
        dropout    : float = 0.1,
    ):
        super().__init__()

        self.seq_len  = seq_len
        self.d_model  = d_model

        # ── embedding layer ───────────────────────────────────────────────────
        # Maps (B, T) integer ids to (B, T, D) float vectors.
        self.embeddings = Embeddings(vocab_size, d_model, seq_len, dropout)

        # ── N transformer blocks ──────────────────────────────────────────────
        # nn.ModuleList ensures PyTorch tracks all parameters.
        self.blocks = nn.ModuleList([
            TransformerBlock(d_model, n_heads, seq_len, dropout)
            for _ in range(n_layers)
        ])

        # ── final layer norm ──────────────────────────────────────────────────
        # Applied after all blocks before the language model head.
        # Stabilises the magnitude of hidden states before the final projection.
        self.ln_final = nn.LayerNorm(d_model)

        # ── language model head ───────────────────────────────────────────────
        # Projects each token's D-dimensional state to a distribution over V.
        #   input  (B, T, D)  →  output (B, T, V)
        # bias=False is conventional when using weight tying.
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)

        # ── weight tying ──────────────────────────────────────────────────────
        # Share the lm_head weights with the token embedding weights.
        # After this, lm_head.weight IS embeddings.token_emb.weight —
        # they are the same tensor in memory.
        self.lm_head.weight = self.embeddings.token_emb.weight

        # ── weight initialisation ─────────────────────────────────────────────
        self._init_weights()

    # ──────────────────────────────────────────────────────────────────────────

    def _init_weights(self):
        """
        Standard GPT-2 initialisation:
          - Linear and Embedding layers: N(0, 0.02)
          - LayerNorm: weight=1, bias=0
          - Residual projections scaled down by 1/sqrt(2 * n_layers)
            to prevent the residual stream from growing in magnitude
            with depth.
        """
        n_layers = len(self.blocks)

        for name, module in self.named_modules():
            if isinstance(module, nn.Linear):
                nn.init.normal_(module.weight, mean=0.0, std=0.02)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)

            elif isinstance(module, nn.Embedding):
                nn.init.normal_(module.weight, mean=0.0, std=0.02)

            elif isinstance(module, nn.LayerNorm):
                nn.init.ones_(module.weight)
                nn.init.zeros_(module.bias)

        # Scale output projections of attention and FFN blocks
        # (the layers that feed directly into residual connections)
        for block in self.blocks:
            nn.init.normal_(
                block.attn.out_proj.weight,
                mean=0.0,
                std=0.02 / math.sqrt(2 * n_layers)
            )
            nn.init.normal_(
                block.ffn.fc2.weight,
                mean=0.0,
                std=0.02 / math.sqrt(2 * n_layers)
            )

    # ──────────────────────────────────────────────────────────────────────────

    def forward(
        self,
        idx     : torch.Tensor,             # (B, T)  integer token IDs
        targets : torch.Tensor | None = None  # (B, T)  integer token IDs (optional)
    ):
        """
        Forward pass.

        idx     : (B, T)  — input token indices
        targets : (B, T)  — next-token targets for computing loss (optional)

        Returns:
            logits : (B, T, V)
            loss   : scalar (cross-entropy) if targets provided, else None
        """
        B, T = idx.shape
        assert T <= self.seq_len, (
            f"Sequence length {T} exceeds maximum {self.seq_len}"
        )

        # ── 1. embeddings ─────────────────────────────────────────────────────
        #   idx : (B, T)
        #   x   : (B, T, D)
        x = self.embeddings(idx)              # (B, T, D)

        # ── 2. transformer blocks ─────────────────────────────────────────────
        # Each block: (B, T, D) → (B, T, D)
        # Residual connections inside each block mean x is incrementally refined.
        for block in self.blocks:
            x = block(x)                      # (B, T, D)

        # ── 3. final layer norm ───────────────────────────────────────────────
        #   (B, T, D) → (B, T, D)
        x = self.ln_final(x)                  # (B, T, D)

        # ── 4. language model head ────────────────────────────────────────────
        # Project each token's D-dimensional state to vocabulary logits.
        #   x      : (B, T, D)
        #   weight : (V, D)   (shared with token embedding, transposed in Linear)
        #   logits : (B, T, V)
        logits = self.lm_head(x)              # (B, T, V)

        # ── 5. loss (training only) ───────────────────────────────────────────
        # Cross-entropy between predicted logit distribution and target tokens.
        # We predict token t+1 from position t, so targets are shifted by 1
        # (that shift is handled in the Dataset, not here — targets[:,i] is
        # the token that should follow input[:,i]).
        loss = None
        if targets is not None:
            # Flatten (B, T, V) → (B*T, V) and (B, T) → (B*T,) for F.cross_entropy
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)),   # (B*T, V)
                targets.view(-1),                   # (B*T,)
            )

        return logits, loss

    # ──────────────────────────────────────────────────────────────────────────

    @torch.no_grad()
    def generate(
        self,
        idx         : torch.Tensor,   # (B, T)  prompt token IDs
        max_new_tokens : int = 100,
        temperature : float = 1.0,    # >1 = more random, <1 = more focused
        top_k       : int   = 50,     # keep only top-k logits before sampling
    ) -> torch.Tensor:
        """
        Autoregressive text generation.

        At each step:
          1. Run the model on the current context
          2. Take the logits at the LAST position only (the next-token prediction)
          3. Apply temperature scaling and top-k filtering
          4. Sample the next token
          5. Append it to the sequence and repeat

        Unlike training (which processes all T positions in parallel via
        teacher forcing), inference must be sequential because each new token
        depends on the previous one.
        """
        self.eval()

        for _ in range(max_new_tokens):

            # Crop context to seq_len if it has grown too long
            idx_cond = idx if idx.size(1) <= self.seq_len else idx[:, -self.seq_len:]

            # Forward pass — we only care about the last position's logits
            #   logits : (B, T_cond, V)
            logits, _ = self(idx_cond)

            # Take logits at the final position: (B, V)
            logits = logits[:, -1, :] / temperature   # (B, V)

            # Top-k filtering: zero out all but the top-k logits
            if top_k is not None:
                # values shape: (B, k); indices not needed here
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                # Set everything below the k-th value to -inf
                logits[logits < v[:, [-1]]] = float('-inf')

            # Softmax → probability distribution over vocabulary
            probs = F.softmax(logits, dim=-1)          # (B, V)

            # Sample one token from the distribution
            idx_next = torch.multinomial(probs, num_samples=1)  # (B, 1)

            # Append to the running sequence
            idx = torch.cat([idx, idx_next], dim=1)    # (B, T+1)

        return idx   # (B, T + max_new_tokens)


# ─────────────────────────────────────────────────────────────────────────────
# 6.  PARAMETER COUNT HELPER
# ─────────────────────────────────────────────────────────────────────────────

def count_parameters(model: nn.Module) -> dict:
    """
    Returns a breakdown of parameter counts by component.
    """
    total = sum(p.numel() for p in model.parameters())
    breakdown = {
        name: sum(p.numel() for p in module.parameters())
        for name, module in model.named_children()
    }
    return {"total": total, **breakdown}


# ─────────────────────────────────────────────────────────────────────────────
# 7.  QUICK SANITY CHECK
# ─────────────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    torch.manual_seed(42)

    # Instantiate GPT-2 small configuration
    model = TransformerLM(
        vocab_size = 50257,
        d_model    = 768,
        n_heads    = 12,
        n_layers   = 12,
        seq_len    = 1024,
        dropout    = 0.1,
    )

    # Dummy batch: 2 sequences of length 16
    B, T = 2, 16
    idx     = torch.randint(0, 50257, (B, T))   # (2, 16)
    targets = torch.randint(0, 50257, (B, T))   # (2, 16)

    logits, loss = model(idx, targets)

    print("── TransformerLM sanity check ──────────────────────────")
    print(f"  Input shape    : {idx.shape}          (B, T)")
    print(f"  Logits shape   : {logits.shape}   (B, T, V)")
    print(f"  Loss           : {loss.item():.4f}")
    print()

    params = count_parameters(model)
    print("── Parameter counts ────────────────────────────────────")
    for k, v in params.items():
        print(f"  {k:20s}: {v:>12,}")