[1/3] feat(dsv4): add tensor-parallel DeepSeek-V4 runtime

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/engine/config.py

@@ -72,7 +72,12 @@ 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
     use_pynccl: bool = True
     max_seq_len_override: int | None = None

python/freetoken/engine/engine.py

@@ -245,8 +245,18 @@ def _make_dummy_weight_state_dict(
 ) -> Dict[str, torch.Tensor]:
     state_dict: Dict[str, torch.Tensor] = {}
     fp8_dtypes = (torch.float8_e4m3fn, torch.float8_e5m2)
+    e8m0 = getattr(torch, "float8_e8m0fnu", None)
     for key, param in model_state.items():
-        if param.dtype in fp8_dtypes:
+        if e8m0 is not None and param.dtype == e8m0:
+            # e8m0 is a bare exponent code (value = 2^(code-127)), so 127 is scale 1.0.
+            # It reports is_floating_point, but there is no normal_ kernel for it -- and a
+            # RANDOM exponent would scale a block by up to 2^127 anyway, so 1.0 is both the
+            # only fillable value and the only sane one. Without this a --dummy-weight run
+            # of any block-scaled model dies with "normal_kernel_cuda not implemented".
+            t = torch.empty(param.shape, dtype=param.dtype, device=device)
+            t.view(torch.uint8).fill_(127)
+            state_dict[key] = t
+        elif param.dtype in fp8_dtypes:
             # torch.randn is not implemented for fp8; fill via a uint8 view with small
             # codes (avoid NaN/inf fp8 encodings). Lets dummy-weight startup work for
             # block-fp8 models (the dense fp8 linears are fp8 regardless of moe_backend).
@@ -289,12 +299,38 @@ class ForwardOutput(NamedTuple):
     copy_done_event: torch.cuda.Event
 
 
