Class LlamaModel

models/llama_7b.py:115–161

Compact LLaMA (160M params) -- fits on any GPU, good for testing AutoKernel. Config: dim=768, n_layers=12, n_heads=12, n_kv_heads=4, hidden_dim=2048

class LlamaModel(nn.Module):
    """
    Compact LLaMA (160M params) -- fits on any GPU, good for testing AutoKernel.

    Config: dim=768, n_layers=12, n_heads=12, n_kv_heads=4, hidden_dim=2048
    """

    def __init__(
        self,
        vocab_size: int = 32000,
        dim: int = 768,
        n_layers: int = 12,
        n_heads: int = 12,
        n_kv_heads: int = 4,
        hidden_dim: int = 2048,
        max_seq_len: int = 2048,
    ):
        super().__init__()
        self.tok_embeddings = nn.Embedding(vocab_size, dim)
        self.layers = nn.ModuleList([
            TransformerBlock(dim, n_heads, n_kv_heads, hidden_dim)
            for _ in range(n_layers)
        ])
        self.norm = RMSNorm(dim)
        self.output = nn.Linear(dim, vocab_size, bias=False)

        # RoPE frequencies
        self.register_buffer(
            "freqs_cis",
            precompute_freqs_cis(dim // n_heads, max_seq_len * 2),
            persistent=False,
        )

        n_params = sum(p.numel() for p in self.parameters())
        print(f"LlamaModel: {n_params / 1e6:.1f}M parameters")

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        B, T = input_ids.shape
        h = self.tok_embeddings(input_ids)
        freqs = self.freqs_cis[:T]

        for layer in self.layers:
            h = layer(h, freqs)

        h = self.norm(h)
        logits = self.output(h)
        return logits


class LlamaModel7B(nn.Module):