Update modeling_steerling.py

modeling_steerling.py

@@ -0,0 +1,1574 @@
+from __future__ import annotations
+# Auto-generated by scripts/build_hf_files_v3.py — do not edit manually.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import logging
+import os
+from functools import partial
+from typing import TYPE_CHECKING, Any
+from dataclasses import dataclass
+from torch import Tensor
+import math
+import warnings
+
+
+# ======================================================================
+# steerling/models/layers/primitives.py
+# ======================================================================
+
+class RMSNorm(nn.Module):
+    """
+    Root Mean Square Layer Normalization.
+    """
+
+    def __init__(self, config, size: int | None=None):
+        super().__init__()
+        self.eps = getattr(config, 'norm_eps', 1e-05)
+        norm_size = size if size is not None else config.n_embd
+        self.weight = nn.Parameter(torch.ones(norm_size))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        og = x.dtype
+        x = x.float()
+        var = x.pow(2).mean(-1, keepdim=True)
+        x = x * torch.rsqrt(var + self.eps)
+        return (self.weight * x).to(og)
+
+class BufferCache:
+    """Simple cache for storing tensors (used by RotaryEmbedding)."""
+
+    def __init__(self):
+        self._cache: dict[str, torch.Tensor] = {}
+
+    def get(self, key: str) -> torch.Tensor | None:
+        return self._cache.get(key)
+
+    def __setitem__(self, key: str, value: torch.Tensor):
+        self._cache[key] = value
+
+    def __getitem__(self, key: str) -> torch.Tensor:
+        return self._cache[key]
+
+class RotaryEmbedding(nn.Module):
+    """
+    Rotary Position Embeddings (RoPE).
+
+    Applies rotary embeddings to queries and keys for position information.
+
+    Args:
+        dim: Dimension of the rotary embeddings (typically head_dim)
+        max_seq_len: Maximum sequence length to cache
+        base: Base for inverse frequency computation (theta)
+        rope_full_precision: Whether to compute RoPE in full precision
+    """
+
+    def __init__(self, dim: int, max_seq_len: int=2048, base: float=10000.0, rope_full_precision: bool=True):
+        super().__init__()
+        self.dim = dim
+        self.max_seq_len = max_seq_len
+        self.rope_theta = base
+        self.rope_full_precision = rope_full_precision
+        self.__cache = BufferCache()
+        self.get_rotary_embedding(max_seq_len, torch.device('cpu'))
+
+    def get_rotary_embedding(self, seq_len: int, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
+        """Get or compute rotary embeddings for given sequence length."""
+        pos_sin = self.__cache.get('rope_pos_sin')
+        pos_cos = self.__cache.get('rope_pos_cos')
+        if pos_sin is not None and pos_cos is not None and (pos_sin.shape[-2] >= seq_len) and (pos_cos.shape[-2] >= seq_len):
+            if pos_sin.device != device:
+                pos_sin = pos_sin.to(device)
+                self.__cache['rope_pos_sin'] = pos_sin
+            if pos_cos.device != device:
+                pos_cos = pos_cos.to(device)
+                self.__cache['rope_pos_cos'] = pos_cos
+            return (pos_sin[:, :, :seq_len, :], pos_cos[:, :, :seq_len, :])
+        with torch.autocast(device.type, enabled=False):
+            inv_freq = 1.0 / self.rope_theta ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float) / self.dim)
+            seq = torch.arange(seq_len, device=device, dtype=torch.float)
+            freqs = torch.outer(seq, inv_freq)
+            positions = torch.cat((freqs, freqs), dim=-1)
+            pos_sin = positions.sin()[None, None, :, :]
+            pos_cos = positions.cos()[None, None, :, :]
+        self.__cache['rope_pos_sin'] = pos_sin
+        self.__cache['rope_pos_cos'] = pos_cos
+        return (pos_sin, pos_cos)
+
+    def rotate_half(self, x: torch.Tensor) -> torch.Tensor:
+        """Rotate half the hidden dims of the input."""
+        B, nh, T, hs = x.size()
+        x = x.view(B, nh, T, 2, hs // 2)
+        x1, x2 = x.unbind(dim=-2)
+        return torch.cat((-x2, x1), dim=-1)
+
+    def apply_rotary_pos_emb(self, pos_sin: torch.Tensor, pos_cos: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+        """Apply rotary position embeddings to input tensor."""
+        return (t * pos_cos + self.rotate_half(t) * pos_sin).to(t.dtype)
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Apply rotary embeddings to queries and keys."""
+        if self.rope_full_precision:
+            q_, k_ = (q.float(), k.float())
+        else:
+            q_, k_ = (q, k)
+        with torch.autocast(q.device.type, enabled=False):
+            query_len, key_len = (q_.shape[-2], k_.shape[-2])
+            pos_sin, pos_cos = self.get_rotary_embedding(key_len, q_.device)
+            pos_sin = pos_sin.type_as(q_)
+            pos_cos = pos_cos.type_as(q_)
+            q_ = self.apply_rotary_pos_emb(pos_sin[:, :, key_len - query_len:key_len, :], pos_cos[:, :, key_len - query_len:key_len, :], q_)
+            k_ = self.apply_rotary_pos_emb(pos_sin, pos_cos, k_)
+        return (q_.type_as(q), k_.type_as(k))
+
+class MLP(nn.Module):
+    """
+    Multi-Layer Perceptron with SwiGLU or standard activation.
+
+    Args:
+        config: Model config with n_embd, mlp_ratio, use_bias, mlp_type, activation
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        if hasattr(config, 'intermediate_size') and config.intermediate_size is not None:
+            intermediate_size = config.intermediate_size
+        else:
+            intermediate_size = getattr(config, 'mlp_ratio', 4) * config.n_embd
+        use_bias = config.use_bias
+        mlp_type = config.mlp_type
+        if mlp_type == 'swiglu':
+            self.c_fc = nn.Linear(config.n_embd, 2 * intermediate_size, bias=use_bias)
+            self.c_proj = nn.Linear(intermediate_size, config.n_embd, bias=use_bias)
+            self.activation = None
+        else:
+            self.c_fc = nn.Linear(config.n_embd, intermediate_size, bias=use_bias)
+            self.c_proj = nn.Linear(intermediate_size, config.n_embd, bias=use_bias)
+            act_map = {'gelu': nn.GELU(approximate='tanh'), 'relu': nn.ReLU(), 'silu': nn.SiLU()}
+            self.activation = act_map[config.activation]
+        self.c_proj.SCALE_INIT = 1
+        self.config = config
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        mlp_type = getattr(self.config, 'mlp_type', 'swiglu')
+        if mlp_type == 'swiglu':
+            gate_up = self.c_fc(x)
+            up, gate = gate_up.chunk(2, dim=-1)
+            intermediate = F.silu(gate) * up
+        else:
+            intermediate = self.c_fc(x)
+            intermediate = self.activation(intermediate)
+        return self.c_proj(intermediate)
+
+# ======================================================================
+# steerling/models/layers/causal_diffusion_layers.py
+# ======================================================================
+
+logger = logging.getLogger(__name__)
+try:
+    from torch.nn.attention.flex_attention import BlockMask, _dense_to_ordered, flex_attention
+    _FLEX_ATTN_AVAILABLE = True
+except ImportError:
+    _FLEX_ATTN_AVAILABLE = False
+    BlockMask: Any = None
+    flex_attention: Any = None
+    _dense_to_ordered: Any = None
+if os.environ.get('STEERLING_USE_FLEX_ATTN', '0') != '1':
+    _FLEX_ATTN_AVAILABLE = False
+if TYPE_CHECKING:
+    from torch.nn.attention.flex_attention import BlockMask as BlockMaskType
+    from steerling.configs.causal_diffusion import CausalDiffusionConfig
+if torch.cuda.is_available() and _FLEX_ATTN_AVAILABLE:
+    compiled_flex_attention = torch.compile(flex_attention, fullgraph=True)
+else:
+    compiled_flex_attention = flex_attention
+
+def block_causal_mask_mod(b: Any, h: Any, q_idx: torch.Tensor, kv_idx: torch.Tensor, *, block_size: int) -> torch.Tensor:
+    """Block-causal mask: causal across blocks, bidirectional within blocks."""
+    return q_idx // block_size >= kv_idx // block_size
+
+def fast_create_block_causal_mask(attn_block_size: int, seq_length: int, mask_block_size: int, device: torch.device) -> BlockMaskType:
+    """
+    Fast block-causal mask creation for flex_attention.
+
+    Analytically computes the sparse block structure instead of evaluating
+    the mask function at every position.
+    """
+    if not _FLEX_ATTN_AVAILABLE or _dense_to_ordered is None or BlockMask is None:
+        raise RuntimeError('flex_attention not available')
+    num_mask_blocks = -(-seq_length // mask_block_size)
+    attn_blocks_per_mask_block, rem = divmod(mask_block_size, attn_block_size)
+    if rem != 0:
+        raise ValueError(f'mask_block_size ({mask_block_size}) must be divisible by attn_block_size ({attn_block_size})')
+    num_attn_blocks = num_mask_blocks * attn_blocks_per_mask_block
+    lowres_attn_mask = torch.tril(torch.ones(num_attn_blocks, num_attn_blocks, dtype=torch.bool, device=device))
+    block_attn_count = lowres_attn_mask.reshape(num_mask_blocks, attn_blocks_per_mask_block, num_mask_blocks, attn_blocks_per_mask_block).permute(0, 2, 1, 3).sum(dim=[-2, -1])
+    max_count = attn_blocks_per_mask_block * attn_blocks_per_mask_block
+    full_block_mask = block_attn_count == max_count
+    if seq_length % mask_block_size > 0:
+        full_block_mask[-1, :] = False
+    normal_block_mask = (block_attn_count > 0) & ~full_block_mask
+    kv_num_blocks, kv_indices = _dense_to_ordered(normal_block_mask)
+    full_kv_num_blocks, full_kv_indices = _dense_to_ordered(full_block_mask)
+    q_num_blocks, q_indices = _dense_to_ordered(normal_block_mask.transpose(-2, -1))
+    full_q_num_blocks, full_q_indices = _dense_to_ordered(full_block_mask.transpose(-2, -1))
+    return BlockMask(seq_lengths=(seq_length, seq_length), kv_num_blocks=kv_num_blocks[None, None, ...], kv_indices=kv_indices[None, None, ...], full_kv_num_blocks=full_kv_num_blocks[None, None, ...], full_kv_indices=full_kv_indices[None, None, ...], q_num_blocks=q_num_blocks[None, None, ...], q_indices=q_indices[None, None, ...], full_q_num_blocks=full_q_num_blocks[None, None, ...], full_q_indices=full_q_indices[None, None, ...], mask_mod=partial(block_causal_mask_mod, block_size=attn_block_size), BLOCK_SIZE=(mask_block_size, mask_block_size))
+
+def sdpa_with_block_causal_mask(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, diff_block_size: int, mask_cache: dict[str, torch.Tensor], enable_gqa: bool=False) -> torch.Tensor:
+    """Fallback using SDPA with dense mask when flex_attention unavailable."""
+    B, H, T, D = q.shape
+    device = q.device
+    dtype = q.dtype
+    cache_key = f'sdpa_{T}_{device}_{dtype}'
+    if cache_key not in mask_cache:
+        q_idx = torch.arange(T, device=device).unsqueeze(1)
+        kv_idx = torch.arange(T, device=device).unsqueeze(0)
+        bool_mask = q_idx // diff_block_size >= kv_idx // diff_block_size
+        attn_mask = torch.zeros(T, T, device=device, dtype=dtype)
+        attn_mask.masked_fill_(~bool_mask, float('-inf'))
+        mask_cache[cache_key] = attn_mask
+    return F.scaled_dot_product_attention(q, k, v, attn_mask=mask_cache[cache_key], dropout_p=0.0, is_causal=False, enable_gqa=enable_gqa)
+
+class BlockCausalAttention(nn.Module):
+    """Block-causal self-attention with FlexAttention and optional GQA."""
+    FLEX_MASK_BLOCK_SIZE = 128
+
+    def __init__(self, config: CausalDiffusionConfig) -> None:
+        super().__init__()
+        if not hasattr(config, 'diff_block_size'):
+            raise ValueError("BlockCausalAttention requires 'diff_block_size' in config.")
+        assert config.n_embd % config.n_head == 0
+        self.config = config
+        self.n_head = config.n_head
+        self.n_embd = config.n_embd
+        self.head_dim = config.n_embd // config.n_head
+        n_kv = getattr(config, 'n_kv_heads', None)
+        self.n_kv_heads = self.n_head if n_kv is None else int(n_kv)
+        if self.n_kv_heads <= 0:
+            raise ValueError(f'n_kv_heads must be >= 1 (got {self.n_kv_heads})')
+        if self.n_head % self.n_kv_heads != 0:
+            raise ValueError(f'n_head ({self.n_head}) must be divisible by n_kv_heads ({self.n_kv_heads})')
+        self.kv_repeat = self.n_head // self.n_kv_heads
+        use_bias = getattr(config, 'use_bias', False)
+        kv_out = self.n_kv_heads * self.head_dim
+        attn_out = self.n_embd + 2 * kv_out
+        self.c_attn = nn.Linear(config.n_embd, attn_out, bias=use_bias)
+        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=use_bias)
+        self.c_proj.SCALE_INIT = 1
+        if getattr(config, 'use_qk_norm', False):
+            if getattr(config, 'use_rms_norm', True):
+                self.q_norm: nn.Module | None = RMSNorm(config, size=self.head_dim)
+                self.k_norm: nn.Module | None = RMSNorm(config, size=self.head_dim)
+            else:
+                self.q_norm = nn.LayerNorm(self.head_dim)
+                self.k_norm = nn.LayerNorm(self.head_dim)
+        else:
+            self.q_norm = None
+            self.k_norm = None
+        if getattr(config, 'use_rope', True):
+            self.rope: RotaryEmbedding | None = RotaryEmbedding(dim=self.head_dim, max_seq_len=config.block_size, base=getattr(config, 'rope_base', 500000.0), rope_full_precision=getattr(config, 'rope_full_precision', True))
+        else:
+            self.rope = None
+        self._mask_cache: dict = {}
+        self._sdpa_mask_cache: dict[str, torch.Tensor] = {}
+        self._logged_attention_mode = False
+
+    def _get_block_mask(self, T: int, device: torch.device):
+        cache_key = f'flex_{T}_{device}'
+        if cache_key not in self._mask_cache:
+            diff_block_size = self.config.diff_block_size
+            mask_block_size = self.FLEX_MASK_BLOCK_SIZE
+            if mask_block_size % diff_block_size != 0:
+                mask_block_size = diff_block_size * (mask_block_size // diff_block_size)
+                if mask_block_size == 0:
+                    mask_block_size = diff_block_size
+            self._mask_cache[cache_key] = fast_create_block_causal_mask(attn_block_size=diff_block_size, seq_length=T, mask_block_size=mask_block_size, device=device)
+        return self._mask_cache[cache_key]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, T, C = x.size()
+        device = x.device
+        use_flex = _FLEX_ATTN_AVAILABLE and x.is_cuda and (flex_attention is not None)
+        if not self._logged_attention_mode:
+            self._logged_attention_mode = True
+            mode = 'flex_attention' if use_flex else 'SDPA fallback'
+            logger.debug(f'[CausalDiffusion] Using {mode} with GQA (n_head={self.n_head}, n_kv_heads={self.n_kv_heads})')
+        qkv = self.c_attn(x)
+        clip_qkv = getattr(self.config, 'clip_qkv', None)
+        if clip_qkv is not None:
+            qkv = qkv.clamp(min=-clip_qkv, max=clip_qkv)
+        kv_dim = self.n_kv_heads * self.head_dim
+        q, k, v = qkv.split([self.n_embd, kv_dim, kv_dim], dim=2)
+        q = q.reshape(B, T, self.n_head, self.head_dim).transpose(1, 2)
+        k = k.reshape(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
+        v = v.reshape(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
+        if self.q_norm is not None and self.k_norm is not None:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+        if self.rope is not None:
+            q, k = self.rope(q, k)
+        if use_flex:
+            block_mask = self._get_block_mask(T, device)
+            assert flex_attention is not None and compiled_flex_attention is not None
+            if q.is_cuda:
+                y = compiled_flex_attention(q, k, v, block_mask=block_mask, enable_gqa=True)
+            else:
+                y = flex_attention(q, k, v, block_mask=block_mask, enable_gqa=True)
+        else:
+            y = sdpa_with_block_causal_mask(q, k, v, diff_block_size=self.config.diff_block_size, mask_cache=self._sdpa_mask_cache, enable_gqa=True)
+        y = y.transpose(1, 2).reshape(B, T, C)
+        y = self.c_proj(y)
+        return y
+
+class CausalDiffusionBlock(nn.Module):
+    """Transformer block for CausalDiffusionLM (block-causal attention + MLP)."""
+
+    def __init__(self, config: CausalDiffusionConfig) -> None:
+        super().__init__()
+        use_rms_norm = getattr(config, 'use_rms_norm', True)
+        if use_rms_norm:
+            self.ln_1: nn.Module = RMSNorm(config)
+            self.ln_2: nn.Module = RMSNorm(config)
+        else:
+            self.ln_1 = nn.LayerNorm(config.n_embd)
+            self.ln_2 = nn.LayerNorm(config.n_embd)
+        self.norm_order = getattr(config, 'norm_order', 'post')
+        self.attn = BlockCausalAttention(config)
+        self.mlp = MLP(config)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.norm_order == 'pre':
+            x = x + self.attn(self.ln_1(x))
+            x = x + self.mlp(self.ln_2(x))
+        else:
+            x = x + self.ln_1(self.attn(x))
+            x = x + self.ln_2(self.mlp(x))
+        return x
+
+# ======================================================================
+# steerling/models/causal_diffusion.py
+# ======================================================================
+
+class CausalDiffusionLM(nn.Module):
+    """
+    CausalDiffusionLM transformer backbone with block-causal attention.
+
+    Pure compute graph — no training code, no loss logic.
+
+    Args:
+        config: CausalDiffusionConfig with model hyperparameters
+        vocab_size: Vocabulary size (including special tokens)
+    """
+
+    def __init__(self, config: CausalDiffusionConfig, vocab_size: int) -> None:
+        super().__init__()
+        self.config = config
+        self.vocab_size = vocab_size
+        self.tok_emb = nn.Embedding(vocab_size, config.n_embd)
+        self.blocks = nn.ModuleList([CausalDiffusionBlock(config) for _ in range(config.n_layers)])
+        if config.use_rms_norm:
+            self.ln_f: nn.Module = RMSNorm(config)
+        else:
+            self.ln_f = nn.LayerNorm(config.n_embd)
+        self.lm_head = nn.Linear(config.n_embd, vocab_size, bias=False)
+        if config.weight_sharing:
+            self.tok_emb.weight = self.lm_head.weight
+
+    def forward(self, input_ids: torch.Tensor, *, input_embeds: torch.Tensor | None=None, return_hidden: bool=False) -> torch.Tensor:
+        """
+        Forward pass.
+
+        Args:
+            input_ids: Token indices [B, T] (may contain mask tokens)
+            input_embeds: Pre-computed embeddings [B, T, D]. If provided, input_ids is ignored.
+            return_hidden: If True, return hidden states before lm_head.
+
+        Returns:
+            logits [B, T, vocab_size] or hidden_states [B, T, n_embd]
+        """
+        if input_embeds is not None:
+            x = input_embeds
+        elif input_ids is not None:
+            x = self.tok_emb(input_ids)
+        else:
+            raise ValueError('Either input_ids or input_embeds must be provided')
+        for block in self.blocks:
+            x = block(x)
+        x = self.ln_f(x)
+        if return_hidden:
+            return x
+        return self.lm_head(x)
+
+    def get_num_params(self, non_embedding: bool=True) -> int:
+        """Return number of parameters."""
+        n_params = sum((p.numel() for p in self.parameters()))
+        if non_embedding:
+            n_params -= self.tok_emb.weight.numel()
+        return n_params
+
+    def _restore_weight_tying(self) -> None:
+        """Re-establish weight tying after to_empty() or device transfer."""
+        if self.config.weight_sharing:
+            self.tok_emb.weight = self.lm_head.weight
+
+    def _init_weights(self, module: nn.Module) -> None:
+        """Initialize model weights (used for fresh models, not loaded checkpoints)."""
+        if isinstance(module, nn.Linear):
+            std = 0.02
+            if hasattr(module, 'SCALE_INIT'):
+                std *= (2 * self.config.n_layers) ** (-0.5)
+            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        elif isinstance(module, RMSNorm):
+            torch.nn.init.ones_(module.weight)
+
+# ======================================================================
+# steerling/models/interpretable/outputs.py
+# ======================================================================
+
+@dataclass
+class InterpretableOutput:
+    """
+    Full output from InterpretableCausalDiffusionLM; it contains all decomposition components for attribution and analysis.
+    """
+    hidden: Tensor
+    known_features: Tensor
+    known_logits: Tensor | None
+    known_gt_features: Tensor | None
+    known_predicted: Tensor
+    known_weights: Tensor | None
+    known_topk_indices: Tensor | None
+    known_topk_logits: Tensor | None
+    unk: Tensor
+    unk_hat: Tensor | None
+    unk_for_lm: Tensor
+    unknown_logits: Tensor | None
+    unknown_weights: Tensor | None
+    unknown_topk_indices: Tensor | None
+    unknown_topk_logits: Tensor | None
+    composed: Tensor
+    epsilon: Tensor | None
+    epsilon_true: Tensor | None
+
+# ======================================================================
+# steerling/models/interpretable/concept_head.py
+# ======================================================================
+
+logger = logging.getLogger(__name__)
+LARGE_CONCEPT_THRESHOLD = 50000
+
+@dataclass
+class ConceptHeadOutput:
+    """Output from ConceptHead forward pass.
+
+    Attributes:
+        features: Final concept features after teacher forcing/intervention (B, T, D)
+        gt_features: Ground truth pooled features. None for unknown heads. (B, T, D) or None
+        logits: Full concept logits (B, T, C). Only set if return_logits=True. Usually None.
+        predicted: Predicted features before teacher forcing mixing (B, T, D)
+        weights: Full concept weights (B, T, C). Only set if return_logits=True. Usually None.
+        topk_indices: Top-k concept indices (B, T, k). Set when using streaming top-k.
+        topk_logits: Logits for top-k concepts (B, T, k). Set when using streaming top-k.
+        hidden: Hidden states passed to this head (B, T, D). Stored for attribution.
+    """
+    features: Tensor
+    gt_features: Tensor | None
+    logits: Tensor | None
+    predicted: Tensor
+    weights: Tensor | None = None
+    topk_indices: Tensor | None = None
+    topk_logits: Tensor | None = None
+    hidden: Tensor | None = None
+
+class ConceptHead(nn.Module):
+    """
+    Concept decomposition head supporting both known and unknown concepts.
+    Memory-efficient implementation that avoids (B, T, C) allocations by default.
+
+    Modes:
+    - Known (is_unknown=False): Supports GT, teacher forcing, top-k, interventions
+    - Unknown (is_unknown=True): No GT, no teacher forcing
+
+    Architectures:
+    - use_attention=False: Linear predictor (n_embd -> n_concepts)
+    - use_attention=True: Query projection + sigmoid attention over embeddings
+
+    Factorization (for large unknown heads):
+    - factorize=False: Dense embeddings (C, D) and predictor (D, C)
+    - factorize=True: Factorized embeddings (C, r) @ (r, D) where r << D
+                      Reduces memory by ~10-20x for large C
+
+    Memory Safety:
+    - Unknown heads with n_concepts > 50k cannot use dense operations
+    - Interventions are only supported for known heads
+    - return_logits=True is forbidden for large unknown heads
+    - All tensor indexing uses F.embedding for DTensor safety
+
+    Args:
+        n_concepts: Number of concepts (C)
+        concept_dim: Dimension of concept embeddings (should equal n_embd)
+        n_embd: Model hidden dimension
+        is_unknown: If True, skip GT pooling and teacher forcing
+        use_attention: If True, use attention; else use linear predictor
+        topk: Top-k sparsity for concept weights. None = no sparsity.
+        block_size: Block size for memory-efficient operations
+        pad_multiple: Pad n_concepts to a multiple of this for efficiency
+        store_unknown_weights: If True and use_attention & is_unknown, store logits/weights
+        apply_topk_to_unknown: If True, also apply top-k to unknown concepts
+        topk_on_logits: If True, apply top-k on logits (then sigmoid). If False, on weights.
+        teacher_force_alpha: If None, hard TF. If in [0,1], soft mixing.
+        factorize: If True, use low-rank factorized embeddings
+        factorize_rank: Rank for factorization (r). Lower = less memory, less expressivity.
+    """
+
+    class ConceptPooling(nn.Module):
+        """Memory-efficient sum pooling using scatter-add."""
+
+        def __init__(self, concept_dim: int):
+            super().__init__()
+            self.concept_dim = concept_dim
+
+        def forward(self, concept_ids: Tensor, concept_mask: Tensor, concept_embeddings: nn.Embedding) -> Tensor:
+            """
+            Pool concept embeddings based on ground truth IDs.
+            Uses scatter-add to avoid (B, T, K, D) allocation when K is sparse.
+
+            Args:
+                concept_ids: (B, T, K) concept indices, -1 for invalid
+                concept_mask: (B, T, K) boolean mask for valid concepts
+                concept_embeddings: Embedding layer to look up
+
+            Returns:
+                Pooled features (B, T, D)
+            """
+            B, T, K = concept_ids.shape
+            D = concept_embeddings.embedding_dim
+            device = concept_ids.device
+            valid_mask = concept_mask & (concept_ids != -1)
+            pooled = torch.zeros(B, T, D, device=device, dtype=concept_embeddings.weight.dtype)
+            if not valid_mask.any():
+                return pooled
+            b_idx, t_idx, k_idx = torch.where(valid_mask)
+            c_ids = concept_ids[b_idx, t_idx, k_idx].long()
+            emb = concept_embeddings(c_ids)
+            flat_idx = b_idx * T + t_idx
+            flat_idx = flat_idx.unsqueeze(-1).expand(-1, D)
+            pooled_flat = pooled.view(B * T, D)
+            pooled_flat.scatter_add_(0, flat_idx, emb)
+            return pooled.view(B, T, D)
+
+    def __init__(self, n_concepts: int, concept_dim: int, n_embd: int, is_unknown: bool=False, use_attention: bool=False, topk: int | None=16, topk_features: int | None=None, block_size: int=8192, *, pad_multiple: int=16, store_unknown_weights: bool=False, apply_topk_to_unknown: bool=False, topk_on_logits: bool=False, factorize: bool=False, factorize_rank: int=256):
+        super().__init__()
+        self.n_concepts = n_concepts
+        self.concept_dim = concept_dim
+        self.n_embd = n_embd
+        self.is_unknown = is_unknown
+        self.use_attention = use_attention
+        self.topk = topk
+        self.topk_features = topk_features if topk_features is not None else topk
+        self.block_size = block_size
+        self.pad_multiple = pad_multiple
+        self.store_unknown_weights = store_unknown_weights
+        self.apply_topk_to_unknown = apply_topk_to_unknown
+        self.topk_on_logits = topk_on_logits
+        self.factorize = factorize
+        self.factorize_rank = factorize_rank
+        self._is_large = n_concepts > LARGE_CONCEPT_THRESHOLD
+        self.n_concepts_padded = (n_concepts + pad_multiple - 1) // pad_multiple * pad_multiple
+        if factorize:
+            self.embedding_coef = nn.Embedding(self.n_concepts_padded, factorize_rank)
+            self.embedding_basis = nn.Linear(factorize_rank, concept_dim, bias=False)
+            self.concept_embedding = None
+            if not use_attention:
+                self.predictor_down = nn.Linear(n_embd, factorize_rank, bias=False)
+                self.predictor_up = nn.Linear(factorize_rank, self.n_concepts_padded, bias=False)
+                self.concept_predictor = None
+            else:
+                self.concept_query_projection = nn.Linear(n_embd, concept_dim, bias=False)
+                self.predictor_down = None
+                self.predictor_up = None
+                self.concept_predictor = None
+            dense_params = n_concepts * concept_dim * 2
+            factorized_params = n_concepts * factorize_rank + factorize_rank * concept_dim + (n_embd * factorize_rank + factorize_rank * n_concepts if not use_attention else 0)
+            logger.info(f'[ConceptHead] Factorized mode: {n_concepts} concepts, rank={factorize_rank}')
+            logger.info(f'[ConceptHead] Memory: {dense_params * 2 / 1000000000.0:.2f} GB (dense) -> {factorized_params * 2 / 1000000000.0:.2f} GB (factorized) = {(1 - factorized_params / dense_params) * 100:.1f}% reduction')
+        else:
+            self.concept_embedding = nn.Embedding(self.n_concepts_padded, concept_dim)
+            self.embedding_coef = None
+            self.embedding_basis = None
+            if use_attention:
+                self.concept_query_projection = nn.Linear(n_embd, concept_dim, bias=False)
+                self.concept_predictor = None
+            else:
+                self.concept_predictor = nn.Linear(n_embd, self.n_concepts_padded, bias=False)
+            self.predictor_down = None
+            self.predictor_up = None
+        self.concept_pooling = self.ConceptPooling(concept_dim)
+        if self.topk_features != self.topk:
+            logger.info(f"[ConceptHead] {('Unknown' if is_unknown else 'Known')} head: topk={self.topk} (loss), topk_features={self.topk_features} (features)")
+        if is_unknown and apply_topk_to_unknown:
+            logger.info(f'[ConceptHead] Unknown head: apply_topk_to_unknown=True, topk={self.topk}')
+        self._init_weights()
+
+    def _init_weights(self):
+        """Initialize weights with small values."""
+        if self.factorize:
+            nn.init.normal_(self.embedding_coef.weight, mean=0.0, std=0.02)
+            nn.init.normal_(self.embedding_basis.weight, mean=0.0, std=0.02)
+            if self.predictor_down is not None:
+                nn.init.normal_(self.predictor_down.weight, mean=0.0, std=0.02)
+            if self.predictor_up is not None:
+                nn.init.normal_(self.predictor_up.weight, mean=0.0, std=0.02)
+        else:
+            if self.concept_embedding is not None:
+                nn.init.normal_(self.concept_embedding.weight, mean=0.0, std=0.02)
+            if self.concept_predictor is not None:
+                nn.init.normal_(self.concept_predictor.weight, mean=0.0, std=0.02)
+        if hasattr(self, 'concept_query_projection') and self.concept_query_projection is not None:
+            nn.init.normal_(self.concept_query_projection.weight, mean=0.0, std=0.02)
+
+    def _check_dense_allowed(self, operation: str) -> None:
+        """Raise error if dense operations are requested for large unknown heads."""
+        if self.is_unknown and self._is_large:
+            raise ValueError(f'{operation} requested for unknown head with {self.n_concepts} concepts. This would allocate multi-GB tensors. Use streaming mode instead. (Threshold: {LARGE_CONCEPT_THRESHOLD})')
+
+    @staticmethod
+    def _safe_index(weight: Tensor, indices: Tensor) -> Tensor:
+        """
+        DTensor-safe indexing using F.embedding.
+
+        Replaces weight[indices] which crashes under FSDP2/DTensor.
+
+        Args:
+            weight: (N, D) weight matrix
+            indices: (...) indices to select
+
+        Returns:
+            (..., D) selected embeddings
+        """
+        original_shape = indices.shape
+        flat_indices = indices.reshape(-1)
+        flat_result = F.embedding(flat_indices, weight)
+        return flat_result.reshape(*original_shape, -1)
+
+    def _get_embedding_weight(self) -> Tensor:
+        """
+        Get full embedding matrix.
+
+        For dense: returns concept_embedding.weight
+        For factorized: computes coef @ basis (materializes full matrix)
+
+        Returns:
+            (C, D) embedding matrix
+        """
+        if self.concept_embedding is not None:
+            return self.concept_embedding.weight
+        else:
+            return self.embedding_basis(self.embedding_coef.weight)
+
+    def _get_embedding(self, indices: Tensor) -> Tensor:
+        """
+        Get embeddings for specific indices (DTensor-safe).
+
+        For dense: uses F.embedding
+        For factorized: looks up coef, then applies basis
+
+        Args:
+            indices: (...) concept indices
+
+        Returns:
+            (..., D) embeddings
+        """
+        if self.concept_embedding is not None:
+            return self.concept_embedding(indices)
+        else:
+            coef = self.embedding_coef(indices)
+            return self.embedding_basis(coef)
+
+    def _get_predictor_weight(self) -> Tensor | None:
+        """
+        Get full predictor weight matrix (for linear path only).
+
+        Returns:
+            (C, D) predictor weight, or None if using attention
+        """
+        if self.concept_predictor is not None:
+            return self.concept_predictor.weight
+        elif self.predictor_down is not None and self.predictor_up is not None:
+            return self.predictor_up.weight @ self.predictor_down.weight
+        else:
+            return None
+
+    @staticmethod
+    def _merge_topk(topv: Tensor, topi: Tensor, v_blk: Tensor, i_blk: Tensor, k: int) -> tuple[Tensor, Tensor]:
+        """Efficient merge of two top-k sets. Memory: O(BT × 2k)."""
+        cand_v = torch.cat([topv, v_blk], dim=1)
+        cand_i = torch.cat([topi, i_blk], dim=1)
+        new_v, sel = torch.topk(cand_v, k, dim=1)
+        new_i = torch.gather(cand_i, 1, sel)
+        return (new_v, new_i)
+
+    @staticmethod
+    def linear_block_features(hidden: Tensor, predictor_weight: Tensor, embeddings: Tensor, block_size: int=4096) -> Tensor:
+        """
+        Memory-efficient linear prediction without materializing (B, T, C).
+
+        Args:
+            hidden: (B, T, D)
+            predictor_weight: (C, D)
+            embeddings: (C, D)
+            block_size: Concepts per block
+
+        Returns:
+            Features (B, T, D)
+        """
+        B, T, D = hidden.shape
+        C = predictor_weight.size(0)
+        output = torch.zeros(B, T, D, dtype=hidden.dtype, device=hidden.device)
+        flat_h = hidden.reshape(-1, D)
+        W_t = predictor_weight.t().contiguous()
+        for start in range(0, C, block_size):
+            end = min(start + block_size, C)
+            logits_block = (flat_h @ W_t[:, start:end]).to(torch.float32)
+            logits_block = logits_block.clamp(-15, 15)
+            weights_block = torch.sigmoid(logits_block)
+            E_block = embeddings[start:end].to(weights_block.dtype)
+            output.add_((weights_block @ E_block).reshape(B, T, D))
+        return output.to(hidden.dtype)
+
+    @staticmethod
+    def attention_block_features(query: Tensor, embeddings: Tensor, block_size: int=4096) -> Tensor:
+        """Memory-efficient attention features without materializing (B, T, C)."""
+        B, T, D = query.shape
+        C = embeddings.shape[0]
+        scale = 1.0 / math.sqrt(D)
+        flat_q = query.reshape(-1, D)
+        emb_T = embeddings.t().contiguous()
+        output = torch.zeros(B * T, D, dtype=query.dtype, device=query.device)
+        for start in range(0, C, block_size):
+            end = min(start + block_size, C)
+            scores = (flat_q @ emb_T[:, start:end]).to(torch.float32) * scale
+            scores = scores.clamp(-15, 15)
+            weights = torch.sigmoid(scores)
+            output.add_(weights @ embeddings[start:end].to(weights.dtype))
+        return output.reshape(B, T, D).to(query.dtype)
+
+    @staticmethod
+    def linear_features_topk_streaming(hidden: Tensor, predictor_weight: Tensor, embeddings: Tensor, k: int, block_size: int=4096, topk_on_logits: bool=False) -> tuple[Tensor, Tensor, Tensor]:
+        """
+        Memory-efficient linear prediction with streaming top-k.
+
+        Uses merge-k-with-k to keep memory O(BT × k), not O(BT × block_size).
+
+        Args:
+            hidden: (B, T, D)
+            predictor_weight: (C, D)
+            embeddings: (C, D)
+            k: Number of top concepts
+            block_size: Concepts per block
+            topk_on_logits: If True, select top-k by logits; else by sigmoid
+
+        Returns:
+            features: (B, T, D) weighted concept features
+            topk_indices: (B, T, k) indices of top-k concepts
+            topk_logits: (B, T, k) logits for top-k concepts
+        """
+        B, T, D = hidden.shape
+        C = predictor_weight.size(0)
+        BT = B * T
+        device = hidden.device
+        k = min(k, C)
+        flat_h = hidden.reshape(BT, D)
+        W_t = predictor_weight.t().contiguous()
+        topv = torch.full((BT, k), float('-inf'), device=device, dtype=hidden.dtype)
+        topi = torch.zeros((BT, k), device=device, dtype=torch.long)
+        for start in range(0, C, block_size):
+            end = min(start + block_size, C)
+            logits_blk = (flat_h @ W_t[:, start:end]).to(torch.float32).clamp_(-15, 15)
+            vals_blk = logits_blk if topk_on_logits else torch.sigmoid(logits_blk)
+            blk_k = min(k, end - start)
+            v_blk, idx_blk = torch.topk(vals_blk, blk_k, dim=1)
+            i_blk = idx_blk + start
+            if blk_k < k:
+                pad_v = torch.full((BT, k - blk_k), float('-inf'), device=device, dtype=torch.float32)
+                pad_i = torch.zeros((BT, k - blk_k), device=device, dtype=torch.long)
+                v_blk = torch.cat([v_blk, pad_v], dim=1)
+                i_blk = torch.cat([i_blk, pad_i], dim=1)
+            topv, topi = ConceptHead._merge_topk(topv, topi, v_blk, i_blk, k)
+        W_sel = ConceptHead._safe_index(predictor_weight, topi)
+        logits_sel = torch.einsum('bd,bkd->bk', flat_h.to(torch.float32), W_sel.to(torch.float32))
+        logits_sel = logits_sel.clamp(-15, 15)
+        del W_sel
+        weights_sel = torch.sigmoid(logits_sel)
+        E_sel = ConceptHead._safe_index(embeddings, topi)
+        features = torch.einsum('bk,bkd->bd', weights_sel.to(E_sel.dtype), E_sel)
+        return (features.reshape(B, T, D).to(hidden.dtype), topi.reshape(B, T, k), logits_sel.reshape(B, T, k))
+
+    @staticmethod
+    def attention_features_topk_streaming(query: Tensor, embeddings: Tensor, k: int, block_size: int=4096, topk_on_logits: bool=False) -> tuple[Tensor, Tensor, Tensor]:
+        """Memory-efficient attention with streaming top-k."""
+        B, T, D = query.shape
+        C = embeddings.shape[0]
+        BT = B * T
+        device = query.device
+        scale = 1.0 / math.sqrt(D)
+        k = min(k, C)
+        flat_q = query.reshape(BT, D)
+        emb_T = embeddings.t().contiguous()
+        topv = torch.full((BT, k), float('-inf'), device=device, dtype=query.dtype)
+        topi = torch.zeros((BT, k), device=device, dtype=torch.long)
+        for start in range(0, C, block_size):
+            end = min(start + block_size, C)
+            logits_blk = (flat_q @ emb_T[:, start:end]).to(torch.float32) * scale
+            logits_blk = logits_blk.clamp(-15, 15)
+            vals_blk = logits_blk if topk_on_logits else torch.sigmoid(logits_blk)
+            blk_k = min(k, end - start)
+            v_blk, idx_blk = torch.topk(vals_blk, blk_k, dim=1)
+            i_blk = idx_blk + start
+            if blk_k < k:
+                pad_v = torch.full((BT, k - blk_k), float('-inf'), device=device, dtype=torch.float32)
+                pad_i = torch.zeros((BT, k - blk_k), device=device, dtype=torch.long)
+                v_blk = torch.cat([v_blk, pad_v], dim=1)
+                i_blk = torch.cat([i_blk, pad_i], dim=1)
+            topv, topi = ConceptHead._merge_topk(topv, topi, v_blk, i_blk, k)
+        E_sel = ConceptHead._safe_index(embeddings, topi)
+        logits_sel = torch.einsum('bd,bkd->bk', flat_q.to(torch.float32), E_sel.to(torch.float32)) * scale
+        logits_sel = logits_sel.clamp(-15, 15)
+        weights_sel = torch.sigmoid(logits_sel)
+        features = torch.einsum('bk,bkd->bd', weights_sel.to(E_sel.dtype), E_sel)
+        return (features.reshape(B, T, D).to(query.dtype), topi.reshape(B, T, k), logits_sel.reshape(B, T, k))
+
+    def attention_block_features_factorized(self, query: Tensor, block_size: int=4096) -> Tensor:
+        """
+        Memory-efficient factorized attention over ALL concepts.
+
+        Uses factorized scoring and feature computation:
+        - Scoring: (query @ basis.T) @ coef.T instead of query @ E.T
+        - Features: (weights @ coef) @ basis instead of weights @ E
+
+        FLOPs: O(BT * r * (D + C)) instead of O(BT * D * C)
+
+        Args:
+            query: (B, T, D) query vectors from concept_query_projection
+            block_size: Concepts per block for chunked processing
+
+        Returns:
+            (B, T, D) weighted concept features
+        """
+        assert self.factorize, 'Only valid for factorized head'
+        B, T, D = query.shape
+        BT = B * T
+        C = self.n_concepts
+        _ = self.factorize_rank
+        device = query.device
+        scale = 1.0 / math.sqrt(D)
+        flat_q = query.reshape(BT, D)
+        coef = self.embedding_coef.weight[:C]
+        basis_weight = self.embedding_basis.weight
+        q_compressed = flat_q @ basis_weight
+        output = torch.zeros(BT, D, dtype=query.dtype, device=device)
+        _ = (C + block_size - 1) // block_size
+        for _block_idx, start in enumerate(range(0, C, block_size)):
+            end = min(start + block_size, C)
+            coef_chunk = coef[start:end]
+            scores_chunk = (q_compressed @ coef_chunk.T).float() * scale
+            scores_chunk = scores_chunk.clamp(-15, 15)
+            weights_chunk = torch.sigmoid(scores_chunk)
+            weighted_coef = weights_chunk @ coef_chunk.float()
+            features_chunk = weighted_coef @ basis_weight.T.to(weighted_coef.dtype)
+            output.add_(features_chunk)
+        return output.reshape(B, T, D).to(query.dtype)
+
+    def attention_features_topk_factorized(self, query: Tensor, k: int, block_size: int=4096) -> tuple[Tensor, Tensor, Tensor]:
+        """
+        Memory-efficient factorized attention with streaming top-k.
+
+        Pass 1: Find top-k concepts using factorized scoring
+        Pass 2: Compute features using only top-k embeddings
+
+        Args:
+            query: (B, T, D) query vectors
+            k: Number of top concepts per token
+            block_size: Concepts per block
+
+        Returns:
+            features: (B, T, D) weighted concept features
+            topk_indices: (B, T, k) top-k concept indices
+            topk_logits: (B, T, k) logits for top-k concepts
+        """
+        assert self.factorize, 'Only valid for factorized head'
+        B, T, D = query.shape
+        BT = B * T
+        C = self.n_concepts
+        _ = self.factorize_rank
+        device = query.device
+        scale = 1.0 / math.sqrt(D)
+        k = min(k, C)
+        flat_q = query.reshape(BT, D)
+        coef = self.embedding_coef.weight[:C]
+        basis_weight = self.embedding_basis.weight
+        q_compressed = flat_q @ basis_weight
+        topv = torch.full((BT, k), float('-inf'), device=device, dtype=query.dtype)
+        topi = torch.zeros((BT, k), device=device, dtype=torch.long)
+        for start in range(0, C, block_size):
+            end = min(start + block_size, C)
+            coef_chunk = coef[start:end]
+            scores_chunk = q_compressed.float() @ coef_chunk.T.float() * scale
+            scores_chunk = scores_chunk.clamp(-15, 15)
+            blk_k = min(k, end - start)
+            v_chunk, idx_chunk = torch.topk(scores_chunk, blk_k, dim=1)
+            i_chunk = idx_chunk + start
+            if blk_k < k:
+                pad_v = torch.full((BT, k - blk_k), float('-inf'), device=device, dtype=torch.float32)
+                pad_i = torch.zeros((BT, k - blk_k), device=device, dtype=torch.long)
+                v_chunk = torch.cat([v_chunk, pad_v], dim=1)
+                i_chunk = torch.cat([i_chunk, pad_i], dim=1)
+            topv, topi = self._merge_topk(topv, topi, v_chunk, i_chunk, k)
+        coef_sel = self.embedding_coef(topi)
+        logits_sel = torch.einsum('br,bkr->bk', q_compressed.float(), coef_sel.float()) * scale
+        logits_sel = logits_sel.clamp(-15, 15)
+        weights_sel = torch.sigmoid(logits_sel)
+        weighted_coef = torch.einsum('bk,bkr->br', weights_sel.to(coef_sel.dtype), coef_sel)
+        features = weighted_coef @ basis_weight.T.to(weighted_coef.dtype)
+        return (features.reshape(B, T, D).to(query.dtype), topi.reshape(B, T, k), logits_sel.reshape(B, T, k))
+
+    def linear_block_features_factorized(self, hidden: Tensor, block_size: int=4096) -> Tensor:
+        """
+        Memory-efficient factorized linear prediction over ALL concepts.
+
+        Uses factorized predictor: logits = hidden @ down @ up.T
+        Uses factorized embeddings: features = weights @ coef @ basis
+
+        Args:
+            hidden: (B, T, D) hidden states
+            block_size: Concepts per block
+
+        Returns:
+            (B, T, D) weighted concept features
+        """
+        assert self.factorize, 'Only valid for factorized head'
+        assert self.predictor_down is not None, 'Linear path requires predictor'
+        B, T, D = hidden.shape
+        BT = B * T
+        C = self.n_concepts
+        _ = self.factorize_rank
+        device = hidden.device
+        flat_h = hidden.reshape(BT, D)
+        coef = self.embedding_coef.weight[:C]
+        basis_weight = self.embedding_basis.weight
+        down_weight = self.predictor_down.weight
+        up_weight = self.predictor_up.weight[:C]
+        h_compressed = flat_h @ down_weight.T
+        output = torch.zeros(BT, D, dtype=hidden.dtype, device=device)
+        for start in range(0, C, block_size):
+            end = min(start + block_size, C)
+            up_chunk = up_weight[start:end]
+            coef_chunk = coef[start:end]
+            logits_chunk = h_compressed.float() @ up_chunk.T.float()
+            logits_chunk = logits_chunk.clamp(-15, 15)
+            weights_chunk = torch.sigmoid(logits_chunk)
+            weighted_coef = weights_chunk @ coef_chunk.float()
+            features_chunk = weighted_coef @ basis_weight.T.to(weighted_coef.dtype)
+            output.add_(features_chunk)
+        return output.reshape(B, T, D).to(hidden.dtype)
+
+    def linear_features_topk_factorized(self, hidden: Tensor, k: int, block_size: int=4096) -> tuple[Tensor, Tensor, Tensor]:
+        """
+        Memory-efficient factorized linear with streaming top-k.
+
+        Args:
+            hidden: (B, T, D) hidden states
+            k: Number of top concepts per token
+            block_size: Concepts per block
+
+        Returns:
+            features: (B, T, D) weighted concept features
+            topk_indices: (B, T, k) top-k concept indices
+            topk_logits: (B, T, k) logits for top-k concepts
+        """
+        assert self.factorize, 'Only valid for factorized head'
+        assert self.predictor_down is not None, 'Linear path requires predictor'
+        B, T, D = hidden.shape
+        BT = B * T
+        C = self.n_concepts
+        _ = self.factorize_rank
+        device = hidden.device
+        k = min(k, C)
+        flat_h = hidden.reshape(BT, D)
+        down_weight = self.predictor_down.weight
+        up_weight = self.predictor_up.weight[:C]
+        basis_weight = self.embedding_basis.weight
+        h_compressed = flat_h @ down_weight.T
+        topv = torch.full((BT, k), float('-inf'), device=device, dtype=hidden.dtype)
+        topi = torch.zeros((BT, k), device=device, dtype=torch.long)
+        for start in range(0, C, block_size):
+            end = min(start + block_size, C)
+            up_chunk = up_weight[start:end]
+            logits_chunk = h_compressed.float() @ up_chunk.T.float()
+            logits_chunk = logits_chunk.clamp(-15, 15)
+            blk_k = min(k, end - start)
+            v_chunk, idx_chunk = torch.topk(logits_chunk, blk_k, dim=1)
+            i_chunk = idx_chunk + start
+            if blk_k < k:
+                pad_v = torch.full((BT, k - blk_k), float('-inf'), device=device, dtype=torch.float32)
+                pad_i = torch.zeros((BT, k - blk_k), device=device, dtype=torch.long)
+                v_chunk = torch.cat([v_chunk, pad_v], dim=1)
+                i_chunk = torch.cat([i_chunk, pad_i], dim=1)
+            topv, topi = self._merge_topk(topv, topi, v_chunk, i_chunk, k)
+        coef_sel = self.embedding_coef(topi)
+        up_sel = self._safe_index(self.predictor_up.weight[:C], topi)
+        logits_sel = torch.einsum('br,bkr->bk', h_compressed.float(), up_sel.float())
+        logits_sel = logits_sel.clamp(-15, 15)
+        weights_sel = torch.sigmoid(logits_sel)
+        weighted_coef = torch.einsum('bk,bkr->br', weights_sel.to(coef_sel.dtype), coef_sel)
+        features = weighted_coef @ basis_weight.T.to(weighted_coef.dtype)
+        return (features.reshape(B, T, D).to(hidden.dtype), topi.reshape(B, T, k), logits_sel.reshape(B, T, k))
+
+    def compute_logits_for_indices(self, hidden: Tensor, indices: Tensor) -> Tensor:
+        """
+        Compute logits for specific concept indices only (sparse).
+
+        Supports both dense and factorized heads.
+
+        IMPORTANT: This function materializes (M, K, D) where M is the number of
+        tokens in hidden. Only call this with small M (e.g., masked tokens only).
+
+        Args:
+            hidden: (M, D) or (B, T, D) hidden states
+            indices: (M, K) or (B, T, K) concept indices
+
+        Returns:
+            logits: Same shape as indices
+        """
+        if hidden.dim() == 2:
+            M, D = hidden.shape
+            K = indices.size(-1)
+            flat_h = hidden
+            flat_idx = indices
+            output_shape = indices.shape
+        else:
+            B, T, D = hidden.shape
+            K = indices.size(-1)
+            M = B * T
+            flat_h = hidden.reshape(M, D)
+            flat_idx = indices.reshape(M, K)
+            output_shape = indices.shape
+        estimated_bytes = M * K * D * 2
+        if estimated_bytes > 1000000000.0:
+            warnings.warn(f'compute_logits_for_indices will allocate ~{estimated_bytes / 1000000000.0:.1f} GB. Consider reducing M={M} (use masked tokens only) or K={K}.')
+        n_valid = self.n_concepts
+        indices_safe = flat_idx.clamp(0, n_valid - 1)
+        if self.use_attention:
+            query = self.concept_query_projection(flat_h.unsqueeze(0)).squeeze(0)
+            scale = 1.0 / math.sqrt(self.concept_dim)
+            E_sel = self._get_embedding(indices_safe)
+            logits = torch.einsum('md,mkd->mk', query.float(), E_sel.float()) * scale
+        else:
+            if self.factorize:
+                W = self._get_predictor_weight()[:n_valid]
+                W_sel = self._safe_index(W, indices_safe)
+            else:
+                W = self.concept_predictor.weight[:n_valid]
+                W_sel = self._safe_index(W, indices_safe)
+            logits = torch.einsum('md,mkd->mk', flat_h.float(), W_sel.float())
+        return logits.clamp(-15, 15).reshape(output_shape)
+
+    def get_concept_weights(self, hidden: Tensor, concept_ids: Tensor) -> Tensor:
+        """
+        Get sigmoid weights for specific concepts (for attribution).
+
+        Args:
+            hidden: (B, T, D) or (M, D) hidden states
+            concept_ids: (B, T, K) or (M, K) or (K,) concept indices
+
+        Returns:
+            weights: Same shape as concept_ids, values in [0, 1]
+        """
+        if concept_ids.dim() == 1:
+            if hidden.dim() == 2:
+                M = hidden.size(0)
+                concept_ids = concept_ids.unsqueeze(0).expand(M, -1)
+            else:
+                B, T, _ = hidden.shape
+                concept_ids = concept_ids.unsqueeze(0).unsqueeze(0).expand(B, T, -1)
+        logits = self.compute_logits_for_indices(hidden, concept_ids)
+        return torch.sigmoid(logits)
+
+    @staticmethod
+    def blocked_logits(query: Tensor, embeddings: Tensor, block_size: int=8192, out_device: torch.device | None=None, out_dtype: torch.dtype=torch.float32) -> Tensor:
+        """
+        Compute concept logits in column blocks for memory efficiency.
+
+        logits = query @ embeddings.T / sqrt(D)
+        """
+        B, T, D = query.shape
+        C = embeddings.size(0)
+        scale = 1.0 / math.sqrt(D)
+        dev = query.device if out_device is None else out_device
+        logits = torch.empty(B, T, C, device=dev, dtype=out_dtype)
+        q = query.reshape(-1, D).to(torch.float32)
+        Et = embeddings.t().contiguous().to(torch.float32)
+        for s in range(0, C, block_size):
+            e = min(s + block_size, C)
+            scores = q @ Et[:, s:e] * scale
+            scores = scores.clamp(-15, 15)
+            logits[:, :, s:e] = scores.reshape(B, T, e - s).to(out_dtype)
+        return logits
+
+    @staticmethod
+    def blocked_mix(weights: Tensor, embeddings: Tensor, block_size: int=8192) -> Tensor:
+        """
+        Compute weighted sum of embeddings in column blocks.
+
+        output = weights @ embeddings
+        """
+        B, T, C = weights.shape
+        D = embeddings.size(1)
+        out = torch.zeros(B, T, D, device=weights.device, dtype=weights.dtype)
+        for s in range(0, C, block_size):
+            e = min(s + block_size, C)
+            w_blk = weights[:, :, s:e].to(torch.float32)
+            V_blk = embeddings[s:e].to(w_blk.dtype)
+            out.add_(w_blk @ V_blk)
+        return out.to(weights.dtype)
+
+    @staticmethod
+    def sigmoid_block_attention(query: Tensor, embeddings: Tensor, block_size: int=8192, return_logits: bool=False) -> Tensor | tuple[Tensor, Tensor]:
+        """Memory-efficient sigmoid attention using block processing."""
+        B, T, D = query.shape
+        C = embeddings.shape[0]
+        scale = 1.0 / math.sqrt(D)
+        flat_q = query.reshape(-1, D)
+        emb_T = embeddings.t().contiguous()
+        output = torch.zeros(B * T, D, dtype=query.dtype, device=query.device)
+        logits: Tensor | None = None
+        if return_logits:
+            logits = torch.empty(B, T, C, dtype=torch.float32, device=query.device)
+        for start in range(0, C, block_size):
+            end = min(start + block_size, C)
+            scores = (flat_q @ emb_T[:, start:end]).to(torch.float32) * scale
+            scores = scores.clamp(-15, 15)
+            if logits is not None:
+                logits[:, :, start:end] = scores.reshape(B, T, end - start)
+            weights = torch.sigmoid(scores)
+            output.add_(weights @ embeddings[start:end].to(weights.dtype))
+        output = output.reshape(B, T, D).to(query.dtype)
+        if return_logits:
+            assert logits is not None
+            return (output, logits)
+        return output
+
+    def _apply_sparse_interventions(self, features: Tensor, hidden: Tensor, intervene_ids: Tensor, intervene_vals: Tensor) -> Tensor:
+        """
+        Apply sparse interventions matching original dense behavior.
+
+        Original dense behavior:
+            weights = sigmoid(logits)  # (B, T, C)
+            weights[..., c] = new_val  # Override
+            features = weights @ embeddings
+
+        Sparse equivalent:
+            features += (new_val - current_weight) * embedding[c]
+        """
+        B, T, D = features.shape
+        valid = intervene_ids != -1
+        if not valid.any():
+            return features
+        ids_safe = intervene_ids.clamp(0, self.n_concepts - 1)
+        current_logits = self.compute_logits_for_indices(hidden, ids_safe)
+        current_weights = torch.sigmoid(current_logits)
+        emb = self._get_embedding(ids_safe)
+        delta = (intervene_vals - current_weights) * valid.float()
+        correction = (delta.unsqueeze(-1) * emb).sum(dim=2)
+        return features + correction
+
+    def _apply_dense_interventions(self, concept_weight: Tensor, intervene_ids: Tensor, intervene_vals: Tensor) -> Tensor:
+        """Apply interventions by overriding concept weights (dense path)."""
+        n_valid = min(self.n_concepts, concept_weight.size(-1))
+        valid_edit = intervene_ids != -1
+        ids = intervene_ids.clamp(0, n_valid - 1).long()
+        vals = intervene_vals.to(concept_weight.dtype)
+        updates = torch.zeros_like(concept_weight)
+        updates.scatter_add_(2, ids, torch.where(valid_edit, vals, torch.zeros_like(vals)))
+        set_mask = torch.zeros_like(concept_weight, dtype=torch.bool)
+        set_mask.scatter_(2, ids, valid_edit)
+        return torch.where(set_mask, updates, concept_weight)
+
+    def topk_with_cutoff(self, tensor: Tensor, dim: int=-1) -> Tensor:
+        """
+        Apply top-k sparsity, zeroing out all but top-k values.
+
+        Args:
+            tensor: Input tensor, typically (B, T, C)
+            dim: Dimension to apply top-k (default: last)
+
+        Returns:
+            Sparse tensor with only top-k values preserved
+        """
+        assert dim == -1 or dim == tensor.dim() - 1
+        if self.topk is None:
+            return tensor
+        padded = tensor.size(dim)
+        n_valid = min(self.n_concepts, padded)
+        if n_valid <= 0:
+            return torch.zeros_like(tensor)
+        x = tensor.narrow(dim, 0, n_valid)
+        kk = min(self.topk, n_valid)
+        topv, topi = torch.topk(x, kk, dim=dim)
+        out = torch.zeros_like(x)
+        out.scatter_(dim, topi, topv)
+        if n_valid < padded:
+            pad_shape = list(out.shape)
+            pad_shape[dim] = padded - n_valid
+            pad_zeros = out.new_zeros(pad_shape)
+            out = torch.cat([out, pad_zeros], dim=dim)
+        return out
+
+    def _compute_weights(self, concept_logits: Tensor, E: Tensor) -> Tensor:
+        """Compute concept weights from logits, with optional top-k sparsity."""
+        apply_topk = self.topk is not None and (not self.is_unknown or self.apply_topk_to_unknown)
+        if apply_topk and self.topk_on_logits:
+            logits_for_weights = self.topk_with_cutoff(concept_logits)
+            weights = torch.sigmoid(logits_for_weights).to(E.dtype)
+            return weights
+        weights = torch.sigmoid(concept_logits).to(E.dtype)
+        if apply_topk and (not self.topk_on_logits):
+            weights = self.topk_with_cutoff(weights)
+        return weights
+
+    @torch.compiler.disable
+    def forward(self, hidden: Tensor, intervene_ids: Tensor | None=None, intervene_vals: Tensor | None=None, return_logits: bool=False, store_hidden: bool=False) -> ConceptHeadOutput:
+        """
+        Forward pass for concept decomposition (inference only, no teacher forcing).
+
+        Args:
+            hidden: Transformer hidden states (B, T, n_embd)
+            intervene_ids: Concept IDs to intervene on (B, T, K_int), -1 = skip
+            intervene_vals: Intervention strength values (B, T, K_int)
+            return_logits: If True, compute full (B, T, C) logits. Forbidden for large heads.
+            store_hidden: If True, store hidden in output for later attribution.
+
+        Returns:
+            ConceptHeadOutput with features, predicted, topk_indices, topk_logits
+        """
+        B, T, _ = hidden.shape
+        has_interventions = intervene_ids is not None and intervene_vals is not None
+        if return_logits:
+            self._check_dense_allowed('return_logits=True')
+        n_valid = self.n_concepts
+        concept_logits: Tensor | None = None
+        concept_weight: Tensor | None = None
+        predicted: Tensor
+        topk_indices: Tensor | None = None
+        topk_logits: Tensor | None = None
+        apply_topk = self.topk is not None and (not self.is_unknown or self.apply_topk_to_unknown)
+        k_features = self.topk_features if self.topk_features is not None else self.topk
+        use_dense_intervention = has_interventions and (not self._is_large)
+        if use_dense_intervention:
+            E = self._get_embedding_weight()[:n_valid]
+            if self.use_attention:
+                query = self.concept_query_projection(hidden)
+                concept_logits = self.blocked_logits(query, E, block_size=self.block_size)
+            else:
+                if self.factorize:
+                    W = self._get_predictor_weight()[:n_valid]
+                    raw_logits = hidden @ W.T
+                else:
+                    raw_logits = self.concept_predictor(hidden)[..., :n_valid]
+                concept_logits = raw_logits.float().clamp(-15, 15)
+            concept_weight = self._compute_weights(concept_logits, E)
+            assert intervene_ids is not None and intervene_vals is not None
+            concept_weight = self._apply_dense_interventions(concept_weight, intervene_ids, intervene_vals)
+            predicted = self.blocked_mix(concept_weight, E, block_size=self.block_size)
+        elif self.factorize:
+            if self.use_attention:
+                query = self.concept_query_projection(hidden)
+                if apply_topk:
+                    predicted, topk_indices, topk_logits = self.attention_features_topk_factorized(query, k=k_features, block_size=self.block_size)
+                else:
+                    predicted = self.attention_block_features_factorized(query, block_size=self.block_size)
+            elif apply_topk:
+                predicted, topk_indices, topk_logits = self.linear_features_topk_factorized(hidden, k=k_features, block_size=self.block_size)
+            else:
+                predicted = self.linear_block_features_factorized(hidden, block_size=self.block_size)
+        elif apply_topk:
+            E = self._get_embedding_weight()[:n_valid]
+            if self.use_attention:
+                query = self.concept_query_projection(hidden)
+                predicted, topk_indices, topk_logits = self.attention_features_topk_streaming(query, E, k=k_features, block_size=self.block_size, topk_on_logits=self.topk_on_logits)
+            else:
+                W = self.concept_predictor.weight[:n_valid]
+                predicted, topk_indices, topk_logits = self.linear_features_topk_streaming(hidden, W, E, k=k_features, block_size=self.block_size, topk_on_logits=self.topk_on_logits)
+        else:
+            E = self._get_embedding_weight()[:n_valid]
+            if self.use_attention:
+                query = self.concept_query_projection(hidden)
+                predicted = self.attention_block_features(query, E, block_size=self.block_size)
+            else:
+                W = self.concept_predictor.weight[:n_valid]
+                predicted = self.linear_block_features(hidden, W, E, block_size=self.block_size)
+        if topk_indices is not None and self.topk is not None and (self.topk_features is not None) and (self.topk_features > self.topk):
+            _, rerank_idx = torch.topk(topk_logits, self.topk, dim=-1)
+            topk_indices = torch.gather(topk_indices, -1, rerank_idx)
+            topk_logits = torch.gather(topk_logits, -1, rerank_idx)
+        if return_logits and (not use_dense_intervention):
+            E = self._get_embedding_weight()[:n_valid]
+            if self.use_attention:
+                query = self.concept_query_projection(hidden)
+                concept_logits = self.blocked_logits(query, E, block_size=self.block_size)
+            else:
+                if self.factorize:
+                    W = self._get_predictor_weight()[:n_valid]
+                    raw_logits = hidden @ W.T
+                else:
+                    raw_logits = self.concept_predictor(hidden)[..., :n_valid]
+                concept_logits = raw_logits.float().clamp(-15, 15)
+            concept_weight = self._compute_weights(concept_logits, E)
+        if not hasattr(self, '_logged_forward_path'):
+            self._logged_forward_path = True
+            path = 'dense_intervention' if use_dense_intervention else 'factorized_topk' if self.factorize and apply_topk else 'factorized_all' if self.factorize else 'streaming_topk' if apply_topk else 'dense_all'
+            logger.info(f"[ConceptHead] {('Unknown' if self.is_unknown else 'Known')} head: path={path}, topk={self.topk}, topk_features={self.topk_features}, n_concepts={self.n_concepts}, factorize={self.factorize}, apply_topk={apply_topk}")
+        if topk_indices is not None and self.topk is not None and (self.topk_features is not None) and (self.topk_features > self.topk):
+            if not hasattr(self, '_logged_topk_slice'):
+                self._logged_topk_slice = True
+                logger.info(f"[ConceptHead] {('Unknown' if self.is_unknown else 'Known')} head: Sliced topk: {self.topk_features} features -> {self.topk} for loss")
+        if has_interventions and (not use_dense_intervention):
+            assert intervene_ids is not None and intervene_vals is not None
+            predicted = self._apply_sparse_interventions(predicted, hidden, intervene_ids, intervene_vals)
+        return ConceptHeadOutput(features=predicted, gt_features=None, logits=concept_logits, predicted=predicted, weights=concept_weight, topk_indices=topk_indices, topk_logits=topk_logits, hidden=hidden.detach() if store_hidden else None)
+
+# ======================================================================
+# steerling/models/interpretable/interpretable_causal_diffusion.py
+# ======================================================================
+
+logger = logging.getLogger(__name__)
+
+class InterpretableCausalDiffusionLM(nn.Module):
+    """
+    Interpretable CausalDiffusionLM with concept decomposition heads.
+
+    Wraps a CausalDiffusionLM and adds:
+    - Known concept head: predicts known concepts from hidden states
+    - Unknown concept head: captures residual features (optional)
+    - Steering via concept interventions
+
+    Args:
+        config: CausalDiffusionConfig (model architecture)
+        concept_config: ConceptConfig (concept decomposition)
+        vocab_size: Vocabulary size
+    """
+
+    def __init__(self, config: CausalDiffusionConfig, concept_config: ConceptConfig, vocab_size: int):
+        super().__init__()
+        self.config = config
+        self.concept_config = concept_config
+        self.vocab_size = vocab_size
+        self.transformer = CausalDiffusionLM(config, vocab_size)
+        self.known_head = ConceptHead(n_concepts=concept_config.n_concepts, concept_dim=concept_config.concept_dim, n_embd=config.n_embd, is_unknown=False, use_attention=concept_config.use_attention_known, topk=concept_config.topk_known, topk_features=concept_config.topk_known_features, block_size=concept_config.block_size, pad_multiple=concept_config.pad_multiple, store_unknown_weights=False, apply_topk_to_unknown=False, topk_on_logits=concept_config.topk_on_logits)
+        if concept_config.use_unknown:
+            if concept_config.n_unknown_concepts is None:
+                raise ValueError('n_unknown_concepts must be set when use_unknown=True')
+            self.unknown_head: ConceptHead | None = ConceptHead(n_concepts=concept_config.n_unknown_concepts, concept_dim=concept_config.concept_dim, n_embd=config.n_embd, is_unknown=True, use_attention=concept_config.use_attention_unknown, topk=concept_config.unknown_topk, block_size=concept_config.block_size, pad_multiple=concept_config.pad_multiple, store_unknown_weights=False, apply_topk_to_unknown=concept_config.apply_topk_to_unknown, topk_on_logits=concept_config.topk_on_logits, factorize=concept_config.factorize_unknown, factorize_rank=concept_config.factorize_rank)
+        else:
+            self.unknown_head = None
+
+    def forward(self, input_ids: Tensor, *, input_embeds: Tensor | None=None, intervene_known_ids: Tensor | None=None, intervene_known_vals: Tensor | None=None, intervene_unknown_ids: Tensor | None=None, intervene_unknown_vals: Tensor | None=None, minimal_output: bool=False, position_injection: Tensor | None=None, steering_inject_layer: int | None=None, steering_inject_alpha: float=1.0, unknown_topk: int=64) -> tuple[Tensor, InterpretableOutput]:
+        """
+        Forward pass with concept decomposition.
+
+        Args:
+            input_ids: Token IDs (B, T). May contain mask tokens.
+            input_embeds: Pre-computed embeddings (B, T, D). Overrides input_ids.
+            intervene_known_ids: Known concept IDs to intervene (B, T, K_int)
+            intervene_known_vals: Intervention values for known (B, T, K_int)
+            intervene_unknown_ids: Unknown concept IDs to intervene (B, T, K_int)
+            intervene_unknown_vals: Intervention values for unknown (B, T, K_int)
+            minimal_output: If True, skip some expensive computations
+            position_injection: Per-position steering injection (B, T, D)
+            steering_inject_layer: Inject at layers >= this
+            steering_inject_alpha: Injection strength
+            unknown_topk: Top-k for unknown head attribution
+
+        Returns:
+            logits: LM logits (B, T, V)
+            outputs: InterpretableOutput with all decomposition components
+        """
+        need_dense_logits = not minimal_output
+        if position_injection is not None and steering_inject_layer is not None:
+            hidden = self._forward_with_injection(input_ids, input_embeds, position_injection, steering_inject_layer, steering_inject_alpha)
+        else:
+            hidden = self.transformer(input_ids, input_embeds=input_embeds, return_hidden=True)
+        known_out: ConceptHeadOutput = self.known_head(hidden, intervene_ids=intervene_known_ids, intervene_vals=intervene_known_vals, return_logits=need_dense_logits)
+        known_features = known_out.features.to(hidden.dtype)
+        unk = hidden - known_features.detach()
+        unk_for_lm: Tensor = unk
+        unknown_out: ConceptHeadOutput | None = None
+        unk_hat: Tensor | None = None
+        if self.unknown_head is not None:
+            unknown_out = self.unknown_head(hidden.detach(), intervene_ids=intervene_unknown_ids, intervene_vals=intervene_unknown_vals, return_logits=not minimal_output and (not self.unknown_head._is_large))
+            assert unknown_out is not None
+            unk_hat = unknown_out.features.to(hidden.dtype)
+            unk_for_lm = unk_hat.detach()
+        epsilon_true = None
+        if self.unknown_head is not None and unk_hat is not None:
+            epsilon_true = hidden.detach() - (known_out.predicted + unk_hat)
+        epsilon = None
+        if self.concept_config.use_epsilon_correction and intervene_known_ids is None:
+            epsilon = hidden - (unk_for_lm + known_features)
+            unk_for_lm = unk_for_lm + epsilon
+        composed = unk_for_lm + known_features
+        logits = self.transformer.lm_head(composed)
+        _unk_topk_indices = unknown_out.topk_indices if unknown_out else None
+        _unk_topk_logits = unknown_out.topk_logits if unknown_out else None
+        if not minimal_output and self.unknown_head is not None and (unknown_out is not None) and (_unk_topk_indices is None) and (unknown_topk > 0):
+            with torch.no_grad():
+                _unk_topk_indices, _unk_topk_logits = self._compute_unknown_topk(hidden, unknown_topk)
+        outputs = InterpretableOutput(hidden=hidden, known_features=known_features, known_logits=known_out.logits, known_gt_features=known_out.gt_features, known_predicted=known_out.predicted, known_weights=known_out.weights, known_topk_indices=known_out.topk_indices, known_topk_logits=known_out.topk_logits, unk=unk, unk_hat=unk_hat, unk_for_lm=unk_for_lm, unknown_logits=unknown_out.logits if unknown_out else None, unknown_weights=unknown_out.weights if unknown_out else None, unknown_topk_indices=_unk_topk_indices, unknown_topk_logits=_unk_topk_logits, composed=composed, epsilon=epsilon, epsilon_true=epsilon_true)
+        return (logits, outputs)
+
+    def _compute_unknown_topk(self, hidden: Tensor, unknown_topk: int) -> tuple[Tensor | None, Tensor | None]:
+        """Compute unknown head top-k indices for attribution."""
+        assert self.unknown_head is not None
+        if self.unknown_head.factorize:
+            if self.unknown_head.use_attention:
+                _query = self.unknown_head.concept_query_projection(hidden.detach())
+                _, indices, logits = self.unknown_head.attention_features_topk_factorized(_query, k=unknown_topk, block_size=self.unknown_head.block_size)
+            else:
+                _, indices, logits = self.unknown_head.linear_features_topk_factorized(hidden.detach(), k=unknown_topk, block_size=self.unknown_head.block_size)
+        else:
+            _E = self.unknown_head._get_embedding_weight()[:self.unknown_head.n_concepts]
+            if self.unknown_head.use_attention:
+                _query = self.unknown_head.concept_query_projection(hidden.detach())
+                _, indices, logits = self.unknown_head.attention_features_topk_streaming(_query, _E, k=unknown_topk, block_size=self.unknown_head.block_size)
+            else:
+                _W = self.unknown_head.concept_predictor.weight[:self.unknown_head.n_concepts]
+                _, indices, logits = self.unknown_head.linear_features_topk_streaming(hidden.detach(), _W, _E, k=unknown_topk, block_size=self.unknown_head.block_size)
+        return (indices, logits)
+
+    def _forward_with_injection(self, input_ids: Tensor, input_embeds: Tensor | None, position_injection: Tensor, inject_layer: int, inject_alpha: float) -> Tensor:
+        """Forward through transformer with steering injection at specified layers."""
+        x = input_embeds if input_embeds is not None else self.transformer.tok_emb(input_ids)
+        for i, block in enumerate(self.transformer.blocks):
+            x = block(x)
+            if i + 1 >= inject_layer:
+                x = x + inject_alpha * position_injection
+        x = self.transformer.ln_f(x)
+        return x
+
+    @torch.no_grad()
+    def intervene(self, input_ids: Tensor, known: dict[int, float] | None=None, unknown: dict[int, float] | None=None, positions: Tensor | None=None) -> tuple[Tensor, InterpretableOutput]:
+        """
+        Run inference with concept interventions.
+
+        Args:
+            input_ids: Input token IDs (B, T)
+            known: Dict mapping known concept IDs to intervention strengths
+            unknown: Dict mapping unknown concept IDs to intervention strengths
+            positions: Bool mask of positions to intervene (B, T). Default: all.
+
+        Returns:
+            logits: LM logits (B, T, V)
+            outputs: InterpretableOutput
+        """
+        B, T = input_ids.shape
+        device = input_ids.device
+        if positions is None:
+            positions = torch.ones(B, T, dtype=torch.bool, device=device)
+        int_known_ids, int_known_vals = (None, None)
+        if known is not None and len(known) > 0:
+            int_known_ids, int_known_vals = self._build_intervention_tensors(known, B, T, positions, device)
+        int_unknown_ids, int_unknown_vals = (None, None)
+        if unknown is not None and len(unknown) > 0:
+            int_unknown_ids, int_unknown_vals = self._build_intervention_tensors(unknown, B, T, positions, device)
+        return self(input_ids, intervene_known_ids=int_known_ids, intervene_known_vals=int_known_vals, intervene_unknown_ids=int_unknown_ids, intervene_unknown_vals=int_unknown_vals, minimal_output=False)
+
+    @staticmethod
+    def _build_intervention_tensors(interventions: dict[int, float], B: int, T: int, positions: Tensor, device: torch.device) -> tuple[Tensor, Tensor]:
+        """Build intervention tensors for concept steering."""
+        K = len(interventions)
+        concept_ids = list(interventions.keys())
+        directions = list(interventions.values())
+        ids = torch.full((B, T, K), -1, dtype=torch.long, device=device)
+        vals = torch.zeros((B, T, K), dtype=torch.float32, device=device)
+        concept_tensor = torch.tensor(concept_ids, device=device)
+        direction_tensor = torch.tensor(directions, dtype=torch.float32, device=device)
+        n_active = int(positions.sum().item())
+        ids[positions] = concept_tensor.unsqueeze(0).expand(n_active, -1)
+        vals[positions] = direction_tensor.unsqueeze(0).expand(n_active, -1)
+        return (ids, vals)
+
+    def get_num_params(self, non_embedding: bool=True) -> int:
+        n_params = sum((p.numel() for p in self.parameters()))
+        if non_embedding and hasattr(self.transformer, 'tok_emb'):
+            n_params -= self.transformer.tok_emb.weight.numel()
+        return n_params
+from transformers import PreTrainedModel
+from .configuration_steerling import SteerlingConfig
+
+
+# CausalDiffusionLM is the backbone — alias to HF-friendly name
+SteerlingBackbone = CausalDiffusionLM
+
+
+class SteerlingForCausalLM(PreTrainedModel):
+    config_class = SteerlingConfig
+    supports_gradient_checkpointing = False
+    _tied_weights_keys = ["transformer.lm_head.weight"]
+
+    def __init__(self, config: SteerlingConfig):
+        super().__init__(config)
+        # SteerlingConfig has all fields from both arch and concept configs
+        self.concept_config = config
+        self.transformer = SteerlingBackbone(config, config.vocab_size)
+        self.known_head = ConceptHead(
+            n_concepts=config.n_concepts,
+            concept_dim=config.concept_dim,
+            n_embd=config.n_embd,
+            is_unknown=False,
+            use_attention=config.use_attention_known,
+            topk=config.topk_known,
+            topk_features=config.topk_known_features,
+            block_size=config.concept_block_size,
+            pad_multiple=config.pad_multiple,
+            store_unknown_weights=False,
+            apply_topk_to_unknown=False,
+            topk_on_logits=config.topk_on_logits,
+            factorize=False,
+        )
+        if config.use_unknown:
+            self.unknown_head = ConceptHead(
+                n_concepts=config.n_unknown_concepts,
+                concept_dim=config.concept_dim,
+                n_embd=config.n_embd,
+                is_unknown=True,
+                use_attention=config.use_attention_unknown,
+                topk=config.unknown_topk,
+                block_size=config.concept_block_size,
+                pad_multiple=config.pad_multiple,
+                store_unknown_weights=config.store_unknown_weights,
+                apply_topk_to_unknown=config.apply_topk_to_unknown,
+                topk_on_logits=config.topk_on_logits,
+                factorize=config.factorize_unknown,
+                factorize_rank=config.factorize_rank,
+            )
+        else:
+            self.unknown_head = None
+        self.post_init()
+
+    def _init_weights(self, module):
+        pass
+
+    def _tie_weights(self):
+        if self.config.weight_sharing:
+            self.transformer.lm_head.weight = self.transformer.tok_emb.weight
+
+    def forward(self, input_ids=None, **kwargs):
+        if self.config.interpretable:
+            return InterpretableCausalDiffusionLM.forward(self, input_ids, **kwargs)
+        else:
+            kwargs.pop('minimal_output', None)
+            return CausalDiffusionLM.forward(self, input_ids, **kwargs)