transformer_lens.model_bridge.supported_architectures.deepseek_v4 module

DeepSeek V4 architecture adapter.

DeepSeek V4 replaces V2/V3’s MLA path with a hybrid local/compressed attention stack and keeps hc_mult residual streams alive between blocks through manifold-constrained Hyper-Connections (mHC). The adapter delegates those architecture-specific calculations to Transformers while exposing the modules that are useful for interpretability: mHC collapse/mix tensors, compressed KV states and masks, Lightning Indexer selections, attention projections, and MoE routing/expert outputs.

class transformer_lens.model_bridge.supported_architectures.deepseek_v4.DeepSeekV4ArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Adapter for DeepseekV4ForCausalLM (Flash and Pro variants).

applicable_phases: list[int] = [2, 4]

prepare_loading(model_name: str, model_kwargs: dict) → None

Force eager attention so the delegated attention path is deterministic.

prepare_model(hf_model: Any) → None

Force eager attention on a pre-loaded model before installing bridges.

class transformer_lens.model_bridge.supported_architectures.deepseek_v4.DeepseekV4BlockBridge(name: str, config: Any | None = None, submodules: Dict[str, GeneralizedComponent] | None = None, hook_alias_overrides: Dict[str, str] | None = None)

Bases: BlockBridge

Block bridge whose input/output hooks carry the full mHC stream stack.

Standard residual aliases are intentionally omitted: V4’s block boundary is four-dimensional, and presenting it as a conventional single residual stream would make otherwise-valid patching code silently target the wrong tensor. The collapsed attention/MLP inputs are available at attn_hc.hook_out and mlp_hc.hook_out respectively.

hook_aliases: dict[str, str | list[str]] = {}

hook_out_is_single_residual_stream: bool = False

maintain_native_attention: bool = True

class transformer_lens.model_bridge.supported_architectures.deepseek_v4.DeepseekV4CompressorBridge(name: str, config: Any | None = None, submodules: Dict[str, GeneralizedComponent] | None = None, optional: bool = False)

Bases: GeneralizedComponent

Bridge CSA/HCA compression and expose compressed KV plus block bias.

forward(*args: Any, **kwargs: Any) → Any

Run the native compressor, preserving and hooking both outputs.

class transformer_lens.model_bridge.supported_architectures.deepseek_v4.DeepseekV4HyperConnectionBridge(name: str, config: Any | None = None, submodules: Dict[str, GeneralizedComponent] | None = None)

Bases: GeneralizedComponent

Bridge an mHC module without discarding its three distinct outputs.

hook_in sees the full [batch, pos, hc_mult, d_model] residual stack. hook_post and hook_comb expose the learned expansion and stream-mix weights, while hook_out exposes the collapsed conventional residual that enters attention or the MLP.

forward(*args: Any, **kwargs: Any) → Any

Run the native mHC module and hook each returned tensor separately.