@@ -3767,6 +3767,148 @@ def from_transformers(cls, config, parent_config=None) -> BambaConfig:
37673767 )
37683768
37693769
3770+ @dataclasses .dataclass
3771+ class Phi4FlashConfig (ArchitectureConfig ):
3772+ """Configuration for Phi-4 Flash's SambaY hybrid decoder.
3773+
3774+ The first half alternates Mamba-1 and local differential GQA. Layer 16
3775+ produces transient Mamba memory, layer 17 produces global shared KV, and
3776+ the final half consumes those two shared values through cross-Mamba and
3777+ cross-differential-attention layers. It deliberately has no RoPE despite
3778+ the inactive ``rope_theta`` field in the remote configuration.
3779+ """
3780+
3781+ layer_norm_eps : float = 1e-5
3782+ attention_dropout : float = 0.0
3783+ mamba_d_state : int = 16
3784+ mamba_d_conv : int = 4
3785+ mamba_expand : int = 2
3786+ mamba_dt_rank : int = 160
3787+ mamba_conv_bias : bool = True
3788+ mamba_proj_bias : bool = False
3789+ mb_per_layer : int = 2
3790+ local_attention_window : int = 512
3791+
3792+ def __post_init__ (self ) -> None :
3793+ if self .num_hidden_layers % 4 :
3794+ raise ValueError ("Phi4FlashConfig num_hidden_layers must be divisible by four" )
3795+ if self .mb_per_layer != 2 :
3796+ raise ValueError (
3797+ "Phi4FlashConfig supports the SambaY mb_per_layer=2 schedule only"
3798+ )
3799+ if self .export_paged_attention :
3800+ raise ValueError (
3801+ "Phi4FlashConfig cannot use paged attention: SambaY has heterogeneous "
3802+ "recurrent, local-KV, and shared-global-KV state"
3803+ )
3804+ if self .local_attention_window <= 0 :
3805+ raise ValueError ("Phi4FlashConfig local_attention_window must be positive" )
3806+ if self .mamba_d_state <= 0 or self .mamba_d_conv <= 1 or self .mamba_expand <= 0 :
3807+ raise ValueError (
3808+ "Phi4FlashConfig requires positive Mamba state/expansion and a convolution width above one"
3809+ )
3810+ if self .mamba_dt_rank <= 0 :
3811+ raise ValueError ("Phi4FlashConfig mamba_dt_rank must be positive" )
3812+ if self .hidden_size % self .num_attention_heads :
3813+ raise ValueError (
3814+ "Phi4FlashConfig hidden_size must be divisible by num_attention_heads"
3815+ )
3816+ if self .head_dim != self .hidden_size // self .num_attention_heads :
3817+ raise ValueError (
3818+ "Phi4FlashConfig head_dim must equal hidden_size / num_attention_heads"
3819+ )
3820+ if self .num_attention_heads % 2 or self .num_key_value_heads % 2 :
3821+ raise ValueError (
3822+ "Phi4FlashConfig differential attention requires even Q and KV head counts"
3823+ )
3824+ if self .num_attention_heads % self .num_key_value_heads :
3825+ raise ValueError (
3826+ "Phi4FlashConfig num_attention_heads must be divisible by num_key_value_heads"
3827+ )
3828+ expected = self ._derive_layer_types (self .num_hidden_layers )
3829+ if self .layer_types is None :
3830+ self .layer_types = expected
3831+ elif self .layer_types != expected :
3832+ raise ValueError (
3833+ "Phi4FlashConfig layer_types is derived from the fixed SambaY schedule; "
3834+ f"expected { expected } { self .layer_types }
3835+ )
3836+ # `rope_theta` is configuration residue. The source has no RoPE call.
3837+ self .rope_type = None
3838+ self .rope_theta = None
3839+ self .rope_scaling = None
3840+ self .partial_rotary_factor = None
3841+
3842+ @staticmethod
3843+ def _derive_layer_types (num_hidden_layers : int ) -> list [str ]:
3844+ midpoint = num_hidden_layers // 2
3845+ global_attention = midpoint + 1
3846+ layer_types = []
3847+ for index in range (num_hidden_layers ):
3848+ if index < midpoint :
3849+ layer_types .append (
3850+ "mamba" if index % 2 == 0 else "local_differential_attention"
3851+ )
3852+ elif index == midpoint :
3853+ layer_types .append ("shared_memory_mamba" )
3854+ elif index == global_attention :
3855+ layer_types .append ("global_differential_attention" )
3856+ else :
3857+ layer_types .append (
3858+ "cross_mamba" if index % 2 == 0 else "cross_differential_attention"
3859+ )
3860+ return layer_types
3861+
3862+ @property
3863+ def cache_slot_count (self ) -> int :
3864+ """The source cache owns layers 0 through the global attention layer."""
3865+ return self .num_hidden_layers // 2 + 2
3866+
3867+ @classmethod
3868+ def from_transformers (cls , config , parent_config = None ) -> Phi4FlashConfig :
3869+ base = ArchitectureConfig .from_transformers (config , parent_config )
3870+ # The pinned remote config exposes its checkpoint precision as
3871+ # ``torch_dtype`` rather than the modern ``dtype`` property.
3872+ checkpoint_dtype = _resolve_dtype_value (getattr (config , "torch_dtype" , None ))
3873+ local_window = getattr (config , "sliding_window" , 512 )
3874+ if isinstance (local_window , (list , tuple )):
3875+ local_window = next ((value for value in local_window if value is not None ), 512 )
3876+ dt_rank = getattr (config , "mamba_dt_rank" , "auto" )
3877+ if dt_rank == "auto" :
3878+ dt_rank = math .ceil (base .hidden_size / 16 )
3879+ excluded = {
3880+ "layer_types" ,
3881+ "sliding_window" ,
3882+ "rope_type" ,
3883+ "rope_theta" ,
3884+ "rope_scaling" ,
3885+ "partial_rotary_factor" ,
3886+ "dtype" ,
3887+ }
3888+ base_fields = {
3889+ key : value for key , value in _shallow_fields (base ).items () if key not in excluded
3890+ }
3891+ return cls (
3892+ ** base_fields ,
3893+ dtype = checkpoint_dtype or base .dtype ,
3894+ layer_types = cls ._derive_layer_types (base .num_hidden_layers ),
3895+ layer_norm_eps = getattr (config , "layer_norm_eps" , 1e-5 ),
3896+ attention_dropout = getattr (config , "attention_dropout" , 0.0 ),
3897+ mamba_d_state = getattr (config , "mamba_d_state" , 16 ),
3898+ mamba_d_conv = getattr (config , "mamba_d_conv" , 4 ),
3899+ mamba_expand = getattr (config , "mamba_expand" , 2 ),
3900+ mamba_dt_rank = int (dt_rank ),
3901+ mamba_conv_bias = getattr (config , "mamba_conv_bias" , True ),
3902+ mamba_proj_bias = getattr (config , "mamba_proj_bias" , False ),
3903+ mb_per_layer = getattr (config , "mb_per_layer" , 2 ),
3904+ local_attention_window = int (local_window ),
3905+ rope_type = None ,
3906+ rope_theta = None ,
3907+ rope_scaling = None ,
3908+ partial_rotary_factor = None ,
3909+ )
3910+
3911+
37703912@dataclasses .dataclass
37713913class FalconH1Config (ArchitectureConfig ):
37723914 """Configuration for Falcon-H1 parallel Attention + Mamba2 decoder layers."""
0 commit comments