[2/3] feat(dsv4): implement exact DSpark speculative decoding

assets/ascii-art.txt

@@ -0,0 +1,18 @@
+
+
+      MMM MMMMMMMMMMMMMMMMMMM                                                                                                                                          
+       M  MMMMMMMMMMMMMMMMMMM                                                MMMMMMMMMMMMMMMMMMM                MMMMMM                                                 
+      MMMMMMMMMMMMMMMMMMMMMM                                                MMMMMMMMMMMMMMMMMMM                 MMMMMM                                                 
+      MMMMMMMMMMMM                                MMMM              MMMP    MMMMMMMMMMMMMMMMMMM     MMM         MMMMMM                 MMM                  MM         
+    MMM  MMMMMMMMMMMMMMMM     MMMMMMMMMMMMM  MMMMMMPMMMMM      MMMMMMPMMMMM       MMMMMMM      MMMMMMMMMMMM    MMMMMMM  MMMMMMMM  MMMMMMMMMMMM    MMMMMMMMMMMMMMMM     
+   MMMM MMMMMMMMMMMMMMMMM     MMMMMMMMMMMPM MPMMMPMMMMMPMMM   MMPMMPMMMMPMMMM     MMMMMM     MMMMMMMMMMMMMMM   MMMMMM MMMMMMMM   MMMMMMMMMMMMMM   MMMMMMMMMMMMMMMMM    
+          MMMMMMMMMMMMMMM     MPMPMPMPMPMMPMMMMPM    PMMMMP  MPMMMM    MMMPMM    MMMMMMM    MMMMMMM   MMMMMMM  MMMMMMMMMMMMM   MMMMMMM    MMMMMM  MMMMMMMM MMMMMMMM    
+     MMMMMMMMMMMMMMMMMMMM    MMMMMMM      MMMPMMMMMMMMMPMMM MMMMPMMMMMMMPMMMM    MMMMMMM   MMMMMMM     MMMMMM MMMMMMMMMMMM     MMMMMMMMMMMMMMMMM MMMMMMM    MMMMMMM    
+          MMMMMM             MMMMMMM      MPMMMMMMMMMPMMMPM MMPMMMMMMMPMMMPMM    MMMMMM    MMMMMMM    MMMMMMM MMMMMMMMMMMM     MMMMMMMMMMMMMMMMM MMMMMM     MMMMMM     
+  MMM MMMMMMMMMM             MPMPMP       MMMPMP            MMMMPM              MMMMMMM    MMMMMM     MMMMMM  MMMMMMMMMMMMM   MMMMMMM            MMMMMM     MMMMMM     
+   M  MMMMMMMMMM            MMMMMMM       MMMMMMMMMMMMMMMM  MMPMMMMMMMMMMMMM    MMMMMMM    MMMMMMMMMMMMMMMM   MMMMMM MMMMMMM   MMMMMMMMMMMMMMM  MMMMMMM    MMMMMMM     
+      MMMMMMMMM             MMMMMM         PMMMMMMMMMMMMM    MMMMMMMMMMMMPM     MMMMMM      MMMMMMMMMMMMMM   MMMMMM   MMMMMMMM  MMMMMMMMMMMMM   MMMMMM     MMMMMM      
+      MMMMMMMMM             MPMPMP           PMMMPMPMM         MPMPMPMPM        MMMMMM        MMMMMMMMMM     MMMMMM    MMMMMMM    MMMMMMMM      MMMMMM     MMMMMM      
+
+
+

docs/cli.md

@@ -40,6 +40,8 @@ parsers all resolve automatically from the checkpoint and the GPU.
 |---|---|---|
 | `--host` | 127.0.0.1 | Bind address |
 | `--port` | 1919 | Bind port |
+| `--tensor-parallel-size`, `--tp-size` | 1 | GPUs to shard the model over; one scheduler process per rank, rank `i` on `cuda:i` |
+| `--distributed-timeout` | 1800 | Seconds a TP collective waits before failing; sized for the per-rank skew in weight-load time |
 | `--max-running-requests` | 4 | Max concurrently running requests |
 | `--max-output-tokens` | 32768 | Default output budget for requests that omit one |
 | `--max-seq-len-override` | from checkpoint | Max sequence length |

docs/models.md

@@ -31,6 +31,27 @@ for them; other checkpoints of the same architectures work too.
   `offload`, upgraded to `hybrid` when a cached `ft bench bw` profile
   recommends it.
 
+## Tensor parallelism
+
+`ft serve --tensor-parallel-size N` shards the model over N GPUs on one host. One
+scheduler process runs per rank, and rank `i` uses `cuda:i`.
+
+DeepSeek-V4 shards:
+
+- MLA query and output heads (`wq_b` column-parallel, `wo_a` by output group,
+  `wo_b` row-parallel with one all-reduce per block)
+- the MoE intermediate dimension, for both the shared expert and the routed FP4
+  experts, so the host expert banks divide by N instead of being replicated
+- the vocabulary, for the embedding table and the output head
+
+It keeps replicated: the MLA latent KV path (`wkv` and the paged KV pools — every
+head reads the same latent KV), the compressors and the Lightning Indexer (every
+rank must select the same blocks), and the router.
+
+Constraints for DeepSeek-V4: N must divide `o_groups` (8), so N is 1, 2, 4 or 8.
+The KV pool is replicated, so its cost per GPU does not fall with N; the weights
+and the expert banks do.
+
 ## Notes
 
 - `ft checkpoint` conversion is optional — it pre-converts a checkpoint into

python/freetoken/attention/dsv4_compress.py

@@ -51,16 +51,25 @@ def compress_rows_of(self, ti: int, block_starts: torch.Tensor, ratio: int) -> t
 
     def decode_compress_rows(
         self, rows: torch.Tensor, pos: torch.Tensor, ratio: int, layer_id: int, tier: str,
-        completed: torch.Tensor,
+        completed: torch.Tensor, ti: int | None = None,
     ) -> torch.Tensor:
         """Per-row decode store destination: the completed block's arithmetic row, or the row's
         OWN scratch row when this step did not finish a block.
 
         Reads the decode SNAPSHOT (not the live map) so a concurrent allocate_paged cannot
         redirect the write; scratch keeps the masked store free of negative indices (which
         ``index_copy_`` would treat as out of bounds) and collision-free across rows.
+
+        ``ti`` addresses the LIVE map by table index instead, for a caller that is not a
+        decode batch. A speculative block walks the compressor per token from inside a
+        prefill, where no decode snapshot exists -- and it is safe there for the same
+        reason the rest of prefill is: the batch owns its rows for the whole forward.
         """
-        row_of_block = self.pool.cmp_rows(self.snapshot()[rows, pos], ratio)
+        source = (
+            self.pool.full_loc_map[ti, pos] if ti is not None
+            else self.snapshot()[rows, pos]
+        )
+        row_of_block = self.pool.cmp_rows(source, ratio)
         scratch = rows + self.compress_scratch_base(layer_id, tier)
         return torch.where(completed, row_of_block, scratch)
 

python/freetoken/attention/dsv4_sparse.py

@@ -221,6 +221,18 @@ def window_slots_of(self, ti: int, lo: int, hi: int) -> torch.Tensor:
         that have slid out of the window translate to -1)."""
         return self.pool.translate_full_to_window(self.pool.full_loc_map[ti, lo:hi])
 
+    def window_slots_at(self, rows: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
+        """Window slots for SCATTERED ``(row, position)`` pairs, one per token.
+
+        ``window_slots_of`` covers one request's contiguous range. A flat batch spans
+        several requests, and its tokens arrive as parallel row/position vectors, so the
+        same lookup is a gather rather than a slice. Positions that have slid out of the
+        window translate to -1, exactly as they do there.
+        """
+        return self.pool.translate_full_to_window(
+            self.pool.full_loc_map[rows.long(), positions.long()]
+        )
+
     def win_cols_to_global(self, win_cols: torch.Tensor, slot_lut: torch.Tensor) -> torch.Tensor:
         """Per-query window columns (or -1) -> GLOBAL window-pool slots via ``slot_lut``."""
         g = slot_lut[win_cols.clamp_min(0)]

python/freetoken/core.py

@@ -89,6 +89,22 @@ def complete_one(self) -> None:
         self.cached_len = self.device_len
         self.device_len += 1
 
+    def complete_n(self, n: int) -> None:
+        """Advance by ``n`` accepted tokens instead of one.
+
+        A speculative step emits a whole block, and each request in a batch keeps a
+        DIFFERENT amount of it -- whatever prefix the target agreed with. So the advance
+        is per request, not per batch.
+
+        ``cached_len`` moves to where the step started, not to ``device_len - 1``: every
+        token of the accepted prefix is now cached history, and the next step extends
+        from the end of it.
+        """
+        if n < 1:
+            raise ValueError(f"a step must advance by at least one token, got {n}")
+        self.cached_len = self.device_len
+        self.device_len += n
+
     def append_host(self, next_token: torch.Tensor) -> None:
         n = self.input_ids.numel()
         m = n + next_token.numel()
@@ -133,6 +149,23 @@ class Batch:
     active_table_idx: "torch.Tensor | None" = None
     # this field should be set by attention backend
     attn_metadata: BaseAttnMetadata = field(init=False)
+    # Speculative block state. A verify step rides the PREFILL path (1+k tokens per
+    # request resuming from its own position), so the phase alone cannot distinguish
+    # it from a real prefill -- these say so explicitly.
+    speculative: bool = field(default=False, init=False)
+    spec_block: int = field(default=0, init=False)
+    draft_confidence: torch.Tensor | None = field(default=None, init=False)
+    # Per-request accepted tokens from a speculative step: the reply path emits one
+    # token per request, so the rest of each block is read from here.
+    spec_emitted: List[torch.Tensor] | None = field(default=None, init=False)
+    # The draft distribution q for this block, kept so acceptance can run the
+    # p/q ratio test rather than an argmax comparison.
+    draft_probs: torch.Tensor | None = field(default=None, init=False)
+    # Compressor carry saved before the block advanced it, restored on rejection.
+    carry_snapshot: object | None = field(default=None, init=False)
+    # Returns a rejected block's unused pages/SWA slots. Supplied by the scheduler,
+    # which owns the cache manager and is what inflated device_len to the block width.
+    release_tail: object | None = field(default=None, init=False)
     # concatenated multimodal soft-token embeddings for a prefill batch (or None)
     mm_embeds: torch.Tensor | None = field(default=None, init=False)
     # Prefill log stats snapshotted at schedule time (before forward's complete_one()

python/freetoken/engine/config.py

@@ -72,8 +72,16 @@ class EngineConfig:
     # ratio default above. A runtime cache rebuild sets this (num_swa_pages) to pin the window
     # regardless of the full anchor; the ratio is the startup default and the fallback.
     swa_num_pages_override: int | None = None
-    distributed_timeout: float = 60.0
+    # Collective timeout for the TP process group. The FIRST collective is the post-load
+    # free-memory all-reduce, and ranks reach it minutes apart on a large MoE checkpoint
+    # (each reads its own expert shards, at its own speed), so a short timeout kills the
+    # server during a healthy startup rather than protecting it. Sized for that load skew;
+    # lower it with --distributed-timeout if a hung rank should fail faster.
+    distributed_timeout: float = 1800.0
     use_dummy_weight: bool = False
+    # DeepSeek-V4 only: build the checkpoint's dSpark drafter for block speculative
+    # decoding. Reaches the model through dsv4_args.dspark_enabled (_adjust_dsv4_config).
+    speculative_dspark: bool = False
     use_pynccl: bool = True
     max_seq_len_override: int | None = None
     num_page_override: int | None = None  # if not None, will override the number of pages