Step-3.7-Flash

self.num_key_value_groups = self.num_heads // config.num_key_value_heads

    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)

self.is_local_attention = config.layer_types[layer_idx] == "sliding_attention"
        self.sliding_window = config.sliding_window if self.is_local_attention else None

causal_mask_mapping = {
                "full_attention": create_causal_mask(**mask_kwargs),
                "sliding_attention": create_sliding_window_causal_mask(**sliding_mask_kwargs),
            }

attn_output = attn_output * gate_states.unsqueeze(-1).sigmoid()

self.gate_per_head = config.gating is True or config.gating == "per-head"
        g_proj_dim = self.num_heads if self.gate_per_head else self.num_heads * self.head_dim
        self.g_proj = nn.Linear(config.hidden_size, g_proj_dim, bias=False)

query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
        key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)

partial_rotary_factor = config.rope_parameters[layer_type].get("partial_rotary_factor", 1.0)
        head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
        dim = int(head_dim * partial_rotary_factor)

_, top_k_index = torch.topk(scores_for_choice, self.top_k, dim=-1, sorted=False)

self.shared_experts = Step3p7MLP(config, layer_idx, is_shared_expert=True)

hidden_states = hidden_states + shared_output

scores_for_choice = routing_weights + self.e_score_correction_bias
        _, top_k_index = torch.topk(scores_for_choice, self.top_k, dim=-1, sorted=False)
        top_k_weights = routing_weights.gather(1, top_k_index)

if config.num_nextn_predict_layers:
            # Checkpoints append `num_nextn_predict_layers` MTP layers; ignore them as unexpected keys
            # on regular load. Matches loosely on `layers.<N>.` (not anchored to this model's module

gate = self.act_fn(self.gate_proj(x)).clamp(max=self.limit)
        up = self.up_proj(x).clamp(min=-self.limit, max=self.limit)
        return self.down_proj(gate * up)