dflash

type Attention struct {
	QProj  nn.LinearLayer
	KVProj nn.LinearLayer
	GProj  nn.LinearLayer
	OProj  nn.LinearLayer
	QNorm  *nn.RMSNorm
	KNorm  *nn.RMSNorm
	// contains filtered or unexported fields
}

type Config struct {
	HiddenSize        int32
	NumHiddenLayers   int32
	NumAttentionHeads int32
	NumKeyValueHeads  int32
	HeadDim           int32
	RMSNormEps        float32
	RopeTheta         float32
	Scale             float32
	SlidingWindow     int32
	LayerTypes        []string

	BlockSize      int
	MaskTokenID    int32
	VocabSize      int32
	TargetLayerIDs []int

	// RopeInterleaved selects the draft's rotary pairing convention:
	// true pairs adjacent dims (torch view_as_complex over pairs, the glimmer
	// publisher convention); false pairs split halves (HF rotate_half, the
	// laguna convention). Defaults to false for backwards compatibility with
	// laguna drafts.
	RopeInterleaved bool

	// Causal, when set, overrides every layer's attention direction;
	// otherwise only sliding layers run causal.
	Causal *bool
}

type Layer struct {
	InputNorm    *nn.RMSNorm
	PostAttnNorm *nn.RMSNorm
	Attention    *Attention
	MLP          *MLP
	IsSliding    bool
	IsCausal     bool
}

type MLP struct {
	// GateUpProj is gate|up stacked at load; SwiGLU splits the halves.
	GateUpProj nn.LinearLayer
	DownProj   nn.LinearLayer
}

type Model struct {
	FC         nn.LinearLayer
	HiddenNorm *nn.RMSNorm
	Norm       *nn.RMSNorm
	Layers     []*Layer

	// AuxNorms, when shipped, normalize each target slice before fusion.
	AuxNorms []*nn.RMSNorm

	*Config

	QuantGroupSize int
	QuantBits      int
	QuantMode      string
	TensorQuant    map[string]*model.TensorQuantInfo
	// contains filtered or unexported fields
}