File size: 1,883 Bytes
ade29d5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 | from typing import Literal
from pydantic import Field
from speculators import SpeculatorModelConfig
from speculators.models.dflash.config import DFlashSpeculatorConfig
__all__ = [
"DSparkSpeculatorConfig",
]
@SpeculatorModelConfig.register("dspark")
class DSparkSpeculatorConfig(DFlashSpeculatorConfig):
"""DFlash config plus a Markov logit-bias head and a confidence head.
The Markov head lets each draft position condition on previously sampled
tokens within the block; the confidence head predicts the per-position
acceptance probability. All DFlash fields are inherited unchanged.
"""
speculators_model_type: Literal["dspark"] = "dspark" # type: ignore[assignment]
architectures: list[str] = Field(
default_factory=lambda: ["DSparkSpeculator"],
description="Model architectures that can load these weights",
)
# Sequential (Markov) head.
markov_rank: int = Field(
default=256,
description=(
"Low-rank dimension of the Markov logit-bias factorization B = W1 @ W2. "
"Set to 0 to disable the sequential head (pure DFlash drafting)."
),
)
markov_head_type: Literal["vanilla", "gated", "rnn"] = Field(
default="vanilla",
description=(
"Sequential head variant: 'vanilla' (first-order Markov bias), 'gated' "
"(hidden-gated bias), or 'rnn' (recurrent state over the block)."
),
)
# Confidence head.
enable_confidence_head: bool = Field(
default=True,
description="Whether to attach the per-position acceptance-probability head.",
)
confidence_head_with_markov: bool = Field(
default=True,
description=(
"Concatenate the Markov previous-token embedding with the backbone "
"hidden state as the confidence-head input."
),
)
|