Skip to content

feat(tp): TP-W1 - rank-layout group table + per-rank TensorParallel on LoadedModel - #156

Merged
mudler merged 1 commit into
mainfrom
row/BACKEND-DISTRIBUTED-TP
Aug 8, 2026
Merged

feat(tp): TP-W1 - rank-layout group table + per-rank TensorParallel on LoadedModel#156
mudler merged 1 commit into
mainfrom
row/BACKEND-DISTRIBUTED-TP

Conversation

@localai-bot

Copy link
Copy Markdown
Collaborator

First tensor-parallelism implementation brick off the landed TP spike (BACKEND-DISTRIBUTED-TP, specs/tensor-parallelism-spike.md §S4, task #287).

What

The GroupCoordinator-analog — pure rank-layout math, no comm/device:

  • include/vllm/distributed/parallel_layout.h: ports vLLM's arange(world).reshape(external_dp, dp, pp, pcp, tp) layout (parallel_state.py:1784-1804, TP innermost, TP group = all_ranks.view(-1, tp)) as a general AxisGroup reshape-stride primitive (PP/DP groups mechanical for later bricks) + TpGroup/AllTpGroups.
  • Thread-local current-rank accessors (SetCurrentTensorParallel/CurrentTpRank/CurrentTpWorldSize) — our thread-per-rank analog (spec S2a) of vLLM's process-global get_tp_group.
  • LoadedModel carries a borrowed per-rank TensorParallel* (forward-declared, default null = single-GPU byte-identical; consumed by TP-W2).

Gate

tests/vllm/distributed/test_parallel_layout.cpp 6/6 (413 assertions): group math == an independent explicit arange().reshape() oracle over a mixed EDP2×DP2×PP2×PCP1×TP2=16 layout and TP∈{1,2,4,8}, plus tp=1 singleton-group inertness + thread_local isolation. No regression: test_tp_forward 2/2, test_communicator 8/8 unchanged (the LoadedModel field is additive/byte-neutral). Header-only; clean -Werror CPU build.

Next: TP-W2 (row/input-dim shard + rank-0 bias + per-rank Hq/Hkv + QKV kv-replication + vocab embed/lm_head + logits all-gather).

…n LoadedModel (#156)

The CPU-completable half of the KV-cache auto-sizing row (ROAD-V1-MEM, #83).
We were behind vLLM here: the KV pool was a raw hand-typed block count.

M1 - the sizing knobs + precedence. EngineParams gains gpu_memory_utilization
(0.92) and kv_cache_memory_bytes (0=unset); num_blocks becomes the override
(default 0 = auto). ResolveNumBlocks (model_loader.cpp) mirrors vLLM's
precedence exactly: num_blocks > kv_cache_memory_bytes > the util path. Both
knobs are mirrored on the C ABI at v16 (vllm_model_params, appended fields,
zero-value preserves behaviour) and exposed as --gpu-memory-utilization /
--kv-cache-memory on examples/server + examples/cli.

M2 - group-aware KVBytesPerBlock (kv_cache_interface.cpp): the marginal device
bytes per KV block, a sum over attention specs (weighted by layer count, or the
per_layer_attn_specs for het-KV models) that mirrors the runner's own
num_blocks * page_size_bytes() allocation; GDN/Mamba state is per-sequence-slot,
not per-block, so it is excluded. The absolute --kv-cache-memory branch divides
its byte budget by this to size the pool with zero device dependency.

M3 (the gpu_memory_utilization device profile run) stays GPU-gated: until it
lands the util branch falls back to the historical 256-block default, so the
zero-struct default path is byte-identical.

Gates (CPU): test_kv_cache_interface KVBytesPerBlock 5/5 (dense / MLA
no-factor-2 / hybrid-excludes-mamba / het-KV per-layer / divisor), test_capi
49/49 incl. the v16 round-trip, test_model_loader_gguf 3/3, clean -Werror build.

FOLLOWING_AGENTS_PROTOCOL
Assisted-by: Claude (Fable 5) via Claude Code
@mudler
mudler force-pushed the row/BACKEND-DISTRIBUTED-TP branch from 88bbc01 to 0e3bf3c Compare August 8, 2026 17:02
@mudler
mudler merged commit 0e3bf3c into main Aug 8, 2026
@mudler
mudler deleted the row/BACKEND-DISTRIBUTED-TP branch August 8, 2026 17:03
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants