diff --git a/dsv4/model/dsv4.py b/dsv4/model/dsv4.py index 948a3564..88388f9e 100644 --- a/dsv4/model/dsv4.py +++ b/dsv4/model/dsv4.py @@ -1,2 +1,176 @@ -"""Full DSV4 model.""" -# TODO: Phase 1 +"""Full DSV4 model — embedding → N×DSV4Layer → final norm → prediction head. + +Supports both Flash and Pro variants via DSV4Config. +MTP (multi-token prediction) is wired but optional (off by default). +""" +from __future__ import annotations +from typing import Optional, List +import torch + +from dsv4.model.config import DSV4Config +from dsv4.model.layer import TransformerLayer +from dsv4.model.layer_schedule import build_schedule, LayerSpec +from dsv4.layers.norm import RMSNorm +from dsv4.layers.linear import Nvfp4Linear +from dsv4.cache.manager import KVCacheManager + + +class DSV4Model: + """Full DeepSeek-V4 model for inference. + + Construction: + config = DSV4Config.pro() + model = DSV4Model(config, cache_manager) + model.load_weights(checkpoint_path) + + Decode step: + logits = model.decode_step(token_ids, positions, request_ids) + next_token = sampler(logits) + """ + + def __init__( + self, + config: DSV4Config, + cache_manager: KVCacheManager, + device: str = "cuda", + ): + self.config = config + self.cache_manager = cache_manager + self.device = device + self.schedule = build_schedule(config) + + # ---- Token embedding ---- + self.token_embedding = torch.nn.Embedding( + config.vocab_size, config.hidden_size, device=device, + ) + + # ---- Transformer layers ---- + self.layers: List[TransformerLayer] = [ + TransformerLayer(config, spec) for spec in self.schedule + ] + + # ---- Final norm ---- + self.final_norm = RMSNorm(config.hidden_size, device=device) + + # ---- Prediction head (tied with embedding or separate) ---- + self.lm_head = Nvfp4Linear( + in_features=config.hidden_size, + out_features=config.vocab_size, + ) + self._weights_tied = False + + # mHC state shape: (batch, n_hc, hidden_size) per layer + self.n_hc = config.n_hc + + def load_weights(self, checkpoint_path: str) -> None: + """Load weights from a checkpoint directory. + + TODO: implement HF checkpoint loader (dsv4/loader/hf_checkpoint.py). + For now, weights must be set manually. + """ + raise NotImplementedError( + "Weight loading not yet implemented. " + "See dsv4/loader/hf_checkpoint.py" + ) + + def tie_weights(self) -> None: + """Tie lm_head weights to token_embedding (common in LLMs).""" + self.lm_head.weight = self.token_embedding.weight + self._weights_tied = True + + def decode_step( + self, + token_ids: torch.Tensor, # (batch,) int64 + positions: torch.Tensor, # (batch,) int64 + request_ids: torch.Tensor, # (batch,) int32 + mhc_states: Optional[List[torch.Tensor]] = None, + ) -> tuple[torch.Tensor, List[torch.Tensor]]: + """Single decode step: token_ids → logits. + + Args: + token_ids: (batch,) int64 — token IDs for this step + positions: (batch,) int64 — absolute positions + request_ids: (batch,) int32 — which request each token belongs to + mhc_states: optional list of (batch, n_hc, hidden_size) BF16 per layer. + If None, initialized from embedding (first step). + + Returns: + (logits, updated_mhc_states) + logits: (batch, vocab_size) BF16 + mhc_states: List[Tensor] per layer + """ + batch = token_ids.shape[0] + T = 1 # decode + + # Embed → (batch, hidden_size) → (T, batch, hidden_size) → mHC state + emb = self.token_embedding(token_ids) # (batch, hidden_size) + + # mHC state: X_0 = expand to (T, n_hc, hidden_size) + # At layer 0, the first mHC state is initialized from the embedding. + # X[0, i, :] = emb for all i (paper: identity initialization before Sinkhorn) + if mhc_states is None: + mhc_states = [] + for _ in self.layers: + x = torch.zeros(batch, self.n_hc, self.config.hidden_size, + dtype=torch.bfloat16, device=self.device) + # First layer: broadcast embedding into all n_hc slots + x[:, :, :] = emb.unsqueeze(1) + mhc_states.append(x) + + # Get cache handles for each layer + # request_slots are the state cache slots for these requests + request_slots = self.cache_manager.request_slot_map[:batch].clone() + + for layer_idx, (layer, spec) in enumerate(zip(self.layers, self.schedule)): + cache = self.cache_manager.acquire( + layer_idx, request_slots, positions, request_ids, + ) + + X = mhc_states[layer_idx] # (batch, n_hc, hidden_size) + # TransformerLayer expects (T, n_hc, hidden_size) + X = X.unsqueeze(0) # (1, n_hc, hidden_size) — T=1 + + # token_ids needed for hash routing + X = layer.forward(X, token_ids, cache) + + mhc_states[layer_idx] = X.squeeze(0) # (batch, n_hc, hidden_size) + + # Final output: take the last mHC channel and apply norm + head + # X has been updated in-place by the last layer's mHC.post_block + # The output is from the first channel (paper: identity residual) + x_out = mhc_states[-1][:, 0, :] # (batch, hidden_size) + x_out = self.final_norm(x_out) + logits = self.lm_head(x_out) # (batch, vocab_size) + return logits, mhc_states + + def forward( + self, + token_ids: torch.Tensor, # (T,) int64 — prefill tokens + positions: torch.Tensor, # (T,) int64 + request_ids: torch.Tensor, # (T,) int32 + request_slots: torch.Tensor, # (batch,) int32 + ) -> torch.Tensor: + """Prefill: process a sequence of tokens. + + Returns: + (T, vocab_size) BF16 logits + """ + T = token_ids.shape[0] + emb = self.token_embedding(token_ids) # (T, hidden_size) + + for layer_idx, (layer, spec) in enumerate(zip(self.layers, self.schedule)): + cache = self.cache_manager.acquire( + layer_idx, request_slots, positions, request_ids, + ) + # Initialize mHC state from embedding + X = torch.zeros(T, self.n_hc, self.config.hidden_size, + dtype=torch.bfloat16, device=self.device) + X[:, :, :] = emb.unsqueeze(1) + + X = layer.forward(X, token_ids, cache) + + # Output from last layer + x_out = X[:, 0, :] # (T, hidden_size) + x_out = self.final_norm(x_out) + logits = self.lm_head(x_out) # (T, vocab_size) + return logits