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

[calvarado2004]

Outcome

DeepSeek-V4-Flash can now run as a real tensor-parallel model in FreeToken instead of replicating model and expert storage on every rank. This is the memory and execution foundation required by the DSpark work in PRs 2 and 3.

On the validation system, the standalone milestone loads and serves DeepSeek-V4-Flash-0731 with TP=4 across four 16 GiB RTX A4000 GPUs while keeping the routed expert banks in host memory.

Stack and review boundary

This is PR 1 of 3:

PR	Milestone	Branch	Commits introduced
1/3 (this PR)	DSV4 tensor-parallel runtime	pr/dsv4-tp-runtime	7f0d8fc
2/3	Exact DSpark draft/verify loop	pr/dsv4-dspark-exact	423edd1..d42be9b
3/3	Adaptive verification and device-resident sampling	pr/dsv4-dspark-adaptive	13d4560..38d7b2f

This PR is intentionally independent of speculative decoding: 13 files, 680 additions, 95 deletions relative to main. Reviewing it first keeps checkpoint ownership and collective placement separate from the much larger draft/verify state machine.

Why this is needed

The original DSV4 runtime had single-rank model shapes and checkpoint-loading assumptions. Merely launching multiple ranks would still replicate the largest tensors and could produce mathematically incomplete rank-local MoE outputs.

The implementation therefore defines an explicit ownership rule for every DSV4 tensor, validates that rule before model construction, loads only the rank-owned checkpoint slices, and places collectives at the points where rank-local partials become complete model values.

The goal is not generic “split every large tensor” TP. DSV4's latent-KV attention, Lightning Indexer, recurrent compressors, routed experts, and quantization metadata have different ownership requirements. Those requirements are preserved rather than reshaped to make a requested TP size fit.

Tensor ownership and collectives

Component	TP layout	Completion operation / reason
Token embedding	Vocabulary rows	Masked local lookup, then one all-reduce so exactly one rank contributes each token row
LM head	Vocabulary rows	Local logits followed by rank-ordered all-gather back to [rows, vocab]
MLA wq_b	Column-parallel over query heads	Each rank owns a deterministic head block
MLA wo_a	Output groups	Each rank owns whole o_groups; partial groups are rejected at configuration time
MLA wo_b	Row-parallel	One all-reduce completes the attention output
Shared SwiGLU expert	moe_inter_dim	Produces a rank-local hidden-width partial
Routed FP4 experts	Same moe_inter_dim split	Added to the shared partial locally; both are completed by one MoE all-reduce
Latent KV path and paged KV pools	Replicated	MLA exposes one latent KV per token that every local head reads
Router, compressors, Lightning Indexer	Replicated	Every rank must make the same expert and sparse-block selection
Hyper-connection mixers	Replicated	Small control tensors whose decisions must remain identical across ranks

The shared and routed MoE branches deliberately do not reduce independently. They are partial sums over the same output dimension, so they are added first and all-reduced once. Reducing each branch separately would add a collective per layer without changing the result.

Checkpoint loading and memory correctness

• Add DSV4 TP helpers for (rank, world_size), exact division, rank slicing, and up-front model validation.
• Slice weights and their quantization scale grids on the same axis. FP8 128×128 blocks and FP4 32-element scale blocks cannot be cut at arbitrary boundaries.
• Require TP to divide o_groups, attention heads, vocabulary, and the quantized MoE intermediate dimension. Incompatible sizes fail with a targeted configuration error before any layer is built.
• Stream each rank's non-expert checkpoint tensors instead of materializing a complete model on every GPU.
• Build each rank's routed FP4 host banks directly from its intermediate-dimension slices.
• Clone a shard into independent storage. A contiguous dim-0 narrow() can still be a view; retaining it kept the complete parent allocation alive and was measured to cost 5.1 GiB per GPU on this checkpoint at TP=4.
• Keep TP=1 as the identity layout and loading path.

For the validation checkpoint, TP=4 distributes the routed host banks across ranks instead of holding four full copies; the ranks retain approximately the same aggregate expert data that a single-rank run requires.

CPU-offload integration

The CPU MoE executor now follows the same rank-local expert layout as the GPU path. It also partitions the process's available physical cores across TP ranks, including affinity handling, so four ranks do not each create a full-machine worker pool and oversubscribe the host.

This matters for the later hybrid path: GPU cache hits/fetches and CPU overflow computation must produce compatible rank-local partials before the TP reduction.

Suggested review order

Area	Files	What to verify
Ownership rules	deepseek_v4/parallel.py, args.py	Exact divisibility, quantization alignment, TP=1 identity
Model math	model.py, attention.py, moe.py	Vocabulary collectives, MLA head/group ownership, single MoE reduction
Checkpoint path	weight.py	Weight/scale co-slicing, independent shard storage, rank-local expert banks
Host execution	moe/cpu_executor.py	Rank-local intermediate width and CPU-core partitioning
Wiring and UX	engine/config.py, engine/engine.py, server/args.py, docs	CLI propagation and early validation
Tests	tests/dsv4/test_dsv4_tensor_parallel.py	Shape tiling, storage ownership, invalid-layout failure cases

Validation

Focused coverage includes:

• every sharded tensor tiles exactly one architectural axis;
• replicated tensors retain their full shape;
• vocabulary ranges cover the complete vocabulary exactly once;
• expert-bank slices tile moe_inter_dim with no gaps or overlap;
• shards own their storage instead of retaining the full parent allocation;
• invalid head-group or quantization-block splits fail loudly;
• model construction and checkpoint ownership remain correct at TP=1/2/4.

The completed three-PR stack passes the full DSV4 suite on the target server:

192 passed in 36.14s

Runtime milestone:

Model:     DeepSeek-V4-Flash-0731
Hardware:  4 × NVIDIA RTX A4000 16 GiB
Topology:  PCIe-only, two NUMA nodes
Mode:      TP=4, hybrid CPU/GPU routed-MoE execution
Result:    model loads, CUDA graphs capture, and OpenAI chat generation serves successfully

Scope and non-goals

• This PR does not add speculative decoding or the checkpoint's mtp.* drafter.
• It does not shard latent KV or sparse-selection state; those values are architecturally shared across local heads.
• It does not silently pad invalid dimensions or quantization blocks.
• It is single-host TP; no multi-node transport policy is introduced here.

The next PR uses these ownership and collective contracts to add exact DSpark decoding without changing the TP model math established here.

[gdevenyi]

Benchmarked this on 2× RTX 6000 Ada (48 GiB each, PCIe Gen4 x16, no NVLink), 2× Xeon Gold 6526Y, 503 GiB host, serving DeepSeek-V4-Flash-0731 with offloaded experts. It does exactly what it says, and it is the largest single win I have measured on this box.

Before this PR, TP replicated the expert banks per rank, so TP=2 was a regression:

	moe_cache_size	host RAM	single stream
TP=1	2559	143 GiB	18.06 tok/s
TP=2 before this PR	2477	291 GiB	15.63 tok/s (−13%)
TP=2 with this PR	5622	156 GiB	39.25 tok/s (2.17×)

The cache more than doubles and host memory nearly halves, which is the mechanism you describe — "the host expert banks divide by N instead of being replicated" — landing precisely as claimed. Residency goes from 23% to 51% of the pool, and the second PCIe link starts carrying half the gather instead of a duplicate of it.

Under load it holds up: 8 concurrent streams give 94.0 tok/s aggregate against 66.5 at TP=1 (1.41×). It also beats the workaround I had been using — two independent single-GPU servers behind a load balancer — which managed 66.9 aggregate while consuming 2× host RAM. This PR obsoletes that recipe on every axis.

I had independently written up the same design as an issue (#62) before finding this PR, having measured the replication problem and traced it to OffloadMoeCache having no tp_size/tp_rank anywhere. Your div_tp(..., multiple_of=128) handles the divisibility constraint I flagged there (block-quantized formats need I/tp to stay a multiple of the scale block), and splitting moe_inter_dim rather than doing expert-parallel all-to-all avoids the CUDA-graph problem entirely, since shapes stay static. That is the right trade.

Test suite on the merged tree: 1606 passed, 10 skipped. The two failures are unrelated to this PR (test_cpu_moe_q4_0, pre-existing on main; and test_e4m3_compat, which fails on #48 in isolation).

Two notes rather than objections:

1. NCCL is a hard prerequisite and the failure is opaque. Without libnccl-dev installed, TP dies at ld: cannot find -lnccl from the JIT build of pynccl.cu, even though torch ships libnccl.so.2 in site-packages. Not this PR's doing, but this PR is what makes people hit it. A fallback to torch's bundled copy, or a clearer message, would save someone an hour.
2. A server binds PORT and PORT+1 (a multiprocessing.spawn child). Unrelated to TP, but it bit me while comparing multi-instance setups.

Strong +1 from me. This is the change that makes a second GPU worth having for offloaded MoE.

[calvarado2004]

Thank you for the unusually thorough independent validation. The 2.17x single-stream gain, roughly 1.41x aggregate gain (94 versus 66.5 tok/s), larger KV cache, and lower host-RAM use on a 2x RTX 6000 Ada / 2-socket Xeon system validate both the TP execution path and the memory-sharding mechanism on hardware very different from the 4x A4000 development box. The full test result is also very helpful.

Both operational issues are actionable:

1. NCCL discovery: agreed that ld: cannot find -lnccl is too opaque. A runtime wheel may contain libnccl.so.2 while the extension link step still lacks the unversioned development link expected by -lnccl. The build should either discover/use PyTorch's bundled NCCL library explicitly or fail early with a precise message naming the missing development link/package. This matters especially because the direct PyNCCL route is the intended fallback on PCIe TP systems.
2. Second port: agreed that the extra rendezvous port must be documented and configurable as a first-class setting. Our deployment script already assigns it explicitly (RDZV_PORT, normally API port + 1), but relying on an implicit adjacent port in the generic multi-instance case is surprising and collision-prone.

I will keep these as focused operational follow-ups rather than mixing them into the TP math/data-path changes in this PR.

Thanks again for testing scaling, concurrency, memory behavior, and the failure modes—not just the happy-path token rate. This is strong confirmation that the TP milestone is real.

[calvarado2004]

Follow-up from the DSpark TP>1 hybrid CPU+GPU integration in #71.

The TP sharding work here remains the foundation of our approach, and the newest measurements strengthen that decision. On the 4-GPU, dual-NUMA workstation we now emit a placement/accounting summary immediately before API readiness:

• NUMA 0: ranks 0/1, 73.31 GiB host experts, 20/20 physical cores;
• NUMA 1: ranks 2/3, 73.31 GiB host experts, 20/20 physical cores;
• CPU-hosted expert storage: 146.62 GiB;
• GPU model allocation: 55.24 GiB across 4 GPUs;
• counted model split: 72.6% CPU / 27.4% GPU.

One nuance: this is logical ownership/accounting plus CPU affinity. Physical page residency still depends on allocator first-touch unless explicit mbind is added; we are keeping that distinction visible rather than overstating the NUMA guarantee.

We also replaced the previous symmetric TP collective in this path with direct standard NCCL. In a comparable 87-call profile, NCCL CUDA time dropped from 878.0 ms / 92.1% to 27.2 ms / 15.8%. The bottleneck is now repeated expert-cache copying (65.9 ms / 38.3%), which is a much healthier place to tune the DSpark paper's hybrid CPU/GPU balance.

Quality/performance remained stable: ~29.6 tok/s sustained on long coding, 33.07 tok/s observed peak on the repeated-route-aware experiment, clean 3,000-token code generation, valid required tool-call plus follow-up, and no repeated-token regression. Explicit reasoning_effort=high coding averaged 30.36 tok/s after warm-up with ~86.9% isolated acceptance.

Our recommendation is therefore to keep this TP-aware sharding design and build adaptive DSpark verification on top, while treating explicit NUMA page binding as a separately measurable follow-up rather than conflating it with rank/core affinity.

Detailed integration measurements: #71 (comment)