+def _share_cpu_threads_across_ranks(tp_size: int) -> None:
+    """Split the machine's CPU threads across the TP ranks, before any tensor work.
+
+    torch sizes its intra-op pool from the WHOLE machine, and every rank does the same,
+    so N ranks each start a machine-sized pool and oversubscribe every core N times.
+    Measured on a 40-core / 80-thread host at TP=4: 117 threads per rank, 468 in total,
+    ~6x the hardware, all of it contending through the weight load -- which is thread-
+    bound (host-side copies into the expert banks), not disk-bound. The NVMe sat at
+    187 MB/s while the CPU ran at 92% user.
+
+    The engine already narrows threads once the pinned CPU MoE pool exists; this covers
+    everything BEFORE that, the weight load included. Honour an explicit OMP_NUM_THREADS.
+    """
+    if tp_size <= 1 or os.environ.get("OMP_NUM_THREADS"):
+        return
+    total = os.cpu_count() or tp_size
+    per_rank = max(1, total // tp_size)
+    if per_rank < torch.get_num_threads():
+        torch.set_num_threads(per_rank)
+        logger.info_rank0(
+            f"torch intra-op threads: {total} -> {per_rank} per rank "
+            f"({tp_size} ranks share this host)"
+        )
+
+
 class Engine:
     def __init__(self, config: EngineConfig):
         assert not torch.cuda.is_initialized()
         set_tp_info(rank=config.tp_info.rank, size=config.tp_info.size)
         _ensure_expandable_segments()  # before the first CUDA allocation below
         _adjust_config(config)
+        _share_cpu_threads_across_ranks(config.tp_info.size)
 
         self.device = torch.device(f"cuda:{config.tp_info.rank}")
         torch.cuda.set_device(self.device)
@@ -323,6 +359,19 @@ def __init__(self, config: EngineConfig):
         self.model.load_state_dict(self._load_weight_state_dict(config))
         post_weights_free = self._sync_get_memory()[0]
         self._weights_bytes = self._baseline_free - post_weights_free
+        # What the parameters declare vs what the load actually cost. Every byte of the gap
+        # is resident-but-unaccounted (staging buffers, per-layer scratch) and comes straight
+        # out of the MoE-cache + KV budget, so surface it instead of leaving it to arithmetic
+        # on the cache-sizing line.
+        declared = sum(t.numel() * t.element_size() for t in self.model.state_dict().values())
+        logger.info_rank0(
+            f"Weights: {mem_GB(declared)} declared by parameters, "
+            f"{mem_GB(self._weights_bytes)} measured on device "
+            f"(torch allocator holds {mem_GB(torch.cuda.memory_allocated(self.device))} live, "
+            f"{mem_GB(torch.cuda.memory_reserved(self.device))} reserved; "
+            f"peak {mem_GB(torch.cuda.max_memory_reserved(self.device))}). "
+            "A gap between 'reserved' and 'measured' is held OUTSIDE the allocator."
+        )
         # Pool-budget baseline for the desktop cache sliders: free VRAM after the weights are
         # resident but before ANY runtime cache pool (MoE expert cache below, KV pages, GDN
         # state) is allocated. This is the stable "if all free VRAM went to one pool" budget —

python/freetoken/models/deepseek_v4/args.py

@@ -69,15 +69,56 @@ class DeepseekV4Args:
     hc_sinkhorn_iters: int = 20
     hc_eps: float = 1e-6
 
+    # ----- dSpark (semi-autoregressive speculative decoding) -----
+    # The checkpoint's ``mtp.{0..n_mtp_layers-1}`` stack. Unlike a one-token-per-step MTP
+    # head, dSpark proposes a whole block: ``dspark_block_size`` tokens per draft pass,
+    # scored by a Markov head and gated by a confidence head. Each mtp layer is a FULL
+    # DSV4 block, so enabling it costs n_mtp_layers x n_routed_experts more expert banks.
+    dspark_block_size: int = 0
+    dspark_markov_rank: int = 0
+    dspark_noise_token_id: int = -1
+    # Which target layers feed the drafter its hidden state.
+    dspark_target_layer_ids: Tuple[int, ...] = ()
+    # Runtime opt-in, set by the server from --speculative-dspark. The fields above only
+    # describe what the CHECKPOINT ships; this says whether to pay for it. Off by default
+    # because the drafter's own routed experts enlarge the host banks and the GPU cache.
+    dspark_enabled: bool = False
+
     def __post_init__(self) -> None:
         # JSON lists -> tuple so the dataclass stays hashable / immutable-ish.
         if isinstance(self.compress_ratios, list):
             self.compress_ratios = tuple(self.compress_ratios)
+        if isinstance(self.dspark_target_layer_ids, list):
+            self.dspark_target_layer_ids = tuple(self.dspark_target_layer_ids)
 
     @property
     def nope_head_dim(self) -> int:
         return self.head_dim - self.rope_head_dim
 
+    @property
+    def has_dspark(self) -> bool:
+        """Does this checkpoint ship a usable dSpark drafter?"""
+        return (
+            self.dspark_block_size > 1
+            and self.n_mtp_layers > 0
+            and self.dspark_markov_rank > 0
+        )
+
+    @property
+    def n_draft_layers(self) -> int:
+        """Drafter layers actually built. Zero unless dSpark is both shipped and enabled."""
+        return self.n_mtp_layers if (self.has_dspark and self.dspark_enabled) else 0
+
+    @property
+    def n_moe_layers(self) -> int:
+        """MoE layers the expert banks and the offload cache must cover.
+
+        The drafter's layers are appended AFTER the target's, so a draft layer keeps the
+        cache-facing id ``n_layers + k``. Every layer-indexed structure (host banks, GPU
+        slot cache, KV pools) then addresses target and draft layers the same way.
+        """
+        return self.n_layers + self.n_draft_layers
+
 
 def _config_path(model_path: str) -> str:
     """Locate the authors' ModelArgs JSON inside the checkpoint directory."""

python/freetoken/models/deepseek_v4/attention.py

@@ -8,13 +8,15 @@
 from torch import nn
 
 from freetoken.core import get_global_ctx
+from freetoken.distributed import DistributedCommunicator
 from freetoken.kernel.triton.dsv4.fp8_linear import act_quant_fp8_inplace
 from freetoken.kernel.triton.dsv4.norm import rms_norm
 
 from .args import DeepseekV4Args
 from .compress import Compressor, Indexer
 from .layers import Linear, RMSNorm, get_compress_topk_idxs, get_window_topk_idxs
 from .ops import apply_rotary_emb, apply_rotary_emb_decode, get_freqs_cis
+from .parallel import div_tp, tp_size
 
 
 class Attention(nn.Module):
@@ -30,26 +32,33 @@ def __init__(self, layer_id: int, args: DeepseekV4Args):
         super().__init__()
         self.layer_id = layer_id
         self.dim = args.dim
-        self.n_heads = args.n_heads
         self.q_lora_rank = args.q_lora_rank
         self.o_lora_rank = args.o_lora_rank
         self.head_dim = args.head_dim
         self.rope_head_dim = args.rope_head_dim
-        self.n_groups = args.o_groups
         self.window_size = args.window_size
         self.compress_ratio = args.compress_ratios[layer_id]
         self.eps = args.norm_eps
 
+        # Head parallelism: a rank owns whole o_groups, and therefore whole heads. The
+        # per-group head width (wo_a_k) is a property of one group, so it does NOT shard.
+        self.tp_size = tp_size()
+        self.n_heads = div_tp(args.n_heads, "n_heads")
+        self.n_groups = div_tp(args.o_groups, "o_groups")
+        self._comm = DistributedCommunicator() if self.tp_size > 1 else None
+
         self.attn_sink = nn.Parameter(torch.empty(self.n_heads, dtype=torch.float32), requires_grad=False)
+        # wq_a / wkv / kv_norm stay replicated: the MLA latent KV is shared by every head.
         self.wq_a = Linear(self.dim, self.q_lora_rank, kind="fp8")
         self.q_norm = RMSNorm(self.q_lora_rank, self.eps)
         self.wq_b = Linear(self.q_lora_rank, self.n_heads * self.head_dim, kind="fp8")
         self.wkv = Linear(self.dim, self.head_dim, kind="fp8")
         self.kv_norm = RMSNorm(self.head_dim, self.eps)
         # wo_a: dequantized to bf16 (reference runs a bf16 grouped-output einsum).
         wo_a_rows = self.n_groups * args.o_lora_rank
-        wo_a_k = self.n_heads * self.head_dim // self.n_groups
+        wo_a_k = args.n_heads * self.head_dim // args.o_groups
         self.wo_a = nn.Parameter(torch.empty(wo_a_rows, wo_a_k, dtype=torch.bfloat16), requires_grad=False)
+        # Row-parallel: each rank consumes its own groups and contributes a partial sum.
         self.wo_b = Linear(self.n_groups * args.o_lora_rank, self.dim, kind="fp8")
         self.softmax_scale = self.head_dim ** -0.5
 
@@ -99,10 +108,16 @@ def reset(self) -> None:
 
 
     def _wo(self, o: torch.Tensor, bsz: int, seqlen: int) -> torch.Tensor:
+        # Under TP, ``o`` holds only this rank's heads, so the grouped einsum and wo_b
+        # produce a PARTIAL sum over the full output dim. One all-reduce completes it --
+        # the single attention-side collective per block.
         o = o.reshape(bsz, seqlen, self.n_groups, -1)
         wo_a = self.wo_a.view(self.n_groups, self.o_lora_rank, -1)
         o = torch.einsum("bsgd,grd->bsgr", o, wo_a).flatten(2)
-        return self.wo_b(o)
+        o = self.wo_b(o)
+        if self._comm is not None:
+            o = self._comm.all_reduce(o)
+        return o
 
     def _prefill_segment(self, x_seg, qr_seg, kv_seg, ti: int, start_pos: int, n: int):
         """One request's prefill work inside a (possibly ragged) batch: persist its window KV,

python/freetoken/models/deepseek_v4/model.py

@@ -26,6 +26,7 @@
 from torch import nn
 
 from freetoken.core import get_global_ctx
+from freetoken.distributed import DistributedCommunicator
 from freetoken.kernel.triton.dsv4.hc import hc_post_combine, hc_pre_combine
 from freetoken.kernel.triton.dsv4.sinkhorn import hc_split_sinkhorn
 from freetoken.models.blocks import BaseLLMModel
@@ -45,6 +46,7 @@
     get_window_topk_idxs,
 )
 from .moe import Expert, Gate  # noqa: F401
+from .parallel import div_tp, tp_info, validate_tp
 
 
 class Block(nn.Module):
@@ -130,15 +132,30 @@ def decode_step(self, x, pos, rows, cmp_stage_cap, input_ids, wctx=None):
 class Transformer(nn.Module):
     def __init__(self, args: DeepseekV4Args):
         super().__init__()
+        # Check every tensor-parallel split before building a single layer, so a bad
+        # --tensor-parallel-size fails with one clear message, not a reshape deep in a
+        # forward. embed / head / norm / hyper-connections stay replicated.
+        validate_tp(args)
         self.args = args
         self.norm_eps = args.norm_eps
         self.hc_eps = args.hc_eps
         self.hc_mult = hc_mult = args.hc_mult
-        self.embed = nn.Embedding(args.vocab_size, args.dim)
+        # Vocabulary parallelism. The embedding table and the output head are the two
+        # largest replicated tensors (3.0 GiB together), so each rank keeps one
+        # contiguous vocabulary block: the embedding all-reduces its masked lookup, the
+        # head all-gathers its logit slice.
+        rank, tp = tp_info()
+        self.tp_size = tp
+        self.vocab_size = args.vocab_size
+        vocab_local = div_tp(args.vocab_size, "vocab_size")
+        self.vocab_start = rank * vocab_local
+        self.vocab_local = vocab_local
+        self._comm = DistributedCommunicator() if tp > 1 else None
+        self.embed = nn.Embedding(vocab_local, args.dim)
         self.embed.weight.requires_grad_(False)
         self.layers = nn.ModuleList([Block(i, args) for i in range(args.n_layers)])
         self.norm = RMSNorm(args.dim, self.norm_eps)
-        self.head = nn.Parameter(torch.empty(args.vocab_size, args.dim, dtype=torch.bfloat16), requires_grad=False)
+        self.head = nn.Parameter(torch.empty(vocab_local, args.dim, dtype=torch.bfloat16), requires_grad=False)
         hc_dim = hc_mult * args.dim
         self.hc_head_fn = nn.Parameter(torch.empty(hc_mult, hc_dim, dtype=torch.float32), requires_grad=False)
         self.hc_head_base = nn.Parameter(torch.empty(hc_mult, dtype=torch.float32), requires_grad=False)
@@ -148,6 +165,33 @@ def bind(self, pool, device: torch.device) -> None:
         for layer in self.layers:
             layer.attn.bind(pool, device)
 
+    def embed_tokens(self, input_ids: torch.Tensor) -> torch.Tensor:
+        """Vocabulary-parallel lookup. Rows outside this rank's block contribute zero, so
+        the all-reduce sums exactly one real row per token. Index arithmetic only -- no
+        host sync, so a captured decode graph replays it."""
+        if self._comm is None:
+            return self.embed(input_ids)
+        local = input_ids - self.vocab_start
+        outside = (local < 0) | (local >= self.vocab_local)
+        y = self.embed(local.masked_fill(outside, 0))
+        y = y.masked_fill(outside.unsqueeze(-1), 0)
+        return self._comm.all_reduce(y)
+
+    def logits(self, h: torch.Tensor) -> torch.Tensor:
+        """Head projection over this rank's vocabulary block, then an all-gather back to
+        the full [rows, vocab]. all_gather concatenates on dim 0, so the gathered tensor
+        is rank-major: one row block per rank, which is the vocabulary order."""
+        local = F.linear(h, self.head)  # [rows, vocab_local]
+        if self._comm is None:
+            return local
+        rows = local.shape[0]
+        gathered = self._comm.all_gather(local)  # [tp*rows, vocab_local]
+        if rows == 1:
+            return gathered.view(1, -1)[:, : self.vocab_size]
+        gathered = gathered.view(self.tp_size, rows, self.vocab_local)
+        gathered = gathered.permute(1, 0, 2).reshape(rows, self.tp_size * self.vocab_local)
+        return gathered[:, : self.vocab_size]
+
     def hc_head(self, x):
         shape, dtype = x.size(), x.dtype
         dim = self.args.dim
@@ -172,13 +216,13 @@ def prefill_batched(
         # metadata; ``flat_positions`` [T] is the scheduler-staged batch.positions (per-token
         # ABSOLUTE position); ``last_indices`` [B] is the flattened index of each request's
         # final token -> its next-token logits row.
-        h = self.embed(input_ids)
+        h = self.embed_tokens(input_ids)
         h = h.unsqueeze(2).repeat(1, 1, self.hc_mult, 1)
         for layer in self.layers:
             h = layer.prefill_batched(h, input_ids, segments, flat_positions)
         h = self.hc_head(h)
         h = self.norm(h)
-        return F.linear(h[0, last_indices], self.head)  # [B, vocab]
+        return self.logits(h[0, last_indices])  # [B, vocab]
 
     def decode(
         self, input_ids: torch.Tensor, pos: torch.Tensor, cmp_stage_cap: int
@@ -194,7 +238,7 @@ def decode(
         # in-flight replay (it mutates only the live map).
         B = input_ids.size(0)
         rows = torch.arange(B, device=input_ids.device)
-        h = self.embed(input_ids)
+        h = self.embed_tokens(input_ids)
         h = h.unsqueeze(2).repeat(1, 1, self.hc_mult, 1)
         # Hoist the layer-invariant per-step decode tensors (shared window-ring global slots):
         # resolved ONCE off the attention metadata (recomputed per call, so a capture records
@@ -206,7 +250,7 @@ def decode(
             h = layer.decode_step(h, pos, rows, cmp_stage_cap, input_ids, wctx)
         h = self.hc_head(h)
         h = self.norm(h)
-        return F.linear(h[:, -1], self.head)
+        return self.logits(h[:, -1])
 
 
 class DeepseekV4ForCausalLM(BaseLLMModel):