nnterp · encyclopedia deepseek_v4 · nnterp.families.deepseek_v4

family deepseek_v4

DeepSeek-V4

Four parallel residual streams, which a learned matrix mixes around each sublayer while the sublayer reads one weighted sum of them and writes into all four; multi-query attention over a sliding window plus compressed keys; and a mixture of experts on every block.

The block

01 · the residual stream

One block of the stack. Hover any part for its nnterp name; click to pin the card. The slider picks a block; this family's blocks differ by layer_types.

embed_tokenstoken_embeddings
→
layers× 43 · 4 streams
→
hc_headstreams → one
→
normfinal_norm_output
→
lm_headlm_head_output
→
logitsthe output
i = 0

identity mlp_combᵀ · (attention_combᵀ · layers[i].input + attention_post ⊗ self_attn.attention_output) + mlp_post ⊗ mlp.mlp_output == layer_output

The API

02 · support()

Values, by host

⚠︎ marks a value this checkpoint has only under some load: hover it for the condition.

model DeepseekV4ForCausalLMlayoutwhat it isread at
logits Logits [batch seq vocab_size] The logits the model returns, after anything it does past lm_head (a softcap, a scale)

the module's own output

token_embeddings Residual [batch seq hidden_size] The embedding module's output; layers[0].input is what enters the first block

the child module embed_tokens's output

final_norm_output Residual [batch seq hidden_size] The residual stream after the final norm, what lm_head reads

the child module norm's output

lm_head_output Logits [batch seq vocab_size] The head's logits before anything the model does past lm_head (a softcap, a scale)

the child module lm_head's output

next_token_probs NextTokenProbs [batch vocab_size] The next-token distribution at the last position; derived, read-only

the module's own output

input_ids Tokens [batch seq] The token ids the model was called with

the module's own input

attention_mask Tokens [batch seq] The attention mask the model was called with; zeros are padding

the module's own input

input_size scalar [batch, seq] of the current call; read-only

the module's own input

model.layers[i] DeepseekV4DecoderLayerlayoutwhat it isread at
layer_input Streams [batch seq streams hidden_size] The residual streams entering the block

the module's own input

⚠︎attention_norm_output scalar Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet

nowhere: the family does not serve it

⚠︎layer_mid scalar Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet

nowhere: the family does not serve it

⚠︎mlp_norm_output scalar Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet

nowhere: the family does not serve it

⚠︎mlp_input scalar Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet

nowhere: the family does not serve it

layer_output Streams [batch seq streams hidden_size] The residual streams leaving the block

the module's own output

attention_post StreamWeights [batch seq streams] How much of the attention's output each stream receives

the child module attn_hc's output, element 0

attention_comb StreamMixing [batch seq streams streams] The matrix mixing the streams around the attention, applied transposed

the child module attn_hc's output, element 1

mlp_post StreamWeights [batch seq streams] How much of the MLP's output each stream receives

the child module ffn_hc's output, element 0

mlp_comb StreamMixing [batch seq streams streams] The matrix mixing the streams around the MLP, applied transposed

the child module ffn_hc's output, element 1

model.layers[i].self_attn DeepseekV4Attentionlayoutwhat it isread at
attention_queries Queries [batch num_heads seq qk_head_dim] The queries entering attention

inside the forward: the arguments of attention_interface_1, element 1

attention_keys Keys [batch num_kv_heads seq qk_head_dim] The keys entering attention

inside the forward: the arguments of attention_interface_1, element 2

attention_values Values [batch num_kv_heads seq head_dim] The values entering attention

inside the forward: the arguments of attention_interface_1, element 3

attention_scores Pattern [batch num_heads query key] The attention scores entering the softmax, masked

inside the forward: the output of attn_weights_1 inside attention_interface_1

⚠︎attention_premix scalar Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side

nowhere: the family does not serve it

attention_output Residual [batch seq hidden_size] What the attention adds to the residual stream

the module's own output

attention_probabilities Pattern [batch num_heads query key] The attention pattern the values are mixed with

inside the forward: the output of nn_functional_dropout_0 inside attention_interface_1

attention_head_outputs HeadOutputs [batch seq num_heads head_dim] The per-head outputs before the output projection

inside the forward: the output of attn_output_0

model.layers[i].mlp DeepseekV4SparseMoeBlocklayoutwhat it isread at
⚠︎mlp_activation Neurons [batch seq intermediate_size] The activation function's output

the child module act_fn's output

⚠︎mlp_neurons Neurons [batch seq intermediate_size] What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise

