@support_torch_compile
class MuseSparkDSparkModel(DFlashQwen3Model):
"""Muse draft backbone with independently cached context K/V."""
def __init__(
self,
*,
vllm_config: VllmConfig,
start_layer_id: int = 0,
prefix: str = "",
) -> None:
nn.Module.__init__(self)
assert vllm_config.speculative_config is not None
config = vllm_config.speculative_config.draft_model_config.hf_config
if not isinstance(config, MuseSparkDSparkConfig):
raise TypeError("MuseSparkDSparkModel requires MuseSparkDSparkConfig")
self.config = config
self.default_dtype = torch.get_default_dtype()
self.vocab_size = config.vocab_size
self.quant_config = get_draft_quant_config(vllm_config)
self.use_aux_hidden_state = True
self.target_layer_ids = tuple(config.target_layer_ids)
# Avoid rerunning VllmConfig validation against the target's shared configs.
draft_config = copy(vllm_config)
draft_config.model_config = vllm_config.speculative_config.draft_model_config
draft_config.quant_config = self.quant_config
self.embed_tokens = VocabParallelEmbedding(
config.vocab_size,
config.hidden_size,
prefix=maybe_prefix(prefix, "embed_tokens"),
)
self.mask_token_id = config.mask_token_id
self.register_buffer(
"mask_embedding",
torch.zeros(config.hidden_size, dtype=vllm_config.model_config.dtype),
persistent=False,
)
self.has_separate_mask_embedding = False
self.fc = ReplicatedLinear(
input_size=config.target_hidden_size * len(self.target_layer_ids),
output_size=config.hidden_size,
bias=False,
params_dtype=vllm_config.model_config.dtype,
quant_config=self.quant_config,
prefix=maybe_prefix(prefix, "fc"),
return_bias=False,
)
self.hidden_norm = MuseSpark1xRMSNorm(
config.hidden_size,
eps=config.rms_norm_eps,
has_weight=config.norm_affine,
zero_centered_gamma=config.norm_zero_centered_gamma,
)
self.layers = nn.ModuleList(
MuseSparkDSparkDecoderLayer(
draft_config,
prefix=maybe_prefix(prefix, f"layers.{layer_idx + start_layer_id}"),
)
for layer_idx in range(config.num_hidden_layers)
)
self.norm = MuseSpark1xRMSNorm(
config.hidden_size,
eps=config.rms_norm_eps,
has_weight=config.norm_affine,
zero_centered_gamma=config.norm_zero_centered_gamma,
)
self.markov_head = DSparkMarkovHead(
config.vocab_size,
config.draft_vocab_size,
config.markov_rank,
prefix=maybe_prefix(prefix, "markov_head"),
quant_config=self.quant_config,
)
self.confidence_head: DSparkConfidenceHead | None = None
if config.enable_confidence_head:
confidence_width = config.hidden_size + (
config.markov_rank if config.confidence_head_with_markov else 0
)
self.confidence_head = DSparkConfidenceHead(
confidence_width,
prefix=maybe_prefix(prefix, "confidence_head"),
bias=True,
with_markov=config.confidence_head_with_markov,
)
def _build_fused_kv_buffers(self) -> None:
layers_attn = [layer.self_attn for layer in self.layers]
attn0 = layers_attn[0]
self._context_pre_norms = tuple(attn.pre_qkv_proj_norm for attn in layers_attn)
self._stacked_kv_weight = torch.stack(
[
attn.qkv_proj.weight[attn.q_size : attn.q_size + 2 * attn.kv_size]
for attn in layers_attn
],
dim=0,
)
self._num_attn_layers = len(layers_attn)
self._kv_size = attn0.kv_size
self._head_dim = attn0.head_dim
self._num_kv_heads = attn0.num_kv_heads
self._rms_norm_eps = attn0.qk_norm.variance_epsilon
self._rotary_emb = attn0.rotary_emb
self._attn_layers = [attn.attn for attn in layers_attn]
for attn in layers_attn[1:]:
if (
attn.kv_size != self._kv_size
or attn.head_dim != self._head_dim
or attn.num_kv_heads != self._num_kv_heads
or attn.qk_norm.variance_epsilon != self._rms_norm_eps
):
raise ValueError("Muse Spark DSpark layers must share KV geometry")
def _project_muse_context_kv(self, context_states: torch.Tensor) -> torch.Tensor:
"""Apply the two trained Muse input norms, then each layer's K/V map."""
fc_output = self.hidden_norm(context_states)
per_layer_context = torch.stack(
[norm(fc_output) for norm in self._context_pre_norms], dim=0
)
return torch.bmm(per_layer_context, self._stacked_kv_weight.transpose(1, 2))
def precompute_and_store_context_kv(
self,
context_states: torch.Tensor,
context_positions: torch.Tensor,
context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None,
) -> None:
if not hasattr(self, "_num_attn_layers"):
self._build_fused_kv_buffers()
num_ctx = context_states.shape[0]
num_layers = self._num_attn_layers
all_kv_flat = self._project_muse_context_kv(context_states)
all_kv = (
all_kv_flat.view(
num_layers,
num_ctx,
2,
self._num_kv_heads,
self._head_dim,
)
.permute(2, 0, 1, 3, 4)
.contiguous()
)
all_k = all_kv[0]
all_v = all_kv[1]
k_fp32 = all_k.float()
variance = k_fp32.pow(2).mean(dim=-1, keepdim=True)
all_k = (k_fp32 * torch.rsqrt(variance + self._rms_norm_eps)).to(all_k.dtype)
all_k_flat = all_k.view(
num_layers * num_ctx, self._num_kv_heads, self._head_dim
)
if self._rotary_emb is not None:
positions = context_positions.repeat(num_layers)
all_k_flat, _ = self._rotary_emb(
positions,
all_k_flat,
# CUDA RoPE updates both operands in place. Aliasing the same
# context-key tensor as Q and K rotates it twice.
None,
)
if context_slot_mapping is None:
return
all_k = all_k_flat.view(num_layers, num_ctx, self._num_kv_heads, self._head_dim)
per_layer = isinstance(context_slot_mapping, (list, tuple))
for layer_idx, attn in enumerate(self._attn_layers):
slot_mapping = (
context_slot_mapping[layer_idx] if per_layer else context_slot_mapping
)
if slot_mapping is None:
continue
attn.impl.do_kv_cache_update(
attn,
all_k[layer_idx],
all_v[layer_idx],
attn.kv_cache,
slot_mapping,
)
def forward(
self,
input_ids: torch.Tensor,
positions: torch.Tensor,
input_embeds: torch.Tensor | None = None,
) -> torch.Tensor:
hidden_states = (
self.embed_input_ids(input_ids) if input_embeds is None else input_embeds
)
hidden_states = hidden_states.to(self.config.residual_dtype)
for layer in self.layers:
hidden_states, _ = layer(positions, hidden_states, None)
return self.norm(hidden_states.to(self.default_dtype))