config.py

from typing import Literal

from pydantic import Field, model_validator

from speculators import SpeculatorModelConfig
from speculators.models.dflash.config import DFlashSpeculatorConfig


@SpeculatorModelConfig.register("dflash2")
class DFlash2SpeculatorConfig(DFlashSpeculatorConfig):
    """DFlash2 draft-model configuration."""

    speculators_model_type: Literal["dflash2"] = "dflash2"  # type: ignore[assignment]

    architectures: list[str] = Field(
        default_factory=lambda: ["DFlash2DraftModel"],
    )

    sliding_window_non_causal: bool = True

    conv_kernel_size: int = Field(default=2, ge=1)
    conv_group_size: int = Field(default=16, ge=1)
    selector_rank: int = Field(default=256, ge=1)
    selector_top_k: int = Field(default=16, ge=1)

    draft_ffn_type: Literal["dense", "moe"] = "dense"
    num_experts: int = Field(default=256, ge=1)
    num_experts_per_tok: int = Field(default=8, ge=1)
    moe_intermediate_size: int = Field(default=512, ge=1)
    shared_expert_intermediate_size: int = Field(default=512, ge=1)

    @model_validator(mode="after")
    def validate_moe_routing(self) -> "DFlash2SpeculatorConfig":
        if (
            self.draft_ffn_type == "moe"
            and self.num_experts_per_tok > self.num_experts
        ):
            raise ValueError(
                "num_experts_per_tok cannot exceed num_experts: "
                f"{self.num_experts_per_tok} > {self.num_experts}."
            )
        return self