the child module down_proj's input

mlp_output Residual [batch seq hidden_size] What the MLP adds to the residual stream

the module's own output

router_logits RouterLogits [batch seq num_experts] The router's logits, one per expert, before the scoring

inside router's forward: the output of F_linear_0

expert_weights ExpertWeights [batch seq top_k] The weight each selected expert's output is scaled by

the child module experts's inputs, element 2

expert_indices ExpertIndices [batch seq top_k] The expert each routing slot sends the token to

the child module experts's inputs, element 1

expert_outputs ExpertOutputs [batch seq top_k hidden_size] Each routing slot's weighted expert output

inside experts's forward: the output of weighted_out_view_0 inside experts_forward_1

routed_output Residual [batch seq hidden_size] The routed experts' combined output

the child module experts's output

shared_expert_output Residual [batch seq hidden_size] The shared expert's output

the child module shared_experts's output

Printout

What print(model) shows on deepseek-ai/DeepSeek-V4-Flash: the native modules, standard names and every standard value with its layout.

DeepseekV4ForCausalLM(
  (lm_head): Linear(in_features=4096, out_features=129280, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(129280, 4096)
  (layers): ModuleList(
    (0-1): 2 x DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4HashRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (2): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4CSACompressor(
          (kv_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
          (indexer): DeepseekV4Indexer(
            (kv_proj): Linear(in_features=4096, out_features=256, bias=False)
            (gate_proj): Linear(in_features=4096, out_features=256, bias=False)
            (kv_norm): DeepseekV4RMSNorm((128,), eps=1e-06)
            (q_b_proj): Linear(in_features=1024, out_features=8192, bias=False)
            (rotary_emb): DeepseekV4RotaryEmbedding()
            (scorer): DeepseekV4IndexerScorer(
              (weights_proj): Linear(in_features=4096, out_features=64, bias=False)
            )
          )
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4HashRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (3): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4HCACompressor(
          (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=512, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (4): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4CSACompressor(
          (kv_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
          (indexer): DeepseekV4Indexer(
            (kv_proj): Linear(in_features=4096, out_features=256, bias=False)
            (gate_proj): Linear(in_features=4096, out_features=256, bias=False)
            (kv_norm): DeepseekV4RMSNorm((128,), eps=1e-06)
            (q_b_proj): Linear(in_features=1024, out_features=8192, bias=False)
            (rotary_emb): DeepseekV4RotaryEmbedding()
            (scorer): DeepseekV4IndexerScorer(
              (weights_proj): Linear(in_features=4096, out_features=64, bias=False)
            )
          )
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (5): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4HCACompressor(
          (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=512, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (6): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4CSACompressor(
          (kv_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
          (indexer): DeepseekV4Indexer(
            (kv_proj): Linear(in_features=4096, out_features=256, bias=False)
            (gate_proj): Linear(in_features=4096, out_features=256, bias=False)
            (kv_norm): DeepseekV4RMSNorm((128,), eps=1e-06)
            (q_b_proj): Linear(in_features=1024, out_features=8192, bias=False)
            (rotary_emb): DeepseekV4RotaryEmbedding()
            (scorer): DeepseekV4IndexerScorer(
              (weights_proj): Linear(in_features=4096, out_features=64, bias=False)
            )
          )
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (7): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4HCACompressor(
          (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=512, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (8): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4CSACompressor(
          (kv_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
          (indexer): DeepseekV4Indexer(
            (kv_proj): Linear(in_features=4096, out_features=256, bias=False)
            (gate_proj): Linear(in_features=4096, out_features=256, bias=False)
            (kv_norm): DeepseekV4RMSNorm((128,), eps=1e-06)
            (q_b_proj): Linear(in_features=1024, out_features=8192, bias=False)
            (rotary_emb): DeepseekV4RotaryEmbedding()
            (scorer): DeepseekV4IndexerScorer(
              (weights_proj): Linear(in_features=4096, out_features=64, bias=False)
            )
          )
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (9): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4HCACompressor(
          (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=512, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (10): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4CSACompressor(
          (kv_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
          (indexer): DeepseekV4Indexer(
            (kv_proj): Linear(in_features=4096, out_features=256, bias=False)
            (gate_proj): Linear(in_features=4096, out_features=256, bias=False)
            (kv_norm): DeepseekV4RMSNorm((128,), eps=1e-06)
            (q_b_proj): Linear(in_features=1024, out_features=8192, bias=False)
            (rotary_emb): DeepseekV4RotaryEmbedding()
            (scorer): DeepseekV4IndexerScorer(
              (weights_proj): Linear(in_features=4096, out_features=64, bias=False)
            )
          )
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (11): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4HCACompressor(
          (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=512, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (12): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4CSACompressor(
          (kv_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
          (indexer): DeepseekV4Indexer(
            (kv_proj): Linear(in_features=4096, out_features=256, bias=False)
            (gate_proj): Linear(in_features=4096, out_features=256, bias=False)
            (kv_norm): DeepseekV4RMSNorm((128,), eps=1e-06)
            (q_b_proj): Linear(in_features=1024, out_features=8192, bias=False)
            (rotary_emb): DeepseekV4RotaryEmbedding()
            (scorer): DeepseekV4IndexerScorer(
              (weights_proj): Linear(in_features=4096, out_features=64, bias=False)
            )
          )
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (13): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4HCACompressor(
          (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=512, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (14): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4CSACompressor(
          (kv_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
          (indexer): DeepseekV4Indexer(
            (kv_proj): Linear(in_features=4096, out_features=256, bias=False)
            (gate_proj): Linear(in_features=4096, out_features=256, bias=False)
            (kv_norm): DeepseekV4RMSNorm((128,), eps=1e-06)
            (q_b_proj): Linear(in_features=1024, out_features=8192, bias=False)
            (rotary_emb): DeepseekV4RotaryEmbedding()
            (scorer): DeepseekV4IndexerScorer(
              (weights_proj): Linear(in_features=4096, out_features=64, bias=False)
            )
          )
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (15): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4HCACompressor(
          (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=512, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (16): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4CSACompressor(
          (kv_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
          (indexer): DeepseekV4Indexer(
            (kv_proj): Linear(in_features=4096, out_features=256, bias=False)
            (gate_proj): Linear(in_features=4096, out_features=256, bias=False)
            (kv_norm): DeepseekV4RMSNorm((128,), eps=1e-06)
            (q_b_proj): Linear(in_features=1024, out_features=8192, bias=False)
            (rotary_emb): DeepseekV4RotaryEmbedding()
            (scorer): DeepseekV4IndexerScorer(
              (weights_proj): Linear(in_features=4096, out_features=64, bias=False)
            )
          )
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (17): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4HCACompressor(
          (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=512, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (18): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4CSACompressor(
          (kv_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
          (indexer): DeepseekV4Indexer(
            (kv_proj): Linear(in_features=4096, out_features=256, bias=False)
            (gate_proj): Linear(in_features=4096, out_features=256, bias=False)
            (kv_norm): DeepseekV4RMSNorm((128,), eps=1e-06)
            (q_b_proj): Linear(in_features=1024, out_features=8192, bias=False)
            (rotary_emb): DeepseekV4RotaryEmbedding()
            (scorer): DeepseekV4IndexerScorer(
              (weights_proj): Linear(in_features=4096, out_features=64, bias=False)
            )
          )
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (19): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4HCACompressor(
          (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=512, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (20): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4CSACompressor(
          (kv_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
          (indexer): DeepseekV4Indexer(
            (kv_proj): Linear(in_features=4096, out_features=256, bias=False)
            (gate_proj): Linear(in_features=4096, out_features=256, bias=False)
            (kv_norm): DeepseekV4RMSNorm((128,), eps=1e-06)
            (q_b_proj): Linear(in_features=1024, out_features=8192, bias=False)
            (rotary_emb): DeepseekV4RotaryEmbedding()
            (scorer): DeepseekV4IndexerScorer(
              (weights_proj): Linear(in_features=4096, out_features=64, bias=False)
            )
          )
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (21): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4HCACompressor(
          (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=512, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (22): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4CSACompressor(
          (kv_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
          (indexer): DeepseekV4Indexer(
            (kv_proj): Linear(in_features=4096, out_features=256, bias=False)
            (gate_proj): Linear(in_features=4096, out_features=256, bias=False)
            (kv_norm): DeepseekV4RMSNorm((128,), eps=1e-06)
            (q_b_proj): Linear(in_features=1024, out_features=8192, bias=False)
            (rotary_emb): DeepseekV4RotaryEmbedding()
            (scorer): DeepseekV4IndexerScorer(
              (weights_proj): Linear(in_features=4096, out_features=64, bias=False)
            )
          )
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (23): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4HCACompressor(
          (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=512, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (24): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4CSACompressor(
          (kv_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
          (indexer): DeepseekV4Indexer(
            (kv_proj): Linear(in_features=4096, out_features=256, bias=False)
            (gate_proj): Linear(in_features=4096, out_features=256, bias=False)
            (kv_norm): DeepseekV4RMSNorm((128,), eps=1e-06)
            (q_b_proj): Linear(in_features=1024, out_features=8192, bias=False)
            (rotary_emb): DeepseekV4RotaryEmbedding()
            (scorer): DeepseekV4IndexerScorer(
              (weights_proj): Linear(in_features=4096, out_features=64, bias=False)
            )
          )
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (25): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4HCACompressor(
          (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=512, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (26): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4CSACompressor(
          (kv_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
          (indexer): DeepseekV4Indexer(
            (kv_proj): Linear(in_features=4096, out_features=256, bias=False)
            (gate_proj): Linear(in_features=4096, out_features=256, bias=False)
            (kv_norm): DeepseekV4RMSNorm((128,), eps=1e-06)
            (q_b_proj): Linear(in_features=1024, out_features=8192, bias=False)
            (rotary_emb): DeepseekV4RotaryEmbedding()
            (scorer): DeepseekV4IndexerScorer(
              (weights_proj): Linear(in_features=4096, out_features=64, bias=False)
            )
          )
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (27): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4HCACompressor(
          (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=512, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (28): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4CSACompressor(
          (kv_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
          (indexer): DeepseekV4Indexer(
            (kv_proj): Linear(in_features=4096, out_features=256, bias=False)
            (gate_proj): Linear(in_features=4096, out_features=256, bias=False)
            (kv_norm): DeepseekV4RMSNorm((128,), eps=1e-06)
            (q_b_proj): Linear(in_features=1024, out_features=8192, bias=False)
            (rotary_emb): DeepseekV4RotaryEmbedding()
            (scorer): DeepseekV4IndexerScorer(
              (weights_proj): Linear(in_features=4096, out_features=64, bias=False)
            )
          )
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (29): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4HCACompressor(
          (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=512, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (30): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4CSACompressor(
          (kv_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
          (indexer): DeepseekV4Indexer(
            (kv_proj): Linear(in_features=4096, out_features=256, bias=False)
            (gate_proj): Linear(in_features=4096, out_features=256, bias=False)
            (kv_norm): DeepseekV4RMSNorm((128,), eps=1e-06)
            (q_b_proj): Linear(in_features=1024, out_features=8192, bias=False)
            (rotary_emb): DeepseekV4RotaryEmbedding()
            (scorer): DeepseekV4IndexerScorer(
              (weights_proj): Linear(in_features=4096, out_features=64, bias=False)
            )
          )
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (31): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4HCACompressor(
          (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=512, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (32): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4CSACompressor(
          (kv_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
          (indexer): DeepseekV4Indexer(
            (kv_proj): Linear(in_features=4096, out_features=256, bias=False)
            (gate_proj): Linear(in_features=4096, out_features=256, bias=False)
            (kv_norm): DeepseekV4RMSNorm((128,), eps=1e-06)
            (q_b_proj): Linear(in_features=1024, out_features=8192, bias=False)
            (rotary_emb): DeepseekV4RotaryEmbedding()
            (scorer): DeepseekV4IndexerScorer(
              (weights_proj): Linear(in_features=4096, out_features=64, bias=False)
            )
          )
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (33): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4HCACompressor(
          (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=512, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (34): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4CSACompressor(
          (kv_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
          (indexer): DeepseekV4Indexer(
            (kv_proj): Linear(in_features=4096, out_features=256, bias=False)
            (gate_proj): Linear(in_features=4096, out_features=256, bias=False)
            (kv_norm): DeepseekV4RMSNorm((128,), eps=1e-06)
            (q_b_proj): Linear(in_features=1024, out_features=8192, bias=False)
            (rotary_emb): DeepseekV4RotaryEmbedding()
            (scorer): DeepseekV4IndexerScorer(
              (weights_proj): Linear(in_features=4096, out_features=64, bias=False)
            )
          )
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (35): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4HCACompressor(
          (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=512, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (36): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4CSACompressor(
          (kv_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
          (indexer): DeepseekV4Indexer(
            (kv_proj): Linear(in_features=4096, out_features=256, bias=False)
            (gate_proj): Linear(in_features=4096, out_features=256, bias=False)
            (kv_norm): DeepseekV4RMSNorm((128,), eps=1e-06)
            (q_b_proj): Linear(in_features=1024, out_features=8192, bias=False)
            (rotary_emb): DeepseekV4RotaryEmbedding()
            (scorer): DeepseekV4IndexerScorer(
              (weights_proj): Linear(in_features=4096, out_features=64, bias=False)
            )
          )
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (37): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4HCACompressor(
          (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=512, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (38): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4CSACompressor(
          (kv_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
          (indexer): DeepseekV4Indexer(
            (kv_proj): Linear(in_features=4096, out_features=256, bias=False)
            (gate_proj): Linear(in_features=4096, out_features=256, bias=False)
            (kv_norm): DeepseekV4RMSNorm((128,), eps=1e-06)
            (q_b_proj): Linear(in_features=1024, out_features=8192, bias=False)
            (rotary_emb): DeepseekV4RotaryEmbedding()
            (scorer): DeepseekV4IndexerScorer(
              (weights_proj): Linear(in_features=4096, out_features=64, bias=False)
            )
          )
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (39): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4HCACompressor(
          (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=512, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (40): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4CSACompressor(
          (kv_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
          (indexer): DeepseekV4Indexer(
            (kv_proj): Linear(in_features=4096, out_features=256, bias=False)
            (gate_proj): Linear(in_features=4096, out_features=256, bias=False)
            (kv_norm): DeepseekV4RMSNorm((128,), eps=1e-06)
            (q_b_proj): Linear(in_features=1024, out_features=8192, bias=False)
            (rotary_emb): DeepseekV4RotaryEmbedding()
            (scorer): DeepseekV4IndexerScorer(
              (weights_proj): Linear(in_features=4096, out_features=64, bias=False)
            )
          )
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (41): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4HCACompressor(
          (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=512, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
    (42): DeepseekV4DecoderLayer(
      (self_attn): DeepseekV4Attention(
        (q_a_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (q_a_norm): DeepseekV4RMSNorm((1024,), eps=1e-06)
        (q_b_proj): Linear(in_features=1024, out_features=32768, bias=False)
        (q_b_norm): DeepseekV4UnweightedRMSNorm()
        (kv_proj): Linear(in_features=4096, out_features=512, bias=False)
        (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
        (o_a_proj): DeepseekV4GroupedLinear(in_features=4096, out_features=8192, bias=False)
        (o_b_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (compressor): DeepseekV4CSACompressor(
          (kv_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (gate_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (kv_norm): DeepseekV4RMSNorm((512,), eps=1e-06)
          (rotary_emb): DeepseekV4RotaryEmbedding()
          (indexer): DeepseekV4Indexer(
            (kv_proj): Linear(in_features=4096, out_features=256, bias=False)
            (gate_proj): Linear(in_features=4096, out_features=256, bias=False)
            (kv_norm): DeepseekV4RMSNorm((128,), eps=1e-06)
            (q_b_proj): Linear(in_features=1024, out_features=8192, bias=False)
            (rotary_emb): DeepseekV4RotaryEmbedding()
            (scorer): DeepseekV4IndexerScorer(
              (weights_proj): Linear(in_features=4096, out_features=64, bias=False)
            )
          )
        )
        (attention_premix): Unavailable: a low-rank output projection: o_a_proj compresses the heads before o_b_proj, so nothing receives them side by side
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
      )
      (mlp): DeepseekV4SparseMoeBlock(
        (router/gate): DeepseekV4TopKRouter(
          (score_fn): SqrtSoftplusActivation()
        )
        (experts): DeepseekV4Experts(
          (act_fn): SiLUActivation()
        )
        (shared_experts): DeepseekV4MLP(
          (gate_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (up_proj): Linear(in_features=4096, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=4096, bias=False)
          (act_fn): SiLUActivation()
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (post_attention_layernorm): DeepseekV4RMSNorm((4096,), eps=1e-06)
      (attn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (ffn_hc): DeepseekV4HyperConnection(
        (input_norm): DeepseekV4UnweightedRMSNorm()
      )
      (attention_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_mid): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_norm_output): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (mlp_input): Unavailable: hyper-connection streams: the sublayers read a mix of the streams, not the stream; the block's inner boundaries are not mapped for this family yet
      (layer_input) -> Streams [batch seq streams hidden_size]: The residual streams entering the block
      (layer_output) -> Streams [batch seq streams hidden_size]: The residual streams leaving the block
      (attention_post) -> StreamWeights [batch seq streams]: How much of the attention's output each stream receives
      (attention_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the attention, applied transposed
      (mlp_post) -> StreamWeights [batch seq streams]: How much of the MLP's output each stream receives
      (mlp_comb) -> StreamMixing [batch seq streams streams]: The matrix mixing the streams around the MLP, applied transposed
    )
  )
  (norm): DeepseekV4RMSNorm((4096,), eps=1e-06)
  (logits) -> Logits [batch seq vocab_size]: The logits the model returns, after anything it does past lm_head (a softcap, a scale)
  (token_embeddings) -> Residual [batch seq hidden_size]: The embedding module's output; layers[0].input is what enters the first block
  (final_norm_output) -> Residual [batch seq hidden_size]: The residual stream after the final norm, what lm_head reads
  (lm_head_output) -> Logits [batch seq vocab_size]: The head's logits before anything the model does past lm_head (a softcap, a scale)
  (next_token_probs) -> NextTokenProbs [batch vocab_size]: The next-token distribution at the last position; derived, read-only
  (input_ids) -> Tokens [batch seq]: The token ids the model was called with
  (attention_mask) -> Tokens [batch seq]: The attention mask the model was called with; zeros are padding
  (input_size): [batch, seq] of the current call; read-only
)

Config

model.num_layers43
model.hidden_size4096
model.num_heads64
model.num_kv_heads1
model.head_dim512
model.qk_head_dim512
model.vocab_size129280
model.intermediate_size2048
architecturesDeepseekV4ForCausalLM
num_hidden_layers43
hidden_size4096
num_attention_heads64
num_key_value_heads1
head_dim512
vocab_size129280
max_position_embeddings1048576
sliding_window128
hidden_actsilu
rms_norm_eps1e-06
rope_theta10000
tie_word_embeddingsFalse
dtypebfloat16
partial_rotary_factor0.125
q_lora_rank1024
qk_rope_head_dim64
moe_intermediate_size2048
n_routed_experts256
num_experts_per_tok6
n_shared_experts1
routed_scaling_factor1.5
norm_topk_probTrue
attention_biasFalse

The Notes

03 · what to know before you run

The block, in order

post_a, comb_a, xa = attn_hc(x)                  # x: [batch, seq, 4, hidden_size]
h   = comb_aᵀ · x + post_a ⊗ self_attn(input_layernorm(xa))
post_f, comb_f, xf = ffn_hc(h)
out = comb_fᵀ · h + post_f ⊗ mlp(post_attention_layernorm(xf))

The residual is hc_mult (4) streams. Before each sublayer a hyper-connection (attn_hc, ffn_hc) reads all four and returns (post, comb, collapsed): collapsed is the sum of the streams, each weighted by a sigmoid computed from all four streams' content, and the sublayer's pre-norm reads it. The sublayer's [batch, seq, hidden_size] output goes into every stream, scaled per stream by post, and the streams themselves are mixed by comb, applied transposed: stream k of the result is Σ_j comb[j, k] · x_j. Llama's norm names hold: input_layernorm before the attention, post_attention_layernorm before the mixture.

layer_output is four streams

layer_output and layers[i].input are [batch, seq, streams, hidden_size] (the Streams layout); attention_output and mlp_output stay [batch, seq, hidden_size]. Block 0 reads the embedding copied into each stream; after it the streams differ. Code written for a rank-3 residual runs here without an error and gives one answer per stream: resid[:, -1] is the last position's four streams, [batch, streams, hidden_size]. Index one stream, or take the stream mean, which is the additive part of the block:

with model.trace(prompt):
    out = model.layers[1].layer_output.save()   # [1, seq, 4, hidden_size]

last = out[:, -1, -1]          # the last position, the last stream: [1, hidden_size]
mean = out[:, -1].mean(-2)     # the last position, the stream mean: [1, hidden_size]

A write to layer_output lands natively; steer adds a [hidden_size] vector to every stream, and skip_layers passes the streams through.

The identity is the block's formula

The block is not a sum. The weights that put each contribution into the streams are the block's own values: attention_post and mlp_post, [batch, seq, streams] (StreamWeights, each in (0, 2)), and attention_comb and mlp_comb, [batch, seq, streams, streams] (StreamMixing), read off attn_hc.output and ffn_hc.output. All four are float32 and writable. The formula holds exactly in float32 on every block of the tiny checkpoint:

layer = model.layers[3]
with model.trace(prompt):
    x = layer.input.save()
    post_a, comb_a = layer.attention_post.save(), layer.attention_comb.save()
    attn = layer.self_attn.attention_output.save()
    post_f, comb_f = layer.mlp_post.save(), layer.mlp_comb.save()
    mlp = layer.mlp.mlp_output.save()
    out = layer.layer_output.save()

h = comb_a.mT @ x + post_a[..., None] * attn[..., None, :]
torch.testing.assert_close(comb_f.mT @ h + post_f[..., None] * mlp[..., None, :], out)

comb is a Sinkhorn projection onto the doubly stochastic matrices: its columns sum to one, so each stream of the result is a weighted average of the streams. Its rows also sum to one, up to the projection's residual, so the stream mean is a sum in which each contribution is weighted by the mean of its post. On the tiny checkpoint it matches to about 2e-6 of the stream's largest value:

mean = x.mean(2) + post_a.mean(-1, keepdim=True) * attn + post_f.mean(-1, keepdim=True) * mlp
torch.testing.assert_close(mean, out.mean(2), rtol=1e-4, atol=1e-4)

Zeroing attention_output removes the attention's write from every stream but leaves the mixing; zeroing attention_post does the same through the weights. A direct logit attribution needs the later blocks' comb and post too.

The readout collapses the streams

The model reads the last block's streams out through model.hc_head, a learned, content-dependent weighted sum over the stream axis, then norm and lm_head. The family's project_on_vocab does the same on a [batch, seq, streams, hidden_size] tensor (or one position's [streams, hidden_size]), so the lens on the last block is logits exactly. A rank-3 tensor (a contribution, one stream) goes through norm and lm_head alone, and one stream's lens is not logits.

with model.trace(prompt):
    resid = model.layers[-1].layer_output.save()
    logits = model.logits.save()

assert torch.equal(model.project_on_vocab(resid), logits)

Multi-query attention whose keys are its values

One key/value head (num_key_value_heads 1) serves all 64 query heads (128 on Pro), each 512 wide. The attention projects one kv per token and passes the same tensor as keys and values, so attention_keys and attention_values are one object: an in-place edit of either edits both, and an assignment gives one a tensor of its own. Queries are normed per head without a weight (q_b_norm) and the scores are scaled by head_dim ** -0.5. Rotary turns the last 64 dimensions of each head (partial_rotary_factor 0.125) in adjacent pairs. The values are the keys, so they carry the rotation too. The module rotates the attention's output back before the grouped output projection (o_a_proj over o_groups groups, 8 on Flash and 16 on Pro, then o_b_proj), and attention_head_outputs is that rotated-back tensor. A learned per-head sink (self_attn.sinks) joins the softmax, so pattern rows sum to less than one. transformers runs this family with eager attention only, so the default load already serves the interior values.

Three attention blocks: a window, and compressed keys after it

Every block attends to the latest sliding_window (128) tokens. config.layer_types adds compressed keys on most blocks: a compressed_sparse_attention block appends one entry per 4 tokens, of which a lightning indexer (self_attn.compressor.indexer) keeps each query's index_topk best (512 on Flash, 1024 on Pro); a heavily_compressed_attention block appends one entry per 128 tokens and keeps all of them. The entries follow the token keys, so the key axis of attention_keys, attention_scores and attention_probabilities is seq + seq // rate long there, and a query sees an entry once it has reached the entry's last token. The indexer runs before the keys are joined, so its selection is read first. On a 13-token prompt the tiny checkpoint's block 2 has 16 keys:

attn = model.layers[2].self_attn                 # compressed_sparse_attention
with model.trace(prompt):
    sel = attn.compressor.indexer.output.save()  # [1, 13, min(index_topk, 3)]; -1: window not closed
    k = attn.attention_keys.save()               # [1, 1, 13 + 13 // 4, 128]

Flash has two sliding_attention blocks (0 and 1), then compressed_sparse_attention and heavily_compressed_attention alternating (21 compressed-sparse, 20 heavily compressed); Pro starts with two heavily compressed blocks and alternates after them (30 compressed-sparse, 31 heavily compressed). The selection has no standard name; within the compressed entries the pattern is zero outside it.

The mixture: square-root-softplus scores, hashed on the first three blocks

Every block's MLP is a mixture: 256 routed experts (384 on Pro), 6 per token, and one shared expert. Every expert, the shared one included, clamps its gate projection to at most swiglu_limit (10) and its up projection to ±10. router_logits are the logits before the scoring, sqrt(softplus(logits)) (scoring_func sqrtsoftplus). On a moe block the token takes the 6 best scores plus a selection bias (router.e_score_correction_bias), with no expert groups; on the hash_moe blocks (0 to 2) the experts are a fixed table of the token id (router.tid2eid[input_ids]), so a run needs input_ids, and a written router_logits changes expert_weights but not expert_indices. On both, expert_weights are the chosen experts' scores without the bias, renormalized to one and multiplied by routed_scaling_factor (1.5 on Flash, 2.5 on Pro).

import torch.nn.functional as F

moe = model.layers[4].mlp
with model.trace(prompt):
    logits = moe.router_logits.save()
    w = moe.expert_weights.save()
    idx = moe.expert_indices.save()
    routed = moe.routed_output.save()
    shared = moe.shared_expert_output.save()
    out = moe.mlp_output.save()

chosen = F.softplus(logits).sqrt().gather(-1, idx)
scale = model.config.routed_scaling_factor
torch.testing.assert_close(w, chosen / chosen.sum(-1, keepdim=True) * scale)
torch.testing.assert_close(routed + shared, out)

moe.SCORING is "hash" on a hash block and "sqrtsoftplus" on the others. Ablating an expert is zeroing the weight of every slot that chose it; its effect on the block is read on the streams, one stream or their mean:

e = idx[0, -1, 0].item()                       # an expert the last token uses
with model.trace(prompt):
    moe.expert_weights = moe.expert_weights.masked_fill(moe.expert_indices == e, 0)
    ablated = model.layers[4].layer_output.save()    # [1, seq, 4, hidden_size]

ablated[:, -1, -1]          # the last stream at the last position
ablated[:, -1].mean(-2)     # the stream mean at the last position

From the family module

04 · nnterp/families/deepseek_v4.py

DeepSeek-V4 (DeepseekV4ForCausalLM): a hyper-connection residual, compressed attention, all-MoE.

model.{embed_tokens, layers[i].{attn_hc, input_layernorm, self_attn, ffn_hc, post_attention_layernorm, mlp}, hc_head, norm} and lm_head.

The residual is hc_mult parallel streams, [batch, seq, streams, hidden_size], between every pair of blocks (the model copies the embedding into each stream before block 0). Each sublayer reads one collapse of the streams and writes back into all of them through a hyper-connection (attn_hc before the attention, ffn_hc before the MLP), which returns (post, comb, collapsed): collapsed is the pre-weighted sum of the streams that the sublayer's norm reads, post [batch, seq, streams] how much of the sublayer's output each stream receives, comb [batch, seq, streams, streams] a doubly stochastic matrix mixing the streams, consumed transposed. The block, per token::

h   = attention_combᵀ · x  + attention_post ⊗ self_attn(input_layernorm(collapsed_a))
out = mlp_combᵀ       · h  + mlp_post       ⊗ mlp(post_attention_layernorm(collapsed_f))
  • layer_output is the block's own stream tensor, under the Streams layout, and so is layers[i].input; writes land natively, and skip_layers and steer work unchanged (a [hidden_size] steering vector broadcasts over the streams).
  • The contributions are the sublayers' own outputs, [batch, seq, hidden_size], unscaled: attention_output is self_attn.output[0], mlp_output the mixture's return (routed plus shared experts). The weights that put them into the streams are the block's attention_post, attention_comb, mlp_post and mlp_comb, read off the hyper-connections' outputs. The identity is the block's formula above, not a sum; its stream mean is additive (comb is doubly stochastic) up to the Sinkhorn projection's residual, about 2e-6 relative in float32: layer_output.mean(2) == input.mean(2) + attention_post.mean(-1, keepdim=True) * attention_output + mlp_post.mean(-1, keepdim=True) * mlp_output.
  • The readout collapses the streams once more: norm(hc_head(streams)), hc_head a learned, content-dependent weighting. The family's project_on_vocab does the same, so the lens on the last block is logits exactly.

The attention runs the shared interface, with a per-head sink (s_aux), so the scores are GPT-OSS's binding. It is multi-query (one key/value head), and the keys and values are one tensor: the same object is passed as both, so an in-place edit of one edits the other. On a compressed block (compressed_sparse_attention, heavily_compressed_attention) the compressor's entries are concatenated after the token keys once the prompt is compress_rates[type] tokens long, so the key axis is longer than the sequence there. The interface's output keeps a rotation on its rotary slice (the values are the rotated keys); the module rotates it back before the grouped output projection, and attention_head_outputs is that de-rotated tensor, what the projection reads. Not latent attention: the plain root sizes hold.

Every MLP is a mixture (DeepseekV4SparseMoeBlock, a Moe; its router, gate, is aliased router); on the hash_moe blocks the router picks experts from the token ids (tid2eid[input_ids]), so a run needs input_ids. There the logits only weight the experts the ids chose: writing router_logits changes expert_weights, not expert_indices.