diff --git a/.gitignore b/.gitignore index 1c56ade71a..85a6167eb6 100644 --- a/.gitignore +++ b/.gitignore @@ -10,6 +10,7 @@ /ds4flash.gguf /TODO.md /gguf/ +/.deps/ /core /core.* /tests/test_engine_mgpu_placement @@ -29,7 +30,9 @@ /tests/test_mxfp4_dot /tests/test_mxfp4_metal /tests/test_q4k_dot +/cuda/mmq/test/test_q8_aligned_rows /tests/test_sampling +/tests/test_tp_protocol *.o *.dSYM/ __pycache__/ diff --git a/Makefile b/Makefile index 089bcd76de..aa17c61272 100644 --- a/Makefile +++ b/Makefile @@ -47,12 +47,18 @@ NVCC_ARCH_FLAGS := -arch=$(CUDA_ARCH) endif endif NVCCFLAGS ?= -O3 -g -lineinfo --use_fast_math $(NVCC_ARCH_FLAGS) -Xcompiler $(NATIVE_CPU_FLAG) -Xcompiler -pthread +# NCCL is optional at build and run time. A repo-local extraction keeps +# cluster development unprivileged; distro installations are detected too. +NCCL_INCLUDE_DIR ?= $(firstword $(dir $(wildcard .deps/nccl/usr/include/nccl.h /usr/include/nccl.h /usr/local/include/nccl.h))) +ifneq ($(strip $(NCCL_INCLUDE_DIR)),) +NVCCFLAGS += -DDS4_CUDA_HAVE_NCCL=1 -I$(NCCL_INCLUDE_DIR) +endif # Vendored llama.cpp mmq prefill tier (cuda/mmq/, see cuda/mmq/VENDOR.md). MMQ_INCLUDES := -Icuda/mmq MMQ_OBJS := cuda/mmq/ds4_ggml_stubs.o cuda/mmq/ds4_mmq.o cuda/mmq/ds4_mmq_d2r.o cuda/mmq/quantize.o cuda/mmq/mmid.o cuda/mmq/mmvq.o cuda/mmq/ds4_repack.o CORE_OBJS = ds4.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_cuda.o ds4_layer_pack.o $(MMQ_OBJS) CPU_CORE_OBJS = ds4_cpu.o ds4_distributed.o ds4_tp.o ds4_ssd.o ds4_layer_pack.o -CUDA_LDLIBS ?= -lm -Xcompiler -pthread -L$(CUDA_HOME)/targets/sbsa-linux/lib -L$(CUDA_HOME)/lib64 -lcudart -lcublas +CUDA_LDLIBS ?= -lm -ldl -Xcompiler -pthread -L$(CUDA_HOME)/targets/sbsa-linux/lib -L$(CUDA_HOME)/lib64 -lcudart -lcublas HIPCC ?= $(shell command -v hipcc 2>/dev/null || echo /opt/rocm/bin/hipcc) ROCM_ARCH ?= gfx1151 ROCM_CFLAGS ?= -O3 -ffast-math -g -fno-finite-math-only -pthread -D__HIP_PLATFORM_AMD__ -Wno-unused-command-line-argument --offload-arch=$(ROCM_ARCH) @@ -62,7 +68,7 @@ DS4_LINK_LIBS ?= $(CUDA_LDLIBS) METAL_LDLIBS := $(LDLIBS) endif -.PHONY: all help clean test test-metal-session-batch test-mxfp4-cuda test-cuda-session-batch test-cuda-mixed-batch dspark-acceptance dspark-verify-depth mtp-verify-depth cpu cuda cuda-spark cuda-generic cuda-regression strix-halo rocm +.PHONY: all help clean test test-metal-session-batch test-mxfp4-cuda test-q8-aligned-rows test-cuda-session-batch test-cuda-mixed-batch dspark-acceptance dspark-verify-depth mtp-verify-depth cpu cuda cuda-spark cuda-generic cuda-regression strix-halo rocm ifeq ($(UNAME_S),Darwin) .PHONY: metal-decode-schedule-bench metal-prefill-variant-bench check-mxfp4-half-lut @@ -97,7 +103,7 @@ ds4-eval: ds4_eval.o ds4_help.o $(CORE_OBJS) ds4-agent: ds4_agent.o ds4_help.o ds4_web.o ds4_kvstore.o linenoise.o ds4_gpu_args.o $(CORE_OBJS) $(CC) $(CFLAGS) -o $@ ds4_agent.o ds4_help.o ds4_web.o ds4_kvstore.o linenoise.o ds4_gpu_args.o $(CORE_OBJS) $(METAL_LDLIBS) -gguf-tools/quality-testing/score_official: gguf-tools/quality-testing/score_official.c ds4.h $(CORE_OBJS) rax.o ds4_gpu_args.o +gguf-tools/quality-testing/score_official: gguf-tools/quality-testing/score_official.c ds4.h ds4_distributed.h ds4_tp.h $(CORE_OBJS) rax.o ds4_gpu_args.o $(CC) $(QUALITY_CFLAGS) -I. -o $@ gguf-tools/quality-testing/score_official.c $(CORE_OBJS) rax.o ds4_gpu_args.o $(METAL_LDLIBS) tests/test_metal_session_batch.o: tests/test_metal_session_batch.c ds4.h @@ -200,7 +206,7 @@ ds4-eval: ds4_eval.o ds4_help.o $(CORE_OBJS) ds4-agent: ds4_agent.o ds4_help.o ds4_web.o ds4_kvstore.o linenoise.o ds4_gpu_args.o $(CORE_OBJS) $(DS4_LINK) -o $@ $^ $(DS4_LINK_LIBS) -gguf-tools/quality-testing/score_official.o: gguf-tools/quality-testing/score_official.c ds4.h +gguf-tools/quality-testing/score_official.o: gguf-tools/quality-testing/score_official.c ds4.h ds4_distributed.h ds4_tp.h $(CC) $(filter-out -ffast-math,$(QUALITY_CFLAGS)) -I. -c -o $@ $< gguf-tools/quality-testing/score_official: gguf-tools/quality-testing/score_official.o $(CORE_OBJS) rax.o ds4_gpu_args.o @@ -221,15 +227,21 @@ tests/test_mxfp4_cuda: tests/test_mxfp4_cuda.cu $(MMQ_OBJS) test-mxfp4-cuda: tests/test_mxfp4_cuda ./tests/test_mxfp4_cuda + +cuda/mmq/test/test_q8_aligned_rows: cuda/mmq/test/test_q8_aligned_rows.cu cuda/mmq/ds4_mmq.h $(MMQ_OBJS) + $(NVCC) $(NVCCFLAGS) -std=c++17 $(MMQ_INCLUDES) -o $@ $< $(MMQ_OBJS) $(CUDA_LDLIBS) + +test-q8-aligned-rows: cuda/mmq/test/test_q8_aligned_rows + ./cuda/mmq/test/test_q8_aligned_rows endif -ds4.o: ds4.c ds4.h ds4_ssd.h ds4_distributed.h ds4_gpu.h +ds4.o: ds4.c ds4.h ds4_ssd.h ds4_distributed.h ds4_tp.h ds4_gpu.h $(CC) $(CFLAGS) -c -o $@ ds4.c ds4_ssd.o: ds4_ssd.c ds4_ssd.h $(CC) $(CFLAGS) -c -o $@ ds4_ssd.c -ds4_cli.o: ds4_cli.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h linenoise.h +ds4_cli.o: ds4_cli.c ds4.h ds4_ssd.h ds4_distributed.h ds4_tp.h ds4_help.h linenoise.h $(CC) $(CFLAGS) -c -o $@ ds4_cli.c ds4_distributed.o: ds4_distributed.c ds4_distributed.h ds4.h ds4_ssd.h @@ -244,16 +256,16 @@ ds4_help.o: ds4_help.c ds4_help.h ds4_gpu_args.o: ds4_gpu_args.c ds4_gpu_args.h ds4_gpu_mgpu.h $(CC) $(CFLAGS) -c -o $@ ds4_gpu_args.c -ds4_server.o: ds4_server.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h ds4_kvstore.h rax.h +ds4_server.o: ds4_server.c ds4.h ds4_ssd.h ds4_distributed.h ds4_tp.h ds4_help.h ds4_kvstore.h rax.h $(CC) $(CFLAGS) -c -o $@ ds4_server.c -ds4_bench.o: ds4_bench.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h +ds4_bench.o: ds4_bench.c ds4.h ds4_ssd.h ds4_distributed.h ds4_tp.h ds4_help.h $(CC) $(CFLAGS) -c -o $@ ds4_bench.c -ds4_eval.o: ds4_eval.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h +ds4_eval.o: ds4_eval.c ds4.h ds4_ssd.h ds4_distributed.h ds4_tp.h ds4_help.h $(CC) $(CFLAGS) -c -o $@ ds4_eval.c -ds4_agent.o: ds4_agent.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h ds4_kvstore.h ds4_web.h linenoise.h +ds4_agent.o: ds4_agent.c ds4.h ds4_ssd.h ds4_distributed.h ds4_tp.h ds4_help.h ds4_kvstore.h ds4_web.h linenoise.h $(CC) $(CFLAGS) -c -o $@ ds4_agent.c ds4_web.o: ds4_web.c ds4_web.h @@ -262,10 +274,10 @@ ds4_web.o: ds4_web.c ds4_web.h ds4_kvstore.o: ds4_kvstore.c ds4_kvstore.h ds4.h ds4_ssd.h $(CC) $(CFLAGS) -c -o $@ ds4_kvstore.c -ds4_test.o: tests/ds4_test.c ds4_server.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h ds4_kvstore.h rax.h +ds4_test.o: tests/ds4_test.c ds4_server.c ds4.h ds4_ssd.h ds4_distributed.h ds4_tp.h ds4_help.h ds4_kvstore.h rax.h $(CC) $(CFLAGS) -Wno-unused-function -c -o $@ tests/ds4_test.c -ds4_agent_test.o: tests/ds4_agent_test.c ds4_agent.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h ds4_kvstore.h ds4_web.h linenoise.h +ds4_agent_test.o: tests/ds4_agent_test.c ds4_agent.c ds4.h ds4_ssd.h ds4_distributed.h ds4_tp.h ds4_help.h ds4_kvstore.h ds4_web.h linenoise.h $(CC) $(CFLAGS) -Wno-unused-function -c -o $@ tests/ds4_agent_test.c tests/cuda_long_context_smoke.o: tests/cuda_long_context_smoke.c ds4_gpu.h @@ -277,25 +289,25 @@ rax.o: rax.c rax.h rax_malloc.h linenoise.o: linenoise.c linenoise.h $(CC) $(CFLAGS) -c -o $@ linenoise.c -ds4_cpu.o: ds4.c ds4.h ds4_ssd.h ds4_distributed.h ds4_gpu.h +ds4_cpu.o: ds4.c ds4.h ds4_ssd.h ds4_distributed.h ds4_tp.h ds4_gpu.h $(CC) $(CFLAGS) -Wno-unused-function -DDS4_NO_GPU -c -o $@ ds4.c -ds4_cli_cpu.o: ds4_cli.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h linenoise.h +ds4_cli_cpu.o: ds4_cli.c ds4.h ds4_ssd.h ds4_distributed.h ds4_tp.h ds4_help.h linenoise.h $(CC) $(CFLAGS) -DDS4_NO_GPU -c -o $@ ds4_cli.c ds4_gpu_args_cpu.o: ds4_gpu_args.c ds4_gpu_args.h ds4_gpu_mgpu.h $(CC) $(CFLAGS) -DDS4_NO_GPU -c -o $@ ds4_gpu_args.c -ds4_server_cpu.o: ds4_server.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h ds4_kvstore.h rax.h +ds4_server_cpu.o: ds4_server.c ds4.h ds4_ssd.h ds4_distributed.h ds4_tp.h ds4_help.h ds4_kvstore.h rax.h $(CC) $(CFLAGS) -DDS4_NO_GPU -c -o $@ ds4_server.c -ds4_bench_cpu.o: ds4_bench.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h +ds4_bench_cpu.o: ds4_bench.c ds4.h ds4_ssd.h ds4_distributed.h ds4_tp.h ds4_help.h $(CC) $(CFLAGS) -DDS4_NO_GPU -c -o $@ ds4_bench.c -ds4_eval_cpu.o: ds4_eval.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h +ds4_eval_cpu.o: ds4_eval.c ds4.h ds4_ssd.h ds4_distributed.h ds4_tp.h ds4_help.h $(CC) $(CFLAGS) -DDS4_NO_GPU -c -o $@ ds4_eval.c -ds4_agent_cpu.o: ds4_agent.c ds4.h ds4_ssd.h ds4_distributed.h ds4_help.h ds4_kvstore.h ds4_web.h linenoise.h +ds4_agent_cpu.o: ds4_agent.c ds4.h ds4_ssd.h ds4_distributed.h ds4_tp.h ds4_help.h ds4_kvstore.h ds4_web.h linenoise.h $(CC) $(CFLAGS) -DDS4_NO_GPU -c -o $@ ds4_agent.c ds4_metal.o: ds4_metal.m ds4_gpu.h $(METAL_SRCS) @@ -352,6 +364,12 @@ tests/test_gpu_args.o: tests/test_gpu_args.c ds4_gpu_args.h ds4_gpu_mgpu.h tests/test_gpu_args: tests/test_gpu_args.o ds4_gpu_args_cpu.o $(CC) $(CFLAGS) -o $@ $^ $(LDLIBS) +tests/test_tp_protocol.o: tests/test_tp_protocol.c ds4_tp.h ds4.h + $(CC) $(CFLAGS) -I. -c -o $@ $< + +tests/test_tp_protocol: tests/test_tp_protocol.o $(CPU_CORE_OBJS) + $(CC) $(CFLAGS) -o $@ $^ $(LDLIBS) + ds4_cpu_test_hooks.o: ds4.c ds4.h ds4_gpu.h ds4_gpu_mgpu.h ds4_layer_pack.h $(CC) $(CFLAGS) -Wno-unused-function -DDS4_NO_GPU -DDS4_TEST_HOOKS -c -o $@ ds4.c @@ -442,6 +460,7 @@ endif test: ds4_test ds4_agent_test ds4-eval q4k-dot-test mxfp4-dot-test \ tests/test_layer_pack tests/test_engine_mgpu_placement tests/test_gpu_args \ + tests/test_tp_protocol gguf-tools/quality-testing/score_official \ $(SAMPLING_TEST) ds4 ds4-server ds4-bench ds4-agent ./ds4-eval --self-test-extractors ./ds4_agent_test @@ -449,6 +468,7 @@ test: ds4_test ds4_agent_test ds4-eval q4k-dot-test mxfp4-dot-test \ ./tests/test_layer_pack ./tests/test_engine_mgpu_placement ./tests/test_gpu_args + ./tests/test_tp_protocol ./tests/test_gpu_args_cli.sh ifneq ($(UNAME_S),Darwin) ./tests/test_sampling @@ -488,4 +508,4 @@ mxfp4-dot-test: tests/test_mxfp4_dot.c ./tests/test_mxfp4_dot clean: - rm -f ds4 ds4-server ds4-bench ds4-eval ds4-agent ds4_cpu ds4_native ds4_server_test ds4_test ds4_agent_test gguf-tools/quality-testing/score_official gguf-tools/quality-testing/score_official.o speed-bench/metal_decode_schedule_bench speed-bench/metal_prefill_variant_bench speed-bench/*.o tests/test_q4k_dot tests/test_mxfp4_dot tests/test_mxfp4_metal tests/test_mxfp4_cuda tests/test_metal_session_batch tests/test_gpu_xdev tests/test_gpu_model_cache tests/test_gpu_lookup_cache_strict tests/test_engine_mgpu_refusal tests/test_engine_mgpu_runtime tests/test_engine_correctness tests/test_sampling tests/test_cuda_session_batch tests/test_cuda_mixed_batch tests/*.o *.o tests/cuda_long_context_smoke tests/cuda_long_context_smoke.o + rm -f ds4 ds4-server ds4-bench ds4-eval ds4-agent ds4_cpu ds4_native ds4_server_test ds4_test ds4_agent_test gguf-tools/quality-testing/score_official gguf-tools/quality-testing/score_official.o speed-bench/metal_decode_schedule_bench speed-bench/metal_prefill_variant_bench speed-bench/*.o tests/test_q4k_dot tests/test_mxfp4_dot tests/test_mxfp4_metal tests/test_mxfp4_cuda tests/test_metal_session_batch tests/test_gpu_xdev tests/test_gpu_model_cache tests/test_gpu_lookup_cache_strict tests/test_engine_mgpu_refusal tests/test_engine_mgpu_runtime tests/test_engine_correctness tests/test_sampling tests/test_cuda_session_batch tests/test_cuda_mixed_batch tests/test_tp_protocol tests/*.o *.o tests/cuda_long_context_smoke tests/cuda_long_context_smoke.o cuda/mmq/test/test_q8_aligned_rows diff --git a/QA_BEFORE_RELEASES.md b/QA_BEFORE_RELEASES.md index a74d50b5ad..9d1a611d0b 100644 --- a/QA_BEFORE_RELEASES.md +++ b/QA_BEFORE_RELEASES.md @@ -178,6 +178,15 @@ top-logprob slices, so do not replace them with one sampled chat answer. about `0.800` unless the quantization changed deliberately. - Run the 100-case DeepSeek V4 PRO fixture for every released PRO GGUF: `gguf-tools/quality-testing/score_official /path/to/deepseek-v4-pro.gguf gguf-tools/quality-testing/data/pro/manifest.tsv /tmp/pro.tsv 4096`. +- After network EP/TP graph, collective, ownership, or numerical-path changes, + run the matching 100-case fixture through `score_official` with + `--role coordinator`, the selected `--expert-parallel` or + `--tensor-parallel` mode, and the real rank count. Score the exact/default + path and every proposed faster path into separate TSVs, then compare them + with `compare_scores.py`. A clear NLL regression, first-token-match drop, or + material API top-1/pair-order drop blocks making the faster path the default. + Set path-selecting environment variables identically on every rank and keep + both raw summaries in the QA record. - For SSD streaming, run the same official-continuation scorer once with full residency and once with `--ssd-streaming` for the release model. The summary and API agreement should stay in the same quality band. @@ -453,6 +462,54 @@ release-ready without this pass. - Verify that any CUDA-only warning fixes are also clean on macOS and do not change Metal behavior. +### CUDA network EP/TP + +Run these gates when NCCL bootstrap, ownership mapping, collective MoE, +DeepSeek/GLM TP, or the shared frontend TP options change: + +- Build the same warning-free commit on every rank and verify the complete GGUF + has the same size and checksum on every host. Start at `--ctx 512`; DGX Spark + uses one unified CPU/GPU memory pool, and an overcommitted two-rank experiment + can make a machine require a manual reboot. +- Run `tests/test_tp_protocol` locally, then complete short two- and four-rank + `--expert-parallel` and `--tensor-parallel` prompts where the guarded model + plan fits. For an oversized two-rank GLM plan, verify preflight rejection + instead of disabling the memory guard. Every worker must report the intended + rank/world and exit after the coordinator instead of reconnecting + indefinitely. +- For DeepSeek Flash Q2, compare a greedy single-rank raw-expert trace with + four-rank EP and TP over a prompt long enough to exercise batch prefill plus + at least 32 decode tokens. Set `DS4_CUDA_MOE_NO_IQ2_ALIGNED=1` and + `DS4_CUDA_MOE_NO_Q2K_ALIGNED=1` on the single-rank reference. EP and TP model + output must be byte-identical to that reference without + `DS4_TP_ORDERED_REDUCE`; a matching first token alone is insufficient. +- When shard-local DeepSeek artifacts or aligned dense row slicing changes, + run `make test-q8-aligned-rows CUDA_ARCH=sm_121` on a Spark. Confirm every + rank reports the same owned expert count, `expert raw residency replaced`, + and the replicated aligned-dense bytes in its guarded memory plan. Repeat + the parity trace once with `--expert-parallel` and once with + `--tensor-parallel`; do not use the replication diagnostic environment + variables to make the default TP gate pass. +- If `DS4_CUDA_TP_FAST_ALIGNED_EXPERTS=1` changes, set it on every rank and + repeat matching EP/TP speed sweeps plus the four-case, 4,096-token DeepSeek + evaluation. This opt-in path may drift from the raw-reference token stream, + but EP and TP must agree with each other, keep finite output, and stay in the + default path's quality band. It does not replace the byte-identical default + parity gate above. +- For GLM 5.2, compare four-rank EP and TP greedy traces over a prompt long + enough to exercise batch prefill plus at least 32 decode tokens. The complete + `--dump-logprobs` JSON, including selected tokens and top-logprob entries, + must be byte-identical. GLM TP must not require ordered reduction to pass. +- Run matching single-rank (where it fits), four-rank EP, and four-rank TP + `ds4-bench` sweeps and preserve the CSVs. Record prefill and generation + separately: both supported model paths intentionally replicate dense prefill + kernels for numerical parity, while model-specific decode work is + partitioned. +- Exercise one real request through `ds4-server`, one non-interactive + `ds4-agent` prompt, and one `ds4-eval` case with network workers. Confirm all + mirrored sessions are destroyed and no rank remains running after rank 0 + exits. + ## 9. ROCm / Strix Halo Use the Strix Halo Framework Desktop via the VPN hostname `strixhalo` diff --git a/README.md b/README.md index 0f5324dec5..5e7edaf86b 100644 --- a/README.md +++ b/README.md @@ -161,9 +161,12 @@ Q8/F32 paths and supports routed expert gate/up tensors in `Q2_K`, `Q4_K`, or are added deliberately and scored against the official 100-case fixture. These formats do not all support the same execution modes. The Q4 files work -for normal Metal and CUDA inference. Two-Mac tensor parallelism currently -requires an ownership-aware IQ2_XXS or Q2_K routed layout; a routed Q4 GLM -must be rejected before evaluation. +for normal Metal and CUDA inference. Two-Mac tensor parallelism accepts the +ownership-aware IQ2_XXS and Q2_K routed layouts. CUDA network expert/tensor +parallelism currently requires IQ2_XXS gate/up experts in executable layers +(IQ2_XXS and Q2_K down projections are supported); a fully routed-Q2_K or routed-Q4 +CUDA network model is rejected before evaluation. CUDA tensor parallelism also +requires sliceable Q8_0 dense and attention projections. GLM's MTP block is part of the main GGUF; it does not use the separate Flash MTP file. Ordinary decode remains the default. `--glm-mtp` enables experimental @@ -614,6 +617,329 @@ encryption or authentication, and is not release-stable yet; coordinator and workers should be built from the same commit and used on trusted machines and trusted networks. +## Network expert and tensor parallelism on DGX Spark + +CUDA network parallelism runs one GLM 5.2 or supported DeepSeek V4 graph in +lockstep across two or four machines using NCCL collectives. It is different +from the layer pipeline above: every rank evaluates the same tokens at the same +time and keeps a full KV cache, while routed weights and selected per-layer work +are partitioned inside the graph. + +There are two modes: + +- `--expert-parallel` partitions the routed MoE experts. Attention, leading + dense FFNs, shared experts, embeddings, routers, and the output head remain + replicated. +- `--tensor-parallel` includes that same expert partition and additionally + partitions model-specific decode work. GLM partitions complete attention + heads and leading dense-FFN lanes, then partitions their projections by + complete output rows. DeepSeek partitions decode attention heads/output + groups, shared-expert lanes and complete output rows, and output vocabulary + rows. Both paths keep dense prefill attention and dense/shared FFNs replicated + to preserve the CUDA batch kernels' numerical path; routed experts remain + partitioned during both prefill and decode. + +Tensor parallelism here is not token, sequence, or context parallelism. Every +rank sees the same token rows. Decode remains autoregressive, so the ranks work +simultaneously on one token but the next token cannot start early. + +The GLM CUDA path supports DSA GGUFs whose executable routed layers use +ownership-aware `IQ2_XXS` gate/up experts and supported `IQ2_XXS` or `Q2_K` +down projections. A fully routed-`Q2_K` layer is not shard-safe yet. Full GLM +tensor parallelism also requires the checkpoint's dense and attention +projections to be `Q8_0`. + +GLM TP does not all-reduce partial dot products. During decode, disjoint +complete heads or FFN lanes are gathered first, and each following projection +is divided into complete output rows. This preserves the EP arithmetic within +each output row. Dense prefill stays replicated for the same numerical-parity +reason; `DS4_GLM_TP_SPLIT_PREFILL=1` restores the old split-prefill path for +diagnosis only and is not a supported production setting. + +The tested DeepSeek path is the 0731 Flash Q2 tensor mix downloaded by +`./download_model.sh ds4f-q2`: routed gate/up tensors must be `IQ2_XXS`, routed +down tensors must be `Q2_K`, and attention output, shared-expert, and output-head +projections must be `Q8_0`. The mixed `ds4f-q2-q4` file is not ownership-shard +compatible yet. Unsupported layouts and dimensions are rejected before the +expert shard is made resident. CUDA uses NCCL only; `--transport auto` selects +it, and `--transport nccl` makes the requirement explicit. External MTP/DSpark +drafting, GLM MTP, directional steering, and SSD streaming are not enabled in +CUDA network mode. Native mixed prefill-plus-decode batches are supported by +`ds4-server` and mirrored to every rank. + +### Build and NCCL setup + +NCCL is optional for ordinary CUDA builds and is loaded at runtime. The build +needs `nccl.h` to enable network collectives. A system install is detected in +`/usr/include` or `/usr/local/include`; an unprivileged install can be selected +explicitly: + +```sh +NCCL_INCLUDE_DIR=/path/to/nccl/include make cuda-spark +``` + +If `libnccl.so.2` is not in the dynamic loader's normal search path, point ds4 +at it on every rank: + +```sh +export DS4_NCCL_LIBRARY=/path/to/libnccl.so.2 +``` + +NCCL should use the high-speed cluster interface rather than Wi-Fi. Interface +and HCA names are machine-specific. For the tested four-Spark cluster on +`192.168.2.30` through `192.168.2.33`, the RoCE settings are: + +```sh +export NCCL_SOCKET_IFNAME=enp1s0f0np0 +export NCCL_IB_HCA=rocep1s0f0,roceP2p1s0f0 +``` + +Use `NCCL_DEBUG=INFO` while diagnosing interface selection and return it to +`WARN` for measurements. All ranks must use the same world size, mode, model +contents, and build. Each worker's `--ctx` must be at least as large as the +coordinator session or benchmark allocation. Rank 0 is always the coordinator; +workers have unique ranks 1 through `world-1`. Do not pass `--layers`: network +expert/tensor parallelism owns model placement. + +### Four-Spark GLM 5.2 example + +Put the full GGUF on local storage on every Spark. Each rank opens the complete +file but makes only its owned routed-expert range plus replicated tensors +resident on the CUDA device. +The tested 196.6 GiB IQ2_XXS GGUF maps about 63.84 GiB of model tensors per rank +with a four-way expert split. + +Start the three workers first: + +```sh +# 192.168.2.31, rank 1 +./ds4 --cuda \ + -m gguf/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \ + --ctx 4096 --role worker --tensor-parallel \ + --tensor-parallel-world 4 --tensor-parallel-rank 1 \ + --coordinator 192.168.2.30 9911 --transport nccl + +# 192.168.2.32, rank 2 +./ds4 --cuda \ + -m gguf/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \ + --ctx 4096 --role worker --tensor-parallel \ + --tensor-parallel-world 4 --tensor-parallel-rank 2 \ + --coordinator 192.168.2.30 9911 --transport nccl + +# 192.168.2.33, rank 3 +./ds4 --cuda \ + -m gguf/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \ + --ctx 4096 --role worker --tensor-parallel \ + --tensor-parallel-world 4 --tensor-parallel-rank 3 \ + --coordinator 192.168.2.30 9911 --transport nccl +``` + +Then start rank 0: + +```sh +# 192.168.2.30, rank 0 +./ds4 --cuda \ + -m gguf/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \ + --ctx 4096 --role coordinator --tensor-parallel \ + --tensor-parallel-world 4 --listen 192.168.2.30 9911 \ + --transport nccl --temp 0 \ + -p "Explain why the sky is blue." +``` + +Replace `--tensor-parallel` with `--expert-parallel` on every rank to partition +only the routed experts. Workers retry until the coordinator is listening, so +connection-refused messages during startup are expected. + +The full GGUF is still required on disk on every host; only resident model +memory is sharded. Startup prints both the mapped shard and the complete planned +memory. The GLM memory guard keeps 32 GiB of system/GPU headroom by default on +128 GiB unified-memory systems. Do not disable that guard on DGX Spark: pushing +the shared CPU/GPU memory pool to exhaustion can make the machine unresponsive. +Start with a conservative context and increase it only after checking the +reported plan on every rank. + +This particular 196.6 GiB GLM file does not fit safely across only two 128 GiB +Sparks. Its two-way model shard plus graph state exceeds the guarded budget and +must be refused. Use all four Sparks for this model. + +### DeepSeek V4 Flash example + +The tested 80.76 GiB 0731 Q2 Flash file maps about 26.34 GiB of model tensors +per rank with four ranks. Every host still needs the complete GGUF on local +storage. On ranks 1 through 3, change both `--tensor-parallel-rank` and the +worker host as appropriate: + +```sh +MODEL=gguf/DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix-0731.gguf + +# 192.168.2.31, rank 1. Repeat on .32/rank 2 and .33/rank 3. +./ds4 --cuda -m "$MODEL" --ctx 4096 \ + --role worker --tensor-parallel \ + --tensor-parallel-world 4 --tensor-parallel-rank 1 \ + --coordinator 192.168.2.30 9911 --transport nccl + +# 192.168.2.30, rank 0, after all workers are loading. +./ds4 --cuda -m "$MODEL" --ctx 4096 \ + --role coordinator --tensor-parallel --tensor-parallel-world 4 \ + --listen 192.168.2.30 9911 --transport nccl --temp 0 \ + -p "Explain why the sky is blue." +``` + +Use `--expert-parallel` on every rank for expert-only execution. DeepSeek TP +uses disjoint expert slots and complete dense output rows during decode, rather +than summing partial dot products. The supported path is therefore designed to +preserve the single-rank arithmetic order and does not require +`DS4_TP_ORDERED_REDUCE`. Prefill attention and the shared expert remain +replicated because their optimized CUDA batch kernels are shape-sensitive. +Decode keeps the row split: sliced dense Q8 projections resolve the complete +aligned artifact and evaluate only the rank's row interval, so their per-row +arithmetic stays identical to a complete single-rank launch. + +On DeepSeek Q2, each network rank builds aligned IQ2_XXS/Q2_K artifacts only +for its owned expert interval and replaces that interval's raw CUDA residency. +Decode consumes the shard-local artifacts directly. Batch prefill reconstructs +only the current layer into reusable device scratch before running the exact +raw-layout arithmetic; it does not restore or pin the complete expert shard. +Dense Q8 aligned artifacts stay replicated because dense/shared weights are +replicated. Startup reports their additive memory separately, for example +`aligned dense artifacts 6.15 GiB` for the tested 0731 file. + +`DS4_CUDA_TP_FAST_ALIGNED_EXPERTS=1` is an opt-in throughput mode for this +DeepSeek layout. Set it identically on every rank. It runs prefill and decode +directly through the shard-local aligned MMQ representation instead of using +the raw-arithmetic parity twins. This changes floating-point rounding and is +therefore not the byte-identical release-parity path. On the four-Spark 0731 +Q2 test at a 4,096-token frontier it measured 535.74/14.92 prefill/generation +t/s in EP and 511.36/19.43 t/s in TP; both modes passed the same four-case, +4,096-token evaluation and produced identical extracted answers. Treat those +numbers as one-cluster measurements, not a general performance guarantee. +GLM currently keeps its existing routed-expert representation; enabling this +environment variable does not change the GLM graph. + +For a strict greedy parity reference, disable only the single-rank routed +expert artifacts. Dense Q8 artifacts remain enabled on both the reference and +the network ranks: + +```sh +DS4_CUDA_MOE_NO_IQ2_ALIGNED=1 \ +DS4_CUDA_MOE_NO_Q2K_ALIGNED=1 \ +./ds4 --cuda -m "$MODEL" --ctx 512 --temp 0 --tokens 32 -p "Hello" +``` + +`DS4_CUDA_TP_NO_ALIGNED_EXPERTS=1` is a diagnostic rollback for a network +rank. It disables the shard-local artifact build, including the replicated Q8 +artifacts, and restores the older raw-resident path; set it identically on all +ranks. It is not the parity reference above and uses substantially more model +residency. + +### Two ranks + +The DeepSeek Q2 file maps about 44.48 GiB of model tensors per rank with a +two-way split. Rank 1 may omit its explicit rank in a two-rank world, but +spelling it out makes launch scripts clearer: + +```sh +MODEL=gguf/DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix-0731.gguf + +# Worker, rank 1. +./ds4 --cuda -m "$MODEL" --ctx 4096 \ + --role worker --tensor-parallel \ + --tensor-parallel-world 2 --tensor-parallel-rank 1 \ + --coordinator 192.168.2.30 9911 --transport nccl + +# Coordinator, rank 0. +./ds4 --cuda -m "$MODEL" --ctx 4096 \ + --role coordinator --tensor-parallel \ + --tensor-parallel-world 2 --listen 192.168.2.30 9911 \ + --transport nccl --temp 0 -p "Hello" +``` + +The same two-way protocol works for a smaller supported GLM model whose +per-rank plan fits with headroom; the 196.6 GiB GLM example above does not. +Two- and four-rank collective setup use the same code path. Before model tests, +`nccl-tests` is useful for confirming that the selected interfaces complete an +in-place all-reduce without errors. + +### Benchmarking + +Start workers with `./ds4` exactly as above, giving them enough `--ctx` for the +benchmark allocation. Run `ds4-bench` as rank 0. This short command is a smoke +benchmark rather than a quality evaluation: + +```sh +./ds4-bench --cuda \ + -m gguf/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \ + --prompt-file speed-bench/promessi_sposi.txt \ + --ctx-start 16 --ctx-max 16 --ctx-alloc 128 \ + --step-incr 16 --gen-tokens 64 \ + --role coordinator --tensor-parallel \ + --tensor-parallel-world 4 --listen 192.168.2.30 9911 \ + --transport nccl --csv /tmp/glm52-tp4.csv +``` + +Supported GLM TP should match its EP reference exactly, and supported DeepSeek +TP should match its raw-expert single-rank/EP reference. Compare a multi-token +`--dump-logprobs` trace, not only the first selected token: selected token bytes +and stored top-logprob entries must be byte-identical. Treat a greedy mismatch +as a correctness failure. All ranks must also remain bit-identical to each other +after every collective. + +### Server, agent, and evaluation frontends + +Network expert/tensor parallelism is wired into `ds4`, `ds4-bench`, +`ds4-server`, `ds4-agent`, `ds4-eval`, and the official-continuation +`score_official` tool. Rank 0 runs the selected frontend; ordinary frontend +workers may run either that same executable or plain `ds4` with the matching +model, mode, world, rank, and context options. Use plain `ds4` workers for the +scorer. Frontend-only flags are ignored after a worker enters its +mirrored-session loop. The scorer's exact commands and comparison workflow are +documented in +`gguf-tools/quality-testing/README.md`. + +For example, after starting the three `ds4` workers from the four-Spark example +above, rank 0 can serve the model over HTTP: + +```sh +./ds4-server --cuda \ + -m gguf/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \ + --ctx 4096 --role coordinator --tensor-parallel \ + --tensor-parallel-world 4 --listen 192.168.2.30 9911 \ + --transport nccl --host 127.0.0.1 --port 8000 +``` + +`--listen` is the rank-control address; `--host` and `--port` remain the HTTP +address. Graceful server shutdown destroys every mirrored session before it +stops the workers. Native `--batched-session N` is supported, but every +resident session has its own KV/cache allocation, so begin with one session +and increase it only when the printed guarded memory plan leaves ample room. + +The agent and evaluation harness use the same coordinator options: + +```sh +./ds4-agent --cuda \ + -m gguf/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \ + --ctx 4096 --non-interactive --nothink --tokens 64 \ + -p "Explain why the sky is blue." \ + --role coordinator --tensor-parallel --tensor-parallel-world 4 \ + --listen 192.168.2.30 9911 --transport nccl + +./ds4-eval --cuda \ + -m gguf/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf \ + --ctx 4096 --questions 1 --tokens 64 --nothink --plain \ + --role coordinator --tensor-parallel --tensor-parallel-world 4 \ + --listen 192.168.2.30 9911 --transport nccl +``` + +`ds4-eval` requires an explicit `--ctx` in network mode; automatic context +sizing cannot establish a safe common capacity before workers connect. Each +frontend sends the same session create/sync/eval/rewind/destroy protocol, and +rank 0 sends a final stop only after its live sessions have been released. +Server and agent KV files may still be written, but direct payload restore is +disabled in network mode because it would update rank 0 alone. A cache hit +therefore falls back to replaying the requested or stored prompt tokens through +every rank in lockstep; this is slower than a local KV restore but preserves +collective and cache alignment. + ## Tensor Parallelism over RDMA Tensor parallelism runs a single decode across two Macs connected with a @@ -678,8 +1004,9 @@ MODEL=gguf/GLM-5.2-UD-IQ2_XXS_RoutedIQ2XXS_blk78Q2K.gguf The active verbs device and IPv4-mapped GID are selected automatically. If that is ambiguous, add `--rdma-device rdma_en6 --rdma-gid-index 1` on the worker and the matching `rdma_en1` flags on the coordinator. Use `--transport tcp` on both -sides to force TCP. Tensor parallel roles are currently exposed by the `ds4` -CLI, not by `ds4-server` or `ds4-agent`. +sides to force TCP. Network tensor-parallel roles are exposed by `ds4`, +`ds4-bench`, `ds4-server`, `ds4-agent`, and `ds4-eval`; workers may use any of +those frontends with matching model, mode, rank, and context options. Startup takes about 9 seconds per machine: each rank pre-faults its ~100 GiB shard from SSD and pins it through a Metal residency set. diff --git a/cuda/mmq/ds4_mmq.cu b/cuda/mmq/ds4_mmq.cu index b424635320..a6544fb0e8 100644 --- a/cuda/mmq/ds4_mmq.cu +++ b/cuda/mmq/ds4_mmq.cu @@ -3677,12 +3677,12 @@ __global__ void q8_0_aligned_dense_vec_kernel( const int4 *qs, // aligned codes, 2 int4 per block const __half *dq, // block scales const block_q8_1 *x8, // [K/32] canonical Q8_1 activation - int M, + int row0, int nb) // blocks per row = K/32 { const int row = blockIdx.x; const int lane = threadIdx.x; - const long long rbase = (long long)row * nb; + const long long rbase = (long long)(row0 + row) * nb; float acc = 0.0f; for (int b0 = 0; b0 < nb; b0 += 32) { @@ -3720,12 +3720,13 @@ __global__ void q8_0_aligned_dense_vec_nc_kernel( const int4 *qs, // aligned codes, 2 int4 per block const __half *dq, // block scales const block_q8_1 *x8, // [NC * nb], col stride nb + int row0, int M, int nb) // blocks per row = K/32 { const int row = blockIdx.x; const int lane = threadIdx.x; - const long long rbase = (long long)row * nb; + const long long rbase = (long long)(row0 + row) * nb; float acc[NC]; #pragma unroll @@ -3768,10 +3769,10 @@ extern "C" uint64_t ds4_mmq_q8_0_aligned_bytes(int M, int K) { return dq_bytes + nblk * 32u; } -extern "C" int ds4_mmq_q8_0_aligned_dense_vec( +extern "C" int ds4_mmq_q8_0_aligned_dense_vec_rows( const void * W_aligned, const float * X_f32, float * out_f32, - int M, int N, int K, cudaStream_t stream) { - const char *tag = "ds4_mmq_q8_0_aligned_dense_vec"; + int M_total, int row0, int M, int N, int K, cudaStream_t stream) { + const char *tag = "ds4_mmq_q8_0_aligned_dense_vec_rows"; if (!W_aligned || !X_f32 || !out_f32) { fprintf(stderr, "%s: null pointer\n", tag); return -1; @@ -3779,7 +3780,9 @@ extern "C" int ds4_mmq_q8_0_aligned_dense_vec( // K % 1024: the kernel's 32-blocks-per-pass loop needs nb % 32 == 0. // N covers the decode/verify-width envelope (mmvq batch bound); K % 1024 // also guarantees ne10_padded == K, so the q8_1 col stride is exactly nb. - if (N < 1 || N > 8 || M <= 0 || K <= 0 || K % 1024 != 0) return -1; + if (N < 1 || N > 8 || M_total <= 0 || row0 < 0 || M <= 0 || + row0 > M_total || M > M_total - row0 || + K <= 0 || K % 1024 != 0) return -1; const int dev = ggml_cuda_get_device(); ggml_backend_cuda_context * ctx = get_ctx_for_device(dev); @@ -3816,7 +3819,7 @@ extern "C" int ds4_mmq_q8_0_aligned_dense_vec( } } - const uint64_t nblk = (uint64_t)M * (uint64_t)(K / 32); + const uint64_t nblk = (uint64_t)M_total * (uint64_t)(K / 32); const uint64_t dq_bytes = (nblk * 2u + 63u) & ~63ull; const int4 *qsp = (const int4 *)((const char *)W_aligned + dq_bytes); const __half *dqp = (const __half *)W_aligned; @@ -3824,15 +3827,15 @@ extern "C" int ds4_mmq_q8_0_aligned_dense_vec( switch (N) { case 1: q8_0_aligned_dense_vec_kernel<<<(unsigned)M, 32, 0, stream>>>( - out_f32, qsp, dqp, x8p, M, K / 32); + out_f32, qsp, dqp, x8p, row0, K / 32); break; - case 2: q8_0_aligned_dense_vec_nc_kernel<2><<<(unsigned)M, 32, 0, stream>>>(out_f32, qsp, dqp, x8p, M, K / 32); break; - case 3: q8_0_aligned_dense_vec_nc_kernel<3><<<(unsigned)M, 32, 0, stream>>>(out_f32, qsp, dqp, x8p, M, K / 32); break; - case 4: q8_0_aligned_dense_vec_nc_kernel<4><<<(unsigned)M, 32, 0, stream>>>(out_f32, qsp, dqp, x8p, M, K / 32); break; - case 5: q8_0_aligned_dense_vec_nc_kernel<5><<<(unsigned)M, 32, 0, stream>>>(out_f32, qsp, dqp, x8p, M, K / 32); break; - case 6: q8_0_aligned_dense_vec_nc_kernel<6><<<(unsigned)M, 32, 0, stream>>>(out_f32, qsp, dqp, x8p, M, K / 32); break; - case 7: q8_0_aligned_dense_vec_nc_kernel<7><<<(unsigned)M, 32, 0, stream>>>(out_f32, qsp, dqp, x8p, M, K / 32); break; - case 8: q8_0_aligned_dense_vec_nc_kernel<8><<<(unsigned)M, 32, 0, stream>>>(out_f32, qsp, dqp, x8p, M, K / 32); break; + case 2: q8_0_aligned_dense_vec_nc_kernel<2><<<(unsigned)M, 32, 0, stream>>>(out_f32, qsp, dqp, x8p, row0, M, K / 32); break; + case 3: q8_0_aligned_dense_vec_nc_kernel<3><<<(unsigned)M, 32, 0, stream>>>(out_f32, qsp, dqp, x8p, row0, M, K / 32); break; + case 4: q8_0_aligned_dense_vec_nc_kernel<4><<<(unsigned)M, 32, 0, stream>>>(out_f32, qsp, dqp, x8p, row0, M, K / 32); break; + case 5: q8_0_aligned_dense_vec_nc_kernel<5><<<(unsigned)M, 32, 0, stream>>>(out_f32, qsp, dqp, x8p, row0, M, K / 32); break; + case 6: q8_0_aligned_dense_vec_nc_kernel<6><<<(unsigned)M, 32, 0, stream>>>(out_f32, qsp, dqp, x8p, row0, M, K / 32); break; + case 7: q8_0_aligned_dense_vec_nc_kernel<7><<<(unsigned)M, 32, 0, stream>>>(out_f32, qsp, dqp, x8p, row0, M, K / 32); break; + case 8: q8_0_aligned_dense_vec_nc_kernel<8><<<(unsigned)M, 32, 0, stream>>>(out_f32, qsp, dqp, x8p, row0, M, K / 32); break; } err = cudaGetLastError(); if (err != cudaSuccess) { @@ -3842,6 +3845,13 @@ extern "C" int ds4_mmq_q8_0_aligned_dense_vec( return 0; } +extern "C" int ds4_mmq_q8_0_aligned_dense_vec( + const void * W_aligned, const float * X_f32, float * out_f32, + int M, int N, int K, cudaStream_t stream) { + return ds4_mmq_q8_0_aligned_dense_vec_rows( + W_aligned, X_f32, out_f32, M, 0, M, N, K, stream); +} + // --------------------------------------------------------------------------- // Aligned row-pair-SoA Q2_K routed-expert decode matvec (megakernel program // M2, moe-down increment). The production down leg runs diff --git a/cuda/mmq/ds4_mmq.h b/cuda/mmq/ds4_mmq.h index fee3530b4b..8b29ee7cf9 100644 --- a/cuda/mmq/ds4_mmq.h +++ b/cuda/mmq/ds4_mmq.h @@ -573,6 +573,21 @@ int ds4_mmq_q8_0_aligned_dense_vec( int K, cudaStream_t stream); +// Row-range form of the aligned Q8_0 decode kernel. W_aligned describes the +// complete M_total-row artifact, while only [row0, row0 + M) is evaluated and +// written densely to out_f32. This lets network TP gather disjoint complete +// rows without changing the single-rank dot-product arithmetic. +int ds4_mmq_q8_0_aligned_dense_vec_rows( + const void * W_aligned, + const float * X_f32, + float * out_f32, + int M_total, + int row0, + int M, + int N, + int K, + cudaStream_t stream); + int ds4_mmq_iq2_xxs_aligned_moe_vec( const void * W_aligned, const float * X_f32, diff --git a/cuda/mmq/test/test_q8_aligned_rows.cu b/cuda/mmq/test/test_q8_aligned_rows.cu new file mode 100644 index 0000000000..7d0c415576 --- /dev/null +++ b/cuda/mmq/test/test_q8_aligned_rows.cu @@ -0,0 +1,155 @@ +// Regression for the aligned Q8_0 output-row slice entry used by CUDA +// network tensor parallelism. A sliced launch must be bit-identical to the +// corresponding rows of a complete launch for both decode and small verify +// widths. + +#include "ds4_mmq.h" + +#include +#include + +#include +#include +#include +#include +#include +#include + +// ds4_mmq.o optionally consumes producer-folded Q8_1 activations supplied by +// ds4_cuda.cu. This standalone kernel test intentionally exercises the +// ordinary quantizer, so report that no folded activation is available. +extern "C" int ds4_cuda_q8_fold_take_q81( + const void *src, uint64_t in_dim, const void **q81) { + (void)src; + (void)in_dim; + (void)q81; + return 0; +} + +#define CUDA_CHECK(expr) do { \ + const cudaError_t err_ = (expr); \ + if (err_ != cudaSuccess) { \ + fprintf(stderr, "%s:%d: %s: %s\n", __FILE__, __LINE__, #expr, \ + cudaGetErrorString(err_)); \ + return 1; \ + } \ +} while (0) + +static int check_width( + const void *artifact, + const float *x, + float *full, + float *slice, + int m_total, + int row0, + int rows, + int n, + int k, + cudaStream_t stream) { + if (ds4_mmq_q8_0_aligned_dense_vec( + artifact, x, full, m_total, n, k, stream) != 0 || + ds4_mmq_q8_0_aligned_dense_vec_rows( + artifact, x, slice, m_total, row0, rows, n, k, stream) != 0) { + fprintf(stderr, "aligned Q8 row launch rejected N=%d\n", n); + return 0; + } + if (cudaStreamSynchronize(stream) != cudaSuccess) return 0; + + std::vector got_full((size_t)n * m_total); + std::vector got_slice((size_t)n * rows); + if (cudaMemcpy(got_full.data(), full, + got_full.size() * sizeof(float), + cudaMemcpyDeviceToHost) != cudaSuccess || + cudaMemcpy(got_slice.data(), slice, + got_slice.size() * sizeof(float), + cudaMemcpyDeviceToHost) != cudaSuccess) { + return 0; + } + for (int col = 0; col < n; col++) { + if (memcmp(got_full.data() + (size_t)col * m_total + row0, + got_slice.data() + (size_t)col * rows, + (size_t)rows * sizeof(float)) != 0) { + fprintf(stderr, + "aligned Q8 row mismatch N=%d column=%d rows=%d:%d\n", + n, col, row0, row0 + rows - 1); + return 0; + } + } + return 1; +} + +int main(void) { + constexpr int m_total = 256; + constexpr int row0 = 64; + constexpr int rows = 64; + constexpr int k = 1024; + constexpr int max_n = 3; + constexpr int blocks = k / 32; + const uint64_t n_blocks = (uint64_t)m_total * blocks; + const uint64_t dq_bytes = (n_blocks * 2u + 63u) & ~63ull; + const uint64_t artifact_bytes = + ds4_mmq_q8_0_aligned_bytes(m_total, k); + if (artifact_bytes != dq_bytes + n_blocks * 32u) { + fprintf(stderr, "aligned Q8 byte geometry mismatch\n"); + return 1; + } + + std::mt19937 rng(0x5148524fu); + std::vector artifact((size_t)artifact_bytes, 0u); + for (uint64_t block = 0; block < n_blocks; block++) { + const __half scale = __float2half_rn( + 0.0025f + (float)(rng() % 2000u) / 100000.0f); + memcpy(artifact.data() + block * 2u, &scale, sizeof(scale)); + int8_t *codes = reinterpret_cast( + artifact.data() + dq_bytes + block * 32u); + for (int i = 0; i < 32; i++) { + codes[i] = (int8_t)((int)(rng() % 255u) - 127); + } + } + std::vector input((size_t)max_n * k); + for (float &value : input) { + value = ((float)(rng() % 20001u) - 10000.0f) / 5000.0f; + } + + if (ds4_mmq_init(0) != 0) { + fprintf(stderr, "ds4_mmq_init failed\n"); + return 1; + } + void *device_artifact = nullptr; + float *device_input = nullptr; + float *device_full = nullptr; + float *device_slice = nullptr; + CUDA_CHECK(cudaMalloc(&device_artifact, artifact.size())); + CUDA_CHECK(cudaMalloc(&device_input, input.size() * sizeof(float))); + CUDA_CHECK(cudaMalloc(&device_full, + (size_t)max_n * m_total * sizeof(float))); + CUDA_CHECK(cudaMalloc(&device_slice, + (size_t)max_n * rows * sizeof(float))); + CUDA_CHECK(cudaMemcpy(device_artifact, artifact.data(), artifact.size(), + cudaMemcpyHostToDevice)); + CUDA_CHECK(cudaMemcpy(device_input, input.data(), + input.size() * sizeof(float), + cudaMemcpyHostToDevice)); + cudaStream_t stream = nullptr; + CUDA_CHECK(cudaStreamCreate(&stream)); + + const int ok = + check_width(device_artifact, device_input, + device_full, device_slice, + m_total, row0, rows, 1, k, stream) && + check_width(device_artifact, device_input, + device_full, device_slice, + m_total, row0, rows, max_n, k, stream) && + ds4_mmq_q8_0_aligned_dense_vec_rows( + device_artifact, device_input, device_slice, + m_total, m_total - rows + 1, rows, 1, k, stream) != 0; + + CUDA_CHECK(cudaStreamDestroy(stream)); + CUDA_CHECK(cudaFree(device_slice)); + CUDA_CHECK(cudaFree(device_full)); + CUDA_CHECK(cudaFree(device_input)); + CUDA_CHECK(cudaFree(device_artifact)); + if (!ok) return 1; + puts("aligned Q8 output-row slices: bit-identical PASS (N=1,3)"); + return 0; +} diff --git a/ds4.c b/ds4.c index 449140b523..de87e2dcdd 100644 --- a/ds4.c +++ b/ds4.c @@ -6410,15 +6410,17 @@ static const uint8_t *tensor_expert_bytes( uint64_t *out_dim, uint64_t *row_bytes); -/* TP sharding keeps full layers but restricts every routed-expert blob to - * one contiguous rank range (rank 0 owns the lower expert ids, matching - * ds4_tp_owns_expert in metal/moe.metal). */ +/* Expert parallelism keeps full layers but restricts every routed-expert + * blob to one contiguous rank range. Boundaries use floor(E*r/world), so + * any indivisible remainder is distributed without overlap or gaps. */ static DS4_MAYBE_UNUSED bool weights_model_map_sharded_spans( const ds4_weights *w, const ds4_model *m, - int rank, + uint32_t rank, + uint32_t world, ds4_model_map_span_vec *spans) { - if (!w || !m || !spans || (rank != 0 && rank != 1)) return false; + if (!w || !m || !spans || world < 2 || world > DS4_TP_MAX_WORLD || + rank >= world) return false; memset(spans, 0, sizeof(*spans)); model_map_span_vec_include_one(spans, w->token_embd); for (uint32_t il = 0; il < (uint32_t)DS4_N_LAYER; il++) { @@ -6434,14 +6436,13 @@ static DS4_MAYBE_UNUSED bool weights_model_map_sharded_spans( uint64_t in_dim = 0, out_dim = 0, row_bytes = 0; (void)tensor_expert_bytes(m, x, 0, &in_dim, &out_dim, &row_bytes); const uint64_t expert_bytes = out_dim * row_bytes; - const uint64_t low_experts = x->dim[2] / 2; - const uint64_t first_expert = rank == 1 ? low_experts : 0; - const uint64_t owned_experts = rank == 1 ? - x->dim[2] - low_experts : low_experts; + const uint64_t first_expert = x->dim[2] * rank / world; + const uint64_t last_expert = x->dim[2] * (rank + 1u) / world; + const uint64_t owned_experts = last_expert - first_expert; const uint64_t owned_bytes = owned_experts * expert_bytes; const uint64_t lo = x->abs_offset + first_expert * expert_bytes; - /* Kernels index experts from the blob base, so the owned range - * must sit in one contiguous view. Rank 1 takes any remainder. */ + /* Ownership-aware kernels remap selected IDs to this compact + * contiguous view before reading it. */ model_map_span_vec_append(spans, lo, lo + owned_bytes, true); if (owned_bytes > spans->max_tensor_bytes) { spans->max_tensor_bytes = owned_bytes; @@ -15257,10 +15258,14 @@ typedef struct { bool streaming_static_decode_map_current; float *cpu_router_norm; - /* Metal network tensor parallelism. These views alias engine-owned - * transport slabs except tp_logits_half, whose view object is session-owned. */ + /* Network parallelism. Metal uses the pair-transport slab views below; + * CUDA collectives keep them NULL and reduce rank-local partials through + * NCCL. tensor_parallel distinguishes full TP (attention/shared/output) + * from expert-only sharding. */ uint32_t tp_world; uint32_t tp_rank; + bool tp_collective; + bool tensor_parallel; ds4_gpu_tensor **tp_out; ds4_gpu_tensor **tp_in; ds4_gpu_tensor **tp_batch_out; @@ -16334,6 +16339,72 @@ static bool metal_graph_needs_ffn_out(const ds4_gpu_graph *g, uint32_t il, uint3 metal_graph_debug_wants("ffn_out", il, pos); } +static bool deepseek_graph_pair_transport(const ds4_gpu_graph *g) { + return g && !g->tp_collective && g->tp_world == 2u && + g->tp_rank < 2u && g->tp_out && g->tp_in; +} + +static bool deepseek_graph_collective(const ds4_gpu_graph *g) { +#if defined(__APPLE__) || defined(DS4_ROCM_BUILD) + (void)g; + return false; +#else + return g && g->tp_collective && g->tp_world >= 2u && + g->tp_world <= DS4_TP_MAX_WORLD && g->tp_rank < g->tp_world; +#endif +} + +static bool deepseek_graph_tensor_parallel(const ds4_gpu_graph *g) { + return deepseek_graph_collective(g) && g->tensor_parallel; +} + +static bool deepseek_graph_tp_component(const ds4_gpu_graph *g, + const char *replicate_env) { + if (!deepseek_graph_tensor_parallel(g)) return false; + const char *env = getenv(replicate_env); + return !env || !env[0] || strcmp(env, "0") == 0; +} + +static bool deepseek_graph_tp_range(const ds4_gpu_graph *g, + uint32_t total, + bool split, + uint32_t *base, + uint32_t *count) { + if (!g || !base || !count || total == 0u) return false; + if (!split) { + *base = 0u; + *count = total; + return true; + } + if (!deepseek_graph_collective(g) || (total % g->tp_world) != 0u) { + return false; + } + *count = total / g->tp_world; + *base = g->tp_rank * *count; + return *count != 0u; +} + +static bool deepseek_graph_expert_range(const ds4_gpu_graph *g, + uint32_t total, + uint32_t *base, + uint32_t *count) { + if (!deepseek_graph_collective(g) || !base || !count || total == 0u) { + return false; + } + return ds4_tp_partition(total, g->tp_rank, g->tp_world, base, count); +} + +static bool deepseek_graph_allreduce(ds4_gpu_tensor *tensor, uint64_t count) { +#if defined(__APPLE__) || defined(DS4_ROCM_BUILD) + (void)tensor; + (void)count; + return false; +#else + return tensor && ds4_gpu_tp_collective_ready() && + ds4_gpu_tp_allreduce_f32(tensor, count) != 0; +#endif +} + /* tier-aware lazy allocator. The Class P ffn_out scratch is * created on demand the first time a layer that materializes ffn_out runs * on a tier; subsequent visits to the same tier reuse the existing slot. @@ -20210,6 +20281,22 @@ static bool metal_graph_matmul_dense_quant_kslice( uint64_t out_dim, const ds4_gpu_tensor *x, uint64_t x_elem_off); +static bool deepseek_graph_collective_dense_output_rows( + ds4_gpu_graph *g, + ds4_gpu_tensor *out, + const ds4_model *model, + const ds4_tensor *weight, + uint64_t in_dim, + uint32_t out_dim, + const ds4_gpu_tensor *x); +static bool deepseek_graph_collective_attention_output_decode( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_layer_weights *layer, + uint32_t group_dim, + uint32_t rank, + uint32_t n_groups, + const ds4_gpu_tensor *heads); static bool metal_graph_attention_output_dense_quant_low( ds4_gpu_tensor *low, ds4_gpu_graph *g, @@ -20220,7 +20307,7 @@ static bool metal_graph_attention_output_dense_quant_low( uint32_t group0, uint32_t group_cnt, const ds4_gpu_tensor *heads); -static bool metal_graph_attention_output_dense_quant_tp( +static bool metal_graph_attention_output_dense_quant_tp_compact_heads( ds4_gpu_tensor *out, ds4_gpu_tensor *low, ds4_gpu_graph *g, @@ -21795,6 +21882,209 @@ static bool metal_graph_cuda_tp_ep_finish_reduce( return ok; } +/* DeepSeek CUDA network parallelism. Every rank evaluates the replicated + * router, runs only its contiguous resident expert range, and gathers the + * resulting expert slots. Full TP additionally shards shared-expert decode + * lanes and complete down-projection rows; expert-only mode leaves that dense + * path replicated. */ +static bool deepseek_graph_collective_ffn_decode( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_layer_weights *layer, + uint32_t il, + uint32_t pos, + bool keep_ffn_out) { + if (!deepseek_graph_collective(g) || !model || !layer) return false; + + const uint32_t shared_dim = (uint32_t)layer->ffn_gate_shexp->dim[1]; + const uint32_t expert_in_dim = (uint32_t)layer->ffn_gate_exps->dim[0]; + const uint32_t expert_mid_dim = (uint32_t)layer->ffn_gate_exps->dim[1]; + const uint32_t down_in_dim = (uint32_t)layer->ffn_down_exps->dim[0]; + const uint32_t routed_out_dim = (uint32_t)layer->ffn_down_exps->dim[1]; + const uint64_t gate_row_bytes = routed_expert_row_bytes(layer->ffn_gate_exps); + const uint64_t gate_expert_bytes = + (uint64_t)expert_mid_dim * gate_row_bytes; + const uint64_t down_row_bytes = routed_expert_row_bytes(layer->ffn_down_exps); + const uint64_t down_expert_bytes = + (uint64_t)routed_out_dim * down_row_bytes; + uint32_t expert_base = 0; + uint32_t expert_count = 0; + uint32_t shared_base = 0; + uint32_t shared_count = 0; + const bool split_shared = deepseek_graph_tp_component( + g, "DS4_DEEPSEEK_TP_REPLICATE_SHARED"); + bool ok = deepseek_graph_expert_range(g, DS4_N_EXPERT, + &expert_base, &expert_count) && + deepseek_graph_tp_range(g, shared_dim, split_shared, + &shared_base, &shared_count) && + shared_count != 0u && (shared_count % 32u) == 0u && + layer->ffn_gate_shexp->type == DS4_TENSOR_Q8_0 && + layer->ffn_up_shexp->type == DS4_TENSOR_Q8_0 && + layer->ffn_down_shexp->type == DS4_TENSOR_Q8_0; + + uint64_t shared_row_bytes = 0; + if (ok) { + ok = metal_graph_dense_quant_row_bytes(layer->ffn_gate_shexp, + DS4_N_EMBD, + &shared_row_bytes); + } + const uint64_t shared_offset = + (uint64_t)shared_base * shared_row_bytes; + ds4_gpu_tensor *shared_gate_view = NULL; + ds4_gpu_tensor *shared_up_view = NULL; + ds4_gpu_tensor *shared_mid_view = NULL; + if (ok && split_shared) { + const uint64_t view_offset = (uint64_t)shared_base * sizeof(float); + const uint64_t view_bytes = (uint64_t)shared_count * sizeof(float); + shared_gate_view = ds4_gpu_tensor_view(metal_graph_shared_gate(g), + view_offset, view_bytes); + shared_up_view = ds4_gpu_tensor_view(metal_graph_shared_up(g), + view_offset, view_bytes); + shared_mid_view = ds4_gpu_tensor_view(metal_graph_shared_mid(g), + view_offset, view_bytes); + ok = shared_gate_view && shared_up_view && shared_mid_view && + ds4_gpu_tensor_fill_f32(metal_graph_shared_mid(g), + 0.0f, + shared_dim) != 0; + } + if (ok) { + ok = ds4_gpu_shared_gate_up_swiglu_q8_0_tensor( + split_shared ? shared_gate_view : metal_graph_shared_gate(g), + split_shared ? shared_up_view : metal_graph_shared_up(g), + split_shared ? shared_mid_view : metal_graph_shared_mid(g), + model->map, + model->size, + layer->ffn_gate_shexp->abs_offset + shared_offset, + layer->ffn_up_shexp->abs_offset + shared_offset, + DS4_N_EMBD, + shared_count, + metal_graph_ffn_norm(g), + DS4_SWIGLU_CLAMP_EXP) != 0; + } + if (ok && split_shared) { + ok = deepseek_graph_allreduce(metal_graph_shared_mid(g), + shared_dim) && + deepseek_graph_collective_dense_output_rows( + g, + metal_graph_shared_out(g), + model, + layer->ffn_down_shexp, + shared_dim, + DS4_N_EMBD, + metal_graph_shared_mid(g)); + } else if (ok) { + ok = metal_graph_matmul_dense_quant_tensor( + metal_graph_shared_out(g), + model, + layer->ffn_down_shexp, + shared_dim, + DS4_N_EMBD, + metal_graph_shared_mid(g), + 1); + } + ds4_gpu_tensor_free(shared_mid_view); + ds4_gpu_tensor_free(shared_up_view); + ds4_gpu_tensor_free(shared_gate_view); + if (ok) { + ok = ds4_gpu_routed_moe_one_owned_tensor( + metal_graph_routed_out(g), + metal_graph_routed_gate(g), + metal_graph_routed_up(g), + metal_graph_routed_mid(g), + metal_graph_routed_down(g), + model->map, + model->size, + layer->ffn_gate_exps->abs_offset, + layer->ffn_up_exps->abs_offset, + layer->ffn_down_exps->abs_offset, + layer->ffn_gate_exps->type, + layer->ffn_down_exps->type, + gate_expert_bytes, + gate_row_bytes, + down_expert_bytes, + down_row_bytes, + expert_in_dim, + down_in_dim, + routed_out_dim, + metal_graph_router_selected(g), + metal_graph_router_weights(g), + DS4_N_EXPERT, + DS4_N_EXPERT_USED, + expert_base, + expert_count, + DS4_SWIGLU_CLAMP_EXP, + metal_graph_ffn_norm(g), + NULL, + false, + NULL) != 0; + } + + if (ok) { + ok = deepseek_graph_allreduce( + metal_graph_routed_down(g), + (uint64_t)DS4_N_EXPERT_USED * DS4_N_EMBD); + } + if (ok) { + ok = ds4_gpu_routed_moe_slots_sum_tensor( + metal_graph_routed_out(g), + metal_graph_routed_down(g), + DS4_N_EMBD, + DS4_N_EXPERT_USED, + 1u) != 0; + } + if (ok) { + metal_graph_debug_dump_tensor( + "ffn_moe_down", metal_graph_routed_down(g), + (uint64_t)DS4_N_EXPERT_USED * DS4_N_EMBD, il, pos); + metal_graph_debug_dump_tensor( + "ffn_moe_out", metal_graph_routed_out(g), + DS4_N_EMBD, il, pos); + } + if (ok && split_shared) { + ok = ds4_gpu_add_tensor(metal_graph_routed_out(g), + metal_graph_routed_out(g), + metal_graph_shared_out(g), + DS4_N_EMBD) != 0; + } + + if (ok && keep_ffn_out) { + ok = metal_graph_ensure_ffn_out(g); + if (ok && split_shared) { + ok = ds4_gpu_tensor_copy(metal_graph_ffn_out(g), 0, + metal_graph_routed_out(g), 0, + (uint64_t)DS4_N_EMBD * sizeof(float)) != 0; + } else if (ok) { + ok = ds4_gpu_add_tensor(metal_graph_ffn_out(g), + metal_graph_routed_out(g), + metal_graph_shared_out(g), + DS4_N_EMBD) != 0; + } + } + if (ok && keep_ffn_out) { + metal_graph_debug_dump_tensor("ffn_out", metal_graph_ffn_out(g), + DS4_N_EMBD, il, pos); + } + if (ok && split_shared) { + ok = ds4_gpu_hc_expand_tensor(metal_graph_after_ffn_hc(g), + metal_graph_routed_out(g), + metal_graph_after_attn_hc(g), + metal_graph_hc_post(g), + metal_graph_hc_comb(g), + DS4_N_EMBD, + DS4_N_HC) != 0; + } else if (ok) { + ok = ds4_gpu_hc_expand_add_split_tensor( + metal_graph_after_ffn_hc(g), + metal_graph_routed_out(g), + metal_graph_shared_out(g), + metal_graph_after_attn_hc(g), + metal_graph_hc_split(g), + DS4_N_EMBD, + DS4_N_HC) != 0; + } + return ok; +} + typedef enum { METAL_DECODE_LAYER_FULL = 0, METAL_DECODE_LAYER_TO_FFN, @@ -21892,9 +22182,13 @@ static bool metal_graph_encode_decode_layer_phase( const int cuda_tp_home_tier = g->active_tier; const int cuda_tp_partner_tier = g->cuda_tp_decode ? metal_graph_cuda_tp_partner_tier(cuda_tp_home_tier) : -1; - const bool tp_split_attn = g->tp_world == 2; + const bool tp_pair_transport = deepseek_graph_pair_transport(g); + const bool tp_collective_tensor = deepseek_graph_tp_component( + g, "DS4_DEEPSEEK_TP_REPLICATE_ATTN"); + const bool tp_split_attn = tp_pair_transport || tp_collective_tensor; + const uint32_t tp_attn_world = tp_collective_tensor ? g->tp_world : 2u; const uint32_t tp_heads = tp_split_attn ? - (uint32_t)DS4_N_HEAD / 2u : (uint32_t)DS4_N_HEAD; + (uint32_t)DS4_N_HEAD / tp_attn_world : (uint32_t)DS4_N_HEAD; const uint32_t tp_head0 = tp_split_attn ? g->tp_rank * tp_heads : 0; bool ok = true; @@ -23518,12 +23812,21 @@ static bool metal_graph_encode_decode_layer_phase( DS4_N_EMBD, DS4_N_HC) != 0; } - } else if (ok && g->tp_world == 2) { + } else if (ok && tp_collective_tensor) { + ok = deepseek_graph_collective_attention_output_decode( + g, + model, + layer, + group_dim, + rank, + n_groups, + metal_graph_heads(g)); + } else if (ok && tp_pair_transport) { /* Group-sliced attention output: this rank computes its half of the * output groups and the matching k-window of the expand projection, * leaving a partial block output in the gate slot. */ const uint32_t tp_groups = n_groups / 2; - ok = metal_graph_attention_output_dense_quant_tp( + ok = metal_graph_attention_output_dense_quant_tp_compact_heads( g->tp_out[il * DS4_TP_GATES_PER_LAYER + DS4_TP_GATE_ATTN], metal_graph_attn_low(g), g, @@ -23536,7 +23839,7 @@ static bool metal_graph_encode_decode_layer_phase( DS4_N_EMBD, metal_graph_heads(g)); } else if (ok && layer->attn_output_a->type != DS4_TENSOR_Q8_0) { - ds4_gpu_tensor *attn_out_dst = g->tp_world == 2 ? + ds4_gpu_tensor *attn_out_dst = tp_pair_transport ? g->tp_out[il * DS4_TP_GATES_PER_LAYER + DS4_TP_GATE_ATTN] : metal_graph_attn_out(g); ok = metal_graph_attention_output_dense_quant_low(metal_graph_attn_low(g), g, @@ -23555,7 +23858,7 @@ static bool metal_graph_encode_decode_layer_phase( metal_graph_attn_low(g), 1); } else if (ok) { - ds4_gpu_tensor *attn_out_dst = g->tp_world == 2 ? + ds4_gpu_tensor *attn_out_dst = tp_pair_transport ? g->tp_out[il * DS4_TP_GATES_PER_LAYER + DS4_TP_GATE_ATTN] : metal_graph_attn_out(g); ok = ds4_gpu_attention_output_q8_batch_tensor(attn_out_dst, metal_graph_attn_low(g), @@ -23569,7 +23872,7 @@ static bool metal_graph_encode_decode_layer_phase( n_groups, DS4_N_EMBD, metal_graph_heads(g), 1) != 0; } - if (ok && g->tp_world == 2) { + if (ok && tp_pair_transport) { /* Gate ATTN: exchange the attention block output with the peer and * rebuild the canonical sum (rank0 first, then rank1) in attn_out * on both ranks — identical expression on both machines keeps them @@ -23866,6 +24169,19 @@ static bool metal_graph_encode_decode_layer_phase( } if (phase == METAL_DECODE_LAYER_TO_ROUTER) return ok; } + if (ok && deepseek_graph_collective(g)) { + const bool keep_collective_ffn = + metal_graph_needs_ffn_out(g, il, pos); + ok = deepseek_graph_collective_ffn_decode( + g, model, layer, il, pos, keep_collective_ffn); + DS4_METAL_PROFILE_DECODE_STAGE("collective_ffn"); + if (ok) { + metal_graph_debug_dump_tensor("hc_ffn_post", + metal_graph_after_ffn_hc(g), + hc_dim, il, pos); + } + return ok; + } const bool external_routed = phase == METAL_DECODE_LAYER_FROM_ROUTER; const bool cuda_tp_shared_requested = g->cuda_tp_shared; const bool cuda_tp_moe_requested = !external_routed && g->cuda_tp_moe; @@ -25122,6 +25438,56 @@ static bool metal_graph_output_logits_head_matmul( uint32_t n_tokens, uint64_t vocab_dim); +static bool deepseek_graph_collective_output_head( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_weights *weights, + uint64_t vocab_dim) { + if (!deepseek_graph_tp_component( + g, "DS4_DEEPSEEK_TP_REPLICATE_OUTPUT") || + !model || !weights || + !weights->output || vocab_dim == 0u || vocab_dim > UINT32_MAX || + weights->output->type != DS4_TENSOR_Q8_0) { + return false; + } + uint32_t vocab_base = 0; + uint32_t vocab_count = 0; + if (!deepseek_graph_tp_range(g, (uint32_t)vocab_dim, true, + &vocab_base, &vocab_count)) { + return false; + } + ds4_gpu_tensor *logits = metal_graph_logits(g); + bool ok = logits && + ds4_gpu_tensor_fill_f32(logits, 0.0f, vocab_dim) != 0; + ds4_gpu_tensor *local_logits = NULL; + uint64_t row_bytes = 0; + if (ok) { + local_logits = ds4_gpu_tensor_view( + logits, + (uint64_t)vocab_base * sizeof(float), + (uint64_t)vocab_count * sizeof(float)); + ok = local_logits && + metal_graph_dense_quant_row_bytes(weights->output, + DS4_N_EMBD, + &row_bytes); + } + if (ok) { + ok = metal_graph_matmul_dense_quant_abs( + local_logits, + model, + weights->output, + weights->output->abs_offset + + (uint64_t)vocab_base * row_bytes, + DS4_N_EMBD, + vocab_count, + metal_graph_output_norm(g), + 1); + } + ds4_gpu_tensor_free(local_logits); + if (ok) ok = deepseek_graph_allreduce(logits, vocab_dim); + return ok; +} + /* Encode the final HC collapse, output norm, and vocab projection on Metal. */ static bool metal_graph_encode_output_head( ds4_gpu_graph *g, @@ -25195,7 +25561,11 @@ static bool metal_graph_encode_output_head( if (ok) { metal_graph_debug_dump_tensor("result_norm", metal_graph_output_norm(g), DS4_N_EMBD, DS4_N_LAYER, 0); } - if (ok && g->tp_world == 2 && g->tp_logits_half) { + if (ok && deepseek_graph_tp_component( + g, "DS4_DEEPSEEK_TP_REPLICATE_OUTPUT")) { + ok = deepseek_graph_collective_output_head(g, model, weights, + vocab_dim); + } else if (ok && g->tp_world == 2 && g->tp_logits_half) { /* Vocab-split: this rank computes its half of the head rows into * its logits view; the halves are bit-identical to the full head * (same kernel, same rows) and the worker ships its half to the @@ -25832,6 +26202,133 @@ static bool metal_graph_matmul_dense_quant_kslice( x_elem_off) != 0; } +/* Preserve the single-rank dot-product order while distributing a dense + * projection. Each rank evaluates complete weight rows, scatters those + * disjoint columns into an otherwise-zero output, then uses the collective + * only as a gather. Unlike an inner-dimension split, no floating-point dot + * product is reassociated across machines. */ +static bool deepseek_graph_collective_dense_output_rows( + ds4_gpu_graph *g, + ds4_gpu_tensor *out, + const ds4_model *model, + const ds4_tensor *weight, + uint64_t in_dim, + uint32_t out_dim, + const ds4_gpu_tensor *x) { + if (!deepseek_graph_tensor_parallel(g) || !out || !model || !weight || + !x || in_dim == 0u || out_dim == 0u) { + return false; + } + uint32_t row0 = 0; + uint32_t row_count = 0; + uint64_t weight_row_bytes = 0; + bool ok = deepseek_graph_tp_range(g, out_dim, true, + &row0, &row_count) && + metal_graph_dense_quant_row_bytes(weight, in_dim, + &weight_row_bytes); + if (!ok || row_count == 0u || weight_row_bytes == 0u || + (uint64_t)row0 > UINT64_MAX / weight_row_bytes) { + return false; + } + const uint64_t weight_offset = + weight->abs_offset + (uint64_t)row0 * weight_row_bytes; + if (weight_offset < weight->abs_offset) { + return false; + } + ok = ds4_gpu_tensor_fill_f32(out, 0.0f, out_dim) != 0; + ds4_gpu_tensor *out_view = ok ? ds4_gpu_tensor_view( + out, + (uint64_t)row0 * sizeof(float), + (uint64_t)row_count * sizeof(float)) : NULL; + ok = ok && out_view && + metal_graph_matmul_dense_quant_abs(out_view, + model, + weight, + weight_offset, + in_dim, + row_count, + x, + 1u); + ds4_gpu_tensor_free(out_view); + return ok && deepseek_graph_allreduce(out, out_dim); +} + +static bool deepseek_graph_collective_attention_output_decode( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_layer_weights *layer, + uint32_t group_dim, + uint32_t rank, + uint32_t n_groups, + const ds4_gpu_tensor *heads) { + if (!g || !model || !layer || !heads || group_dim == 0u || rank == 0u || + n_groups == 0u || !layer->attn_output_a || !layer->attn_output_b) { + return false; + } + uint32_t group0 = 0; + uint32_t group_count = 0; + uint64_t a_row_bytes = 0; + bool ok = deepseek_graph_tp_range(g, n_groups, true, + &group0, &group_count) && + metal_graph_dense_quant_row_bytes(layer->attn_output_a, + group_dim, + &a_row_bytes) && + a_row_bytes != 0u; + const uint64_t low_dim = (uint64_t)n_groups * rank; + const uint64_t local_low_dim = (uint64_t)group_count * rank; + const uint64_t low0 = (uint64_t)group0 * rank; + ds4_gpu_tensor *low_view = NULL; + ds4_tensor out_a_slice; + memset(&out_a_slice, 0, sizeof(out_a_slice)); + if (ok && low0 <= UINT64_MAX / sizeof(float) && + local_low_dim <= UINT64_MAX / sizeof(float)) { + low_view = ds4_gpu_tensor_view( + metal_graph_attn_low(g), + low0 * sizeof(float), + local_low_dim * sizeof(float)); + out_a_slice = *layer->attn_output_a; + const uint64_t a_row0 = (uint64_t)group0 * rank; + ok = low_view && a_row0 <= UINT64_MAX / a_row_bytes; + if (ok) { + const uint64_t a_byte0 = a_row0 * a_row_bytes; + ok = a_byte0 <= UINT64_MAX - out_a_slice.abs_offset; + if (ok) out_a_slice.abs_offset += a_byte0; + } + } else { + ok = false; + } + if (ok) { + ok = ds4_gpu_tensor_fill_f32(metal_graph_attn_low(g), + 0.0f, + low_dim) != 0 && + metal_graph_attention_output_dense_quant_low( + low_view, + g, + model, + &out_a_slice, + group_dim, + rank, + 0u, + group_count, + heads); + } + ds4_gpu_tensor_free(low_view); + if (ok) { + ok = deepseek_graph_allreduce(metal_graph_attn_low(g), low_dim); + } + if (ok) { + ok = deepseek_graph_collective_dense_output_rows( + g, + metal_graph_attn_out(g), + model, + layer->attn_output_b, + low_dim, + DS4_N_EMBD, + metal_graph_attn_low(g)); + } + return ok; +} + static bool metal_graph_attention_output_dense_quant_low( ds4_gpu_tensor *low, ds4_gpu_graph *g, @@ -25897,7 +26394,10 @@ static bool metal_graph_attention_output_dense_quant_low( return ok; } -static bool metal_graph_attention_output_dense_quant_tp( +/* Decode attention writes only this rank's heads, starting at element zero. + * Keep that local buffer coordinate independent from the global output-group + * coordinate used to select the matching A rows and B columns. */ +static bool metal_graph_attention_output_dense_quant_tp_compact_heads( ds4_gpu_tensor *out, ds4_gpu_tensor *low, ds4_gpu_graph *g, @@ -25912,33 +26412,25 @@ static bool metal_graph_attention_output_dense_quant_tp( uint64_t out_dim, const ds4_gpu_tensor *heads) { if (!out || !low || !g || !model || !out_a || !out_b || !heads || - group0 + group_cnt > n_groups_total) { + group0 + group_cnt > n_groups_total || + heads->bytes < (uint64_t)group_cnt * group_dim * sizeof(float)) { return false; } - if (out_a->type == DS4_TENSOR_Q8_0 && out_b->type == DS4_TENSOR_Q8_0) { - return ds4_gpu_attention_output_q8_tp_tensor(out, - low, - model->map, - model->size, - out_a->abs_offset, - out_b->abs_offset, - group_dim, - rank, - n_groups_total, - group0, - group_cnt, - out_dim, - heads) != 0; + uint64_t row_bytes = 0; + if (!metal_graph_dense_quant_row_bytes(out_a, group_dim, &row_bytes)) { + return false; } + ds4_tensor out_a_slice = *out_a; + out_a_slice.abs_offset += (uint64_t)group0 * rank * row_bytes; if (!metal_graph_attention_output_dense_quant_low(low, - g, - model, - out_a, - group_dim, - rank, - group0, - group_cnt, - heads)) { + g, + model, + &out_a_slice, + group_dim, + rank, + 0, + group_cnt, + heads)) { return false; } return metal_graph_matmul_dense_quant_kslice(out, @@ -26045,6 +26537,46 @@ static bool metal_graph_attention_output_dense_quant_batch( return ok; } +static bool deepseek_graph_collective_attention_output_batch( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_layer_weights *layer, + ds4_gpu_tensor *out, + ds4_gpu_tensor *low, + const ds4_gpu_tensor *heads, + uint32_t n_tokens) { + if (!deepseek_graph_tp_component( + g, "DS4_DEEPSEEK_TP_REPLICATE_ATTN") || + !model || !layer || + !out || !low || !heads || n_tokens == 0u || + layer->attn_output_a->type != DS4_TENSOR_Q8_0 || + layer->attn_output_b->type != DS4_TENSOR_Q8_0) { + return false; + } + const uint32_t n_groups = DS4_N_OUT_GROUP; + const uint32_t group_heads = DS4_N_HEAD / n_groups; + const uint32_t group_dim = DS4_N_HEAD_DIM * group_heads; + const uint32_t rank = DS4_N_LORA_O; + /* The optimized Q8 A projection chooses different MMA/token-pair kernels + * as the batch geometry changes. Splitting its groups therefore changes + * rounding before the B projection. Keep prefill attention replicated + * so TP starts decode from the exact single-rank state; decode itself uses + * disjoint A lanes and complete B rows above. */ + return metal_graph_attention_output_dense_quant_batch( + out, + low, + g, + model, + layer->attn_output_a, + layer->attn_output_b, + group_dim, + rank, + n_groups, + DS4_N_EMBD, + heads, + n_tokens); +} + static bool metal_graph_matmul_q8_0_named_tensor( const char *module, uint32_t il, @@ -26851,7 +27383,7 @@ static uint32_t metal_graph_token_adaptive_split_after_layers( second_split_after_layers == 12u && allow_split_flush && pos >= 128u && pos < 2048u && g && !g->quality && !g->ssd_streaming && !g->ssd_streaming_cold && - g->tp_world != 2u && mxfp4_routed && + g->tp_world <= 1u && mxfp4_routed && ds4_gpu_device_is_pre_m5_apple_silicon() && getenv("DS4_METAL_DISABLE_PRE_M5_DECODE_EARLY_SPLIT3") == NULL) { return 3u; @@ -26863,7 +27395,7 @@ static uint32_t metal_graph_token_adaptive_split_after_layers( second_split_after_layers == 0u && allow_split_flush && pos >= 2048u && pos < 2816u && g && !g->quality && !g->ssd_streaming && !g->ssd_streaming_cold && - g->tp_world != 2u && mxfp4_routed && + g->tp_world <= 1u && mxfp4_routed && ds4_gpu_device_is_pre_m5_apple_silicon() && getenv("DS4_METAL_DISABLE_PRE_M5_DECODE_EARLY_SPLIT5") == NULL) { return 5u; @@ -26968,7 +27500,7 @@ static uint32_t metal_graph_token_adaptive_second_split_after_layers( const bool eligible = allow_split_flush && pos < 3328u && g && !g->quality && !g->ssd_streaming && !g->ssd_streaming_cold && - g->tp_world != 2u && mxfp4_routed && + g->tp_world <= 1u && mxfp4_routed && ds4_gpu_device_is_pre_m5_apple_silicon(); /* Once the raw SWA window is full, short eval prompts leave less GPU work * in the four-layer prefix while the host still has the same remaining @@ -27474,7 +28006,7 @@ static bool metal_graph_encode_token_raw_swa( return false; } /* Under the vocab split both ranks materialize their logits half. */ - if (g->tp_world == 2 && g->tp_rank == 1 && + if (deepseek_graph_pair_transport(g) && g->tp_rank == 1 && !g->tp_logits_half) need_logits = false; const uint32_t raw_row = pos % g->raw_cap; const uint32_t n_raw = metal_graph_raw_span_for_batch(g, pos, 1); @@ -27550,7 +28082,7 @@ static bool metal_graph_encode_token_raw_swa( * blocked at an earlier gate, making the transport consume a slab * slot before its payload is ready. Keep each TP token in one command * buffer; non-TP decode retains the encode/execute overlap. */ - if (ok && allow_split_flush && g->tp_world != 2 && + if (ok && allow_split_flush && g->tp_world <= 1 && ((split_after_layers != 0 && il + 1u == split_after_layers) || (second_split_after_layers != 0 && il + 1u == second_split_after_layers))) { @@ -28125,7 +28657,7 @@ static bool metal_graph_encode_layer_attention_batch( const bool tp_attn_indexed = zero_prefix && ratio == 4 && tp_attn_n_comp > DS4_N_INDEXER_TOP_K; const bool tp_row_split_attn = - g->tp_world == 2 && + deepseek_graph_pair_transport(g) && g->tp_batch_rows != n_tokens && (tp_attn_full_raw || tp_attn_static_mixed || tp_attn_indexed) && !metal_graph_directional_steering_attn_enabled(g) && @@ -29697,6 +30229,8 @@ static bool metal_graph_encode_layer_attention_batch( if (ok && !attn_out_debug && !tp_row_split_attn && + !deepseek_graph_tp_component( + g, "DS4_DEEPSEEK_TP_REPLICATE_ATTN") && layer->attn_output_a->type == DS4_TENSOR_Q8_0 && layer->attn_output_b->type == DS4_TENSOR_Q8_0 && !metal_graph_directional_steering_attn_enabled(g)) { @@ -29728,7 +30262,15 @@ static bool metal_graph_encode_layer_attention_batch( tp_row_split_attn && (n_tokens % 256u) == 0u && metal_graph_tp_subgate_pipeline(); if (!attn_out_f16) { - if (ok && tp_attn_pipeline) { + if (ok && deepseek_graph_tp_component( + g, "DS4_DEEPSEEK_TP_REPLICATE_ATTN")) { + ok = deepseek_graph_collective_attention_output_batch( + g, model, layer, + metal_graph_batch_attn_out(g), + metal_graph_batch_attn_low(g), + metal_graph_batch_heads(g), + n_tokens); + } else if (ok && tp_attn_pipeline) { /* Sub-chunk pipelined swap: the output projection runs in two * sub-halves of this rank's rows and each sub-half's row swap * is kicked as soon as its rows land in batch_attn_out, so the @@ -29880,6 +30422,165 @@ static bool metal_graph_encode_mixed_routed_rows( uint32_t il, uint32_t prefill_rows); +static bool deepseek_graph_collective_ffn_batch( + ds4_gpu_graph *g, + const ds4_model *model, + const ds4_layer_weights *layer, + uint32_t il, + uint32_t pos0, + uint32_t n_tokens, + ds4_gpu_tensor *next_hc, + ds4_gpu_tensor *hc_split, + bool keep_ffn_out) { + if (!deepseek_graph_collective(g) || !model || !layer || + !next_hc || !hc_split || n_tokens == 0u) { + return false; + } + const uint32_t shared_dim = (uint32_t)layer->ffn_gate_shexp->dim[1]; + const uint32_t expert_in_dim = (uint32_t)layer->ffn_gate_exps->dim[0]; + const uint32_t expert_mid_dim = (uint32_t)layer->ffn_gate_exps->dim[1]; + const uint32_t down_in_dim = (uint32_t)layer->ffn_down_exps->dim[0]; + const uint32_t routed_out_dim = (uint32_t)layer->ffn_down_exps->dim[1]; + const uint64_t gate_row_bytes = routed_expert_row_bytes(layer->ffn_gate_exps); + const uint64_t gate_expert_bytes = + (uint64_t)expert_mid_dim * gate_row_bytes; + const uint64_t down_row_bytes = routed_expert_row_bytes(layer->ffn_down_exps); + const uint64_t down_expert_bytes = + (uint64_t)routed_out_dim * down_row_bytes; + const uint64_t row_elems = (uint64_t)n_tokens * DS4_N_EMBD; + if (row_elems > UINT32_MAX) return false; + + uint32_t expert_base = 0; + uint32_t expert_count = 0; + const bool split_shared = deepseek_graph_tp_component( + g, "DS4_DEEPSEEK_TP_REPLICATE_SHARED"); + bool ok = deepseek_graph_expert_range(g, DS4_N_EXPERT, + &expert_base, &expert_count) && + shared_dim != 0u && (shared_dim % 32u) == 0u && + layer->ffn_gate_shexp->type == DS4_TENSOR_Q8_0 && + layer->ffn_up_shexp->type == DS4_TENSOR_Q8_0 && + layer->ffn_down_shexp->type == DS4_TENSOR_Q8_0; + if (ok) { + /* The complete batch shared expert is deliberately replicated. Its + * optimized kernels are shape-sensitive; decode uses the stable lane + * and complete-output-row partition instead. */ + ok = ds4_gpu_shared_gate_up_swiglu_q8_0_rows_tensor( + metal_graph_batch_shared_gate(g), + metal_graph_batch_shared_up(g), + metal_graph_batch_shared_mid(g), + model->map, + model->size, + layer->ffn_gate_shexp->abs_offset, + layer->ffn_up_shexp->abs_offset, + DS4_N_EMBD, + shared_dim, + metal_graph_batch_ffn_norm(g), + n_tokens, + DS4_SWIGLU_CLAMP_EXP) != 0; + } + if (ok) { + /* Keep the complete shared expert replicated during prefill. The + * CUDA batch matmul selects tiles using the output shape, so a row + * shard is not bit-equivalent even though its mathematical rows are + * independent. Decode below uses the stable row-parallel path. */ + ok = metal_graph_matmul_q8_0_named_tensor( + "shared_down", il, pos0, + metal_graph_batch_shared_out(g), + model, + layer->ffn_down_shexp, + shared_dim, + DS4_N_EMBD, + metal_graph_batch_shared_mid(g), + n_tokens); + } + + if (ok) { + ok = ds4_gpu_routed_moe_batch_owned_tensor( + metal_graph_batch_routed_out(g), + metal_graph_batch_routed_gate(g), + metal_graph_batch_routed_up(g), + metal_graph_batch_routed_mid(g), + metal_graph_batch_routed_down(g), + model->map, + model->size, + layer->ffn_gate_exps->abs_offset, + layer->ffn_up_exps->abs_offset, + layer->ffn_down_exps->abs_offset, + layer->ffn_gate_exps->type, + layer->ffn_down_exps->type, + gate_expert_bytes, + gate_row_bytes, + down_expert_bytes, + down_row_bytes, + expert_in_dim, + down_in_dim, + routed_out_dim, + metal_graph_batch_router_selected(g), + metal_graph_batch_router_weights(g), + DS4_N_EXPERT, + DS4_N_EXPERT_USED, + expert_base, + expert_count, + DS4_SWIGLU_CLAMP_EXP, + metal_graph_batch_ffn_norm(g), + il, + n_tokens, + &g->batch_routed_mid_is_f16, + true) != 0; + } + if (ok) { + ok = deepseek_graph_allreduce( + metal_graph_batch_routed_down(g), + (uint64_t)n_tokens * DS4_N_EXPERT_USED * DS4_N_EMBD); + } + if (ok) { + ok = ds4_gpu_routed_moe_slots_sum_tensor( + metal_graph_batch_routed_out(g), + metal_graph_batch_routed_down(g), + DS4_N_EMBD, + DS4_N_EXPERT_USED, + n_tokens) != 0; + } + if (ok && split_shared) { + ok = ds4_gpu_add_tensor(metal_graph_batch_routed_out(g), + metal_graph_batch_routed_out(g), + metal_graph_batch_shared_out(g), + (uint32_t)row_elems) != 0; + } + if (ok && keep_ffn_out) { + ok = metal_graph_ensure_batch_ffn_out(g); + if (ok && split_shared) { + ok = ds4_gpu_tensor_copy(metal_graph_batch_ffn_out(g), 0, + metal_graph_batch_routed_out(g), 0, + row_elems * sizeof(float)) != 0; + } else if (ok) { + ok = ds4_gpu_add_tensor(metal_graph_batch_ffn_out(g), + metal_graph_batch_routed_out(g), + metal_graph_batch_shared_out(g), + (uint32_t)row_elems) != 0; + } + } + if (ok && split_shared) { + ok = ds4_gpu_hc_expand_split_tensor( + next_hc, + metal_graph_batch_routed_out(g), + metal_graph_batch_after_attn_hc(g), + hc_split, + DS4_N_EMBD, + DS4_N_HC) != 0; + } else if (ok) { + ok = ds4_gpu_hc_expand_add_split_tensor( + next_hc, + metal_graph_batch_routed_out(g), + metal_graph_batch_shared_out(g), + metal_graph_batch_after_attn_hc(g), + hc_split, + DS4_N_EMBD, + DS4_N_HC) != 0; + } + return ok; +} + /* Encode the batched prefill FFN half: HC pre/norm, shared expert, routed * experts, sum, and HC post. A non-empty decode tail has already been * prepared in rows [n_tokens, n_tokens + decode_count); only the routed @@ -30050,6 +30751,34 @@ static bool metal_graph_encode_layer_ffn_batch( } DS4_METAL_PROFILE_FFN_STAGE("router"); + if (deepseek_graph_collective(g)) { + const bool keep_collective_ffn = + metal_graph_needs_ffn_out(g, il, pos0); + if (decode_count != 0) { + fprintf(stderr, + "ds4: DeepSeek network collectives do not support mixed " + "prefill/decode rows yet\n"); + ok = false; + } + if (ok) { + ok = deepseek_graph_collective_ffn_batch( + g, model, layer, il, pos0, n_tokens, + next_hc_view, hc_split_view, keep_collective_ffn); + } + DS4_METAL_PROFILE_FFN_STAGE("collective_ffn"); + if (ok) { + metal_graph_debug_dump_tensor("hc_ffn_post", + metal_graph_batch_next_hc(g), + (uint64_t)n_tokens * hc_dim, + il, pos0); + } + ds4_gpu_tensor_free(next_hc_view); + ds4_gpu_tensor_free(ffn_cur_view); + ds4_gpu_tensor_free(hc_split_view); + ds4_gpu_tensor_free(hc_mix_view); + return ok; + } + if (ok) { ok = metal_graph_cuda_stream_prefill_batch_selected_load(g, model, @@ -35459,7 +36188,7 @@ static bool metal_graph_verify_suffix_tops_impl( * the host encodes the next layers. Submit short prefixes to overlap * that work without changing the verifier's kernels or arithmetic. */ if (ok && capture_dspark_hidden && !verify_profile && - !selected_profile && g->tp_world != 2 && + !selected_profile && g->tp_world <= 1 && ((il + 1u) % 4u) == 0u) { ok = ds4_gpu_flush_commands() != 0; } @@ -36834,6 +37563,9 @@ typedef struct { uint64_t eval_seq; /* leader: mirrored eval counter */ uint64_t next_session_id; /* leader: stable worker-session handle */ int rank; + uint32_t world; + bool collective; + bool tensor_parallel; /* attention/dense split in addition to experts */ bool vocab_split; /* DS4-only: logits halves cross the wire */ bool active; } ds4_engine_tp_state; @@ -36865,6 +37597,7 @@ struct ds4_engine { uint32_t ssd_streaming_full_layers; uint32_t ssd_streaming_preload_experts; uint64_t startup_model_span_bytes; + uint64_t startup_model_additive_artifact_bytes; ds4_ssd_memory_lock simulated_memory; bool quality; bool glm_mtp; @@ -36977,6 +37710,8 @@ static void ds4_engine_print_startup_memory( uint64_t total = kv_bytes; total = ds4_add_sat_u64(total, mem.scratch_bytes); total = ds4_add_sat_u64(total, e->startup_model_span_bytes); + total = ds4_add_sat_u64(total, + e->startup_model_additive_artifact_bytes); total = ds4_add_sat_u64(total, dynamic_expert_cache_bytes); total = ds4_add_sat_u64(total, e->ssd_streaming_full_layer_bytes); total = ds4_add_sat_u64(total, expert_reserved_bytes); @@ -36995,6 +37730,12 @@ static void ds4_engine_print_startup_memory( ds4_bytes_to_gib(mem.compressed_bytes), ds4_bytes_to_gib(mem.scratch_bytes), ds4_bytes_to_gib(e->startup_model_span_bytes)); + if (e->startup_model_additive_artifact_bytes != 0) { + fprintf(stderr, + " + aligned dense artifacts %.2f GiB", + ds4_bytes_to_gib( + e->startup_model_additive_artifact_bytes)); + } if (e->ssd_streaming_full_layer_bytes != 0) { fprintf(stderr, " + full-layer experts %.2f GiB", @@ -38981,12 +39722,13 @@ typedef struct { bool ssd_streaming_cold; bool generic_routed_moe; bool streaming_static_decode_map_current; - /* Tensor parallelism (50/50 expert sharding): tp_world 2 means - * this rank computes only its contiguous half of the routed experts - * and exchanges the 24KB routed-FFN partial at one gate per sparse - * layer. Views alias the engine's TP slab slots [layer*2 + FFN]. */ + /* Network parallelism. Every mode owns a contiguous routed-expert + * range. tensor_parallel keeps dense prefill replicated for exact batch + * arithmetic. Decode splits attention heads and dense FFN lanes, gathers + * those disjoint intermediates, then evaluates complete output rows. */ uint32_t tp_world; uint32_t tp_rank; + bool tensor_parallel; ds4_gpu_tensor **tp_out; ds4_gpu_tensor **tp_in; /* Prefill batch gate bounce buffers (shared storage; grow on demand). */ @@ -39731,20 +40473,165 @@ static bool glm_tp_validate_ownership_kernels( uint32_t *bad_layer, uint32_t *bad_type) { if (!weights) return false; - for (uint32_t il = 0; il < DS4_N_LAYER; il++) { + uint32_t checked_layers = DS4_N_LAYER; +#if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) + /* CUDA's rank-local path currently covers IQ2_XXS gate/up experts with + * IQ2_XXS or Q2_K down projections. A fully Q2_K routed layer still uses the + * full-expert CUDA kernel, so it must never run against a sharded map. + * Optional next-token-predict layers are not part of ordinary inference; + * --glm-mtp is rejected for CUDA network mode until they are shard-safe. */ + checked_layers = DS4_N_LAYER > DS4_N_NEXTN_PREDICT + ? DS4_N_LAYER - DS4_N_NEXTN_PREDICT : 0u; +#endif + if (checked_layers == 0) return false; + for (uint32_t il = 0; il < checked_layers; il++) { const ds4_layer_weights *l = &weights->layer[il]; if (!l->ffn_gate_exps) continue; - if (glm_graph_layer_uses_generic_routed_moe(l) || - l->ffn_gate_exps->type == DS4_TENSOR_Q2_K) { +#if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) + if (glm_graph_layer_uses_generic_routed_moe(l) && + l->ffn_up_exps->type == DS4_TENSOR_IQ2_XXS && + (l->ffn_down_exps->type == DS4_TENSOR_IQ2_XXS || + l->ffn_down_exps->type == DS4_TENSOR_Q2_K)) { continue; } +#else + if (glm_graph_layer_uses_generic_routed_moe(l)) continue; + if (l->ffn_gate_exps->type == DS4_TENSOR_Q2_K) continue; +#endif if (bad_layer) *bad_layer = il; - if (bad_type) *bad_type = l->ffn_gate_exps->type; + if (bad_type) { + if (l->ffn_gate_exps->type != DS4_TENSOR_IQ2_XXS) { + *bad_type = l->ffn_gate_exps->type; + } else if (!l->ffn_up_exps || + l->ffn_up_exps->type != DS4_TENSOR_IQ2_XXS) { + *bad_type = l->ffn_up_exps ? l->ffn_up_exps->type : UINT32_MAX; + } else { + *bad_type = l->ffn_down_exps ? l->ffn_down_exps->type + : UINT32_MAX; + } + } return false; } return true; } +/* The first CUDA tensor-parallel implementation deliberately targets the + * Q8_0 dense projections used by GLM-5.2. Output-row slices are contiguous; + * down/output projections use the existing Q8 K-slice kernels. Reject a + * different checkpoint layout before any rank commits model memory. */ +static bool glm_tp_validate_tensor_kernels( + const ds4_weights *weights, + uint32_t world, + uint32_t *bad_layer, + const char **bad_tensor) { + if (!weights || (world != 2u && world != 4u) || + DS4_N_HEAD == 0 || (DS4_N_HEAD % world) != 0u || + DS4_N_EMBD == 0 || (DS4_N_EMBD % world) != 0u || + DS4_N_FF_EXP == 0 || (DS4_N_FF_EXP % world) != 0u) { + if (bad_layer) *bad_layer = 0; + if (bad_tensor) *bad_tensor = "model dimensions"; + return false; + } + const uint32_t normal_layers = + DS4_N_LAYER > DS4_N_NEXTN_PREDICT ? + DS4_N_LAYER - DS4_N_NEXTN_PREDICT : 0u; + for (uint32_t il = 0; il < normal_layers; il++) { + const ds4_layer_weights *l = &weights->layer[il]; +#define DS4_GLM_TP_REQUIRE_Q8(field_) do { \ + if (!l->field_ || l->field_->type != DS4_TENSOR_Q8_0) { \ + if (bad_layer) *bad_layer = il; \ + if (bad_tensor) *bad_tensor = #field_; \ + return false; \ + } \ + } while (0) + DS4_GLM_TP_REQUIRE_Q8(attn_q_b); + DS4_GLM_TP_REQUIRE_Q8(attn_k_b); + DS4_GLM_TP_REQUIRE_Q8(attn_v_b); + DS4_GLM_TP_REQUIRE_Q8(attn_output); + if (il < DS4_N_LEADING_DENSE) { + DS4_GLM_TP_REQUIRE_Q8(ffn_gate); + DS4_GLM_TP_REQUIRE_Q8(ffn_up); + DS4_GLM_TP_REQUIRE_Q8(ffn_down); + if (l->ffn_gate->dim[1] == 0 || + (l->ffn_gate->dim[1] % world) != 0u) { + if (bad_layer) *bad_layer = il; + if (bad_tensor) *bad_tensor = "ffn hidden width"; + return false; + } + } else { + DS4_GLM_TP_REQUIRE_Q8(ffn_gate_shexp); + DS4_GLM_TP_REQUIRE_Q8(ffn_up_shexp); + DS4_GLM_TP_REQUIRE_Q8(ffn_down_shexp); + } +#undef DS4_GLM_TP_REQUIRE_Q8 + } + return normal_layers != 0u; +} + +static bool deepseek_tp_validate_ownership_kernels( + const ds4_weights *weights, + uint32_t *bad_layer, + const char **bad_tensor) { + if (!weights || DS4_N_EXPERT_USED != 6u) { + if (bad_layer) *bad_layer = 0u; + if (bad_tensor) *bad_tensor = "active expert count"; + return false; + } + for (uint32_t il = 0; il < (uint32_t)DS4_N_LAYER; il++) { + const ds4_layer_weights *l = &weights->layer[il]; + if (!l->ffn_gate_exps || !l->ffn_up_exps || !l->ffn_down_exps || + l->ffn_gate_exps->type != DS4_TENSOR_IQ2_XXS || + l->ffn_up_exps->type != DS4_TENSOR_IQ2_XXS || + l->ffn_down_exps->type != DS4_TENSOR_Q2_K) { + if (bad_layer) *bad_layer = il; + if (bad_tensor) *bad_tensor = "IQ2_XXS gate/up plus Q2_K down"; + return false; + } + } + return DS4_N_LAYER != 0u; +} + +static bool deepseek_tp_validate_tensor_kernels( + const ds4_weights *weights, + uint32_t world, + uint32_t *bad_layer, + const char **bad_tensor) { + if (!weights || (world != 2u && world != 4u) || + DS4_N_HEAD == 0u || (DS4_N_HEAD % world) != 0u || + DS4_N_OUT_GROUP == 0u || (DS4_N_OUT_GROUP % world) != 0u || + DS4_N_VOCAB == 0u || (DS4_N_VOCAB % world) != 0u || + !weights->output || weights->output->type != DS4_TENSOR_Q8_0) { + if (bad_layer) *bad_layer = 0u; + if (bad_tensor) *bad_tensor = "model/head dimensions"; + return false; + } + for (uint32_t il = 0; il < (uint32_t)DS4_N_LAYER; il++) { + const ds4_layer_weights *l = &weights->layer[il]; +#define DS4_DEEPSEEK_TP_REQUIRE_Q8(field_) do { \ + if (!l->field_ || l->field_->type != DS4_TENSOR_Q8_0) { \ + if (bad_layer) *bad_layer = il; \ + if (bad_tensor) *bad_tensor = #field_; \ + return false; \ + } \ + } while (0) + DS4_DEEPSEEK_TP_REQUIRE_Q8(attn_q_b); + DS4_DEEPSEEK_TP_REQUIRE_Q8(attn_output_a); + DS4_DEEPSEEK_TP_REQUIRE_Q8(attn_output_b); + DS4_DEEPSEEK_TP_REQUIRE_Q8(ffn_gate_shexp); + DS4_DEEPSEEK_TP_REQUIRE_Q8(ffn_up_shexp); + DS4_DEEPSEEK_TP_REQUIRE_Q8(ffn_down_shexp); +#undef DS4_DEEPSEEK_TP_REQUIRE_Q8 + const uint64_t shared_dim = l->ffn_gate_shexp->dim[1]; + if (shared_dim == 0u || (shared_dim % world) != 0u || + ((shared_dim / world) % 32u) != 0u) { + if (bad_layer) *bad_layer = il; + if (bad_tensor) *bad_tensor = "shared expert width"; + return false; + } + } + return DS4_N_LAYER != 0u; +} + static bool glm_graph_stream_map_token( ds4_glm_gpu_graph *g, const ds4_model *model, @@ -41239,6 +42126,247 @@ static bool glm_graph_prefill_stage_boundary( return true; } +#if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) +static bool glm_graph_tp_expert_range(const ds4_glm_gpu_graph *g, + uint32_t total, + uint32_t *base, + uint32_t *count) { + if (!g || g->tp_world < 2 || g->tp_world > DS4_TP_MAX_WORLD || + g->tp_rank >= g->tp_world || total == 0) { + return false; + } + return ds4_tp_partition(total, g->tp_rank, g->tp_world, base, count); +} +#endif + +static bool glm_graph_tensor_parallel(const ds4_glm_gpu_graph *g) { +#if defined(__APPLE__) || defined(DS4_ROCM_BUILD) + (void)g; + return false; +#else + return g && g->tensor_parallel && g->tp_world > 1u && + g->tp_world <= DS4_TP_MAX_WORLD && g->tp_rank < g->tp_world; +#endif +} + +/* GLM's optimized batch kernels select different numerical paths for + * different output shapes. Keep dense prefill replicated by default so every + * rank starts decode from the single-rank state. This diagnostic switch keeps + * the old split-prefill path available for performance isolation only. */ +static bool glm_graph_tp_split_prefill(const ds4_glm_gpu_graph *g) { + return glm_graph_tensor_parallel(g) && + glm_graph_env_truthy(getenv("DS4_GLM_TP_SPLIT_PREFILL")); +} + +static bool glm_graph_tp_range(const ds4_glm_gpu_graph *g, + uint32_t total, + uint32_t *base, + uint32_t *count) { + if (!base || !count || total == 0u) return false; + if (!glm_graph_tensor_parallel(g)) { + *base = 0u; + *count = total; + return true; + } + if ((total % g->tp_world) != 0u) return false; + *count = total / g->tp_world; + *base = g->tp_rank * *count; + return true; +} + +static bool glm_graph_tp_rowslice_offset(const ds4_tensor *w, + uint64_t in_dim, + uint64_t row0, + uint64_t *offset) { + uint64_t row_bytes = 0; + if (!w || !offset || w->type != DS4_TENSOR_Q8_0 || + !tensor_nbytes(w->type, in_dim, &row_bytes) || + row0 > (UINT64_MAX - w->abs_offset) / row_bytes) { + return false; + } + *offset = w->abs_offset + row0 * row_bytes; + return true; +} + +/* Column-parallel projection: each rank evaluates contiguous output rows and + * packs them at the start of every activation row. */ +static bool glm_graph_tp_rowslice_matmul( + ds4_gpu_tensor *out, + const ds4_model *model, + const ds4_tensor *w, + uint64_t in_dim, + uint64_t row0, + uint64_t rows, + const ds4_gpu_tensor *x, + uint32_t n_tokens) { + uint64_t offset = 0; + if (!out || !model || !w || !x || rows == 0u || n_tokens == 0u || + !glm_graph_tp_rowslice_offset(w, in_dim, row0, &offset)) { + return false; + } + if (n_tokens == 1u) { + return ds4_gpu_matmul_quant_decode_mpp_model_view_tensor( + out, model->map, model->size, offset, w->type, + in_dim, rows, x, 1u) != 0; + } + return ds4_gpu_matmul_quant_tensor(out, + model->map, + model->size, + offset, + w->type, + in_dim, + rows, + x, + n_tokens) != 0; +} + +/* Preserve the reference dense-FFN arithmetic while sharding its hidden + * rows. In particular, the single-token CUDA fast path consumes the f32 + * activation directly; routing these projections through the generic Q8 + * matmul would requantize the activation and introduce a much larger error + * than the K-slice reduction itself. */ +static bool glm_graph_tp_gate_up_swiglu( + ds4_gpu_tensor *gate, + ds4_gpu_tensor *up, + ds4_gpu_tensor *mid, + const ds4_model *model, + const ds4_tensor *gate_w, + const ds4_tensor *up_w, + uint64_t in_dim, + uint64_t row0, + uint64_t rows, + const ds4_gpu_tensor *x, + uint32_t n_tokens) { + uint64_t gate_offset = 0; + uint64_t up_offset = 0; + if (!gate || !up || !mid || !model || !gate_w || !up_w || !x || + gate_w->type != DS4_TENSOR_Q8_0 || + up_w->type != DS4_TENSOR_Q8_0 || rows == 0u || n_tokens == 0u || + !glm_graph_tp_rowslice_offset(gate_w, in_dim, row0, + &gate_offset) || + !glm_graph_tp_rowslice_offset(up_w, in_dim, row0, + &up_offset)) { + return false; + } + if (n_tokens == 1u) { + return ds4_gpu_shared_gate_up_swiglu_q8_0_model_view_tensor( + gate, up, mid, + model->map, model->size, + gate_offset, up_offset, + in_dim, rows, x, 0.0f) != 0; + } + return ds4_gpu_shared_gate_up_swiglu_q8_0_rows_tensor( + gate, up, mid, + model->map, model->size, + gate_offset, up_offset, + in_dim, rows, x, n_tokens, 0.0f) != 0; +} + +/* Row-parallel projection: input rows are packed local K slices and every + * rank produces a full-width partial that is subsequently all-reduced. */ +static bool glm_graph_tp_kslice_matmul( + ds4_gpu_tensor *out, + const ds4_model *model, + const ds4_tensor *w, + uint64_t full_in_dim, + uint64_t k0, + uint64_t k_count, + uint64_t out_dim, + const ds4_gpu_tensor *x, + uint32_t n_tokens) { + if (!out || !model || !w || !x || w->type != DS4_TENSOR_Q8_0 || + k_count == 0u || n_tokens == 0u) { + return false; + } + return ds4_gpu_matmul_q8_0_kslice_rows_tensor(out, + model->map, + model->size, + w->abs_offset, + full_in_dim, + out_dim, + k0, + k_count, + x, + n_tokens) != 0; +} + +static bool glm_graph_tp_allreduce(ds4_gpu_tensor *tensor, uint64_t count) { +#if defined(__APPLE__) || defined(DS4_ROCM_BUILD) + (void)tensor; + (void)count; + return false; +#else + return tensor && ds4_gpu_tp_collective_ready() && + ds4_gpu_tp_allreduce_f32(tensor, count) != 0; +#endif +} + +/* Gather one rank's packed contiguous slice into its global coordinates. + * Every other element is zero, so the all-reduce is a gather rather than a + * floating-point reassociation. */ +static bool glm_graph_tp_gather_vector( + const ds4_glm_gpu_graph *g, + ds4_gpu_tensor *dst, + const ds4_gpu_tensor *local, + uint64_t total, + uint64_t local0, + uint64_t local_count) { + if (!glm_graph_tensor_parallel(g) || !dst || !local || total == 0u || + local_count == 0u || local0 > total || local_count > total - local0 || + total > UINT64_MAX / sizeof(float) || + local0 > UINT64_MAX / sizeof(float) || + local_count > UINT64_MAX / sizeof(float)) { + return false; + } + bool ok = ds4_gpu_tensor_fill_f32(dst, 0.0f, total) != 0; + if (ok) { + ok = ds4_gpu_tensor_copy(dst, + local0 * sizeof(float), + local, + 0, + local_count * sizeof(float)) != 0; + } + return ok && glm_graph_tp_allreduce(dst, total); +} + +/* Preserve each Q8 dot product's single-rank order. Each rank evaluates a + * disjoint set of complete output rows into global coordinates; reducing the + * otherwise-zero tensor only gathers those rows. */ +static bool glm_graph_tp_complete_output_rows( + const ds4_glm_gpu_graph *g, + ds4_gpu_tensor *out, + const ds4_model *model, + const ds4_tensor *w, + uint64_t in_dim, + uint32_t out_dim, + const ds4_gpu_tensor *x) { + if (!glm_graph_tensor_parallel(g) || !out || !model || !w || !x || + w->type != DS4_TENSOR_Q8_0 || in_dim == 0u || out_dim == 0u) { + return false; + } + uint32_t row0 = 0; + uint32_t rows = 0; + if (!glm_graph_tp_range(g, out_dim, &row0, &rows) || rows == 0u) { + return false; + } + bool ok = ds4_gpu_tensor_fill_f32(out, 0.0f, out_dim) != 0; + ds4_gpu_tensor *out_view = ok ? ds4_gpu_tensor_view( + out, + (uint64_t)row0 * sizeof(float), + (uint64_t)rows * sizeof(float)) : NULL; + ok = ok && out_view && + glm_graph_tp_rowslice_matmul(out_view, + model, + w, + in_dim, + row0, + rows, + x, + 1u); + ds4_gpu_tensor_free(out_view); + return ok && glm_graph_tp_allreduce(out, out_dim); +} + static int glm_graph_routed_moe_one_dispatch( const ds4_glm_gpu_graph *g, const ds4_model *model, @@ -41260,7 +42388,7 @@ static int glm_graph_routed_moe_one_dispatch( /* Under the TP expert split only the ownership-aware kernels may run: * the generic mul_mv_id family and the GLM q2_K resident pair/down. * Anything else would silently compute the full expert set. */ - if (g->tp_world == 2 && + if (g->tp_world > 1 && !glm_graph_layer_uses_generic_routed_moe(l) && l->ffn_gate_exps->type != DS4_TENSOR_Q2_K) { fprintf(stderr, @@ -41280,6 +42408,47 @@ static int glm_graph_routed_moe_one_dispatch( } return 0; } +#if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) + if (g->tp_world > 1) { + uint32_t expert_base = 0, expert_count = 0; + if (!glm_graph_tp_expert_range(g, DS4_N_EXPERT, + &expert_base, &expert_count)) { + return 0; + } + return ds4_gpu_routed_moe_batch_owned_tensor( + out, + g->routed_gate, + g->routed_up, + mid, + g->routed_down, + model->map, + model->size, + l->ffn_gate_exps->abs_offset, + l->ffn_up_exps->abs_offset, + l->ffn_down_exps->abs_offset, + l->ffn_gate_exps->type, + l->ffn_down_exps->type, + gate_expert_bytes, + gate_row_bytes, + down_expert_bytes, + down_row_bytes, + DS4_N_EMBD, + DS4_N_FF_EXP, + DS4_N_EMBD, + (ds4_gpu_tensor *)selected, + (ds4_gpu_tensor *)weights, + DS4_N_EXPERT, + DS4_N_EXPERT_USED, + expert_base, + expert_count, + 0.0f, + x, + il, + 1u, + NULL, + false); + } +#endif return ds4_gpu_routed_moe_one_tensor(out, g->routed_gate, g->routed_up, @@ -41388,6 +42557,25 @@ static bool glm_graph_tp_batch_ffn_combine( (uint32_t)((uint64_t)n_tokens * DS4_N_EMBD)) != 0; } +static bool glm_graph_tp_batch_ffn_reduce(ds4_glm_gpu_graph *g, + uint32_t il, + ds4_gpu_tensor *ffn_out, + uint32_t n_tokens) { + if (!g || g->tp_world <= 1) return true; + const bool pair_transport = + g->tp_world == 2 && g->tp_out && g->tp_in; + if (pair_transport) { + return glm_graph_tp_batch_ffn_combine(g, il, ffn_out, n_tokens); + } +#if defined(__APPLE__) || defined(DS4_ROCM_BUILD) + return false; +#else + const uint64_t count = (uint64_t)n_tokens * DS4_N_EMBD; + return ds4_gpu_tp_collective_ready() && + ds4_gpu_tp_allreduce_f32(ffn_out, count) != 0; +#endif +} + static int glm_graph_routed_moe_batch_dispatch( ds4_glm_gpu_graph *g, const ds4_model *model, @@ -41416,6 +42604,47 @@ static int glm_graph_routed_moe_batch_dispatch( l->ffn_gate_exps->type != l->ffn_up_exps->type) { return 0; } +#if !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) + if (g->tp_world > 1) { + uint32_t expert_base = 0, expert_count = 0; + if (!glm_graph_tp_expert_range(g, DS4_N_EXPERT, + &expert_base, &expert_count)) { + return 0; + } + return ds4_gpu_routed_moe_batch_owned_tensor( + out, + g->batch_routed_gate, + g->batch_routed_up, + mid, + g->batch_routed_down, + model->map, + model->size, + l->ffn_gate_exps->abs_offset, + l->ffn_up_exps->abs_offset, + l->ffn_down_exps->abs_offset, + l->ffn_gate_exps->type, + l->ffn_down_exps->type, + gate_expert_bytes, + gate_row_bytes, + down_expert_bytes, + down_row_bytes, + DS4_N_EMBD, + DS4_N_FF_EXP, + DS4_N_EMBD, + (ds4_gpu_tensor *)selected, + (ds4_gpu_tensor *)weights, + DS4_N_EXPERT, + DS4_N_EXPERT_USED, + expert_base, + expert_count, + 0.0f, + x, + il, + n_tokens, + &g->batch_routed_mid_is_f16, + false); + } +#endif return ds4_gpu_routed_moe_batch_tensor(out, g->batch_routed_gate, g->batch_routed_up, @@ -41926,14 +43155,25 @@ static bool glm_graph_encode_sparse_ffn_one( stream_t0 = glm_graph_streaming_async_profile_ms(); } } - /* 50/50 TP: this rank's routed partial goes straight into the - * slab out slot, the gate exchanges it with the peer's half, and the - * commutative add rebuilds the full routed output on both ranks - * bit-identically. The shared expert and everything else stay - * replicated, so no other exchange is needed. */ - const bool tp_split_ffn = g->tp_world == 2 && g->tp_out && g->tp_in; + /* Each rank computes its owned routed experts. Metal's original pair + * path exchanges slab vectors; CUDA reduces the local vector in place + * across two or four ranks with NCCL. Shared experts remain replicated. */ + const bool tp_pair_ffn = + g->tp_world == 2 && g->tp_out && g->tp_in; +#if defined(__APPLE__) || defined(DS4_ROCM_BUILD) + const bool tp_collective_ffn = false; +#else + const bool tp_collective_ffn = + g->tp_world > 1 && !tp_pair_ffn && ds4_gpu_tp_collective_ready(); +#endif + const bool tp_split_ffn = tp_pair_ffn || tp_collective_ffn; const uint32_t tp_ffn_slot = il * DS4_TP_GATES_PER_LAYER + DS4_TP_GATE_FFN; - ds4_gpu_tensor *routed_dst = tp_split_ffn ? g->tp_out[tp_ffn_slot] : ffn_out; + ds4_gpu_tensor *routed_dst = tp_pair_ffn ? + g->tp_out[tp_ffn_slot] : ffn_out; + if (ok && g->tp_world > 1 && !tp_split_ffn) { + fprintf(stderr, "ds4: GLM expert collective is not initialized\n"); + ok = false; + } if (ok && tp_split_ffn && g->ssd_streaming) { fprintf(stderr, "ds4: GLM tensor parallelism requires resident weights\n"); ok = false; @@ -41960,13 +43200,23 @@ static bool glm_graph_encode_sparse_ffn_one( ffn_norm, resident_decode_layer) != 0; } - if (ok && tp_split_ffn) { + if (ok && tp_pair_ffn) { ok = ds4_gpu_tp_gate_encode(il, DS4_TP_GATE_FFN) != 0; if (ok) ok = ds4_gpu_add_tensor(ffn_out, g->tp_out[tp_ffn_slot], g->tp_in[tp_ffn_slot], DS4_N_EMBD) != 0; if (!ok) fprintf(stderr, "ds4: GLM TP gate/combine failed (layer %u)\n", il); + } else if (ok && tp_collective_ffn) { +#if defined(__APPLE__) || defined(DS4_ROCM_BUILD) + ok = false; +#else + ok = ds4_gpu_tp_allreduce_f32(ffn_out, DS4_N_EMBD) != 0; +#endif + if (!ok) { + fprintf(stderr, + "ds4: GLM expert collective failed (layer %u)\n", il); + } } else if (!ok && tp_split_ffn) { fprintf(stderr, "ds4: GLM TP routed dispatch failed before the gate (layer %u)\n", il); } @@ -42084,12 +43334,32 @@ static bool glm_graph_encode_ffn_one_normed_from( if (il < DS4_N_LEADING_DENSE) { const uint64_t hidden = l->ffn_gate->dim[1]; + uint32_t tp_hidden0 = 0; + uint32_t tp_hidden = 0; + if (hidden > UINT32_MAX || + !glm_graph_tp_range(g, (uint32_t)hidden, + &tp_hidden0, &tp_hidden)) { + return false; + } + const bool split_dense = glm_graph_tensor_parallel(g); const bool can_fuse_gate_up = glm_graph_weights_are_q8_0(model, l->ffn_gate->abs_offset, l->ffn_up->abs_offset); - const bool fused_gate_up = can_fuse_gate_up && - ds4_gpu_shared_gate_up_swiglu_q8_0_model_view_tensor( + const bool fused_gate_up = split_dense ? + glm_graph_tp_gate_up_swiglu(ffn_gate, + ffn_up, + ffn_mid, + model, + l->ffn_gate, + l->ffn_up, + DS4_N_EMBD, + tp_hidden0, + tp_hidden, + ffn_norm, + 1u) : + (can_fuse_gate_up && + ds4_gpu_shared_gate_up_swiglu_q8_0_model_view_tensor( ffn_gate, ffn_up, ffn_mid, @@ -42100,7 +43370,8 @@ static bool glm_graph_encode_ffn_one_normed_from( DS4_N_EMBD, hidden, ffn_norm, - 0.0f) != 0; + 0.0f) != 0); + if (split_dense && !fused_gate_up) return false; bool ok = fused_gate_up; if (fused_gate_up) { ok = glm_graph_profile_stage(stage_profile, @@ -42112,25 +43383,25 @@ static bool glm_graph_encode_ffn_one_normed_from( stage_t0); } else { ok = glm_graph_matmul_q8_0_decode_profiled_tensor(ffn_gate, - model, - l->ffn_gate->abs_offset, - DS4_N_EMBD, - hidden, - ffn_norm, - il, - pos, - "dense_gate", - g->ssd_streaming) != 0; - if (ok) ok = glm_graph_matmul_q8_0_decode_profiled_tensor(ffn_up, - model, - l->ffn_up->abs_offset, - DS4_N_EMBD, - hidden, - ffn_norm, - il, - pos, - "dense_up", - g->ssd_streaming) != 0; + model, + l->ffn_gate->abs_offset, + DS4_N_EMBD, + hidden, + ffn_norm, + il, + pos, + "dense_gate", + g->ssd_streaming) != 0; + if (ok) ok = glm_graph_matmul_q8_0_decode_profiled_tensor(ffn_up, + model, + l->ffn_up->abs_offset, + DS4_N_EMBD, + hidden, + ffn_norm, + il, + pos, + "dense_up", + g->ssd_streaming) != 0; if (ok) ok = glm_graph_profile_stage(stage_profile, "glm_decode_ffn", "dense_gate_up", @@ -42141,7 +43412,7 @@ static bool glm_graph_encode_ffn_one_normed_from( if (ok) ok = ds4_gpu_swiglu_tensor(ffn_mid, ffn_gate, ffn_up, - (uint32_t)hidden, + tp_hidden, 0.0f, 1.0f) != 0; if (ok) ok = glm_graph_profile_stage(stage_profile, @@ -42152,16 +43423,37 @@ static bool glm_graph_encode_ffn_one_normed_from( 1, stage_t0); } - if (ok) ok = glm_graph_matmul_q8_0_decode_profiled_tensor(ffn_out, - model, - l->ffn_down->abs_offset, - hidden, - DS4_N_EMBD, - ffn_mid, - il, - pos, - "dense_down", - g->ssd_streaming) != 0; + if (ok && split_dense) { + /* Local gate/up rows are complete and retain reference order. + * Gather the disjoint SwiGLU lanes, then shard complete down + * rows so no dot product is summed across ranks. */ + ok = glm_graph_tp_gather_vector(g, + ffn_gate, + ffn_mid, + hidden, + tp_hidden0, + tp_hidden); + if (ok) { + ok = glm_graph_tp_complete_output_rows(g, + ffn_out, + model, + l->ffn_down, + hidden, + DS4_N_EMBD, + ffn_gate); + } + } else if (ok) { + ok = glm_graph_matmul_q8_0_decode_profiled_tensor(ffn_out, + model, + l->ffn_down->abs_offset, + hidden, + DS4_N_EMBD, + ffn_mid, + il, + pos, + "dense_down", + g->ssd_streaming) != 0; + } if (ok) ok = glm_graph_profile_stage(stage_profile, "glm_decode_ffn", "dense_down", @@ -43048,8 +44340,9 @@ static bool glm_graph_encode_sparse_ffn_indexed_batch_routed_moe( il, pos0); - const bool tp_batch_split_ffn2 = g->tp_world == 2; - if (ok && tp_batch_split_ffn2) { + const bool tp_batch_pair_ffn2 = + g->tp_world == 2 && g->tp_out && g->tp_in; + if (ok && tp_batch_pair_ffn2) { ok = glm_graph_tp_batch_bounce_ready(g, n_tokens); } if (ok) { @@ -43060,7 +44353,7 @@ static bool glm_graph_encode_sparse_ffn_indexed_batch_routed_moe( model, l, il, - tp_batch_split_ffn2 ? g->tp_bounce_out : g->batch_ffn_out, + tp_batch_pair_ffn2 ? g->tp_bounce_out : g->batch_ffn_out, g->batch_ffn_mid, gate_out * gate_row_bytes, gate_row_bytes, @@ -43076,6 +44369,15 @@ static bool glm_graph_encode_sparse_ffn_indexed_batch_routed_moe( false, !use_grouped_moe) != 0; } + if (ok && g->tp_world > 1) { + ok = glm_graph_tp_batch_ffn_reduce(g, il, g->batch_ffn_out, + n_tokens); + if (!ok) { + fprintf(stderr, + "ds4: GLM indexed expert collective failed (layer %u)\n", + il); + } + } if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, stage_sync, "glm_indexed_ffn", @@ -43333,12 +44635,39 @@ static bool glm_graph_encode_ffn_batch( if (il < DS4_N_LEADING_DENSE) { const uint64_t hidden = l->ffn_gate->dim[1]; + const bool split_dense = glm_graph_tp_split_prefill(g); if (hidden == 0 || hidden > UINT32_MAX / n_tokens) return false; - const uint32_t mid_elems = (uint32_t)(hidden * n_tokens); + uint32_t tp_hidden0 = 0; + uint32_t tp_hidden = 0; + if (hidden > UINT32_MAX) return false; + if (split_dense) { + if (!glm_graph_tp_range(g, (uint32_t)hidden, + &tp_hidden0, &tp_hidden)) { + return false; + } + } else { + tp_hidden = (uint32_t)hidden; + } + if (tp_hidden > UINT32_MAX / n_tokens) { + return false; + } + const uint32_t mid_elems = tp_hidden * n_tokens; const uint64_t residual_elems = (uint64_t)n_tokens * DS4_N_EMBD; if (residual_elems > UINT32_MAX) return false; - const bool fused_gate_up = glm_graph_shared_gate_up_swiglu_q8_0_tensor( + const bool fused_gate_up = split_dense ? + glm_graph_tp_gate_up_swiglu(g->batch_ffn_gate, + g->batch_ffn_up, + g->batch_ffn_mid, + model, + l->ffn_gate, + l->ffn_up, + DS4_N_EMBD, + tp_hidden0, + tp_hidden, + g->batch_ffn_norm, + n_tokens) : + glm_graph_shared_gate_up_swiglu_q8_0_tensor( g->batch_ffn_gate, g->batch_ffn_up, g->batch_ffn_mid, @@ -43350,6 +44679,7 @@ static bool glm_graph_encode_ffn_batch( g->batch_ffn_norm, n_tokens, 0.0f); + if (split_dense && !fused_gate_up) return false; if (fused_gate_up) { ok = glm_graph_prefill_stage_boundary(stage_profile, stage_sync, @@ -43361,19 +44691,19 @@ static bool glm_graph_encode_ffn_batch( stage_t0); } else { ok = glm_graph_matmul_q8_0_tensor(g->batch_ffn_gate, - model, - l->ffn_gate->abs_offset, - DS4_N_EMBD, - hidden, - g->batch_ffn_norm, - n_tokens); - if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_ffn_up, - model, - l->ffn_up->abs_offset, - DS4_N_EMBD, - hidden, - g->batch_ffn_norm, - n_tokens); + model, + l->ffn_gate->abs_offset, + DS4_N_EMBD, + hidden, + g->batch_ffn_norm, + n_tokens); + if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_ffn_up, + model, + l->ffn_up->abs_offset, + DS4_N_EMBD, + hidden, + g->batch_ffn_norm, + n_tokens); if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, stage_sync, "glm_ffn", @@ -43397,13 +44727,27 @@ static bool glm_graph_encode_ffn_batch( n_tokens, stage_t0); } - if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_ffn_out, - model, - l->ffn_down->abs_offset, - hidden, - DS4_N_EMBD, - g->batch_ffn_mid, - n_tokens); + if (ok && split_dense) { + ok = glm_graph_tp_kslice_matmul(g->batch_ffn_out, + model, + l->ffn_down, + hidden, + tp_hidden0, + tp_hidden, + DS4_N_EMBD, + g->batch_ffn_mid, + n_tokens); + if (ok) ok = glm_graph_tp_allreduce(g->batch_ffn_out, + residual_elems); + } else if (ok) { + ok = glm_graph_matmul_q8_0_tensor(g->batch_ffn_out, + model, + l->ffn_down->abs_offset, + hidden, + DS4_N_EMBD, + g->batch_ffn_mid, + n_tokens); + } if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, stage_sync, "glm_ffn", @@ -43499,8 +44843,9 @@ static bool glm_graph_encode_ffn_batch( il, pos0, n_tokens); - const bool tp_batch_split_ffn = g->tp_world == 2; - if (ok && tp_batch_split_ffn) { + const bool tp_batch_pair_ffn = + g->tp_world == 2 && g->tp_out && g->tp_in; + if (ok && tp_batch_pair_ffn) { ok = glm_graph_tp_batch_bounce_ready(g, n_tokens); } if (ok) ok = glm_graph_capture_prefill_seed_router_selected(g, @@ -43641,7 +44986,7 @@ static bool glm_graph_encode_ffn_batch( model, l, il, - tp_batch_split_ffn ? g->tp_bounce_out : g->batch_ffn_out, + tp_batch_pair_ffn ? g->tp_bounce_out : g->batch_ffn_out, g->batch_ffn_mid, gate_out * gate_row_bytes, gate_row_bytes, @@ -43656,9 +45001,12 @@ static bool glm_graph_encode_ffn_batch( (uint32_t)g->ffn_mid_elems, full_layer_prefill, false) != 0; - if (ok && g->tp_world == 2) { - ok = glm_graph_tp_batch_ffn_combine(g, il, g->batch_ffn_out, n_tokens); - if (!ok) fprintf(stderr, "ds4: GLM TP batch gate failed (layer %u)\n", il); + if (ok && g->tp_world > 1) { + ok = glm_graph_tp_batch_ffn_reduce(g, il, g->batch_ffn_out, n_tokens); + if (!ok) { + fprintf(stderr, + "ds4: GLM expert batch collective failed (layer %u)\n", il); + } } if (ok) ok = glm_graph_prefill_stage_boundary(stage_profile, stage_sync, @@ -44543,6 +45891,15 @@ static bool glm_graph_forward_tokens( const bool stage_sync = glm_graph_small_prefill_stage_sync(n_tokens, logits_out != NULL); + const bool split_dense_prefill = glm_graph_tp_split_prefill(g); + uint32_t tp_head0 = 0; + uint32_t tp_heads = DS4_N_HEAD; + if (split_dense_prefill && + !glm_graph_tp_range(g, DS4_N_HEAD, &tp_head0, &tp_heads)) { + return false; + } + const uint64_t tp_q_dim = (uint64_t)tp_heads * DS4_N_KEY_MLA; + const uint64_t tp_heads_dim = (uint64_t)tp_heads * DS4_N_VALUE_MLA; const uint32_t layer_flush_interval = stage_sync ? 0u : glm_graph_full_prefill_layer_flush_interval(n_tokens, n_rows, @@ -44854,16 +46211,27 @@ static bool glm_graph_forward_tokens( DS4_N_LORA_Q, n_tokens, DS4_RMS_EPS) != 0; - if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_q, - model, - l->attn_q_b->abs_offset, - DS4_N_LORA_Q, - g->q_dim, - g->batch_q_rank_norm, - n_tokens); + if (ok && split_dense_prefill) { + ok = glm_graph_tp_rowslice_matmul(g->batch_q, + model, + l->attn_q_b, + DS4_N_LORA_Q, + (uint64_t)tp_head0 * DS4_N_KEY_MLA, + tp_q_dim, + g->batch_q_rank_norm, + n_tokens); + } else if (ok) { + ok = glm_graph_matmul_q8_0_tensor(g->batch_q, + model, + l->attn_q_b->abs_offset, + DS4_N_LORA_Q, + g->q_dim, + g->batch_q_rank_norm, + n_tokens); + } if (ok) ok = ds4_gpu_rope_tail_tensor(g->batch_q, n_tokens, - DS4_N_HEAD, + tp_heads, DS4_N_KEY_MLA, DS4_N_ROT, pos0, @@ -44938,23 +46306,43 @@ static bool glm_graph_forward_tokens( DS4_N_ROT, glm_graph_compact_cache_is_f16()) != 0; } + uint64_t tp_k_b_offset = l->attn_k_b->abs_offset; + if (ok && split_dense_prefill) { + ok = glm_graph_tp_rowslice_offset( + l->attn_k_b, + DS4_N_KV_LORA, + (uint64_t)tp_head0 * g->q_nope, + &tp_k_b_offset); + } if (ok) ok = ds4_gpu_glm_k_b_project_typed_tensor(g->batch_k_nope, g->batch_kv_norm, model->map, model->size, - l->attn_k_b->abs_offset, + tp_k_b_offset, l->attn_k_b->type, n_tokens, DS4_N_KV_LORA, (uint32_t)g->q_nope, - DS4_N_HEAD) != 0; - if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_value, - model, - l->attn_v_b->abs_offset, - DS4_N_KV_LORA, - g->heads_dim, - g->batch_kv_norm, - n_tokens); + tp_heads) != 0; + if (ok && split_dense_prefill) { + ok = glm_graph_tp_rowslice_matmul( + g->batch_value, + model, + l->attn_v_b, + DS4_N_KV_LORA, + (uint64_t)tp_head0 * DS4_N_VALUE_MLA, + tp_heads_dim, + g->batch_kv_norm, + n_tokens); + } else if (ok) { + ok = glm_graph_matmul_q8_0_tensor(g->batch_value, + model, + l->attn_v_b->abs_offset, + DS4_N_KV_LORA, + g->heads_dim, + g->batch_kv_norm, + n_tokens); + } DS4_GLM_PROFILE_PREFILL_STAGE("glm_attn", "kv_path"); const bool flash_requested = glm_graph_use_flash_attention_prefill(n_tokens); const bool use_staged_flash_kv = @@ -44970,7 +46358,7 @@ static bool glm_graph_forward_tokens( pos0, n_tokens, g->ctx_cap, - DS4_N_HEAD, + tp_heads, kv_raw_dim, DS4_N_KV_LORA, (uint32_t)g->q_nope, @@ -44993,7 +46381,7 @@ static bool glm_graph_forward_tokens( pos0, n_tokens, g->ctx_cap, - DS4_N_HEAD, + tp_heads, kv_raw_dim, DS4_N_KV_LORA, (uint32_t)g->q_nope, @@ -45021,7 +46409,7 @@ static bool glm_graph_forward_tokens( n_tokens, cache_len, g->ctx_cap, - DS4_N_HEAD, + tp_heads, DS4_N_KEY_MLA, DS4_N_VALUE_MLA, true) != 0; @@ -45034,7 +46422,7 @@ static bool glm_graph_forward_tokens( n_tokens, cache_len, g->ctx_cap, - DS4_N_HEAD, + tp_heads, DS4_N_KEY_MLA, DS4_N_VALUE_MLA, true) != 0; @@ -45048,20 +46436,35 @@ static bool glm_graph_forward_tokens( n_tokens, cache_len, g->ctx_cap, - DS4_N_HEAD, + tp_heads, DS4_N_KEY_MLA, DS4_N_VALUE_MLA, true) != 0; } } DS4_GLM_PROFILE_PREFILL_STAGE("glm_attn", "attention"); - if (ok) ok = glm_graph_matmul_q8_0_tensor(g->batch_attn_out, - model, - l->attn_output->abs_offset, - g->heads_dim, - DS4_N_EMBD, - g->batch_heads, - n_tokens); + if (ok && split_dense_prefill) { + ok = glm_graph_tp_kslice_matmul( + g->batch_attn_out, + model, + l->attn_output, + g->heads_dim, + (uint64_t)tp_head0 * DS4_N_VALUE_MLA, + tp_heads_dim, + DS4_N_EMBD, + g->batch_heads, + n_tokens); + if (ok) ok = glm_graph_tp_allreduce(g->batch_attn_out, + residual_elems); + } else if (ok) { + ok = glm_graph_matmul_q8_0_tensor(g->batch_attn_out, + model, + l->attn_output->abs_offset, + g->heads_dim, + DS4_N_EMBD, + g->batch_heads, + n_tokens); + } if (ok) ok = ds4_gpu_add_tensor(g->batch_after_attn, cur, g->batch_attn_out, @@ -45467,14 +46870,15 @@ static bool glm_graph_forward_indexed_tokens( const bool use_split_value_proj = use_batch_attn_kernel && g->batch_attn_lora; - /* Tensor-parallel attention head split: each rank computes half the - * heads in the qk-low / attention-lora / value-project kernels, the - * unowned half of batch_heads stays zero, and the full-width attn - * output projection yields partials combined over the big-gate - * exchange (same commutative add as the routed-FFN combine). Only the - * split-value-proj batch chain has head ownership. */ + /* Metal's two-rank attention split computes one owned head range in the + * qk-low / attention-lora / value-project kernels. The other range in + * batch_heads stays zero, so the full-width output projection produces + * a partial for the big-gate exchange. CUDA collective TP instead packs + * each rank's head range at the start of the scratch rows and applies a + * matching K-slice output projection below. */ const bool tp_attn_head_split = g->tp_world == 2 && + g->tp_out && g->tp_in && use_batch_attn_kernel && use_split_value_proj && (DS4_N_HEAD % 16u) == 0u && @@ -45489,6 +46893,20 @@ static bool glm_graph_forward_indexed_tokens( const bool use_batch_indexer_weights_proj = true; const bool use_batch_attn_out_proj = true; const bool use_batch_ffn = glm_graph_indexed_prefill_batch_ffn(); + const bool collective_head_split = glm_graph_tp_split_prefill(g); + uint32_t collective_head0 = 0; + uint32_t collective_heads = DS4_N_HEAD; + if (collective_head_split && + !glm_graph_tp_range(g, DS4_N_HEAD, + &collective_head0, &collective_heads)) { + return false; + } + const uint64_t collective_q_dim = + (uint64_t)collective_heads * DS4_N_KEY_MLA; + const uint64_t collective_qk_low_dim = + (uint64_t)collective_heads * DS4_N_KV_LORA; + const uint64_t collective_heads_dim = + (uint64_t)collective_heads * DS4_N_VALUE_MLA; const bool stage_sync = glm_graph_small_prefill_stage_sync(n_tokens, logits_out != NULL); const uint32_t layer_flush_interval = stage_sync ? 0u : @@ -45621,6 +47039,21 @@ static bool glm_graph_forward_indexed_tokens( const float rope_base = layer_rope_freq_base(il); const float rope_scale = layer_rope_freq_scale(il); const uint64_t residual_elems = (uint64_t)n_tokens * DS4_N_EMBD; + uint64_t collective_k_b_offset = l->attn_k_b->abs_offset; + uint64_t collective_v_b_offset = l->attn_v_b->abs_offset; + if (collective_head_split) { + ok = glm_graph_tp_rowslice_offset( + l->attn_k_b, + DS4_N_KV_LORA, + (uint64_t)collective_head0 * g->q_nope, + &collective_k_b_offset) && + glm_graph_tp_rowslice_offset( + l->attn_v_b, + DS4_N_KV_LORA, + (uint64_t)collective_head0 * DS4_N_VALUE_MLA, + &collective_v_b_offset); + if (!ok) break; + } const bool layer_stage_profile = metal_graph_layer_stage_profile_enabled(il); double layer_stage_t0 = layer_stage_profile ? now_sec() : 0.0; double trace_stage_t0 = trace ? now_sec() : 0.0; @@ -45686,7 +47119,18 @@ static bool glm_graph_forward_indexed_tokens( n_tokens, DS4_RMS_EPS) != 0; if (ok) { - ok = (use_batch_q_proj ? + if (collective_head_split) { + ok = glm_graph_tp_rowslice_matmul( + g->batch_q, + model, + l->attn_q_b, + DS4_N_LORA_Q, + (uint64_t)collective_head0 * DS4_N_KEY_MLA, + collective_q_dim, + g->batch_q_rank_norm, + n_tokens); + } else { + ok = (use_batch_q_proj ? glm_graph_matmul_q8_0_tensor(g->batch_q, model, l->attn_q_b->abs_offset, @@ -45701,10 +47145,11 @@ static bool glm_graph_forward_indexed_tokens( g->q_dim, g->batch_q_rank_norm, n_tokens)); + } } if (ok) ok = ds4_gpu_rope_tail_tensor(g->batch_q, n_tokens, - DS4_N_HEAD, + collective_heads, DS4_N_KEY_MLA, DS4_N_ROT, pos0, @@ -46003,7 +47448,7 @@ static bool glm_graph_forward_indexed_tokens( DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_attn", "indexer_select"); metal_graph_debug_dump_tensor("glm_indexed_q", g->batch_q, - (uint64_t)n_tokens * DS4_N_HEAD * DS4_N_KEY_MLA, + (uint64_t)n_tokens * collective_q_dim, il, pos0); if (use_batch_qk_low) { @@ -46012,10 +47457,10 @@ static bool glm_graph_forward_indexed_tokens( g->batch_q, model->map, model->size, - l->attn_k_b->abs_offset, + collective_k_b_offset, l->attn_k_b->type, n_tokens, - DS4_N_HEAD, + collective_heads, DS4_N_KV_LORA, (uint32_t)g->q_nope, DS4_N_KEY_MLA) != 0; @@ -46024,14 +47469,14 @@ static bool glm_graph_forward_indexed_tokens( ds4_gpu_tensor *q_view = glm_graph_tensor_row_view_strided(g->batch_q, t, - g->q_dim, - g->q_dim); + collective_q_dim, + collective_q_dim); ds4_gpu_tensor *qk_low_view = glm_graph_tensor_row_view_strided( g->batch_qk_low, t, - (uint64_t)DS4_N_HEAD * DS4_N_KV_LORA, - (uint64_t)DS4_N_HEAD * DS4_N_KV_LORA); + collective_qk_low_dim, + collective_qk_low_dim); ok = q_view && qk_low_view; if (!ok) fprintf(stderr, "ds4: GLM scalar indexed prefill failed to create qk-low row views at layer %u token %u\n", il, t); if (ok) { @@ -46039,9 +47484,9 @@ static bool glm_graph_forward_indexed_tokens( q_view, model->map, model->size, - l->attn_k_b->abs_offset, + collective_k_b_offset, l->attn_k_b->type, - DS4_N_HEAD, + collective_heads, DS4_N_KV_LORA, (uint32_t)g->q_nope, DS4_N_KEY_MLA); @@ -46055,7 +47500,7 @@ static bool glm_graph_forward_indexed_tokens( DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_attn", "qk_low"); metal_graph_debug_dump_tensor("glm_indexed_qk_low", g->batch_qk_low, - (uint64_t)n_tokens * DS4_N_HEAD * DS4_N_KV_LORA, + (uint64_t)n_tokens * collective_qk_low_dim, il, pos0); if (ok && use_batch_attn_kernel) ok = glm_graph_indexed_prefill_attention_boundary(); @@ -46069,20 +47514,20 @@ static bool glm_graph_forward_indexed_tokens( ds4_gpu_tensor *q_view = ds4_gpu_tensor_view(g->batch_q, - (uint64_t)t0 * g->q_dim * sizeof(float), - (uint64_t)slice * g->q_dim * sizeof(float)); + (uint64_t)t0 * collective_q_dim * sizeof(float), + (uint64_t)slice * collective_q_dim * sizeof(float)); ds4_gpu_tensor *qk_low_view = ds4_gpu_tensor_view(g->batch_qk_low, - (uint64_t)t0 * DS4_N_HEAD * DS4_N_KV_LORA * sizeof(float), - (uint64_t)slice * DS4_N_HEAD * DS4_N_KV_LORA * sizeof(float)); + (uint64_t)t0 * collective_qk_low_dim * sizeof(float), + (uint64_t)slice * collective_qk_low_dim * sizeof(float)); ds4_gpu_tensor *heads_view = ds4_gpu_tensor_view(g->batch_heads, - (uint64_t)t0 * g->heads_dim * sizeof(float), - (uint64_t)slice * g->heads_dim * sizeof(float)); + (uint64_t)t0 * collective_heads_dim * sizeof(float), + (uint64_t)slice * collective_heads_dim * sizeof(float)); ds4_gpu_tensor *attn_lora_view = use_split_value_proj ? ds4_gpu_tensor_view(g->batch_attn_lora, - (uint64_t)t0 * DS4_N_HEAD * DS4_N_KV_LORA * sizeof(float), - (uint64_t)slice * DS4_N_HEAD * DS4_N_KV_LORA * sizeof(float)) : + (uint64_t)t0 * collective_qk_low_dim * sizeof(float), + (uint64_t)slice * collective_qk_low_dim * sizeof(float)) : NULL; ds4_gpu_tensor *selected_view = ds4_gpu_tensor_view(last_indexer_selected, @@ -46107,7 +47552,7 @@ static bool glm_graph_forward_indexed_tokens( last_indexer_selected_count, g->compact_cache_cap, glm_graph_compact_cache_is_f16(), - DS4_N_HEAD, + collective_heads, DS4_N_KV_LORA, (uint32_t)g->q_nope, DS4_N_ROT, @@ -46130,7 +47575,7 @@ static bool glm_graph_forward_indexed_tokens( last_indexer_selected_count, g->compact_cache_cap, glm_graph_compact_cache_is_f16(), - DS4_N_HEAD, + collective_heads, DS4_N_KV_LORA, (uint32_t)g->q_nope, DS4_N_ROT, @@ -46158,10 +47603,10 @@ static bool glm_graph_forward_indexed_tokens( attn_lora_view, model->map, model->size, - l->attn_v_b->abs_offset, + collective_v_b_offset, l->attn_v_b->type, slice, - DS4_N_HEAD, + collective_heads, DS4_N_KV_LORA, DS4_N_VALUE_MLA); ok = rc != 0; @@ -46183,14 +47628,14 @@ static bool glm_graph_forward_indexed_tokens( g->layer_k_rope_cache[il], model->map, model->size, - l->attn_v_b->abs_offset, + collective_v_b_offset, l->attn_v_b->type, selected_view, slice, last_indexer_selected_count, g->compact_cache_cap, glm_graph_compact_cache_is_f16(), - DS4_N_HEAD, + collective_heads, DS4_N_KV_LORA, (uint32_t)g->q_nope, DS4_N_ROT, @@ -46226,19 +47671,19 @@ static bool glm_graph_forward_indexed_tokens( ds4_gpu_tensor *q_view = glm_graph_tensor_row_view_strided(g->batch_q, t, - g->q_dim, - g->q_dim); + collective_q_dim, + collective_q_dim); ds4_gpu_tensor *qk_low_view = glm_graph_tensor_row_view_strided( g->batch_qk_low, t, - (uint64_t)DS4_N_HEAD * DS4_N_KV_LORA, - (uint64_t)DS4_N_HEAD * DS4_N_KV_LORA); + collective_qk_low_dim, + collective_qk_low_dim); ds4_gpu_tensor *heads_view = glm_graph_tensor_row_view_strided(g->batch_heads, t, - g->heads_dim, - g->heads_dim); + collective_heads_dim, + collective_heads_dim); ds4_gpu_tensor *selected_view = ds4_gpu_tensor_view(last_indexer_selected, (uint64_t)t * last_indexer_selected_count * sizeof(uint32_t), @@ -46253,13 +47698,13 @@ static bool glm_graph_forward_indexed_tokens( g->layer_k_rope_cache[il], model->map, model->size, - l->attn_v_b->abs_offset, + collective_v_b_offset, l->attn_v_b->type, selected_view, last_indexer_selected_count, g->compact_cache_cap, glm_graph_compact_cache_is_f16(), - DS4_N_HEAD, + collective_heads, DS4_N_KV_LORA, (uint32_t)g->q_nope, DS4_N_ROT, @@ -46283,13 +47728,29 @@ static bool glm_graph_forward_indexed_tokens( DS4_GLM_PROFILE_INDEXED_STAGE("glm_indexed_attn", "attention"); metal_graph_debug_dump_tensor("glm_indexed_heads", g->batch_heads, - (uint64_t)n_tokens * g->heads_dim, + (uint64_t)n_tokens * collective_heads_dim, il, pos0); if (ok && tp_attn_head_split) { ok = glm_graph_tp_batch_bounce_ready(g, n_tokens); } - if (ok) { + if (ok && collective_head_split) { + if (!(n_tokens <= 8u && + (glm_decode_ablate_mask() & DS4_GLM_ABLATE_ATTN_OUT))) { + ok = glm_graph_tp_kslice_matmul( + g->batch_attn_out, + model, + l->attn_output, + g->heads_dim, + (uint64_t)collective_head0 * DS4_N_VALUE_MLA, + collective_heads_dim, + DS4_N_EMBD, + g->batch_heads, + n_tokens); + if (ok) ok = glm_graph_tp_allreduce(g->batch_attn_out, + residual_elems); + } + } else if (ok) { /* Under the head split the projection input has zeros in the * unowned head columns, so the result is this rank's partial; * it must land in the shared bounce tensor for the exchange. */ @@ -47078,6 +48539,15 @@ static bool glm_graph_forward_token( ok = ds4_gpu_end_commands() != 0; } const uint32_t indexer_top_k = glm_graph_indexer_top_k_limit(); + uint32_t tp_head0 = 0; + uint32_t tp_heads = 0; + if (!glm_graph_tp_range(g, DS4_N_HEAD, &tp_head0, &tp_heads)) { + DS4_GLM_FT_FAIL("tensor-parallel head range"); + return false; + } + const uint64_t tp_q_dim = (uint64_t)tp_heads * DS4_N_KEY_MLA; + const uint64_t tp_qk_low_dim = (uint64_t)tp_heads * DS4_N_KV_LORA; + const uint64_t tp_heads_dim = (uint64_t)tp_heads * DS4_N_VALUE_MLA; ds4_gpu_tensor *last_indexer_selected = NULL; uint32_t last_indexer_selected_count = 0; #define DS4_GLM_PROFILE_DECODE_STAGE(part_, name_) do { \ @@ -47104,6 +48574,21 @@ static bool glm_graph_forward_token( const float rope_base = layer_rope_freq_base(il); const float rope_scale = layer_rope_freq_scale(il); const bool decode_stage_profile = metal_graph_decode_stage_profile_enabled(il); + uint64_t tp_k_b_offset = l->attn_k_b->abs_offset; + uint64_t tp_v_b_offset = l->attn_v_b->abs_offset; + if (glm_graph_tensor_parallel(g)) { + ok = glm_graph_tp_rowslice_offset( + l->attn_k_b, + DS4_N_KV_LORA, + (uint64_t)tp_head0 * g->q_nope, + &tp_k_b_offset) && + glm_graph_tp_rowslice_offset( + l->attn_v_b, + DS4_N_KV_LORA, + (uint64_t)tp_head0 * DS4_N_VALUE_MLA, + &tp_v_b_offset); + if (!ok) break; + } double decode_stage_t0 = decode_stage_profile ? now_sec() : 0.0; if (decode_stage_profile) { ok = metal_graph_layer_stage_profile_boundary("glm_decode_attn", @@ -47206,19 +48691,32 @@ static bool glm_graph_forward_token( DS4_RMS_EPS) != 0; } if (ok && !(decode_ablate & DS4_GLM_ABLATE_QPATH)) { - ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->q, - model, - l->attn_q_b->abs_offset, - DS4_N_LORA_Q, - g->q_dim, - g->q_rank_norm, - il, - pos, - "attn_q_b", - g->ssd_streaming) != 0; + if (glm_graph_tensor_parallel(g)) { + ok = glm_graph_tp_rowslice_matmul( + g->q, + model, + l->attn_q_b, + DS4_N_LORA_Q, + (uint64_t)tp_head0 * DS4_N_KEY_MLA, + tp_q_dim, + g->q_rank_norm, + 1u); + } else { + ok = glm_graph_matmul_q8_0_decode_profiled_tensor( + g->q, + model, + l->attn_q_b->abs_offset, + DS4_N_LORA_Q, + g->q_dim, + g->q_rank_norm, + il, + pos, + "attn_q_b", + g->ssd_streaming) != 0; + } if (ok) ok = ds4_gpu_glm_rope_tail_tensor(g->q, 1, - DS4_N_HEAD, + tp_heads, DS4_N_KEY_MLA, DS4_N_ROT, pos, @@ -47233,7 +48731,7 @@ static bool glm_graph_forward_token( DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "q_path"); if (ok) metal_graph_debug_dump_tensor("glm_decode_q", g->q, - (uint64_t)DS4_N_HEAD * DS4_N_KEY_MLA, + tp_q_dim, il, pos); if (ok && g->compact_cache_cap != 0 && glm_graph_layer_uses_full_indexer(il) && @@ -47392,15 +48890,15 @@ static bool glm_graph_forward_token( g->q, model->map, model->size, - l->attn_k_b->abs_offset, + tp_k_b_offset, l->attn_k_b->type, - DS4_N_HEAD, + tp_heads, DS4_N_KV_LORA, (uint32_t)g->q_nope, DS4_N_KEY_MLA) != 0; if (ok) metal_graph_debug_dump_tensor("glm_decode_qk_low", g->qk_low, - (uint64_t)DS4_N_HEAD * DS4_N_KV_LORA, + tp_qk_low_dim, il, pos); } @@ -47409,7 +48907,7 @@ static bool glm_graph_forward_token( /* Skip the indexed attention kernels; zero heads so the * rest of the layer stays finite (timing-only). */ ok = ds4_gpu_tensor_fill_f32(g->heads, 0.0f, - (uint64_t)g->heads_dim) != 0; + tp_heads_dim) != 0; } else if (ok && glm_graph_indexed_decode_split_group8_available(last_indexer_selected_count)) { const uint32_t split_block_rows = glm_graph_indexed_decode_split_block_rows_for(last_indexer_selected_count); @@ -47424,14 +48922,14 @@ static bool glm_graph_forward_token( g->layer_k_rope_cache[il], model->map, model->size, - l->attn_v_b->abs_offset, + tp_v_b_offset, l->attn_v_b->type, last_indexer_selected, last_indexer_selected_count, true, g->compact_cache_cap, glm_graph_compact_cache_is_f16(), - DS4_N_HEAD, + tp_heads, DS4_N_KV_LORA, (uint32_t)g->q_nope, DS4_N_ROT, @@ -47453,13 +48951,13 @@ static bool glm_graph_forward_token( g->layer_k_rope_cache[il], model->map, model->size, - l->attn_v_b->abs_offset, + tp_v_b_offset, l->attn_v_b->type, last_indexer_selected, last_indexer_selected_count, g->compact_cache_cap, glm_graph_compact_cache_is_f16(), - DS4_N_HEAD, + tp_heads, DS4_N_KV_LORA, (uint32_t)g->q_nope, DS4_N_ROT, @@ -47477,22 +48975,35 @@ static bool glm_graph_forward_token( g->kv_norm, model->map, model->size, - l->attn_k_b->abs_offset, + tp_k_b_offset, l->attn_k_b->type, 1, DS4_N_KV_LORA, (uint32_t)g->q_nope, - DS4_N_HEAD) != 0; - if (ok) ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->value, - model, - l->attn_v_b->abs_offset, - DS4_N_KV_LORA, - g->heads_dim, - g->kv_norm, - il, - pos, - "attn_v_b", - g->ssd_streaming) != 0; + tp_heads) != 0; + if (ok && glm_graph_tensor_parallel(g)) { + ok = glm_graph_tp_rowslice_matmul( + g->value, + model, + l->attn_v_b, + DS4_N_KV_LORA, + (uint64_t)tp_head0 * DS4_N_VALUE_MLA, + tp_heads_dim, + g->kv_norm, + 1u); + } else if (ok) { + ok = glm_graph_matmul_q8_0_decode_profiled_tensor( + g->value, + model, + l->attn_v_b->abs_offset, + DS4_N_KV_LORA, + g->heads_dim, + g->kv_norm, + il, + pos, + "attn_v_b", + g->ssd_streaming) != 0; + } DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "kv_path"); if (ok) ok = ds4_gpu_glm_build_kv_cache_tensor(g->layer_key_cache[il], g->layer_value_cache[il], @@ -47502,7 +49013,7 @@ static bool glm_graph_forward_token( pos, 1, g->ctx_cap, - DS4_N_HEAD, + tp_heads, kv_raw_dim, DS4_N_KV_LORA, (uint32_t)g->q_nope, @@ -47525,7 +49036,7 @@ static bool glm_graph_forward_token( 1, pos + 1u, g->ctx_cap, - DS4_N_HEAD, + tp_heads, DS4_N_KEY_MLA, DS4_N_VALUE_MLA, true) != 0; @@ -47533,20 +49044,44 @@ static bool glm_graph_forward_token( DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "attention"); if (ok) metal_graph_debug_dump_tensor("glm_decode_heads", g->heads, - g->heads_dim, + tp_heads_dim, il, pos); if (ok && !(decode_ablate & DS4_GLM_ABLATE_ATTN_OUT)) { - ok = glm_graph_matmul_q8_0_decode_profiled_tensor(g->attn_out, - model, - l->attn_output->abs_offset, - g->heads_dim, - DS4_N_EMBD, - g->heads, - il, - pos, - "attn_o", - g->ssd_streaming) != 0; + if (glm_graph_tensor_parallel(g)) { + /* Each local attention head is already a complete reference + * head. Gather heads in model order, then evaluate complete + * output rows instead of reducing K-slice partial sums. */ + ok = glm_graph_tp_gather_vector( + g, + g->value, + g->heads, + g->heads_dim, + (uint64_t)tp_head0 * DS4_N_VALUE_MLA, + tp_heads_dim); + if (ok) { + ok = glm_graph_tp_complete_output_rows( + g, + g->attn_out, + model, + l->attn_output, + g->heads_dim, + DS4_N_EMBD, + g->value); + } + } else { + ok = glm_graph_matmul_q8_0_decode_profiled_tensor( + g->attn_out, + model, + l->attn_output->abs_offset, + g->heads_dim, + DS4_N_EMBD, + g->heads, + il, + pos, + "attn_o", + g->ssd_streaming) != 0; + } } DS4_GLM_PROFILE_DECODE_STAGE("glm_decode_attn", "attn_output"); if (ok) ok = ds4_gpu_add_rms_norm_weight_tensor(g->ffn_norm, @@ -51567,6 +53102,11 @@ int ds4_session_load_payload(ds4_session *s, FILE *fp, uint64_t payload_bytes, c if (s->distributed) { return ds4_dist_session_load_payload(s->distributed, s, fp, payload_bytes, err, errlen); } + if (s->engine && s->engine->tp.active) { + payload_set_err(err, errlen, + "network-parallel KV restore requires mirrored token replay"); + return 1; + } uint64_t remaining = payload_bytes; uint32_t h[DS4_SESSION_PAYLOAD_U32_FIELDS]; for (uint32_t i = 0; i < DS4_SESSION_PAYLOAD_U32_FIELDS; i++) { @@ -55832,80 +57372,37 @@ static bool ds4_engine_preload_pro_q4_expert_tables( #endif } -/* TP sharding: touch the dense weights and only this rank's contiguous range - * of every routed-expert blob, so the other range is never faulted in. +/* Network sharding: touch the dense weights and only this rank's contiguous + * range of every routed-expert blob, so unowned ranges are never faulted in. * Replaces the whole-file residency request; the pages * fault through the same view buffers, streaming-style, but everything a * rank will ever read is pre-faulted here. */ static void model_warm_weights_sharded(const ds4_model *m, const ds4_weights *w, - int rank) { - typedef struct { uint64_t off, len; } skip_range; - if (rank != 0 && rank != 1) return; - skip_range *skips = xmalloc((size_t)DS4_N_LAYER * 3 * sizeof(*skips)); - uint32_t n_skips = 0; - uint64_t skip_bytes = 0; - for (uint32_t il = 0; il < (uint32_t)DS4_N_LAYER; il++) { - const ds4_layer_weights *l = &w->layer[il]; - const ds4_tensor *exps[3] = { l->ffn_gate_exps, l->ffn_up_exps, - l->ffn_down_exps }; - for (int t = 0; t < 3; t++) { - const ds4_tensor *x = exps[t]; - if (!x || x->ndim != 3 || x->dim[2] < 2) continue; - uint64_t in_dim = 0, out_dim = 0, row_bytes = 0; - (void)tensor_expert_bytes(m, x, 0, &in_dim, &out_dim, &row_bytes); - const uint64_t expert_bytes = out_dim * row_bytes; - const uint64_t total_bytes = x->dim[2] * expert_bytes; - const uint64_t low_bytes = (x->dim[2] / 2) * expert_bytes; - /* Unowned range: rank 0 owns the low ids; rank 1 owns the high - * ids and takes any odd-count remainder. */ - skips[n_skips].off = x->abs_offset + (rank == 0 ? low_bytes : 0); - skips[n_skips].len = rank == 0 ? total_bytes - low_bytes : low_bytes; - skip_bytes += skips[n_skips].len; - n_skips++; - } - } - /* File order should already ascend, but do not rely on it. */ - for (uint32_t i = 1; i < n_skips; i++) { - skip_range key = skips[i]; - uint32_t j = i; - while (j > 0 && skips[j - 1].off > key.off) { - skips[j] = skips[j - 1]; - j--; - } - skips[j] = key; - } + uint32_t rank, + uint32_t world) { + ds4_model_map_span_vec spans; + if (!weights_model_map_sharded_spans(w, m, rank, world, &spans)) return; const uint64_t page = (uint64_t)sysconf(_SC_PAGESIZE); const uint8_t *p = m->map; - const uint64_t start = m->tensor_data_pos, end = m->size; + const uint64_t warm_bytes = model_map_span_vec_total_bytes(&spans); fprintf(stderr, - "ds4: warming sharded tensor pages (rank %d): %.2f of %.2f GiB\n", - rank, - (double)(end - start - skip_bytes) / 1073741824.0, - (double)(end - start) / 1073741824.0); + "ds4: warming sharded tensor pages (rank %u/%u): %.2f of %.2f GiB\n", + rank, world, + (double)warm_bytes / 1073741824.0, + (double)(m->size - m->tensor_data_pos) / 1073741824.0); const double t0 = now_sec(); volatile uint64_t checksum = 0; - uint64_t off = start; - uint32_t si = 0; - /* skips are naturally sorted: layers ascend and tensors within a layer - * are laid out in file order. */ - while (off < end) { - uint64_t stop = end; - while (si < n_skips && skips[si].off + skips[si].len <= off) si++; - if (si < n_skips && skips[si].off <= off) { - off = skips[si].off + skips[si].len; - si++; - continue; - } - if (si < n_skips && skips[si].off < stop) stop = skips[si].off; + for (uint32_t i = 0; i < spans.len; i++) { + uint64_t off = spans.v[i].off; + const uint64_t stop = spans.v[i].end; #if defined(POSIX_MADV_WILLNEED) (void)posix_madvise((void *)(p + off), (size_t)(stop - off), POSIX_MADV_WILLNEED); #endif for (; off < stop; off += page) checksum += p[off]; - off = stop; } - free(skips); + free(spans.v); fprintf(stderr, "ds4: sharded warm done in %.1fs (checksum=%llu)\n", now_sec() - t0, (unsigned long long)checksum); } @@ -57291,14 +58788,27 @@ static int ds4_engine_open_internal(ds4_engine **out, load_output, load_output_optional); - /* TP always maps one contiguous routed-expert half per rank. Decide + /* Network parallelism maps one contiguous routed-expert range per rank. Decide * immediately after binding so memory guards account only the bytes this * rank owns (replicated dense weights plus its expert shard). */ #ifndef DS4_NO_GPU + if (opt->tp.role != DS4_TP_NONE && + DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK4 && + opt->backend == DS4_BACKEND_CUDA && e->ssd_streaming) { + fprintf(stderr, + "ds4: DeepSeek CUDA network EP/TP cannot be combined with " + "--ssd-streaming\n"); + ds4_engine_close(e); + *out = NULL; + return 1; + } const bool tp_shard = opt->tp.role != DS4_TP_NONE && !e->ssd_streaming; - const int tp_shard_rank = opt->tp.role == DS4_TP_WORKER ? 1 : 0; + const uint32_t tp_shard_world = + opt->tp.world_size ? opt->tp.world_size : 2u; + const uint32_t tp_shard_rank = opt->tp.role == DS4_TP_WORKER ? + (opt->tp.rank_set ? opt->tp.rank : 1u) : 0u; if (tp_shard && DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA) { uint32_t bad_layer = 0; uint32_t bad_type = 0; @@ -57306,7 +58816,7 @@ static int ds4_engine_open_internal(ds4_engine **out, &bad_layer, &bad_type)) { fprintf(stderr, - "ds4: GLM tensor parallelism lacks ownership-aware " + "ds4: GLM network parallelism lacks ownership-aware " "kernels for routed expert type %u in layer %u\n", bad_type, bad_layer); @@ -57314,12 +58824,69 @@ static int ds4_engine_open_internal(ds4_engine **out, *out = NULL; return 1; } + if (!opt->tp.expert_only) { + const char *bad_tensor = NULL; + if (!glm_tp_validate_tensor_kernels(&e->weights, + tp_shard_world, + &bad_layer, + &bad_tensor)) { + fprintf(stderr, + "ds4: GLM tensor parallelism requires sliceable Q8_0 " + "dense projections; %s is unsupported in layer %u\n", + bad_tensor ? bad_tensor : "tensor layout", + bad_layer); + ds4_engine_close(e); + *out = NULL; + return 1; + } + } + } else if (tp_shard && + DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK4 && + opt->backend == DS4_BACKEND_CUDA) { + uint32_t bad_layer = 0; + const char *bad_tensor = NULL; + if (!deepseek_tp_validate_ownership_kernels(&e->weights, + &bad_layer, + &bad_tensor)) { + fprintf(stderr, + "ds4: DeepSeek CUDA network EP requires %s in layer %u\n", + bad_tensor ? bad_tensor : "an ownership-aware expert layout", + bad_layer); + ds4_engine_close(e); + *out = NULL; + return 1; + } + if (!opt->tp.expert_only && + !deepseek_tp_validate_tensor_kernels(&e->weights, + tp_shard_world, + &bad_layer, + &bad_tensor)) { + fprintf(stderr, + "ds4: DeepSeek CUDA network TP cannot shard %s in layer %u\n", + bad_tensor ? bad_tensor : "tensor layout", + bad_layer); + ds4_engine_close(e); + *out = NULL; + return 1; + } + if (e->ssd_streaming || + (opt->mtp_path && opt->mtp_path[0]) || opt->dspark || + (opt->directional_steering_file && + opt->directional_steering_file[0])) { + fprintf(stderr, + "ds4: DeepSeek CUDA network EP/TP does not currently " + "combine with SSD streaming, MTP/DSpark, or directional steering\n"); + ds4_engine_close(e); + *out = NULL; + return 1; + } } g_tp_shard_model_bytes = 0; if (tp_shard) { ds4_model_map_span_vec shard_spans; if (weights_model_map_sharded_spans(&e->weights, &e->model, - tp_shard_rank, &shard_spans)) { + tp_shard_rank, tp_shard_world, + &shard_spans)) { g_tp_shard_model_bytes = model_map_span_vec_total_bytes(&shard_spans); free(shard_spans.v); @@ -57626,10 +59193,11 @@ static int ds4_engine_open_internal(ds4_engine **out, return 1; #endif } - /* With a raised wired limit the sharded span views (~97 GiB) fit the + /* With a raised wired limit the sharded span views fit the * GPU budget, so let the residency set pin them — that is what makes * the shard actually resident. Without the sysctl, fall back to lazy * faulting (slow but functional). */ +#ifdef __APPLE__ if (graph_backend && tp_shard && glm_graph_wired_limit_bytes() == 0) { fprintf(stderr, "ds4: iogpu.wired_limit_mb is 0 -- TP expert shard will page " @@ -57637,6 +59205,7 @@ static int ds4_engine_open_internal(ds4_engine **out, "for full residency\n"); ds4_gpu_model_residency_skip(1); } +#endif if (graph_backend) { if (e->multi_tier) { /* Wave-2 multi-tier branch. @@ -57794,6 +59363,41 @@ static int ds4_engine_open_internal(ds4_engine **out, (void)ds4_gpu_build_derived_artifacts(e->model.map, e->model.size, opt->model_path); + } else if (e->backend == DS4_BACKEND_CUDA && + !load_slice && tp_shard && !e->ssd_streaming && + DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_DEEPSEEK4) { + uint32_t expert_base = 0u; + uint32_t expert_count = 0u; + if (!ds4_tp_partition(DS4_N_EXPERT, + tp_shard_rank, + tp_shard_world, + &expert_base, + &expert_count)) { + fprintf(stderr, + "ds4: unable to partition routed experts for " + "shard-local CUDA artifacts\n"); + ds4_engine_close(e); + *out = NULL; + return 1; + } + const int artifact_count = + ds4_gpu_build_sharded_derived_artifacts(e->model.map, + e->model.size, + opt->model_path, + expert_base, + expert_count); + if (artifact_count < 0) { + fprintf(stderr, + "ds4: shard-local CUDA artifact build failed; " + "aborting network rank startup\n"); + ds4_engine_close(e); + *out = NULL; + return 1; + } + } + if (e->backend == DS4_BACKEND_CUDA && !e->ssd_streaming) { + e->startup_model_additive_artifact_bytes = + ds4_gpu_derived_artifact_additive_bytes(e->model.map); } #endif int model_map_ok = 0; @@ -57925,7 +59529,9 @@ static int ds4_engine_open_internal(ds4_engine **out, } else if (tp_shard) { ds4_model_map_span_vec spans; if (!weights_model_map_sharded_spans(&e->weights, &e->model, - tp_shard_rank, &spans)) { + tp_shard_rank, + tp_shard_world, + &spans)) { fprintf(stderr, "ds4: sharded model span build failed\n"); ds4_engine_close(e); *out = NULL; @@ -57944,8 +59550,8 @@ static int ds4_engine_open_internal(ds4_engine **out, load_span_count = spans.len; e->startup_model_span_bytes = span_bytes; fprintf(stderr, - "ds4: TP expert shard (rank %d): mapping %u spans, %.2f GiB of %.2f GiB\n", - tp_shard_rank, + "ds4: expert shard (rank %u/%u): mapping %u spans, %.2f GiB of %.2f GiB\n", + tp_shard_rank, tp_shard_world, spans.len, (double)span_bytes / 1073741824.0, (double)(e->model.size - e->model.tensor_data_pos) / 1073741824.0); @@ -57978,7 +59584,7 @@ static int ds4_engine_open_internal(ds4_engine **out, } if (tp_shard) { model_warm_weights_sharded(&e->model, &e->weights, - tp_shard_rank); + tp_shard_rank, tp_shard_world); } const bool support_model_runtime_ready = e->mtp_ready || @@ -58196,6 +59802,10 @@ int ds4_engine_tp_vocab_split(ds4_engine *e) { return e && e->tp.active && e->tp.vocab_split; } +bool ds4_engine_network_parallel(ds4_engine *e) { + return e && e->tp.active; +} + #if !defined(DS4_NO_GPU) && defined(__APPLE__) static int ds4_engine_tp_exchange(void *ud, uint32_t layer, uint32_t gate, uint64_t seq) { ds4_tp *tp = ud; @@ -58231,12 +59841,27 @@ static int ds4_engine_tp_big_exchange(void *ud, uint32_t layer, uint64_t seq, } #endif +#if !defined(DS4_NO_GPU) && !defined(__APPLE__) && !defined(DS4_ROCM_BUILD) +static int ds4_engine_tp_collective_bootstrap(void *ud, void *data, + uint32_t bytes) { + ds4_tp *tp = ud; + char err[256] = ""; + const int ok = ds4_tp_broadcast_blob(tp, data, bytes, err, sizeof(err)); + if (!ok) { + fprintf(stderr, "ds4: %s\n", + err[0] ? err : "collective bootstrap failed"); + ds4_tp_mark_failed(tp); + } + return ok; +} +#endif + int ds4_engine_tp_bind(ds4_engine *e, struct ds4_tp *tp, char *err, size_t errlen) { -#if defined(DS4_NO_GPU) || !defined(__APPLE__) +#if defined(DS4_NO_GPU) (void)e; (void)tp; - snprintf(err, errlen, "tensor parallelism requires the Metal backend"); + snprintf(err, errlen, "network parallelism requires a GPU backend"); return 0; -#else +#elif defined(__APPLE__) if (e->backend != DS4_BACKEND_METAL) { snprintf(err, errlen, "tensor parallelism requires the Metal backend"); return 0; @@ -58303,12 +59928,62 @@ int ds4_engine_tp_bind(ds4_engine *e, struct ds4_tp *tp, char *err, size_t errle e->tp.vocab_split = DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_GLM_DSA; e->tp.ctx = tp; e->tp.rank = ds4_tp_rank(tp); + e->tp.world = ds4_tp_world(tp); + e->tp.collective = false; + e->tp.tensor_parallel = true; e->tp.eval_seq = 0; e->tp.active = true; ds4_log(stderr, DS4_LOG_OK, - "tensor parallelism bound: rank %d, 50/50 expert split, %s transport", + "tensor parallelism bound: rank %d, 50/50 expert split, %s transport\n", e->tp.rank, ds4_tp_is_rdma(tp) ? "rdma" : "tcp"); return 1; +#elif defined(DS4_ROCM_BUILD) + (void)e; (void)tp; + snprintf(err, errlen, "network parallelism is not implemented for ROCm"); + return 0; +#else + if (!e || !tp || e->backend != DS4_BACKEND_CUDA || + !ds4_tp_is_collective(tp)) { + snprintf(err, errlen, + "CUDA network parallelism requires the NCCL transport"); + return 0; + } + if (DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_GLM_DSA && + DS4_MODEL_FAMILY != DS4_MODEL_FAMILY_DEEPSEEK4) { + snprintf(err, errlen, + "CUDA network parallelism supports GLM DSA and DeepSeek-V4 models"); + return 0; + } + if (e->tp.active) { + snprintf(err, errlen, "network parallelism already bound"); + return 0; + } + const uint32_t world = ds4_tp_world(tp); + const uint32_t rank = (uint32_t)ds4_tp_rank(tp); + const int local_collective_ready = ds4_gpu_tp_collective_probe(); + if (!ds4_tp_collective_preflight(tp, local_collective_ready, + err, errlen)) { + ds4_gpu_tp_collective_shutdown(); + return 0; + } + if (!ds4_gpu_tp_collective_init(rank, world, + ds4_engine_tp_collective_bootstrap, tp)) { + snprintf(err, errlen, "NCCL collective initialization failed"); + return 0; + } + e->tp.vocab_split = false; + e->tp.ctx = tp; + e->tp.rank = (int)rank; + e->tp.world = world; + e->tp.collective = true; + e->tp.tensor_parallel = !ds4_tp_is_expert_only(tp); + e->tp.eval_seq = 0; + e->tp.active = true; + ds4_log(stderr, DS4_LOG_OK, + "%s parallelism bound: rank %u/%u, NCCL transport\n", + e->tp.tensor_parallel ? "tensor" : "expert", + rank, world); + return 1; #endif } @@ -58317,11 +59992,12 @@ bool ds4_engine_is_glm_dsa(ds4_engine *e) { return DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA; } -void ds4_engine_close(ds4_engine *e) { +void ds4_engine_tp_unbind(ds4_engine *e) { if (!e) return; #if !defined(DS4_NO_GPU) && defined(__APPLE__) - if (e->tp.active) { - ds4_gpu_tp_shutdown(); + if (e->tp.active) ds4_gpu_tp_shutdown(); + if (e->tp.active || e->tp.slab || e->tp.zero_vec || e->tp.out_views || + e->tp.in_views || e->tp.batch_out_views || e->tp.batch_in_views) { const uint32_t slots = (uint32_t)DS4_N_LAYER * DS4_TP_GATES_PER_LAYER; for (uint32_t i = 0; i < slots; i++) { if (e->tp.out_views) ds4_gpu_tensor_free(e->tp.out_views[i]); @@ -58339,7 +60015,17 @@ void ds4_engine_close(ds4_engine *e) { ds4_gpu_tensor_free(e->tp.slab); memset(&e->tp, 0, sizeof(e->tp)); } +#elif !defined(DS4_NO_GPU) && !defined(DS4_ROCM_BUILD) + if (e->tp.active) { + if (e->tp.collective) ds4_gpu_tp_collective_shutdown(); + memset(&e->tp, 0, sizeof(e->tp)); + } #endif +} + +void ds4_engine_close(ds4_engine *e) { + if (!e) return; + ds4_engine_tp_unbind(e); ds4_expert_profile_close(); weights_free(&e->weights); vocab_free(&e->vocab); @@ -58478,6 +60164,14 @@ static bool ds4_session_tp_leader(const ds4_session *s) { static int ds4_session_tp_register(ds4_session *s) { if (!ds4_session_tp_leader(s)) return 1; ds4_engine *e = s->engine; + const uint32_t worker_ctx = ds4_tp_peer_ctx(e->tp.ctx); + if ((uint32_t)s->ctx_size > worker_ctx) { + fprintf(stderr, + "ds4: network session context %d exceeds the smallest worker " + "capacity %u; restart every worker with --ctx at least %d\n", + s->ctx_size, worker_ctx, s->ctx_size); + return 0; + } uint64_t id = ++e->tp.next_session_id; if (id == 0) id = ++e->tp.next_session_id; char err[256] = ""; @@ -58567,14 +60261,15 @@ int ds4_session_create(ds4_session **out, ds4_engine *e, int ctx_size) { s->glm_graph.quality = e->quality; s->glm_graph.ssd_streaming = e->ssd_streaming; s->glm_graph.ssd_streaming_cold = e->ssd_streaming_cold; -#if !defined(DS4_NO_GPU) && defined(__APPLE__) if (e->tp.active) { - s->glm_graph.tp_world = 2; + s->glm_graph.tp_world = e->tp.world; s->glm_graph.tp_rank = (uint32_t)e->tp.rank; - s->glm_graph.tp_out = e->tp.out_views; - s->glm_graph.tp_in = e->tp.in_views; + s->glm_graph.tensor_parallel = e->tp.tensor_parallel; + if (!e->tp.collective) { + s->glm_graph.tp_out = e->tp.out_views; + s->glm_graph.tp_in = e->tp.in_views; + } } -#endif if (e->ssd_streaming && !glm_graph_env_present("DS4_ROCM_GLM_DISABLE_STREAMING_SEED_BEFORE_PREFILL", "DS4_METAL_GLM_DISABLE_STREAMING_SEED_BEFORE_PREFILL")) { @@ -58686,22 +60381,26 @@ int ds4_session_create(ds4_session **out, ds4_engine *e, int ctx_size) { s->graph.ssd_streaming_cold = e->ssd_streaming_cold; s->graph.streaming_preload_experts = e->ssd_streaming_preload_experts; if (e->tp.active) { - s->graph.tp_world = 2; + s->graph.tp_world = e->tp.world; s->graph.tp_rank = (uint32_t)e->tp.rank; - s->graph.tp_out = e->tp.out_views; - s->graph.tp_in = e->tp.in_views; - s->graph.tp_batch_out = e->tp.batch_out_views; - s->graph.tp_batch_in = e->tp.batch_in_views; - s->graph.tp_zero = e->tp.zero_vec; - const uint64_t half = (uint64_t)DS4_N_VOCAB / 2u; - s->graph.tp_logits_half = ds4_gpu_tensor_view( - metal_graph_logits(&s->graph), - (uint64_t)e->tp.rank * half * sizeof(float), - half * sizeof(float)); - if (!s->graph.tp_logits_half) { - metal_graph_free(&s->graph); - free(s); - return 1; + s->graph.tp_collective = e->tp.collective; + s->graph.tensor_parallel = e->tp.tensor_parallel; + if (!e->tp.collective) { + s->graph.tp_out = e->tp.out_views; + s->graph.tp_in = e->tp.in_views; + s->graph.tp_batch_out = e->tp.batch_out_views; + s->graph.tp_batch_in = e->tp.batch_in_views; + s->graph.tp_zero = e->tp.zero_vec; + const uint64_t half = (uint64_t)DS4_N_VOCAB / 2u; + s->graph.tp_logits_half = ds4_gpu_tensor_view( + metal_graph_logits(&s->graph), + (uint64_t)e->tp.rank * half * sizeof(float), + half * sizeof(float)); + if (!s->graph.tp_logits_half) { + metal_graph_free(&s->graph); + free(s); + return 1; + } } } s->graph.power_percent = (uint32_t)e->power_percent; @@ -59754,11 +61453,11 @@ static void ds4_session_note_prefill_progress(void *ud, const char *event, int c */ static int ds4_session_sync_internal(ds4_session *s, const ds4_tokens *prompt, char *err, size_t errlen); -/* Under tensor parallelism the leader mirrors every public sync/eval to the - * worker before doing the work itself, so both engines execute the same - * graph sequence and the per-layer gates pair up. The worker acks a sync - * once its matching prefill completes, surfacing worker-side failures - * here instead of as a gate timeout mid-decode. */ +/* Under network parallelism the leader mirrors every public sync/eval to all + * workers before doing the work itself, so every engine executes the same + * graph sequence and the per-layer collectives pair up. Every worker acks a + * sync once its matching prefill completes, surfacing worker-side failures + * here instead of as a collective timeout mid-decode. */ int ds4_session_sync(ds4_session *s, const ds4_tokens *prompt, char *err, size_t errlen) { const bool mirror = ds4_session_tp_leader(s); if (mirror && prompt && prompt->len > 0) { @@ -59956,7 +61655,7 @@ static int ds4_session_sync_internal(ds4_session *s, const ds4_tokens *prompt, c } if (s->engine->glm_tp_token_prefill || s->glm_graph.placement != NULL || - (s->glm_graph.tp_world == 2 && + (s->glm_graph.tp_world > 1 && (uint32_t)(prompt->len - start) <= glm_exact_prefill_max)) { /* Multi-tier GLM has per-tier decode workspaces and KV caches; * its large batch-prefill workspace is not mirrored across @@ -64576,7 +66275,8 @@ static bool metal_graph_encode_routed_session_batch( DS4_N_EXPERT - DS4_N_EXPERT / 2u, DS4_SWIGLU_CLAMP_EXP, &peer_norm, il, (uint32_t)count, - &peer_mid_is_f16) != 0; + &peer_mid_is_f16, + false) != 0; } if (ds4_gpu_set_current_device(home_tier) != 0) ok = false; @@ -64603,7 +66303,8 @@ static bool metal_graph_encode_routed_session_batch( 0, DS4_N_EXPERT / 2u, DS4_SWIGLU_CLAMP_EXP, &local_norm, il, (uint32_t)count, - &g->batch_routed_mid_is_f16) != 0; + &g->batch_routed_mid_is_f16, + false) != 0; } if (ok) { ok = ds4_gpu_tensor_wait_xdev_default(&peer_down, home_tier) != 0; @@ -64777,7 +66478,8 @@ static bool metal_graph_encode_mixed_routed_rows( DS4_N_EXPERT - DS4_N_EXPERT / 2u, DS4_SWIGLU_CLAMP_EXP, &peer_norm, il, total_rows, - &peer_mid_is_f16) != 0; + &peer_mid_is_f16, + false) != 0; } if (ds4_gpu_set_current_device(home_tier) != 0) ok = false; @@ -64804,7 +66506,8 @@ static bool metal_graph_encode_mixed_routed_rows( 0, DS4_N_EXPERT / 2u, DS4_SWIGLU_CLAMP_EXP, &local_norm, il, total_rows, - &g->batch_routed_mid_is_f16) != 0; + &g->batch_routed_mid_is_f16, + false) != 0; } if (ok) { diff --git a/ds4.h b/ds4.h index d9c7e2f62c..7aed397ebe 100644 --- a/ds4.h +++ b/ds4.h @@ -91,12 +91,9 @@ typedef struct { bool debug; } ds4_distributed_options; -/* Tensor parallelism: two identical machines run the model in lockstep and - * split the heavy per-layer matvecs, exchanging partial sums at gates inside - * the graph (see misc/METAL_TENSOR_PARALLELISM.md). Each rank keeps one - * contiguous half of the routed experts resident; dense and shared weights - * remain replicated. The leader owns prompt/sampling and listens; the worker - * dials in and mirrors every session sync/eval. */ +/* Network expert/tensor parallelism. Ranks run the model in lockstep and + * keep one contiguous routed-expert range resident. The leader owns + * prompt/sampling and workers mirror every session operation. */ typedef enum { DS4_TP_NONE = 0, DS4_TP_LEADER, @@ -107,15 +104,20 @@ typedef enum { DS4_TP_TRANSPORT_AUTO = 0, DS4_TP_TRANSPORT_RDMA, DS4_TP_TRANSPORT_TCP, + DS4_TP_TRANSPORT_NCCL, } ds4_tp_transport; typedef struct { ds4_tp_role role; - bool requested; /* --tensor-parallel with shared role options */ - const char *listen_host; /* leader listens here for the worker */ + bool requested; /* network parallel mode with shared role options */ + bool expert_only; /* --expert-parallel; false is full tensor parallel */ + const char *listen_host; /* leader listens here for workers */ int listen_port; const char *leader_host; /* worker dials the leader */ int leader_port; + uint32_t world_size; /* 2 or 4; zero means the default 2 */ + uint32_t rank; /* leader=0, workers=1..world_size-1 */ + bool rank_set; ds4_tp_transport transport; const char *rdma_device; int rdma_gid_index; @@ -236,6 +238,7 @@ uint64_t ds4_engine_hidden_f32_values(ds4_engine *e); int ds4_engine_embd_dim(ds4_engine *e); uint64_t ds4_engine_model_bytes(ds4_engine *e); int ds4_engine_tp_vocab_split(ds4_engine *e); +bool ds4_engine_network_parallel(ds4_engine *e); bool ds4_engine_glm_layer_payload_bytes(ds4_engine *e, uint32_t layer, uint32_t full_live, @@ -326,6 +329,9 @@ int ds4_token_assistant(ds4_engine *e); * with the caller. */ struct ds4_tp; int ds4_engine_tp_bind(ds4_engine *e, struct ds4_tp *tp, char *err, size_t errlen); +/* Detach GPU collective/gate state before its frontend-owned transport is + * released. Safe to call after a partial bind or more than once. */ +void ds4_engine_tp_unbind(ds4_engine *e); int ds4_session_create(ds4_session **out, ds4_engine *e, int ctx_size); void ds4_session_free(ds4_session *s); diff --git a/ds4_agent.c b/ds4_agent.c index 72f91b9c7f..b2a2cfce40 100644 --- a/ds4_agent.c +++ b/ds4_agent.c @@ -3,6 +3,7 @@ #include "ds4_gpu_args.h" #include "ds4_help.h" #include "ds4_kvstore.h" +#include "ds4_tp.h" #include "ds4_web.h" #include "linenoise.h" @@ -605,6 +606,18 @@ static agent_config parse_options(int argc, char **argv) { } if (dist_parse == DS4_DIST_CLI_MATCHED) continue; + char tp_parse_err[256] = {0}; + ds4_tp_cli_parse_result tp_parse = + ds4_tp_parse_cli_arg(arg, &i, argc, argv, &c.engine.tp, + tp_parse_err, sizeof(tp_parse_err)); + if (tp_parse == DS4_TP_CLI_ERROR) { + fprintf(stderr, "ds4-agent: %s\n", + tp_parse_err[0] ? tp_parse_err : + "invalid network parallel option"); + exit(2); + } + if (tp_parse == DS4_TP_CLI_MATCHED) continue; + if (!strcmp(arg, "-p") || !strcmp(arg, "--prompt")) { c.gen.prompt = need_arg(&i, argc, argv, arg); } else if (!strcmp(arg, "--non-interactive")) { @@ -746,6 +759,13 @@ static agent_config parse_options(int argc, char **argv) { if (c.engine.directional_steering_file && !steering_scale_set) c.engine.directional_steering_ffn = 1.0f; + char tp_err[256] = {0}; + if (!ds4_tp_adopt_distributed_options(&c.engine.tp, + &c.engine.distributed, + tp_err, sizeof(tp_err))) { + fprintf(stderr, "ds4-agent: %s\n", tp_err); + exit(2); + } char dist_err[256]; if (ds4_dist_prepare_engine_options(&c.engine.distributed, &c.engine, @@ -754,6 +774,10 @@ static agent_config parse_options(int argc, char **argv) { fprintf(stderr, "ds4-agent: %s\n", dist_err); exit(2); } + if (!ds4_tp_validate_engine_options(&c.engine, tp_err, sizeof(tp_err))) { + fprintf(stderr, "ds4-agent: %s\n", tp_err); + exit(2); + } if (c.engine.distributed.role == DS4_DISTRIBUTED_WORKER) { fprintf(stderr, "ds4-agent: --role worker is a serving mode; start workers with ./ds4\n"); exit(2); @@ -4265,7 +4289,10 @@ static bool agent_kv_load_path(agent_worker *w, const char *path, } char load_err[160] = {0}; - if (ok && hdr.payload_bytes == 0) { + if (ok && (hdr.payload_bytes == 0 || + ds4_engine_network_parallel(w->engine))) { + /* Network ranks cannot restore a leader-local KV payload. Rebuild the + * exact rendered transcript through mirrored session sync instead. */ ds4_tokens rebuilt = {0}; ds4_tokenize_rendered_chat(w->engine, text, &rebuilt); expected_tokens = (uint32_t)rebuilt.len; @@ -11288,6 +11315,27 @@ int main(int argc, char **argv) { } agent_apply_model_sampling_defaults(engine, &cfg.gen); + if (cfg.engine.tp.role == DS4_TP_WORKER) { + int rc = ds4_tp_worker_run(engine, &cfg.engine.tp, + (uint32_t)cfg.gen.ctx_size); + ds4_engine_close(engine); + return rc; + } + + ds4_tp *tp_leader = NULL; + if (cfg.engine.tp.role == DS4_TP_LEADER) { + char tp_err[256] = {0}; + if (!ds4_tp_leader_bind(engine, &cfg.engine.tp, + (uint32_t)cfg.gen.ctx_size, + &tp_leader, tp_err, sizeof(tp_err))) { + fprintf(stderr, "ds4-agent: %s\n", + tp_err[0] ? tp_err : + "network parallel initialization failed"); + ds4_engine_close(engine); + return 1; + } + } + struct sigaction old_int; struct sigaction sa; memset(&sa, 0, sizeof(sa)); @@ -11301,6 +11349,7 @@ int main(int argc, char **argv) { run_agent(engine, &cfg); if (sigint_installed) sigaction(SIGINT, &old_int, NULL); + ds4_tp_leader_shutdown(engine, &tp_leader); ds4_engine_close(engine); return rc; } diff --git a/ds4_bench.c b/ds4_bench.c index 4b2b0d3bd2..0920a6961a 100644 --- a/ds4_bench.c +++ b/ds4_bench.c @@ -2,6 +2,7 @@ #include "ds4_distributed.h" #include "ds4_gpu_args.h" #include "ds4_help.h" +#include "ds4_tp.h" /* Purpose-built throughput benchmark. * @@ -51,6 +52,7 @@ typedef struct { double step_mul; const char *dump_frontier_logits_dir; ds4_dist_options dist; + ds4_tp_options tp; bool warm_weights; bool quality; bool ssd_streaming; @@ -234,6 +236,24 @@ static bench_config parse_options(int argc, char **argv) { } if (dist_parse == DS4_DIST_CLI_MATCHED) continue; + char tp_parse_err[256] = {0}; + ds4_tp_cli_parse_result tp_parse = + ds4_tp_parse_cli_arg(arg, + &i, + argc, + argv, + &c.tp, + tp_parse_err, + sizeof(tp_parse_err)); + if (tp_parse == DS4_TP_CLI_ERROR) { + fprintf(stderr, + "ds4-bench: %s\n", + tp_parse_err[0] ? tp_parse_err : + "invalid network parallel option"); + exit(2); + } + if (tp_parse == DS4_TP_CLI_MATCHED) continue; + if (!strcmp(arg, "-m") || !strcmp(arg, "--model")) { c.model_path = need_arg(&i, argc, argv, arg); } else if (!strcmp(arg, "--prompt-file")) { @@ -360,6 +380,19 @@ static bench_config parse_options(int argc, char **argv) { fprintf(stderr, "ds4-bench: --ctx-alloc must be greater than ctx-max + gen-tokens\n"); exit(2); } + if (c.tp.requested && c.dist.role == DS4_DISTRIBUTED_WORKER) { + fprintf(stderr, + "ds4-bench: network workers are a serving mode; " + "start them with ./ds4\n"); + exit(2); + } + char tp_err[256] = {0}; + if (!ds4_tp_adopt_distributed_options(&c.tp, &c.dist, + tp_err, sizeof(tp_err))) { + fprintf(stderr, "ds4-bench: %s\n", + tp_err[0] ? tp_err : "invalid network parallel options"); + exit(2); + } char dist_err[256]; if (ds4_dist_prepare_engine_options(&c.dist, NULL, dist_err, sizeof(dist_err)) != 0) { fprintf(stderr, "ds4-bench: %s\n", dist_err); @@ -594,12 +627,19 @@ int main(int argc, char **argv) { .ssd_streaming_full_layers_set = cfg.ssd_streaming_full_layers_set, .expert_profile_path = cfg.expert_profile_path, .distributed = cfg.dist, + .tp = cfg.tp, }; char dist_err[256]; if (ds4_dist_prepare_engine_options(&cfg.dist, &opt, dist_err, sizeof(dist_err)) != 0) { fprintf(stderr, "ds4-bench: %s\n", dist_err); return 2; } + char tp_err[256] = {0}; + if (!ds4_tp_validate_engine_options(&opt, tp_err, sizeof(tp_err))) { + fprintf(stderr, "ds4-bench: %s\n", + tp_err[0] ? tp_err : "invalid network parallel options"); + return 2; + } ds4_engine *engine = NULL; if (have_gpu_config && !skip_cuda) { const bool was_auto = @@ -640,10 +680,26 @@ int main(int argc, char **argv) { return 1; } + ds4_tp *tp_leader = NULL; + if (cfg.tp.role == DS4_TP_LEADER) { + char tp_err[256] = ""; + if (!ds4_tp_leader_bind(engine, &cfg.tp, + (uint32_t)cfg.ctx_alloc, + &tp_leader, tp_err, sizeof(tp_err))) { + fprintf(stderr, "ds4-bench: %s\n", + tp_err[0] ? tp_err : + "network parallel initialization failed"); + ds4_tokens_free(&prompt); + ds4_engine_close(engine); + return 1; + } + } + ds4_session *session = NULL; if (ds4_session_create(&session, engine, cfg.ctx_alloc) != 0) { fprintf(stderr, "ds4-bench: failed to create session\n"); ds4_tokens_free(&prompt); + ds4_tp_leader_shutdown(engine, &tp_leader); ds4_engine_close(engine); return 1; } @@ -652,6 +708,7 @@ int main(int argc, char **argv) { { ds4_session_free(session); ds4_tokens_free(&prompt); + ds4_tp_leader_shutdown(engine, &tp_leader); ds4_engine_close(engine); return 1; } @@ -664,6 +721,7 @@ int main(int argc, char **argv) { fprintf(stderr, "ds4-bench: failed to open %s: %s\n", cfg.csv_path, strerror(errno)); ds4_session_free(session); ds4_tokens_free(&prompt); + ds4_tp_leader_shutdown(engine, &tp_leader); ds4_engine_close(engine); return 1; } @@ -672,7 +730,8 @@ int main(int argc, char **argv) { fflush(out); const int eos = ds4_token_eos(engine); - const bool distributed = cfg.dist.role == DS4_DISTRIBUTED_COORDINATOR; + const bool distributed = + cfg.dist.role == DS4_DISTRIBUTED_COORDINATOR || tp_leader != NULL; ds4_session_snapshot snap = {0}; const uint64_t snapshot_max_bytes = bench_snapshot_max_bytes(); bool warned_large_snapshot = false; @@ -817,6 +876,7 @@ int main(int argc, char **argv) { ds4_session_snapshot_free(&snap); ds4_session_free(session); ds4_tokens_free(&prompt); + ds4_tp_leader_shutdown(engine, &tp_leader); ds4_engine_close(engine); return rc; } diff --git a/ds4_cli.c b/ds4_cli.c index 811c56e6b1..5e0b0a8b3a 100644 --- a/ds4_cli.c +++ b/ds4_cli.c @@ -2119,7 +2119,8 @@ int main(int argc, char **argv) { } cli_apply_model_sampling_defaults(engine, &cfg.gen); if (cfg.engine.tp.role == DS4_TP_WORKER) { - int rc = ds4_tp_worker_run(engine, &cfg.engine.tp); + int rc = ds4_tp_worker_run(engine, &cfg.engine.tp, + (uint32_t)cfg.gen.ctx_size); ds4_engine_close(engine); ds4_dist_options_free(cfg.dist); free(cfg.prompt_owned); @@ -2128,23 +2129,10 @@ int main(int argc, char **argv) { ds4_tp *tp_leader = NULL; if (cfg.engine.tp.role == DS4_TP_LEADER) { char tp_err[256] = ""; - ds4_tp_identity tp_id = { - .gguf_bytes = ds4_engine_model_bytes(engine), - .model_id = (uint32_t)ds4_engine_model_id(engine), - .n_layer = (uint32_t)ds4_engine_layer_count(engine), - .n_embd = (uint32_t)ds4_engine_embd_dim(engine), - .n_vocab = (uint32_t)ds4_engine_vocab_size(engine), - .quant_bits = (uint32_t)ds4_engine_routed_quant_bits(engine), - .ctx_size = (uint32_t)cfg.gen.ctx_size, - }; - ds4_engine_tp_gate_schedule(engine, - &tp_id.gate_slot_start, - &tp_id.gate_slot_step, - &tp_id.gates_per_token); - if (!ds4_tp_create(&tp_leader, &cfg.engine.tp, &tp_id, tp_err, sizeof(tp_err)) || - !ds4_engine_tp_bind(engine, tp_leader, tp_err, sizeof(tp_err))) { + if (!ds4_tp_leader_bind(engine, &cfg.engine.tp, + (uint32_t)cfg.gen.ctx_size, + &tp_leader, tp_err, sizeof(tp_err))) { fprintf(stderr, "ds4: %s\n", tp_err); - ds4_tp_free(tp_leader); ds4_engine_close(engine); ds4_dist_options_free(cfg.dist); free(cfg.prompt_owned); @@ -2196,9 +2184,8 @@ int main(int argc, char **argv) { } else { rc = run_generation(engine, &cfg); } - if (tp_leader) ds4_tp_send_stop(tp_leader); + ds4_tp_leader_shutdown(engine, &tp_leader); ds4_engine_close(engine); - ds4_tp_free(tp_leader); ds4_dist_options_free(cfg.dist); free(cfg.prompt_owned); return rc; diff --git a/ds4_cuda.cu b/ds4_cuda.cu index 9eb100e536..f616995aa5 100644 --- a/ds4_cuda.cu +++ b/ds4_cuda.cu @@ -20,6 +20,11 @@ #include #include +#if defined(DS4_CUDA_HAVE_NCCL) +#include +#include +#endif + #include "cuda/mmq/ds4_mmq.h" #include "cuda/mmq/ds4_repack.h" @@ -416,6 +421,7 @@ static std::unordered_map g_q8_f32_by_offset; static std::vector g_derived_ranges; static const void *g_derived_replace_map; static uint64_t g_derived_artifact_bytes; +static uint64_t g_derived_additive_artifact_bytes; static double g_derived_artifact_build_secs; static int g_derived_replaces_complete; static void *g_aligned_q81_scratch; @@ -516,6 +522,60 @@ static const char *cuda_derived_weight_ptr( return NULL; } +/* Resolve a contiguous output-row slice inside a complete aligned Q8 + * artifact. Network TP partitions dense projections by complete rows, so + * every local dot product can use the same artifact and arithmetic as the + * single-rank kernel; only the launch's row interval changes. */ +static const char *cuda_derived_q8_rows( + const void *model_map, + uint64_t source_offset, + uint64_t source_bytes, + uint64_t in_dim, + uint64_t out_dim, + uint64_t *full_out_dim, + uint64_t *row0) { + if (!full_out_dim || !row0 || + getenv("DS4_CUDA_NO_DERIVED_WEIGHTS") != NULL || + in_dim == 0u || (in_dim % 32u) != 0u || out_dim == 0u || + in_dim > INT_MAX || out_dim > INT_MAX) { + return NULL; + } + const uint64_t blocks = in_dim / 32u; + if (blocks > UINT64_MAX / 34u) return NULL; + const uint64_t row_bytes = blocks * 34u; + if (out_dim > UINT64_MAX / row_bytes || + source_bytes != out_dim * row_bytes) { + return NULL; + } + for (const cuda_derived_range &r : g_derived_ranges) { + if (r.host_base != model_map || + r.kind != CUDA_DERIVED_Q8_0_ALIGNED_DENSE || + r.in_dim != in_dim || r.group_count != 1u || + r.out_dim == 0u || r.out_dim > INT_MAX || + r.out_dim > UINT64_MAX / row_bytes || + r.source_bytes != r.out_dim * row_bytes || + source_offset < r.source_offset) { + continue; + } + const uint64_t delta = source_offset - r.source_offset; + if ((delta % row_bytes) != 0u || delta > r.source_bytes || + source_bytes > r.source_bytes - delta) { + continue; + } + const uint64_t first_row = delta / row_bytes; + if (first_row > r.out_dim || out_dim > r.out_dim - first_row) { + continue; + } + const uint64_t artifact_bytes = ds4_mmq_q8_0_aligned_bytes( + (int)r.out_dim, (int)in_dim); + if (artifact_bytes == 0u || artifact_bytes > r.bytes) continue; + *full_out_dim = r.out_dim; + *row0 = first_row; + return r.device_ptr; + } + return NULL; +} + static int cuda_model_map_replaces_complete(const void *model_map) { return g_derived_replaces_complete && model_map == g_derived_replace_map; } @@ -2565,6 +2625,7 @@ static void cuda_derived_range_release_all(void) { g_derived_ranges.clear(); g_derived_replace_map = NULL; g_derived_artifact_bytes = 0; + g_derived_additive_artifact_bytes = 0; g_derived_artifact_build_secs = 0.0; g_derived_replaces_complete = 0; } @@ -4346,13 +4407,19 @@ extern "C" int ds4_gpu_set_model_fd(int fd) { return 1; } -extern "C" int ds4_gpu_build_derived_artifacts( +static int cuda_build_derived_artifacts( const void *model_map, uint64_t model_size, - const char *model_path) { + const char *model_path, + bool shard_local, + uint32_t expert_base, + uint32_t expert_count) { if (!model_map || model_size == 0 || !model_path || !model_path[0]) return 0; if (!g_derived_ranges.empty()) return (int)g_derived_ranges.size(); if (getenv("DS4_CUDA_NO_DERIVED_WEIGHTS") != NULL) return 0; + if (shard_local && getenv("DS4_CUDA_TP_NO_ALIGNED_EXPERTS") != NULL) { + return 0; + } const char *build = getenv("DS4_CUDA_BUILD_ARTIFACTS"); if (build && strcmp(build, "0") == 0) return 0; @@ -4375,7 +4442,65 @@ extern "C" int ds4_gpu_build_derived_artifacts( const bool catalog_ok = ds4_repack_collect_catalog("ds4", mapped, nullptr, &records); ds4_repack_unmap_file(mapped); - if (!catalog_ok) return 0; + if (!catalog_ok) return shard_local ? -1 : 0; + + /* Network EP/TP owns one contiguous expert interval. Present each MoE + * tensor's owned byte interval as an ordinary smaller 3-D tensor to the + * shared repacker: the aligned layouts preserve expert-major ordering, so + * no builder or kernel format fork is needed. Dense Q8 artifacts are + * replicated unchanged on every rank. Besides matching the single-rank + * execution path, this follows the distributed numerical-parity rule that + * dense/shared work stays replicated. */ + if (shard_local) { + if (expert_count == 0u) return -1; + std::vector shard_records; + shard_records.reserve(records.size()); + for (const ds4_repack_tensor &tensor : records) { + if (ds4_repack_q8_candidate(tensor)) { + shard_records.push_back(tensor); + continue; + } + if (!ds4_repack_iq2_candidate(tensor) && + !ds4_repack_q2k_candidate(tensor)) { + continue; + } + const uint64_t total_experts = tensor.dims[2]; + if (total_experts == 0u || expert_base >= total_experts || + expert_count > total_experts - expert_base || + tensor.bytes % total_experts != 0u || + tensor.elements % total_experts != 0u) { + fprintf(stderr, + "ds4: shard-local aligned artifact geometry invalid for %s " + "(base=%u count=%u total=%llu)\n", + tensor.name.c_str(), expert_base, expert_count, + (unsigned long long)total_experts); + return -1; + } + const uint64_t expert_bytes = tensor.bytes / total_experts; + if ((uint64_t)expert_base > UINT64_MAX / expert_bytes || + (uint64_t)expert_count > UINT64_MAX / expert_bytes) { + return -1; + } + ds4_repack_tensor sliced = tensor; + sliced.off += (uint64_t)expert_base * expert_bytes; + sliced.bytes = (uint64_t)expert_count * expert_bytes; + sliced.elements = + (tensor.elements / total_experts) * expert_count; + sliced.dims[2] = expert_count; + if (sliced.off > model_size || + sliced.bytes > model_size - sliced.off) { + return -1; + } + shard_records.push_back(std::move(sliced)); + } + if (shard_records.empty()) { + fprintf(stderr, + "ds4: shard-local aligned artifact build found no routed " + "IQ2_XXS/Q2_K tensors; using raw weights\n"); + return 0; + } + records.swap(shard_records); + } const bool build_moe = cuda_aligned_iq2_enabled() && cuda_aligned_q2k_enabled(); @@ -4384,7 +4509,14 @@ extern "C" int ds4_gpu_build_derived_artifacts( ds4_repack_build_args args; args.log_prefix = "ds4"; - args.model_id = "base"; + char shard_label[64]; + if (shard_local) { + snprintf(shard_label, sizeof(shard_label), "experts-%u-%u", + expert_base, expert_base + expert_count - 1u); + args.model_id = shard_label; + } else { + args.model_id = "base"; + } args.path = model_path; args.records = &records; args.device = device; @@ -4411,8 +4543,10 @@ extern "C" int ds4_gpu_build_derived_artifacts( for (ds4_repack_artifact &artifact : artifacts) { if (artifact.dev) (void)cudaFree(artifact.dev); } - fprintf(stderr, "ds4: aligned artifact build failed; using raw weights\n"); - return 0; + fprintf(stderr, "ds4: %saligned artifact build failed%s\n", + shard_local ? "shard-local " : "", + shard_local ? "" : "; using raw weights"); + return shard_local ? -1 : 0; } for (const ds4_repack_artifact &artifact : artifacts) { @@ -4450,7 +4584,14 @@ extern "C" int ds4_gpu_build_derived_artifacts( g_derived_replace_map = model_map; g_derived_replaces_complete = replaces_complete; + uint64_t additive_bytes = 0; + for (const ds4_repack_artifact &artifact : artifacts) { + if (artifact.kind == DS4_REPACK_Q8_0_ALIGNED_DENSE) { + additive_bytes += artifact.bytes; + } + } g_derived_artifact_bytes = built_bytes; + g_derived_additive_artifact_bytes = additive_bytes; g_derived_artifact_build_secs = cuda_wall_sec() - t0; if (!g_aligned_q81_scratch) { cudaError_t scratch_err = cudaMalloc(&g_aligned_q81_scratch, 256u * 1024u); @@ -4466,14 +4607,40 @@ extern "C" int ds4_gpu_build_derived_artifacts( } } fprintf(stderr, - "ds4: built %llu aligned CUDA artifacts (%.2f GiB) in %.1fs%s\n", + "ds4: built %llu %saligned CUDA artifacts (%.2f GiB) in %.1fs%s\n", (unsigned long long)g_derived_ranges.size(), + shard_local ? "shard-local " : "", (double)g_derived_artifact_bytes / 1073741824.0, g_derived_artifact_build_secs, replaces_complete ? "; expert raw residency replaced" : ""); return (int)g_derived_ranges.size(); } +extern "C" int ds4_gpu_build_derived_artifacts( + const void *model_map, + uint64_t model_size, + const char *model_path) { + return cuda_build_derived_artifacts(model_map, model_size, model_path, + false, 0u, 0u); +} + +extern "C" int ds4_gpu_build_sharded_derived_artifacts( + const void *model_map, + uint64_t model_size, + const char *model_path, + uint32_t expert_base, + uint32_t expert_count) { + return cuda_build_derived_artifacts(model_map, model_size, model_path, + true, expert_base, expert_count); +} + +extern "C" uint64_t ds4_gpu_derived_artifact_additive_bytes( + const void *model_map) { + return model_map == g_derived_replace_map + ? g_derived_additive_artifact_bytes + : 0u; +} + extern "C" int ds4_gpu_model_range_replaced( const void *model_map, uint64_t offset, @@ -14448,23 +14615,24 @@ static int cuda_matmul_q8_0_tensor_labeled(ds4_gpu_tensor *out, const void *mode const int physical_device = (g_n_gpus > 1 && logical_tier >= 0 && logical_tier < g_n_gpus) ? g_gpu[logical_tier].device_id : 0; - const uint64_t aligned_bytes = - (in_dim % 1024u) == 0 && (out_dim % 128u) == 0 - ? ds4_mmq_q8_0_aligned_bytes((int)out_dim, (int)in_dim) - : 0; - const char *aligned = aligned_bytes && cuda_aligned_q8_enabled() - ? cuda_derived_weight_ptr( - model_map, weight_offset, weight_bytes, - CUDA_DERIVED_Q8_0_ALIGNED_DENSE, - in_dim, out_dim, 1u, aligned_bytes) - : NULL; + uint64_t aligned_full_out_dim = 0u; + uint64_t aligned_row0 = 0u; + const char *aligned = + (in_dim % 1024u) == 0u && cuda_aligned_q8_enabled() + ? cuda_derived_q8_rows( + model_map, weight_offset, weight_bytes, + in_dim, out_dim, + &aligned_full_out_dim, &aligned_row0) + : NULL; if (aligned && n_tok == 1u) { - const int rc = ds4_mmq_q8_0_aligned_dense_vec( + const int rc = ds4_mmq_q8_0_aligned_dense_vec_rows( aligned, (const float *)x->ptr, (float *)out->ptr, + (int)aligned_full_out_dim, (int)aligned_row0, (int)out_dim, 1, (int)in_dim, cuda_decode_stream()); if (rc == 0) return 1; } - if (aligned && n_tok >= 512u && out_dim >= 2048u && + if (aligned && aligned_row0 == 0u && aligned_full_out_dim == out_dim && + n_tok >= 512u && out_dim >= 2048u && in_dim <= 4096u && cuda_use_mmq()) { const char *d2r = getenv("DS4_MMQ_DENSE_D2R"); if (!d2r || strcmp(d2r, "0") != 0) { @@ -18905,6 +19073,56 @@ __device__ static float dev_dot_iq2_xxs_q8_K_block_lut( return 0.125f * xd * y->d * (float)bsum; } +/* Exact-Q8_K twin for the byte-neutral aligned IQ2 artifact. The code + * payload of each block is unchanged by the repack; only the half scale is + * split into dq[] and the 64 code bytes are placed in an aligned uint2 + * stream. Keep the integer loop and float finish identical to the raw + * ownership helper above so distributed decode retains its parity contract. */ +__device__ static float dev_dot_iq2_xxs_q8_K_block_lut_aligned( + const uint16_t *dq, + const uint2 *qs, + uint64_t block, + const cuda_block_q8_K *y, + const uint64_t *grid, + const uint8_t *signs) { + const float xd = dev_f16_to_f32(dq[block]); + const uint16_t *q2 = (const uint16_t *)(qs + block * 8ull); + const int8_t *q8 = y->qs; + int32_t bsum = 0; + for (int ib32 = 0; ib32 < CUDA_QK_K / 32; ib32++) { + const uint32_t aux0 = + (uint32_t)q2[0] | ((uint32_t)q2[1] << 16); + const uint32_t aux1 = + (uint32_t)q2[2] | ((uint32_t)q2[3] << 16); + q2 += 4; + const int32_t ls = (int32_t)(2u * (aux1 >> 28) + 1u); + int32_t w[8]; + dev_iq2_i8x8_lut(grid, signs, + (uint8_t)(aux0 & 0xffu), + (aux1 >> 0) & 127u, &w[0], &w[1]); + dev_iq2_i8x8_lut(grid, signs, + (uint8_t)((aux0 >> 8) & 0xffu), + (aux1 >> 7) & 127u, &w[2], &w[3]); + dev_iq2_i8x8_lut(grid, signs, + (uint8_t)((aux0 >> 16) & 0xffu), + (aux1 >> 14) & 127u, &w[4], &w[5]); + dev_iq2_i8x8_lut(grid, signs, + (uint8_t)((aux0 >> 24) & 0xffu), + (aux1 >> 21) & 127u, &w[6], &w[7]); + int32_t sumi = 0; + sumi = __dp4a(w[0], *(const int32_t *)(q8 + ib32 * 32u + 0), sumi); + sumi = __dp4a(w[1], *(const int32_t *)(q8 + ib32 * 32u + 4), sumi); + sumi = __dp4a(w[2], *(const int32_t *)(q8 + ib32 * 32u + 8), sumi); + sumi = __dp4a(w[3], *(const int32_t *)(q8 + ib32 * 32u + 12), sumi); + sumi = __dp4a(w[4], *(const int32_t *)(q8 + ib32 * 32u + 16), sumi); + sumi = __dp4a(w[5], *(const int32_t *)(q8 + ib32 * 32u + 20), sumi); + sumi = __dp4a(w[6], *(const int32_t *)(q8 + ib32 * 32u + 24), sumi); + sumi = __dp4a(w[7], *(const int32_t *)(q8 + ib32 * 32u + 28), sumi); + bsum += sumi * ls; + } + return 0.125f * xd * y->d * (float)bsum; +} + __device__ static float dev_dot_iq2_xxs_q8_K_block(const cuda_block_iq2_xxs *x, const cuda_block_q8_K *y) { const float d = dev_f16_to_f32(x->d) * y->d; const uint16_t *q2 = x->qs; @@ -19207,6 +19425,91 @@ __device__ static float dev_dot_q2_K_q8_K_block(const cuda_block_q2_K *x, const return dall * (float)isum - dmin * (float)summs; } +__device__ __forceinline__ static uint32_t dev_q2_K_aligned_word( + const uint2 *words, + uint64_t index, + uint32_t parity) { + const uint2 word = words[index]; + return parity ? word.y : word.x; +} + +__device__ __forceinline__ static uint8_t dev_q2_K_aligned_scale( + const int4 *sc4, + uint64_t pair_block, + uint32_t parity, + uint32_t index) { + const int4 window = sc4[pair_block * 2ull + index / 8u]; + const uint32_t word_index = (index & 7u) / 4u; + uint32_t word; + if (parity) { + word = word_index ? (uint32_t)window.w : (uint32_t)window.z; + } else { + word = word_index ? (uint32_t)window.y : (uint32_t)window.x; + } + return (uint8_t)(word >> (8u * (index & 3u))); +} + +__device__ __forceinline__ static int32_t dev_dot_q2_16_aligned( + const uint2 *qs2, + uint64_t pair_block, + uint32_t parity, + uint32_t word_base, + const int8_t *q8, + int shift) { + int32_t sum = 0; +#pragma unroll + for (uint32_t word = 0; word < 4u; word++) { + const int32_t q2 = (int32_t)dev_q2_K_aligned_word( + qs2, pair_block * 16ull + word_base + word, parity); + const int32_t v = (q2 >> shift) & 0x03030303; + sum = __dp4a(v, *(const int32_t *)(q8 + word * 4u), sum); + } + return sum; +} + +/* Exact-Q8_K twin for one row of the aligned row-pair Q2_K artifact. Field + * selection reconstructs the raw block logically while preserving the raw + * helper's sum order and float operations. */ +__device__ static float dev_dot_q2_K_q8_K_block_aligned( + const uint2 *dm2, + const int4 *sc4, + const uint2 *qs2, + uint64_t pair_block, + uint32_t parity, + const cuda_block_q8_K *y) { + int summs = 0; + for (int j = 0; j < 16; j++) { + const uint8_t sc = dev_q2_K_aligned_scale( + sc4, pair_block, parity, (uint32_t)j); + summs += y->bsums[j] * (sc >> 4); + } + const uint32_t dm = dev_q2_K_aligned_word( + dm2, pair_block, parity); + const float dall = y->d * dev_f16_to_f32((uint16_t)(dm & 0xffffu)); + const float dmin = y->d * dev_f16_to_f32((uint16_t)(dm >> 16u)); + int isum = 0; + int is = 0; + for (int k = 0; k < CUDA_QK_K / 128; k++) { + int shift = 0; + for (int j = 0; j < 4; j++) { + int d = dev_q2_K_aligned_scale( + sc4, pair_block, parity, (uint32_t)is++) & + 0x0f; + isum += d * dev_dot_q2_16_aligned( + qs2, pair_block, parity, (uint32_t)k * 8u, + y->qs + k * 128 + j * 32, shift); + d = dev_q2_K_aligned_scale( + sc4, pair_block, parity, (uint32_t)is++) & + 0x0f; + isum += d * dev_dot_q2_16_aligned( + qs2, pair_block, parity, (uint32_t)k * 8u + 4u, + y->qs + k * 128 + j * 32 + 16, shift); + shift += 2; + } + } + return dall * (float)isum - dmin * (float)summs; +} + __device__ static void dev_dot_q2_K_q8_K_block4( const cuda_block_q2_K *x, const cuda_block_q8_K *y0, @@ -19953,6 +20256,90 @@ __global__ static void moe_gate_up_mid_decode_lut_owned_qwarp32_kernel( } } +/* Shard-local aligned-artifact twin of the ownership decode kernel above. + * Expert IDs remain global at this boundary, matching the network protocol; + * only the resolved weight index becomes rank-local. Activation + * quantization, lane/block assignment, reduction, clamp, and epilogue are + * deliberately identical to the raw reference. */ +__global__ static void +moe_gate_up_mid_decode_lut_owned_qwarp32_aligned_kernel( + float *gate_out, + float *up_out, + float *mid_out, + const uint16_t *gate_dq, + const uint2 *gate_qs, + const uint16_t *up_dq, + const uint2 *up_qs, + const cuda_block_q8_K *xq, + const int32_t *selected, + const float *weights, + uint32_t xq_blocks, + uint32_t expert_mid_dim, + uint32_t n_expert, + uint32_t expert_base, + uint32_t expert_count, + uint32_t write_aux, + float clamp) { + const uint32_t lane = threadIdx.x & 7u; + const uint32_t row_lane = threadIdx.x >> 3u; + const uint32_t pair = blockIdx.y; + uint32_t expert = 0u; + if (!moe_owned_local_expert(selected[pair], expert_base, expert_count, + &expert)) { + return; + } + const cuda_block_q8_K *xqb = xq; + __shared__ cuda_block_q8_K sxq[16]; + __shared__ uint64_t s_iq2_grid[256]; + __shared__ uint8_t s_iq2_signs[128]; + if (xq_blocks <= 16u) { + for (uint32_t i = threadIdx.x; i < xq_blocks; i += blockDim.x) { + sxq[i] = xqb[i]; + } + for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) { + s_iq2_grid[i] = cuda_iq2xxs_grid[i]; + } + for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) { + s_iq2_signs[i] = cuda_ksigns_iq2xs[i]; + } + __syncthreads(); + xqb = sxq; + } + for (uint32_t rr = 0; rr < MOE_DECODE_ROW_TILES; rr++) { + const uint32_t row = blockIdx.x * MOE_DECODE_ROWS_PER_BLOCK + + row_lane + rr * 32u; + if (row >= expert_mid_dim) continue; + const uint64_t row_block = + ((uint64_t)expert * expert_mid_dim + row) * xq_blocks; + float gate = 0.0f; + float up = 0.0f; + for (uint32_t b = lane; b < xq_blocks; b += 8u) { + gate += dev_dot_iq2_xxs_q8_K_block_lut_aligned( + gate_dq, gate_qs, row_block + b, xqb + b, + s_iq2_grid, s_iq2_signs); + up += dev_dot_iq2_xxs_q8_K_block_lut_aligned( + up_dq, up_qs, row_block + b, xqb + b, + s_iq2_grid, s_iq2_signs); + } + gate = quarter_warp_sum_f32(gate, lane); + up = quarter_warp_sum_f32(up, lane); + if (lane == 0u) { + if (clamp > 1.0e-6f) { + if (gate > clamp) gate = clamp; + if (up > clamp) up = clamp; + if (up < -clamp) up = -clamp; + } + const uint64_t off = (uint64_t)pair * expert_mid_dim + row; + if (write_aux) { + gate_out[off] = gate; + up_out[off] = up; + } + mid_out[off] = (gate / (1.0f + expf(-gate))) * up * + weights[pair]; + } + } +} + __global__ static void moe_count_sorted_pairs_kernel( uint32_t *counts, const int32_t *selected, @@ -21653,6 +22040,100 @@ __global__ static void moe_down_owned_packed_qwarp32_kernel( if (lane == 0u) packed_out[(uint64_t)packed_slot * out_dim + row] = packed; } +__global__ static void moe_down_owned_slots_qwarp32_aligned_kernel( + float *down_out, + const uint2 *dm2, + const int4 *sc4, + const uint2 *qs2, + const cuda_block_q8_K *midq, + const int32_t *selected, + uint32_t midq_blocks, + uint32_t out_dim, + uint32_t expert_base, + uint32_t expert_count) { + const uint32_t lane = threadIdx.x & 7u; + const uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); + const uint32_t slot = blockIdx.y; + if (row >= out_dim || slot >= 6u) return; + uint32_t expert = 0u; + if (!moe_owned_local_expert(selected[slot], expert_base, expert_count, + &expert)) { + return; + } + const cuda_block_q8_K *xq = midq + (uint64_t)slot * midq_blocks; + float acc = 0.0f; + for (uint32_t b = lane; b < midq_blocks; b += 8u) { + const uint64_t pair_block = + ((uint64_t)expert * (out_dim / 2u) + row / 2u) * + midq_blocks + + b; + acc += dev_dot_q2_K_q8_K_block_aligned( + dm2, sc4, qs2, pair_block, row & 1u, xq + b); + } + acc = quarter_warp_sum_f32(acc, lane); + if (lane == 0u) { + down_out[(uint64_t)slot * out_dim + row] = acc; + } +} + +__global__ static void moe_down_owned_packed_qwarp32_aligned_kernel( + float *packed_out, + const uint2 *dm2, + const int4 *sc4, + const uint2 *qs2, + const cuda_block_q8_K *midq, + const int32_t *selected, + uint32_t midq_blocks, + uint32_t out_dim, + uint32_t expert_base, + uint32_t expert_count) { + const uint32_t lane = threadIdx.x & 7u; + const uint32_t row = blockIdx.x * 32u + (threadIdx.x >> 3u); + const uint32_t packed_slot = blockIdx.y; + if (row >= out_dim || packed_slot >= 4u) return; + bool prefix_pair = false; + const int first_slot = moe_owned_packed_component( + selected, packed_slot / 2u, packed_slot & 1u, + expert_base, expert_count, &prefix_pair); + if (first_slot < 0) { + if (lane == 0u) { + packed_out[(uint64_t)packed_slot * out_dim + row] = 0.0f; + } + return; + } + + float packed = 0.0f; + const uint32_t n_slots = prefix_pair ? 2u : 1u; +#pragma unroll + for (uint32_t i = 0; i < 2u; i++) { + if (i >= n_slots) break; + const uint32_t slot = (uint32_t)first_slot + i; + uint32_t expert = 0u; + if (!moe_owned_local_expert(selected[slot], expert_base, + expert_count, &expert)) { + continue; + } + const cuda_block_q8_K *xq = + midq + (uint64_t)slot * midq_blocks; + float acc = 0.0f; + for (uint32_t b = lane; b < midq_blocks; b += 8u) { + const uint64_t pair_block = + ((uint64_t)expert * (out_dim / 2u) + row / 2u) * + midq_blocks + + b; + acc += dev_dot_q2_K_q8_K_block_aligned( + dm2, sc4, qs2, pair_block, row & 1u, xq + b); + } + acc = quarter_warp_sum_f32(acc, lane); + if (lane == 0u) { + packed = prefix_pair ? __fadd_rn(packed, acc) : acc; + } + } + if (lane == 0u) { + packed_out[(uint64_t)packed_slot * out_dim + row] = packed; + } +} + __global__ static void moe_down_sum3_qwarp32_kernel( float *out, const char *down_base, @@ -23569,7 +24050,8 @@ static int routed_moe_launch( uint32_t layer_index, uint32_t n_tokens, int allow_streaming, - int owned_filtered) { + int owned_filtered, + int leave_slots) { if (!out || !gate || !up || !mid || !down || !model_map || !selected || !weights || !x || n_tokens == 0 || n_total_expert == 0 || n_expert == 0 || expert_in_dim % CUDA_QK_K != 0 || expert_mid_dim % CUDA_QK_K != 0 || @@ -23586,13 +24068,14 @@ static int routed_moe_launch( } const int q4k_path = (gate_type == 12u && down_type == 12u); const int iq2_path = (gate_type == 16u && down_type == 10u); + const int iq2_all_path = (gate_type == 16u && down_type == 16u); const int mxfp4_path = (gate_type == 39u && down_type == 39u); - if (!q4k_path && !iq2_path && !mxfp4_path) return 0; + if (!q4k_path && !iq2_path && !iq2_all_path && !mxfp4_path) return 0; /* The aligned artifacts replace the raw expert tensors on integrated * CUDA systems. Route both prefill and decode before resolving a raw * pointer, otherwise the fallback cache would duplicate tens of GiB. */ - if (iq2_path && !owned_filtered && !g_ssd_streaming_mode && + if (iq2_path && !g_ssd_streaming_mode && cuda_aligned_iq2_enabled() && cuda_aligned_q2k_enabled()) { const uint64_t gate_total = (uint64_t)n_total_expert * gate_expert_bytes; const uint64_t down_total = (uint64_t)n_total_expert * down_expert_bytes; @@ -23618,8 +24101,20 @@ static int routed_moe_launch( if (gate_aligned && up_aligned && down_aligned) { const cudaStream_t aligned_stream = n_tokens == 1u ? cuda_decode_stream() : (cudaStream_t)0; - int rc; - if (n_tokens == 1u) { + int rc = 0; + /* MMQ writes only valid assignments. An ownership rank leaves + * the six canonical slot rows materialized for NCCL, so clear + * non-owned/-1 rows before the aligned kernels fill their local + * subset. The ordinary summed path never observes those rows. */ + if (leave_slots) { + const uint64_t slot_bytes = + (uint64_t)n_tokens * n_expert * out_dim * sizeof(float); + rc = cuda_ok(cudaMemsetAsync( + down->ptr, 0, (size_t)slot_bytes, + aligned_stream), + "aligned owned moe slots clear") ? 0 : -1; + } + if (rc == 0 && n_tokens == 1u) { rc = ds4_mmq_iq2_xxs_aligned_moe_gate_up_mid_vec( gate_aligned, up_aligned, (const float *)x->ptr, @@ -23640,7 +24135,100 @@ static int routed_moe_launch( /*n_expert_used=*/1, aligned_stream); } - } else { + } else if (rc == 0 && owned_filtered && + (getenv("DS4_CUDA_TP_FAST_ALIGNED_EXPERTS") == NULL || + getenv("DS4_CUDA_TP_EXACT_RAW_EXPERTS") != NULL)) { + /* Network ranks must match the raw-expert reference's exact + * arithmetic. The fused and materialized SoA MMQ variants + * change rounding for this filtered distributed shape, so + * restore one layer's byte-neutral artifacts into a reusable + * device scratch and run the same raw pair -> SwiGLU -> down + * sequence as the reference. This avoids restoring/pinning + * the complete 18 GiB expert shard. */ + const uint64_t gate_raw_off = 0u; + const uint64_t up_raw_off = + gate_total <= UINT64_MAX - 255u + ? (gate_total + 255u) & ~255ull + : UINT64_MAX; + const uint64_t down_raw_off = + up_raw_off != UINT64_MAX && + gate_total <= UINT64_MAX - up_raw_off && + up_raw_off + gate_total <= UINT64_MAX - 255u + ? (up_raw_off + gate_total + 255u) & ~255ull + : UINT64_MAX; + const uint64_t scratch_bytes = + down_raw_off != UINT64_MAX && + down_total <= UINT64_MAX - down_raw_off + ? down_raw_off + down_total + : 0u; + char *raw = scratch_bytes + ? (char *)cuda_tmp_alloc_on( + ds4_tensor_device_idx(out), scratch_bytes, + "exact aligned expert prefill derepack") + : NULL; + char *gate_raw = raw ? raw + gate_raw_off : NULL; + char *up_raw = raw ? raw + up_raw_off : NULL; + char *down_raw = raw ? raw + down_raw_off : NULL; + if (!raw || + ds4_mmq_iq2_xxs_aligned_derepack( + gate_aligned, gate_raw, + (int)expert_mid_dim, (int)expert_in_dim, + (int)n_total_expert, aligned_stream) != 0 || + ds4_mmq_iq2_xxs_aligned_derepack( + up_aligned, up_raw, + (int)expert_mid_dim, (int)expert_in_dim, + (int)n_total_expert, aligned_stream) != 0 || + ds4_mmq_q2_K_aligned_derepack( + down_aligned, down_raw, + (int)out_dim, (int)expert_mid_dim, + (int)n_total_expert, aligned_stream) != 0) { + rc = -1; + } + if (rc == 0) { + rc = ds4_mmq_iq2_xxs_moe_pair( + gate_raw, up_raw, + (const float *)x->ptr, + (const int32_t *)selected->ptr, + (float *)gate->ptr, (float *)up->ptr, + (int)expert_mid_dim, (int)expert_in_dim, + (int)n_tokens, (int)n_total_expert, (int)n_expert, + aligned_stream); + } + if (rc == 0) { + const uint64_t mid_floats = + (uint64_t)n_tokens * n_expert * expert_mid_dim; + moe_mmq_swiglu_weighted_clamp_kernel<<< + (uint32_t)((mid_floats + 255u) / 256u), + 256, 0, aligned_stream>>>( + (float *)mid->ptr, + (const float *)gate->ptr, + (const float *)up->ptr, + (const float *)weights->ptr, + expert_mid_dim, n_tokens, n_expert, clamp); + rc = cuda_ok(cudaGetLastError(), + "aligned owned moe swiglu launch") ? 0 : -1; + } + if (rc == 0) { + const uint64_t assignments = + (uint64_t)n_tokens * n_expert; + rc = ds4_mmq_q2_K_moe( + down_raw, (const float *)mid->ptr, + (const int32_t *)selected->ptr, + (float *)down->ptr, + (int)out_dim, (int)expert_mid_dim, + (int)assignments, (int)n_total_expert, + /*n_expert_used=*/1, aligned_stream); + } + if (rc == 0) { + static int logged_exact_owned_prefill = 0; + if (!logged_exact_owned_prefill) { + logged_exact_owned_prefill = 1; + fprintf(stderr, + "ds4: routed MoE prefill using exact " + "layer-local aligned-artifact derepack\n"); + } + } + } else if (rc == 0) { rc = ds4_mmq_iq2_xxs_q2_K_moe_fused_soa( gate_aligned, up_aligned, down_aligned, (const float *)x->ptr, @@ -23652,12 +24240,23 @@ static int routed_moe_launch( (int)n_tokens, (int)n_total_expert, (int)n_expert, clamp, aligned_stream); } + if (rc == 0 && leave_slots) { + static int logged_owned_slots = 0; + if (!logged_owned_slots) { + logged_owned_slots = 1; + fprintf(stderr, + "ds4: routed MoE using shard-local aligned CUDA " + "artifacts (owned slots)\n"); + } + return 1; + } if (rc == 0) { const uint64_t n = (uint64_t)n_tokens * out_dim; moe_mmq_sum_kernel<<< (uint32_t)((n + 255u) / 256u), 256, 0, aligned_stream>>>( (float *)out->ptr, (const float *)down->ptr, - NULL, out_dim, n_expert, n_tokens, + owned_filtered ? (const int32_t *)selected->ptr : NULL, + out_dim, n_expert, n_tokens, /*guard_nonfinite=*/1); if (cuda_ok(cudaGetLastError(), "aligned moe sum launch")) { static int logged = 0; @@ -23678,6 +24277,148 @@ static int routed_moe_launch( } } + /* Some routed-IQ2 models store the down projection in IQ2_XXS too. + * On integrated CUDA the gate/up source tensors may have been replaced + * by aligned artifacts, while the IQ2 down stack remains in raw form. */ + if (iq2_all_path && !g_ssd_streaming_mode && cuda_use_mmq()) { + const uint64_t gate_total = + (uint64_t)n_total_expert * gate_expert_bytes; + const uint64_t down_total = + (uint64_t)n_total_expert * down_expert_bytes; + const uint64_t gate_aligned_bytes = + ds4_mmq_iq2_xxs_aligned_bytes( + (int)expert_mid_dim, (int)expert_in_dim, + (int)n_total_expert); + /* A derived artifact describes the complete expert tensor. An EP + * rank maps only its contiguous expert slice and passes shifted raw + * offsets here, so it must use that resident raw slice instead. */ + const char *gate_aligned = owned_filtered ? NULL : + cuda_derived_weight_ptr( + model_map, gate_offset, gate_total, + CUDA_DERIVED_IQ2_XXS_ALIGNED_MOE, + expert_in_dim, expert_mid_dim, n_total_expert, + gate_aligned_bytes); + const char *up_aligned = owned_filtered ? NULL : + cuda_derived_weight_ptr( + model_map, up_offset, gate_total, + CUDA_DERIVED_IQ2_XXS_ALIGNED_MOE, + expert_in_dim, expert_mid_dim, n_total_expert, + gate_aligned_bytes); + const int mmq_tier = ds4_tensor_device_idx(out); + const char *down_w = cuda_resolve_weight_ptr( + model_map, down_offset, down_total, mmq_tier, + "iq2-all moe down"); + const char *gate_w = NULL; + const char *up_w = NULL; + if (!gate_aligned || !up_aligned) { + gate_w = cuda_resolve_weight_ptr( + model_map, gate_offset, gate_total, mmq_tier, + "iq2-all moe gate"); + up_w = gate_w ? cuda_resolve_weight_ptr( + model_map, up_offset, gate_total, mmq_tier, + "iq2-all moe up") : NULL; + } + if (down_w && ((gate_aligned && up_aligned) || + (gate_w && up_w))) { + const cudaStream_t stream = + n_tokens == 1u ? cuda_decode_stream() : (cudaStream_t)0; + int rc = -1; + bool mid_ready = false; + if (gate_aligned && up_aligned && n_tokens <= 16u) { + rc = ds4_mmq_iq2_xxs_aligned_moe_gate_up_mid_vec( + gate_aligned, up_aligned, + (const float *)x->ptr, + (const int32_t *)selected->ptr, + (const float *)weights->ptr, + (float *)mid->ptr, + (int)expert_mid_dim, (int)expert_in_dim, + (int)n_tokens, (int)n_total_expert, + (int)n_expert, clamp, stream); + mid_ready = rc == 0; + } + if (!mid_ready) { + if (gate_aligned && up_aligned) { + rc = ds4_mmq_iq2_xxs_moe_pair_soa( + gate_aligned, up_aligned, + (const float *)x->ptr, + (const int32_t *)selected->ptr, + (float *)gate->ptr, (float *)up->ptr, + (int)expert_mid_dim, (int)expert_in_dim, + (int)n_tokens, (int)n_total_expert, + (int)n_expert, stream); + } else { + rc = ds4_mmq_iq2_xxs_moe_pair( + gate_w, up_w, + (const float *)x->ptr, + (const int32_t *)selected->ptr, + (float *)gate->ptr, (float *)up->ptr, + (int)expert_mid_dim, (int)expert_in_dim, + (int)n_tokens, (int)n_total_expert, + (int)n_expert, stream); + } + if (rc == 0) { + const uint64_t mid_floats = + (uint64_t)n_tokens * n_expert * expert_mid_dim; + moe_mmq_swiglu_weighted_clamp_kernel<<< + (uint32_t)((mid_floats + 255u) / 256u), + 256, 0, stream>>>( + (float *)mid->ptr, + (const float *)gate->ptr, + (const float *)up->ptr, + (const float *)weights->ptr, + expert_mid_dim, n_tokens, n_expert, clamp); + rc = cuda_ok(cudaGetLastError(), + "iq2-all moe swiglu launch") ? 0 : -1; + } + } + if (rc == 0) { + const uint64_t assignments = + (uint64_t)n_tokens * n_expert; + rc = ds4_mmq_iq2_xxs_moe( + down_w, (const float *)mid->ptr, + (const int32_t *)selected->ptr, + (float *)down->ptr, + (int)out_dim, (int)expert_mid_dim, + (int)assignments, (int)n_total_expert, + /*n_expert_used=*/1, stream); + } + if (rc == 0) { + const uint64_t n = (uint64_t)n_tokens * out_dim; + moe_mmq_sum_kernel<<< + (uint32_t)((n + 255u) / 256u), 256, 0, stream>>>( + (float *)out->ptr, (const float *)down->ptr, + owned_filtered ? (const int32_t *)selected->ptr : NULL, + out_dim, n_expert, n_tokens, + /*guard_nonfinite=*/1); + if (cuda_ok(cudaGetLastError(), + "iq2-all moe sum launch")) { + static int logged_iq2_all = 0; + if (!logged_iq2_all) { + logged_iq2_all = 1; + fprintf(stderr, + "ds4: routed MoE using %s IQ2 gate/up " + "with IQ2 down%s\n", + gate_aligned ? "aligned" : "raw", + owned_filtered ? " (owned expert slice)" : ""); + } + return 1; + } + rc = -1; + } + fprintf(stderr, + "ds4: CUDA all-IQ2 routed-MoE returned %d " + "(layer=%u n_tokens=%u)\n", + rc, layer_index, n_tokens); + } + } + if (iq2_all_path) { + fprintf(stderr, + "ds4: CUDA all-IQ2 routed-MoE requires the resident MMQ path " + "(layer=%u n_tokens=%u)\n", + layer_index, n_tokens); + return 0; + } + /* Native MXFP4 routed experts use the vendored MMVQ decode kernels and * MMQ matrix kernels. On Blackwell the latter dispatch to FP4 MMA; older * CUDA devices use the mathematically equivalent DP4A implementation. @@ -23822,7 +24563,8 @@ static int routed_moe_launch( * [n_tokens, n_expert, *] by the validation above. Any entry * failure falls through to the legacy sorted-pairs path (the * buffers are scratch there too). */ - if (iq2_path && n_tokens > 1u && !owned_filtered && cuda_use_mmq()) { + if (iq2_path && n_tokens > 1u && + (!owned_filtered || leave_slots) && cuda_use_mmq()) { const uint64_t gate_total = (uint64_t)n_total_expert * gate_expert_bytes; const uint64_t down_total = (uint64_t)n_total_expert * down_expert_bytes; const int mmq_tier = ds4_tensor_device_idx(out); @@ -23847,6 +24589,16 @@ static int routed_moe_launch( expert_mid_dim, n_tokens, n_expert, clamp); rc = cuda_ok(cudaGetLastError(), "mmq moe swiglu launch") ? 0 : -1; } + if (rc == 0) { + if (leave_slots) { + const uint64_t slot_bytes = + n_assignments * out_dim * sizeof(float); + rc = cuda_ok(cudaMemsetAsync( + down->ptr, 0, (size_t)slot_bytes, + (cudaStream_t)0), + "mmq owned moe slots clear") ? 0 : -1; + } + } if (rc == 0) { rc = ds4_mmq_q2_K_moe( down_w, (const float *)mid->ptr, @@ -23857,6 +24609,7 @@ static int routed_moe_launch( /*n_expert_used=*/1, (cudaStream_t)0); } + if (rc == 0 && leave_slots) return 1; if (rc == 0) { const uint64_t n = (uint64_t)n_tokens * out_dim; moe_mmq_sum_kernel<<<(uint32_t)((n + 255) / 256), 256>>>( @@ -23988,7 +24741,8 @@ static int routed_moe_launch( const uint32_t use_p2_sorted = use_sorted_pairs && !owned_filtered && getenv("DS4_CUDA_MOE_NO_P2") == NULL; - const uint32_t use_atomic_down = !q4k_path && use_expert_tiles && + const uint32_t use_atomic_down = !leave_slots && !q4k_path && + use_expert_tiles && (getenv("DS4_CUDA_MOE_ATOMIC_DOWN") != NULL || (n_tokens >= 128u && getenv("DS4_CUDA_MOE_NO_ATOMIC_DOWN") == NULL)); const uint32_t use_owned_sparse_buffers = owned_filtered && @@ -24627,8 +25381,15 @@ static int routed_moe_launch( } } if (prof_ev[4]) (void)cudaEventRecord(prof_ev[4], 0); - if (ok && owned_filtered && use_sorted_pairs && !use_atomic_down && - !use_owned_sparse_buffers) { + if (ok && owned_filtered && leave_slots) { + const uint64_t down_clear_bytes = + (uint64_t)n_tokens * n_expert * out_dim * sizeof(float); + ok = cuda_ok(cudaMemsetAsync(down->ptr, 0, + (size_t)down_clear_bytes, + cuda_decode_stream()), + "owned routed_moe collective slots clear"); + } else if (ok && owned_filtered && use_sorted_pairs && + !use_atomic_down && !use_owned_sparse_buffers) { const uint64_t down_clear_bytes = (uint64_t)n_tokens * n_expert * out_dim * sizeof(float); ok = cuda_ok(cudaMemset(down->ptr, 0, (size_t)down_clear_bytes), @@ -24903,7 +25664,7 @@ static int routed_moe_launch( ok = cuda_ok(cudaGetLastError(), "routed_moe down launch"); } if (prof_ev[5]) (void)cudaEventRecord(prof_ev[5], 0); - if (ok && !use_atomic_down && !use_direct_down_sum) { + if (ok && !leave_slots && !use_atomic_down && !use_direct_down_sum) { uint64_t n = (uint64_t)n_tokens * out_dim; if (use_owned_sparse_buffers) { moe_sum_owned_kernel<<<(n + 255) / 256, 256, 0, cuda_decode_stream()>>>( @@ -25071,19 +25832,140 @@ extern "C" int ds4_gpu_routed_moe_one_owned_tensor( } const int logical_tier = ds4_tensor_device_idx(out); - const char *gate_w = (const char *)cuda_resolve_weight_ptr( + const uint64_t down_output_bytes = + (uint64_t)(pack_fixed3 ? 4u : 6u) * out_dim * sizeof(float); + + /* Sharded derived artifacts are keyed by the owned source interval. The + * aligned ownership kernels below consume global router IDs and preserve + * the raw path's Q8_K arithmetic exactly; this is distinct from the + * single-rank aligned Q8_1 vector kernels. */ + const char *gate_aligned = NULL; + const char *up_aligned = NULL; + const char *down_aligned = NULL; + bool aligned_owned = false; + if (gate_type == 16u && down_type == 10u && + cuda_aligned_iq2_enabled() && cuda_aligned_q2k_enabled()) { + const uint64_t gate_aligned_bytes = + ds4_mmq_iq2_xxs_aligned_bytes( + (int)expert_mid_dim, (int)expert_in_dim, + (int)resident_expert_count); + const uint64_t down_aligned_bytes = + ds4_mmq_q2_k_aligned_bytes( + (int)out_dim, (int)expert_mid_dim, + (int)resident_expert_count); + gate_aligned = cuda_derived_weight_ptr( + model_map, gate_offset + gate_shift, gate_bytes, + CUDA_DERIVED_IQ2_XXS_ALIGNED_MOE, + expert_in_dim, expert_mid_dim, resident_expert_count, + gate_aligned_bytes); + up_aligned = cuda_derived_weight_ptr( + model_map, up_offset + gate_shift, gate_bytes, + CUDA_DERIVED_IQ2_XXS_ALIGNED_MOE, + expert_in_dim, expert_mid_dim, resident_expert_count, + gate_aligned_bytes); + down_aligned = cuda_derived_weight_ptr( + model_map, down_offset + down_shift, down_bytes, + CUDA_DERIVED_Q2_K_ALIGNED_MOE, + expert_mid_dim, out_dim, resident_expert_count, + down_aligned_bytes); + const bool any_aligned = + gate_aligned || up_aligned || down_aligned; + if (gate_aligned && up_aligned && down_aligned) { + aligned_owned = true; + } + if (any_aligned || cuda_model_map_replaces_complete(model_map)) { + if (!aligned_owned) { + fprintf(stderr, + "ds4: incomplete shard-local aligned expert artifacts " + "for owned decode\n"); + return 0; + } + } + } + + /* Optional throughput path: use the ordinary single-rank aligned + * canonical-Q8_1 kernels over this rank's local artifact interval. + * Translating global router IDs leaves non-owned slots as -1/zero; NCCL + * gathers the canonical six slot rows before their ordered sum. This is + * quality-tested but not byte-identical to the raw-expert parity + * reference, so the exact raw-Q8_K artifact kernels remain the default. */ + if (aligned_owned && !pack_fixed3 && !write_aux && + getenv("DS4_CUDA_TP_FAST_ALIGNED_EXPERTS") != NULL && + getenv("DS4_CUDA_TP_EXACT_RAW_EXPERTS") == NULL) { + if (gate->bytes < 6u * sizeof(int32_t) || + up->bytes < 6u * sizeof(float) || + mid->bytes < 6ull * expert_mid_dim * sizeof(float) || + (down_output ? down_output->bytes : down->bytes) < + down_output_bytes) { + return 0; + } + cudaStream_t stream = cuda_decode_stream(); + int32_t *local_ids = (int32_t *)gate->ptr; + float *local_weights = (float *)up->ptr; + float *slot_out = (float *)(down_output + ? down_output->ptr : down->ptr); + mxfp4_prepare_owned_assignments_kernel<<<1, 32, 0, stream>>>( + local_ids, + local_weights, + (const int32_t *)selected->ptr, + (const float *)weights->ptr, + resident_expert_base, + resident_expert_count); + if (!cuda_ok(cudaGetLastError(), + "owned aligned assignment prepare launch")) { + return 0; + } + if (ds4_mmq_iq2_xxs_aligned_moe_gate_up_mid_vec( + gate_aligned, + up_aligned, + (const float *)x->ptr, + local_ids, + local_weights, + (float *)mid->ptr, + (int)expert_mid_dim, + (int)expert_in_dim, + 1, + (int)resident_expert_count, + 6, + clamp, + stream) != 0 || + ds4_mmq_q2_K_aligned_moe_vec( + down_aligned, + (const float *)mid->ptr, + local_ids, + slot_out, + (int)out_dim, + (int)expert_mid_dim, + 6, + (int)resident_expert_count, + 1, + stream) != 0) { + return 0; + } + static int logged_aligned_owned_mmq = 0; + if (!logged_aligned_owned_mmq) { + logged_aligned_owned_mmq = 1; + fprintf(stderr, + "ds4: routed MoE decode using shard-local aligned " + "canonical-Q8_1 MMQ\n"); + } + return 1; + } + + const char *gate_w = aligned_owned ? NULL : + (const char *)cuda_resolve_weight_ptr( model_map, gate_offset + gate_shift, gate_bytes, logical_tier, "moe_owned_gate"); - const char *up_w = (const char *)cuda_resolve_weight_ptr( + const char *up_w = aligned_owned ? NULL : + (const char *)cuda_resolve_weight_ptr( model_map, up_offset + gate_shift, gate_bytes, logical_tier, "moe_owned_up"); - const char *down_w = (const char *)cuda_resolve_weight_ptr( + const char *down_w = aligned_owned ? NULL : + (const char *)cuda_resolve_weight_ptr( model_map, down_offset + down_shift, down_bytes, logical_tier, "moe_owned_down"); - if (!gate_w || !up_w || !down_w) return 0; + if (!aligned_owned && (!gate_w || !up_w || !down_w)) return 0; - const uint64_t down_output_bytes = - (uint64_t)(pack_fixed3 ? 4u : 6u) * out_dim * sizeof(float); if (mxfp4_path) { const uint64_t slot_bytes = 6ull * out_dim * sizeof(float); if (!cuda_use_mxfp4_mmq() || @@ -25192,6 +26074,37 @@ extern "C" int ds4_gpu_routed_moe_one_owned_tensor( float *down_dst = (float *)(down_output ? down_output->ptr : down->ptr); cuda_block_q8_K *xq = (cuda_block_q8_K *)down->ptr; cuda_block_q8_K *midq = (cuda_block_q8_K *)gate->ptr; + const uint16_t *gate_aligned_dq = NULL; + const uint2 *gate_aligned_qs = NULL; + const uint16_t *up_aligned_dq = NULL; + const uint2 *up_aligned_qs = NULL; + const uint2 *down_aligned_dm = NULL; + const int4 *down_aligned_sc = NULL; + const uint2 *down_aligned_qs = NULL; + if (aligned_owned) { + const uint64_t iq2_blocks = + (uint64_t)resident_expert_count * expert_mid_dim * xq_blocks; + const uint64_t iq2_dq_bytes = + (iq2_blocks * sizeof(uint16_t) + 63u) & ~63ull; + gate_aligned_dq = (const uint16_t *)gate_aligned; + gate_aligned_qs = + (const uint2 *)(gate_aligned + iq2_dq_bytes); + up_aligned_dq = (const uint16_t *)up_aligned; + up_aligned_qs = (const uint2 *)(up_aligned + iq2_dq_bytes); + + const uint64_t q2_pairs = + (uint64_t)resident_expert_count * (out_dim / 2u) * + midq_blocks; + const uint64_t q2_dm_bytes = + (q2_pairs * sizeof(uint2) + 63u) & ~63ull; + const uint64_t q2_sc_bytes = + (q2_pairs * 2u * sizeof(int4) + 63u) & ~63ull; + down_aligned_dm = (const uint2 *)down_aligned; + down_aligned_sc = + (const int4 *)(down_aligned + q2_dm_bytes); + down_aligned_qs = + (const uint2 *)(down_aligned + q2_dm_bytes + q2_sc_bytes); + } dim3 xq_grid(xq_blocks, 1, 1); if (shared_prequant) { @@ -25224,6 +26137,27 @@ extern "C" int ds4_gpu_routed_moe_one_owned_tensor( resident_expert_base, resident_expert_count, clamp); + } else if (aligned_owned) { + dim3 gate_grid((expert_mid_dim + 31u) / 32u, 6u, 1u); + moe_gate_up_mid_decode_lut_owned_qwarp32_aligned_kernel<<< + gate_grid, 256>>>( + (float *)gate->ptr, + (float *)up->ptr, + (float *)mid->ptr, + gate_aligned_dq, + gate_aligned_qs, + up_aligned_dq, + up_aligned_qs, + xq, + (const int32_t *)selected->ptr, + (const float *)weights->ptr, + xq_blocks, + expert_mid_dim, + 6u, + resident_expert_base, + resident_expert_count, + write_aux, + clamp); } else { dim3 gate_grid((expert_mid_dim + 31u) / 32u, 6u, 1u); moe_gate_up_mid_decode_lut_owned_qwarp32_kernel<<>>( @@ -25258,6 +26192,15 @@ extern "C" int ds4_gpu_routed_moe_one_owned_tensor( resident_expert_count); if (!cuda_ok(cudaGetLastError(), "owned routed_moe mid quantize launch")) return 0; + /* Non-resident slots are intentionally skipped by the owned down kernels. + * Clear their shared destination after xq is no longer needed so network + * collectives can reduce all six slot rows without carrying stale scratch. */ + if (!cuda_ok(cudaMemsetAsync(down_dst, 0, down_output_bytes, + cuda_decode_stream()), + "owned routed_moe down slots clear")) { + return 0; + } + dim3 down_grid((out_dim + 31u) / 32u, pack_fixed3 ? 4u : 6u, 1u); if (q4k_path && pack_fixed3) { moe_down_q4K_owned_packed_qwarp32_kernel<<>>( @@ -25283,6 +26226,30 @@ extern "C" int ds4_gpu_routed_moe_one_owned_tensor( out_dim, resident_expert_base, resident_expert_count); + } else if (aligned_owned && pack_fixed3) { + moe_down_owned_packed_qwarp32_aligned_kernel<<>>( + down_dst, + down_aligned_dm, + down_aligned_sc, + down_aligned_qs, + midq, + (const int32_t *)selected->ptr, + midq_blocks, + out_dim, + resident_expert_base, + resident_expert_count); + } else if (aligned_owned) { + moe_down_owned_slots_qwarp32_aligned_kernel<<>>( + down_dst, + down_aligned_dm, + down_aligned_sc, + down_aligned_qs, + midq, + (const int32_t *)selected->ptr, + midq_blocks, + out_dim, + resident_expert_base, + resident_expert_count); } else if (pack_fixed3) { moe_down_owned_packed_qwarp32_kernel<<>>( down_dst, @@ -25308,7 +26275,18 @@ extern "C" int ds4_gpu_routed_moe_one_owned_tensor( resident_expert_base, resident_expert_count); } - return cuda_ok(cudaGetLastError(), "owned routed_moe down launch"); + const int ok = cuda_ok(cudaGetLastError(), + "owned routed_moe down launch"); + if (ok && aligned_owned) { + static int logged_aligned_owned = 0; + if (!logged_aligned_owned) { + logged_aligned_owned = 1; + fprintf(stderr, + "ds4: routed MoE decode using exact-Q8_K shard-local " + "aligned CUDA artifacts\n"); + } + } + return ok; } extern "C" int ds4_gpu_routed_moe_owned_slots_combine_rows_tensor( @@ -25356,6 +26334,30 @@ extern "C" int ds4_gpu_routed_moe_owned_slots_combine_tensor( out_dim, expert_split, 1u); } +extern "C" int ds4_gpu_routed_moe_slots_sum_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *slots, + uint32_t out_dim, + uint32_t n_expert, + uint32_t n_tokens) { + if (!out || !slots || out_dim == 0u || n_expert == 0u || + n_tokens == 0u || + out->bytes < (uint64_t)n_tokens * out_dim * sizeof(float) || + slots->bytes < (uint64_t)n_tokens * n_expert * out_dim * + sizeof(float)) { + return 0; + } + const uint64_t n = (uint64_t)n_tokens * out_dim; + moe_sum_kernel<<<(n + 255u) / 256u, 256, 0, + cuda_decode_stream()>>>( + (float *)out->ptr, + (const float *)slots->ptr, + out_dim, + n_expert, + n_tokens); + return cuda_ok(cudaGetLastError(), "routed_moe slots sum launch"); +} + extern "C" int ds4_gpu_routed_moe_owned_packed_combine_tensor( ds4_gpu_tensor *out, const ds4_gpu_tensor *home_slots, @@ -25398,7 +26400,7 @@ extern "C" int ds4_gpu_routed_moe_one_tensor(ds4_gpu_tensor *out, ds4_gpu_tensor down_expert_bytes, down_row_bytes, expert_in_dim, expert_mid_dim, out_dim, selected, weights, n_total_expert, n_expert, clamp, x, - layer_index, 1, force_resident ? 0 : 1, 0); + layer_index, 1, force_resident ? 0 : 1, 0, 0); } extern "C" int ds4_gpu_routed_moe_batch_tensor(ds4_gpu_tensor *out, ds4_gpu_tensor *gate, ds4_gpu_tensor *up, ds4_gpu_tensor *mid, ds4_gpu_tensor *down, const void *model_map, uint64_t model_size, uint64_t gate_offset, uint64_t up_offset, uint64_t down_offset, uint32_t gate_type, uint32_t down_type, uint64_t gate_expert_bytes, uint64_t gate_row_bytes, uint64_t down_expert_bytes, uint64_t down_row_bytes, uint32_t expert_in_dim, uint32_t expert_mid_dim, uint32_t out_dim, const ds4_gpu_tensor *selected, const ds4_gpu_tensor *weights, uint32_t n_total_expert, uint32_t n_expert, float clamp, const ds4_gpu_tensor *x, uint32_t layer_index, uint32_t n_tokens, bool *mid_is_f16, bool force_resident) { (void)force_resident; @@ -25410,7 +26412,7 @@ extern "C" int ds4_gpu_routed_moe_batch_tensor(ds4_gpu_tensor *out, ds4_gpu_tens down_expert_bytes, down_row_bytes, expert_in_dim, expert_mid_dim, out_dim, selected, weights, n_total_expert, n_expert, clamp, x, - layer_index, n_tokens, 1, 0); + layer_index, n_tokens, 1, 0, 0); } extern "C" int ds4_gpu_routed_moe_batch_owned_tensor( @@ -25443,7 +26445,8 @@ extern "C" int ds4_gpu_routed_moe_batch_owned_tensor( const ds4_gpu_tensor *x, uint32_t layer_index, uint32_t n_tokens, - bool *mid_is_f16) { + bool *mid_is_f16, + bool leave_slots) { if (mid_is_f16) *mid_is_f16 = false; if (!selected || !weights || n_tokens == 0u || n_expert == 0u || n_total_expert == 0u || resident_expert_count == 0u || @@ -25489,7 +26492,8 @@ extern "C" int ds4_gpu_routed_moe_batch_owned_tensor( down_expert_bytes, down_row_bytes, expert_in_dim, expert_mid_dim, out_dim, selected, weights, resident_expert_count, n_expert, - clamp, x, layer_index, n_tokens, 0, 1); + clamp, x, layer_index, n_tokens, 0, 1, + leave_slots ? 1 : 0); } extern "C" int ds4_gpu_hc_split_sinkhorn_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *mix, const void *model_map, uint64_t model_size, uint64_t scale_offset, uint64_t base_offset, uint32_t n_hc, uint32_t sinkhorn_iters, float eps) { if (!out || !mix || !model_map || n_hc != 4) return 0; @@ -30092,7 +31096,10 @@ extern "C" int ds4_gpu_set_model_map_spans(const void *model_map, uint64_t model return 0; } } - if (!ds4_gpu_set_model_map(model_map, model_size)) return 0; + /* A span-restricted caller must never inherit DS4_CUDA_COPY_MODEL's + * whole-file allocation. Register the host map, then cache only the + * requested spans when chunked copying is enabled. */ + if (!ds4_gpu_register_model_map_no_copy(model_map, model_size)) return 0; if (getenv("DS4_CUDA_COPY_MODEL_CHUNKED") != NULL) { for (uint32_t i = 0; i < count; i++) { (void)cuda_model_prefetch_range(model_map, model_size, @@ -30447,6 +31454,291 @@ extern "C" int ds4_gpu_tensor_read_after_selected_event(const ds4_gpu_tensor *te "selected tensor read"); } +/* ------------------------------------------------------------------------ + * Network collectives (runtime-loaded NCCL). + * --------------------------------------------------------------------- */ + +typedef int (*ds4_cuda_tp_bootstrap_fn)(void *ud, void *data, uint32_t bytes); + +#if defined(DS4_CUDA_HAVE_NCCL) +typedef struct { + void *handle; + ncclComm_t comm; + int device; + uint32_t rank; + uint32_t world; + ncclResult_t (*get_unique_id)(ncclUniqueId *); + ncclResult_t (*comm_init_rank)(ncclComm_t *, int, ncclUniqueId, int); + ncclResult_t (*comm_destroy)(ncclComm_t); + ncclResult_t (*comm_abort)(ncclComm_t); + ncclResult_t (*all_reduce)(const void *, void *, size_t, + ncclDataType_t, ncclRedOp_t, + ncclComm_t, cudaStream_t); + ncclResult_t (*all_gather)(const void *, void *, size_t, + ncclDataType_t, ncclComm_t, cudaStream_t); + const char *(*get_error_string)(ncclResult_t); + ncclResult_t (*get_version)(int *); + void *ordered_scratch; + uint64_t ordered_scratch_bytes; + int ordered_logged; +} ds4_cuda_nccl_state; + +static ds4_cuda_nccl_state g_cuda_nccl; + +static int ds4_cuda_nccl_symbol(void *handle, const char *name, void **out) { + dlerror(); + *out = dlsym(handle, name); + const char *error = dlerror(); + if (error || !*out) { + fprintf(stderr, "ds4: NCCL missing symbol %s: %s\n", + name, error ? error : "not found"); + return 0; + } + return 1; +} + +static const char *ds4_cuda_nccl_error(ncclResult_t rc) { + return g_cuda_nccl.get_error_string ? + g_cuda_nccl.get_error_string(rc) : "unknown NCCL error"; +} + +static int ds4_cuda_nccl_load(void) { + if (g_cuda_nccl.handle) return 1; + const char *override_path = getenv("DS4_NCCL_LIBRARY"); + const char *path = override_path && override_path[0] ? + override_path : "libnccl.so.2"; + void *handle = dlopen(path, RTLD_NOW | RTLD_LOCAL); + if (!handle) { + fprintf(stderr, + "ds4: unable to load NCCL from %s: %s\n" + "ds4: install libnccl2 or set DS4_NCCL_LIBRARY to libnccl.so.2\n", + path, dlerror()); + return 0; + } +#define DS4_NCCL_LOAD(field, symbol) do { \ + if (!ds4_cuda_nccl_symbol(handle, symbol, \ + reinterpret_cast(&g_cuda_nccl.field))) { \ + dlclose(handle); \ + memset(&g_cuda_nccl, 0, sizeof(g_cuda_nccl)); \ + return 0; \ + } \ +} while (0) + g_cuda_nccl.handle = handle; + DS4_NCCL_LOAD(get_unique_id, "ncclGetUniqueId"); + DS4_NCCL_LOAD(comm_init_rank, "ncclCommInitRank"); + DS4_NCCL_LOAD(comm_destroy, "ncclCommDestroy"); + DS4_NCCL_LOAD(comm_abort, "ncclCommAbort"); + DS4_NCCL_LOAD(all_reduce, "ncclAllReduce"); + DS4_NCCL_LOAD(all_gather, "ncclAllGather"); + DS4_NCCL_LOAD(get_error_string, "ncclGetErrorString"); + DS4_NCCL_LOAD(get_version, "ncclGetVersion"); +#undef DS4_NCCL_LOAD + return 1; +} +#endif + +extern "C" void ds4_gpu_tp_collective_shutdown(void); + +extern "C" int ds4_gpu_tp_collective_probe(void) { +#if !defined(DS4_CUDA_HAVE_NCCL) + fprintf(stderr, + "ds4: CUDA network parallelism was built without NCCL headers\n"); + return 0; +#else + if (!ds4_cuda_nccl_load()) return 0; + int device = 0; + return cuda_ok(cudaGetDevice(&device), "NCCL preflight current device"); +#endif +} + +extern "C" int ds4_gpu_tp_collective_init( + uint32_t rank, uint32_t world, + ds4_cuda_tp_bootstrap_fn bootstrap, void *ud) { +#if !defined(DS4_CUDA_HAVE_NCCL) + (void)rank; (void)world; (void)bootstrap; (void)ud; + fprintf(stderr, + "ds4: CUDA network parallelism was built without NCCL headers\n"); + return 0; +#else + if (!bootstrap || world < 2u || world > 4u || rank >= world || + g_cuda_nccl.comm) { + return 0; + } + if (!ds4_cuda_nccl_load()) return 0; + struct { + uint32_t ready; + ncclUniqueId id; + } boot; + memset(&boot, 0, sizeof(boot)); + if (rank == 0) { + const ncclResult_t rc = g_cuda_nccl.get_unique_id(&boot.id); + if (rc != ncclSuccess) { + fprintf(stderr, "ds4: ncclGetUniqueId failed: %s\n", + ds4_cuda_nccl_error(rc)); + } else { + boot.ready = 1u; + } + } + if (!bootstrap(ud, &boot, (uint32_t)sizeof(boot))) { + fprintf(stderr, "ds4: NCCL communicator bootstrap failed on rank %u\n", + rank); + ds4_gpu_tp_collective_shutdown(); + return 0; + } + if (boot.ready != 1u) { + if (rank != 0) { + fprintf(stderr, + "ds4: NCCL communicator ID creation failed on rank 0\n"); + } + ds4_gpu_tp_collective_shutdown(); + return 0; + } + int device = 0; + if (!cuda_ok(cudaGetDevice(&device), "NCCL current device")) { + ds4_gpu_tp_collective_shutdown(); + return 0; + } + ncclComm_t comm = NULL; + const ncclResult_t rc = g_cuda_nccl.comm_init_rank( + &comm, (int)world, boot.id, (int)rank); + if (rc != ncclSuccess || !comm) { + fprintf(stderr, "ds4: ncclCommInitRank rank %u/%u failed: %s\n", + rank, world, ds4_cuda_nccl_error(rc)); + if (comm) (void)g_cuda_nccl.comm_abort(comm); + ds4_gpu_tp_collective_shutdown(); + return 0; + } + g_cuda_nccl.comm = comm; + g_cuda_nccl.device = device; + g_cuda_nccl.rank = rank; + g_cuda_nccl.world = world; + int version = 0; + if (g_cuda_nccl.get_version(&version) != ncclSuccess) version = 0; + fprintf(stderr, + "ds4: NCCL collective ready: rank %u/%u device=%d version=%d.%d.%d\n", + rank, world, device, + version / 10000, (version / 100) % 100, version % 100); + return 1; +#endif +} + +extern "C" void ds4_gpu_tp_collective_shutdown(void) { +#if defined(DS4_CUDA_HAVE_NCCL) + if (g_cuda_nccl.ordered_scratch) { + (void)cudaFree(g_cuda_nccl.ordered_scratch); + g_cuda_nccl.ordered_scratch = NULL; + g_cuda_nccl.ordered_scratch_bytes = 0; + } + if (g_cuda_nccl.comm) { + const ncclResult_t rc = g_cuda_nccl.comm_destroy(g_cuda_nccl.comm); + if (rc != ncclSuccess) { + fprintf(stderr, "ds4: ncclCommDestroy failed: %s\n", + ds4_cuda_nccl_error(rc)); + } + g_cuda_nccl.comm = NULL; + } + if (g_cuda_nccl.handle) dlclose(g_cuda_nccl.handle); + memset(&g_cuda_nccl, 0, sizeof(g_cuda_nccl)); +#endif +} + +#if defined(DS4_CUDA_HAVE_NCCL) +__global__ static void tp_ordered_reduce_f32_kernel( + float *out, const float *rank_rows, uint64_t count, uint32_t world) { + const uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (i >= count) return; + float sum = rank_rows[i]; + for (uint32_t rank = 1u; rank < world; rank++) { + sum += rank_rows[(uint64_t)rank * count + i]; + } + out[i] = sum; +} +#endif + +extern "C" int ds4_gpu_tp_collective_ready(void) { +#if defined(DS4_CUDA_HAVE_NCCL) + return g_cuda_nccl.comm != NULL; +#else + return 0; +#endif +} + +extern "C" int ds4_gpu_tp_allreduce_f32(ds4_gpu_tensor *tensor, + uint64_t count) { +#if !defined(DS4_CUDA_HAVE_NCCL) + (void)tensor; (void)count; + return 0; +#else + if (!g_cuda_nccl.comm || !tensor || !tensor->ptr || + count > SIZE_MAX || count > tensor->bytes / sizeof(float)) { + return 0; + } + if (ds4_tensor_device_idx(tensor) != g_cuda_nccl.device || + ds4_gpu_set_current_device(g_cuda_nccl.device) != 0) { + fprintf(stderr, "ds4: NCCL allreduce tensor is on the wrong device\n"); + return 0; + } + const char *ordered_env = getenv("DS4_TP_ORDERED_REDUCE"); + const bool ordered = ordered_env && ordered_env[0] && + strcmp(ordered_env, "0") != 0; + ncclResult_t rc = ncclSuccess; + if (ordered) { + if (count > UINT64_MAX / g_cuda_nccl.world / sizeof(float)) return 0; + const uint64_t bytes = count * g_cuda_nccl.world * sizeof(float); + if (bytes > g_cuda_nccl.ordered_scratch_bytes) { + if (g_cuda_nccl.ordered_scratch) { + if (!cuda_ok(cudaFree(g_cuda_nccl.ordered_scratch), + "ordered reduce scratch free")) { + return 0; + } + g_cuda_nccl.ordered_scratch = NULL; + g_cuda_nccl.ordered_scratch_bytes = 0; + } + if (!cuda_ok(cudaMalloc(&g_cuda_nccl.ordered_scratch, + (size_t)bytes), + "ordered reduce scratch allocation")) { + return 0; + } + g_cuda_nccl.ordered_scratch_bytes = bytes; + } + rc = g_cuda_nccl.all_gather( + tensor->ptr, g_cuda_nccl.ordered_scratch, (size_t)count, + ncclFloat32, g_cuda_nccl.comm, cuda_decode_stream()); + if (rc == ncclSuccess) { + const uint64_t blocks = (count + 255u) / 256u; + if (blocks > UINT32_MAX) return 0; + tp_ordered_reduce_f32_kernel<<<(uint32_t)blocks, 256, 0, + cuda_decode_stream()>>>( + (float *)tensor->ptr, + (const float *)g_cuda_nccl.ordered_scratch, + count, + g_cuda_nccl.world); + if (!cuda_ok(cudaGetLastError(), + "ordered rank reduction launch")) { + return 0; + } + if (!g_cuda_nccl.ordered_logged) { + fprintf(stderr, + "ds4: NCCL ordered-reduction diagnostic enabled " + "(rank order 0..%u)\n", + g_cuda_nccl.world - 1u); + g_cuda_nccl.ordered_logged = 1; + } + } + } else { + rc = g_cuda_nccl.all_reduce( + tensor->ptr, tensor->ptr, (size_t)count, ncclFloat32, + ncclSum, g_cuda_nccl.comm, cuda_decode_stream()); + } + if (rc != ncclSuccess) { + fprintf(stderr, "ds4: NCCL allreduce rank %u failed: %s\n", + g_cuda_nccl.rank, ds4_cuda_nccl_error(rc)); + return 0; + } + return 1; +#endif +} + extern "C" int ds4_gpu_tp_big_gate_encode(uint32_t layer, uint32_t rows, const ds4_gpu_tensor *out_t, ds4_gpu_tensor *in_t, diff --git a/ds4_eval.c b/ds4_eval.c index 7aed5d5c5b..c165aee17a 100644 --- a/ds4_eval.c +++ b/ds4_eval.c @@ -1,6 +1,7 @@ #include "ds4.h" #include "ds4_distributed.h" #include "ds4_help.h" +#include "ds4_tp.h" /* ds4-eval: small built-in benchmark integration test. * @@ -1218,6 +1219,7 @@ typedef struct { int soft_limit_think_close_rank; ds4_think_mode think_mode; ds4_dist_options dist; + ds4_tp_options tp; bool plain; bool warm_weights; bool quality; @@ -1550,6 +1552,18 @@ static eval_config parse_options(int argc, char **argv) { } if (dist_parse == DS4_DIST_CLI_MATCHED) continue; + char tp_parse_err[256] = {0}; + ds4_tp_cli_parse_result tp_parse = + ds4_tp_parse_cli_arg(arg, &i, argc, argv, &c.tp, + tp_parse_err, sizeof(tp_parse_err)); + if (tp_parse == DS4_TP_CLI_ERROR) { + fprintf(stderr, "ds4-eval: %s\n", + tp_parse_err[0] ? tp_parse_err : + "invalid network parallel option"); + exit(2); + } + if (tp_parse == DS4_TP_CLI_MATCHED) continue; + if (!strcmp(arg, "-m") || !strcmp(arg, "--model")) { c.model_path = need_arg(&i, argc, argv, arg); } else if (!strcmp(arg, "--mtp")) { @@ -1665,6 +1679,17 @@ static eval_config parse_options(int argc, char **argv) { } if (c.self_test_extractors || c.regrade_trace_path) return c; + char tp_err[256] = {0}; + if (!ds4_tp_adopt_distributed_options(&c.tp, &c.dist, + tp_err, sizeof(tp_err))) { + fprintf(stderr, "ds4-eval: %s\n", tp_err); + exit(2); + } + if (ds4_tp_enabled(&c.tp) && c.ctx_size <= 0) { + fprintf(stderr, + "ds4-eval: network parallelism requires an explicit --ctx so every rank advertises the same capacity\n"); + exit(2); + } char dist_err[256]; if (ds4_dist_prepare_engine_options(&c.dist, NULL, dist_err, sizeof(dist_err)) != 0) { fprintf(stderr, "ds4-eval: %s\n", dist_err); @@ -4172,6 +4197,7 @@ int main(int argc, char **argv) { .ssd_streaming_cold = cfg.ssd_streaming_cold, .ssd_streaming_full_layers_set = cfg.ssd_streaming_full_layers_set, .distributed = cfg.dist, + .tp = cfg.tp, }; char dist_err[256]; if (ds4_dist_prepare_engine_options(&cfg.dist, &opt, dist_err, sizeof(dist_err)) != 0) { @@ -4180,6 +4206,14 @@ int main(int argc, char **argv) { free(case_sequence); return 2; } + char tp_err[256] = {0}; + if (!ds4_tp_validate_engine_options(&opt, tp_err, sizeof(tp_err))) { + fprintf(stderr, "ds4-eval: %s\n", + tp_err[0] ? tp_err : "invalid network parallel options"); + if (trace) fclose(trace); + free(case_sequence); + return 2; + } ds4_engine *engine = NULL; if (ds4_engine_open(&engine, &opt) != 0) { @@ -4188,6 +4222,30 @@ int main(int argc, char **argv) { return 1; } + if (cfg.tp.role == DS4_TP_WORKER) { + int rc = ds4_tp_worker_run(engine, &cfg.tp, + (uint32_t)cfg.ctx_size); + ds4_engine_close(engine); + if (trace) fclose(trace); + free(case_sequence); + return rc; + } + + ds4_tp *tp_leader = NULL; + if (cfg.tp.role == DS4_TP_LEADER) { + if (!ds4_tp_leader_bind(engine, &cfg.tp, + (uint32_t)cfg.ctx_size, + &tp_leader, tp_err, sizeof(tp_err))) { + fprintf(stderr, "ds4-eval: %s\n", + tp_err[0] ? tp_err : + "network parallel initialization failed"); + ds4_engine_close(engine); + if (trace) fclose(trace); + free(case_sequence); + return 1; + } + } + int max_prompt_tokens = 0; int max_prompt_case = -1; const bool auto_ctx = cfg.ctx_size <= 0; @@ -4219,6 +4277,7 @@ int main(int argc, char **argv) { if (ds4_session_create(&session, engine, cfg.ctx_size) != 0) { fprintf(stderr, "ds4-eval: failed to create session\n"); if (trace) fclose(trace); + ds4_tp_leader_shutdown(engine, &tp_leader); ds4_engine_close(engine); free(case_sequence); return 1; @@ -4228,6 +4287,7 @@ int main(int argc, char **argv) { { ds4_session_free(session); if (trace) fclose(trace); + ds4_tp_leader_shutdown(engine, &tp_leader); ds4_engine_close(engine); free(case_sequence); return 1; @@ -4305,6 +4365,7 @@ int main(int argc, char **argv) { tui_free(&ui); ds4_session_free(session); + ds4_tp_leader_shutdown(engine, &tp_leader); ds4_engine_close(engine); if (trace) fclose(trace); free(case_sequence); diff --git a/ds4_gpu.h b/ds4_gpu.h index 9fb03ecee2..1d5e7bab47 100644 --- a/ds4_gpu.h +++ b/ds4_gpu.h @@ -115,6 +115,12 @@ int ds4_gpu_set_model_fd(int fd); int ds4_gpu_set_model_fd_for_map(int fd, const void *model_map); int ds4_gpu_build_derived_artifacts(const void *model_map, uint64_t model_size, const char *model_path); +int ds4_gpu_build_sharded_derived_artifacts(const void *model_map, + uint64_t model_size, + const char *model_path, + uint32_t expert_base, + uint32_t expert_count); +uint64_t ds4_gpu_derived_artifact_additive_bytes(const void *model_map); int ds4_gpu_model_range_replaced(const void *model_map, uint64_t offset, uint64_t bytes); int ds4_gpu_set_model_map_range(const void *model_map, uint64_t model_size, uint64_t map_offset, uint64_t map_size, uint64_t max_tensor_bytes); @@ -263,6 +269,17 @@ int ds4_gpu_stream_expert_cache_seed_experts_gpu_copy( #endif void ds4_gpu_print_memory_report(const char *label); +/* CUDA network collectives. NCCL is loaded dynamically, so ordinary CUDA, + * CPU, ROCm and Metal builds do not acquire a mandatory runtime dependency. + * The bootstrap callback broadcasts an opaque startup payload from rank 0. */ +typedef int (*ds4_gpu_tp_bootstrap_fn)(void *ud, void *data, uint32_t bytes); +int ds4_gpu_tp_collective_probe(void); +int ds4_gpu_tp_collective_init(uint32_t rank, uint32_t world, + ds4_gpu_tp_bootstrap_fn bootstrap, void *ud); +void ds4_gpu_tp_collective_shutdown(void); +int ds4_gpu_tp_allreduce_f32(ds4_gpu_tensor *tensor, uint64_t count); +int ds4_gpu_tp_collective_ready(void); + /* Tensor-parallel per-layer gates (Metal only). The encoder calls * ds4_gpu_tp_gate_encode() right after the kernels that produce a partial * block output in the TP slab: it closes the current encoder, makes the GPU @@ -273,8 +290,8 @@ void ds4_gpu_print_memory_report(const char *label); * construction. The exchange callback runs on the service thread and must * return nonzero on success. */ typedef int (*ds4_gpu_tp_exchange_fn)(void *ud, uint32_t layer, uint32_t gate, uint64_t seq); -/* Bind one rank of the two-way split. slab is the transport slab tensor and - * gpu_flags_off is the offset of its GPU-written gate-ready flag words. */ +/* Bind one rank of the Metal two-way split. slab is the transport slab tensor + * and gpu_flags_off is the offset of its GPU-written gate-ready flag words. */ int ds4_gpu_tp_init(uint32_t rank, ds4_gpu_tensor *slab, uint64_t gpu_flags_off, ds4_gpu_tp_exchange_fn fn, void *ud); @@ -324,7 +341,7 @@ void ds4_gpu_tp_keepalive_pause(int paused); * zeroes the unowned head range of the heads buffer and combines the * attn-output partials over the TP big-gate exchange. */ void ds4_gpu_tp_set_attn_head_split(int enabled); -/* Skip the whole-file model residency set (TP sharding: only the +/* Skip the whole-file model residency set (network sharding: only the * owned ranges are warmed; the rest must never be paged in). Call before * the model is mapped. */ void ds4_gpu_model_residency_skip(int skip); @@ -2495,7 +2512,8 @@ int ds4_gpu_routed_moe_batch_owned_tensor( const ds4_gpu_tensor *x, uint32_t layer_index, uint32_t n_tokens, - bool *mid_is_f16); + bool *mid_is_f16, + bool leave_slots); int ds4_gpu_routed_moe_owned_slots_combine_tensor( ds4_gpu_tensor *out, @@ -2514,6 +2532,13 @@ int ds4_gpu_routed_moe_owned_slots_combine_rows_tensor( uint32_t expert_split, uint32_t rows); +int ds4_gpu_routed_moe_slots_sum_tensor( + ds4_gpu_tensor *out, + const ds4_gpu_tensor *slots, + uint32_t out_dim, + uint32_t n_expert, + uint32_t n_tokens); + int ds4_gpu_routed_moe_owned_packed_combine_tensor( ds4_gpu_tensor *out, const ds4_gpu_tensor *home_slots, diff --git a/ds4_help.c b/ds4_help.c index adc8e634e8..fff6d7a157 100644 --- a/ds4_help.c +++ b/ds4_help.c @@ -239,12 +239,15 @@ static void print_distributed(FILE *fp, const help_colors *c) { opt(fp, c, "--dist-replay-check", "Diagnostic: reset and replay prompt, then compare logits."); opt(fp, c, "--debug", "Print coordinator route/debug logs."); fputc('\n', fp); - title(fp, c, "Tensor Parallelism"); + title(fp, c, "Network Expert/Tensor Parallelism"); fputc('\n', fp); - para(fp, c, "Tensor parallelism uses the same coordinator/worker addresses as distributed mode, but always runs one 50/50 worker. Add --tensor-parallel, omit --layers, start the worker, then start the coordinator."); + para(fp, c, "Network parallelism uses the same coordinator/worker addresses as distributed mode. It runs two or four full-model ranks in lockstep and owns weight placement, so omit --layers. Start every explicitly numbered worker, then start rank 0 as the coordinator. CUDA uses NCCL; Metal retains its two-rank RDMA/TCP path."); fputc('\n', fp); - opt(fp, c, "--tensor-parallel", "Switch --role/--listen/--coordinator to two-machine tensor parallelism."); - opt(fp, c, "--transport auto|rdma|tcp", "Tensor gate transport. Default: auto"); + opt(fp, c, "--expert-parallel", "Shard routed experts across CUDA network ranks."); + opt(fp, c, "--tensor-parallel", "Shard routed experts plus supported model-specific attention/dense decode work."); + opt(fp, c, "--tensor-parallel-world 2|4", "Number of ranks. Default: 2"); + opt(fp, c, "--tensor-parallel-rank N", "Worker rank. Required for ranks 1-3 in four-rank mode."); + opt(fp, c, "--transport auto|nccl|rdma|tcp", "Collective/gate transport. CUDA auto selects NCCL."); opt(fp, c, "--rdma-device NAME", "Select a verbs device when auto-detection is ambiguous."); opt(fp, c, "--rdma-gid-index N", "Select the local verbs GID index."); opt(fp, c, "--tensor-parallel-token-prefill", "GLM diagnostic: prefill one token at a time for exact arithmetic."); @@ -465,6 +468,17 @@ static void print_examples(FILE *fp, const help_colors *c, ds4_help_tool tool, c if (topic_is(topic, "distributed")) { opt(fp, c, "worker", "./ds4 --role worker --layers 21:output --coordinator 192.168.0.181 9000 -m ds4flash.gguf"); opt(fp, c, "coordinator", "./ds4 --role coordinator --layers 0:20 --listen 0.0.0.0 9000 -p \"Hello\" -m ds4flash.gguf"); + if (tool == DS4_HELP_SERVER) { + opt(fp, c, "TP server", "./ds4-server --cuda --ctx 4096 --tensor-parallel --tensor-parallel-world 4 --role coordinator --listen 192.168.2.30 9911 --host 127.0.0.1 --port 8000 -m model.gguf"); + } else if (tool == DS4_HELP_AGENT) { + opt(fp, c, "TP agent", "./ds4-agent --cuda --ctx 4096 --tensor-parallel --tensor-parallel-world 4 --role coordinator --listen 192.168.2.30 9911 --non-interactive -p \"Hello\" -m model.gguf"); + } else if (tool == DS4_HELP_EVAL) { + opt(fp, c, "TP eval", "./ds4-eval --cuda --ctx 4096 --tensor-parallel --tensor-parallel-world 4 --role coordinator --listen 192.168.2.30 9911 --questions 1 --plain -m model.gguf"); + } else if (tool == DS4_HELP_BENCH) { + opt(fp, c, "TP benchmark", "./ds4-bench --cuda --tensor-parallel --tensor-parallel-world 4 --role coordinator --listen 192.168.2.30 9911 --prompt-file long.txt -m model.gguf"); + } else { + opt(fp, c, "TP coordinator", "./ds4 --cuda --ctx 4096 --tensor-parallel --tensor-parallel-world 4 --role coordinator --listen 192.168.2.30 9911 -p \"Hello\" -m model.gguf"); + } } else if (topic_is(topic, "runtime")) { if (tool == DS4_HELP_SERVER) { opt(fp, c, "Metal API", "./ds4-server -m ds4flash.gguf --metal --ctx 100000"); diff --git a/ds4_rocm_unavailable.cu b/ds4_rocm_unavailable.cu index 311c8004da..a135204a47 100644 --- a/ds4_rocm_unavailable.cu +++ b/ds4_rocm_unavailable.cu @@ -31,6 +31,7 @@ ROCM_UNAVAILABLE_INT(ds4_gpu_routed_moe_one_owned_tensor) ROCM_UNAVAILABLE_INT(ds4_gpu_routed_moe_owned_packed_combine_tensor) ROCM_UNAVAILABLE_INT(ds4_gpu_routed_moe_owned_slots_combine_rows_tensor) ROCM_UNAVAILABLE_INT(ds4_gpu_routed_moe_owned_slots_combine_tensor) +ROCM_UNAVAILABLE_INT(ds4_gpu_routed_moe_slots_sum_tensor) ROCM_UNAVAILABLE_INT(ds4_gpu_shared_down_hc_expand_add_q8_0_tensor) ROCM_UNAVAILABLE_INT(ds4_gpu_shared_down_hc_expand_owned_q8_0_tensor) ROCM_UNAVAILABLE_INT(ds4_gpu_shared_mid_swiglu_q8_0_decode_exact_tensor) diff --git a/ds4_server.c b/ds4_server.c index 0ed35d4d39..6bacb35247 100644 --- a/ds4_server.c +++ b/ds4_server.c @@ -3,6 +3,7 @@ #include "ds4_gpu_args.h" #include "ds4_help.h" #include "ds4_kvstore.h" +#include "ds4_tp.h" #include "rax.h" /* OpenAI/Anthropic compatible local server. @@ -8444,6 +8445,7 @@ static void id_list_push_unique(stop_list *ids, const char *id); struct server { ds4_engine *engine; + ds4_tp *tp_leader; server_slot *slots; int slot_count; int ctx_size; @@ -9657,6 +9659,9 @@ static int kv_cache_try_load_text(server *s, server_slot *slot, if (!s || !slot) return 0; if (loaded_path_out) *loaded_path_out = NULL; if (loaded_ext_flags_out) *loaded_ext_flags_out = 0; + /* A payload restore would update rank 0 only. Normal prompt sync below + * mirrors replay to every network-parallel rank. */ + if (ds4_engine_network_parallel(s->engine)) return 0; ds4_kvstore_load_result lr = {0}; ds4_kvstore_trailer_hooks hooks = kv_cache_tool_map_hooks(s, NULL); pthread_mutex_lock(&s->inference_mu); @@ -13039,6 +13044,7 @@ static void server_close_resources(server *s) { pthread_cond_destroy(&s->clients_cv); pthread_cond_destroy(&s->cv); pthread_mutex_destroy(&s->mu); + ds4_tp_leader_shutdown(s->engine, &s->tp_leader); ds4_engine_close(s->engine); memset(s, 0, sizeof(*s)); } @@ -13117,6 +13123,18 @@ static server_config parse_options(int argc, char **argv) { } if (dist_parse == DS4_DIST_CLI_MATCHED) continue; + char tp_parse_err[256] = {0}; + ds4_tp_cli_parse_result tp_parse = + ds4_tp_parse_cli_arg(arg, &i, argc, argv, &c.engine.tp, + tp_parse_err, sizeof(tp_parse_err)); + if (tp_parse == DS4_TP_CLI_ERROR) { + server_log(DS4_LOG_DEFAULT, "ds4-server: %s", + tp_parse_err[0] ? tp_parse_err : + "invalid network parallel option"); + exit(2); + } + if (tp_parse == DS4_TP_CLI_MATCHED) continue; + if (!strcmp(arg, "-m") || !strcmp(arg, "--model")) { c.engine.model_path = need_arg(&i, argc, argv, arg); } else if (!strcmp(arg, "--mtp")) { @@ -13276,6 +13294,13 @@ static server_config parse_options(int argc, char **argv) { if (c.engine.directional_steering_file && !directional_steering_scale_set) { c.engine.directional_steering_ffn = 1.0f; } + char tp_err[256] = {0}; + if (!ds4_tp_adopt_distributed_options(&c.engine.tp, + &c.engine.distributed, + tp_err, sizeof(tp_err))) { + server_log(DS4_LOG_DEFAULT, "ds4-server: %s", tp_err); + exit(2); + } char dist_err[256]; if (ds4_dist_prepare_engine_options(&c.engine.distributed, &c.engine, @@ -13284,6 +13309,10 @@ static server_config parse_options(int argc, char **argv) { server_log(DS4_LOG_DEFAULT, "ds4-server: %s", dist_err); exit(2); } + if (!ds4_tp_validate_engine_options(&c.engine, tp_err, sizeof(tp_err))) { + server_log(DS4_LOG_DEFAULT, "ds4-server: %s", tp_err); + exit(2); + } return c; } @@ -13356,7 +13385,18 @@ int main(int argc, char **argv) { return 1; } + if (cfg.engine.tp.role == DS4_TP_WORKER) { + signal(SIGINT, SIG_DFL); + signal(SIGTERM, SIG_DFL); + int rc = ds4_tp_worker_run(engine, &cfg.engine.tp, + (uint32_t)cfg.ctx_size); + ds4_engine_close(engine); + return rc; + } + if (cfg.engine.distributed.role == DS4_DISTRIBUTED_WORKER) { + signal(SIGINT, SIG_DFL); + signal(SIGTERM, SIG_DFL); ds4_dist_generation_options gen = { .ctx_size = cfg.ctx_size, }; @@ -13365,6 +13405,20 @@ int main(int argc, char **argv) { return rc; } + ds4_tp *tp_leader = NULL; + if (cfg.engine.tp.role == DS4_TP_LEADER) { + char tp_err[256] = {0}; + if (!ds4_tp_leader_bind(engine, &cfg.engine.tp, + (uint32_t)cfg.ctx_size, + &tp_leader, tp_err, sizeof(tp_err))) { + server_log(DS4_LOG_DEFAULT, "ds4-server: %s", + tp_err[0] ? tp_err : + "network parallel initialization failed"); + ds4_engine_close(engine); + return 1; + } + } + const int slot_count = cfg.batched_sessions > 0 ? cfg.batched_sessions : 1; log_context_memory(cfg.engine.backend, cfg.ctx_size, @@ -13374,6 +13428,7 @@ int main(int argc, char **argv) { server s = {0}; s.engine = engine; + s.tp_leader = tp_leader; s.ctx_size = cfg.ctx_size; s.slot_count = slot_count; s.batched_mode = cfg.batched_sessions > 0; diff --git a/ds4_tp.c b/ds4_tp.c index 34ec68ba20..2756064174 100644 --- a/ds4_tp.c +++ b/ds4_tp.c @@ -1,13 +1,11 @@ -/* Tensor-parallel transport and lockstep protocol. See ds4_tp.h and - * misc/METAL_TENSOR_PARALLELISM.md for the design. +/* Network-parallel transport and lockstep protocol. See ds4_tp.h and + * misc/METAL_TENSOR_PARALLELISM.md for the original Metal design. * - * Wire notes: both ranks are identical Apple Silicon machines by - * definition, so the wire format is host little-endian; the hello magic - * doubles as a byte-order check. The control socket is a plain blocking - * TCP stream carrying framed commands. Gate traffic goes over RDMA - * (Thunderbolt UC queue pair, two-sided send/recv — see the driver quirks - * note at ds4_tp_rdma) or over a dedicated full-duplex TCP socket at 16KB - * per direction as the fallback. */ + * Wire notes: ranks are homogeneous machines and the wire format is host + * little-endian; the hello magic doubles as a byte-order check. The control + * path is a plain blocking TCP stream carrying framed commands. Metal gate + * traffic uses RDMA (Thunderbolt UC queue pair, two-sided send/recv — see + * ds4_tp_rdma) or a dedicated TCP socket. CUDA gate traffic uses NCCL. */ #include #include @@ -39,7 +37,7 @@ #define DS4_TP_MAGIC UINT32_C(0x44533454) /* "DS4T" */ #define DS4_TP_BATCH_MAGIC UINT32_C(0x44533442) /* "DS4B" */ -#define DS4_TP_PROTOCOL_VERSION 7u +#define DS4_TP_PROTOCOL_VERSION 10u /* Default gate timeout is generous: the first gate after a sync waits for * the peer's whole (possibly cold page cache) prefill. */ @@ -55,6 +53,10 @@ typedef struct { uint32_t magic; /* also detects byte-order mismatch */ uint32_t version; uint32_t role; + uint32_t rank; + uint32_t world; + uint32_t collective; + uint32_t expert_only; uint32_t rdma_ok; /* this side has a usable verbs device */ uint64_t gguf_bytes; uint32_t model_id; @@ -66,7 +68,6 @@ typedef struct { uint32_t gate_slot_start; uint32_t gate_slot_step; uint32_t gates_per_token; - uint32_t pad; } ds4_tp_hello_fixed; typedef struct { @@ -152,10 +153,13 @@ typedef struct { struct ds4_tp { ds4_tp_options opt; - int rank; /* 0 leader, 1 worker */ - int control_fd; + int rank; /* 0 leader, workers 1..world-1 */ + uint32_t world; + int control_fd; /* rank 0 peer alias for worker; rank 1 for leader */ + int control_fds[DS4_TP_MAX_WORLD]; int data_fd; /* TCP fallback, headers, and verify gates */ bool rdma_active; + bool collective_active; uint32_t peer_ctx; uint32_t n_layer; uint32_t n_embd; @@ -280,7 +284,8 @@ static int tp_listen(const char *host, int port, char *err, size_t errlen) { if (fd < 0) continue; int one = 1; setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one)); - if (bind(fd, ai->ai_addr, ai->ai_addrlen) == 0 && listen(fd, 2) == 0) break; + if (bind(fd, ai->ai_addr, ai->ai_addrlen) == 0 && + listen(fd, (int)(DS4_TP_MAX_WORLD * 2u)) == 0) break; close(fd); fd = -1; } @@ -334,6 +339,18 @@ static int tp_send_frame(int fd, uint32_t type, const void *payload, uint32_t by return 1; } +static int tp_send_frame_workers(ds4_tp *tp, uint32_t type, + const void *payload, uint32_t bytes) { + if (!tp || tp->rank != 0) return 0; + for (uint32_t rank = 1; rank < tp->world; rank++) { + if (tp->control_fds[rank] < 0 || + !tp_send_frame(tp->control_fds[rank], type, payload, bytes)) { + return 0; + } + } + return 1; +} + static int tp_read_frame_header(int fd, uint32_t *type, uint32_t *bytes) { ds4_tp_frame_header h; if (!tp_read_full(fd, &h, sizeof(h))) return 0; @@ -347,15 +364,47 @@ static int tp_read_frame_header(int fd, uint32_t *type, uint32_t *bytes) { * Options and CLI. * --------------------------------------------------------------------- */ +static int tp_parse_u32_arg(const char *name, const char *value, + uint32_t min_value, uint32_t max_value, + uint32_t *out, char *err, size_t errlen) { + char *end = NULL; + errno = 0; + unsigned long parsed = strtoul(value, &end, 10); + if (errno != 0 || !end || *end != '\0' || + parsed < min_value || parsed > max_value) { + tp_set_err(err, errlen, "invalid %s %s", name, value); + return 0; + } + *out = (uint32_t)parsed; + return 1; +} + +bool ds4_tp_partition(uint32_t total, uint32_t rank, uint32_t world, + uint32_t *base, uint32_t *count) { + if (!base || !count || total == 0u || world < 2u || + world > DS4_TP_MAX_WORLD || rank >= world) { + return false; + } + const uint32_t first = (uint32_t)((uint64_t)total * rank / world); + const uint32_t last = (uint32_t)((uint64_t)total * (rank + 1u) / world); + if (last <= first) return false; + *base = first; + *count = last - first; + return true; +} + bool ds4_tp_enabled(const ds4_tp_options *opt) { return opt && opt->role != DS4_TP_NONE; } void ds4_tp_usage(FILE *fp) { fprintf(fp, - "Tensor parallelism (two identical machines):\n" - " --tensor-parallel Use --role/--listen/--coordinator for a 50/50 TP pair.\n" - " --transport Gate transport (default auto).\n" + "Network expert/tensor parallelism (two or four machines):\n" + " --expert-parallel Shard routed experts across network ranks.\n" + " --tensor-parallel Shard attention, dense FFNs, and routed experts.\n" + " --tensor-parallel-world <2|4> Number of ranks (default 2).\n" + " --tensor-parallel-rank Worker rank; required for four-rank mode.\n" + " --transport Collective/gate transport (default auto).\n" " --rdma-device Select a verbs device such as rdma_en1.\n" " --rdma-gid-index Select the local verbs GID index.\n" " --tensor-parallel-token-prefill\n" @@ -373,18 +422,50 @@ int ds4_tp_parse_cli_arg( size_t errlen) { int i = *index; - if (!strcmp(arg, "--tensor-parallel")) { + if (!strcmp(arg, "--expert-parallel")) { + if (opt->requested && !opt->expert_only) { + tp_set_err(err, errlen, + "--expert-parallel and --tensor-parallel are exclusive"); + return DS4_TP_CLI_ERROR; + } opt->requested = true; + opt->expert_only = true; + } else if (!strcmp(arg, "--tensor-parallel")) { + if (opt->requested && opt->expert_only) { + tp_set_err(err, errlen, + "--expert-parallel and --tensor-parallel are exclusive"); + return DS4_TP_CLI_ERROR; + } + opt->requested = true; + opt->expert_only = false; } else if (!strcmp(arg, "--transport")) { if (i + 1 >= argc) goto missing; const char *v = argv[++i]; if (!strcmp(v, "auto")) opt->transport = DS4_TP_TRANSPORT_AUTO; + else if (!strcmp(v, "nccl")) opt->transport = DS4_TP_TRANSPORT_NCCL; else if (!strcmp(v, "rdma")) opt->transport = DS4_TP_TRANSPORT_RDMA; else if (!strcmp(v, "tcp")) opt->transport = DS4_TP_TRANSPORT_TCP; else { tp_set_err(err, errlen, "invalid %s value: %s", arg, v); return DS4_TP_CLI_ERROR; } + } else if (!strcmp(arg, "--tensor-parallel-world")) { + if (i + 1 >= argc) goto missing; + if (!tp_parse_u32_arg(arg, argv[++i], 2, DS4_TP_MAX_WORLD, + &opt->world_size, err, errlen) || + (opt->world_size != 2 && opt->world_size != 4)) { + if (!err || !err[0]) { + tp_set_err(err, errlen, "%s must be 2 or 4", arg); + } + return DS4_TP_CLI_ERROR; + } + } else if (!strcmp(arg, "--tensor-parallel-rank")) { + if (i + 1 >= argc) goto missing; + if (!tp_parse_u32_arg(arg, argv[++i], 0, DS4_TP_MAX_WORLD - 1, + &opt->rank, err, errlen)) { + return DS4_TP_CLI_ERROR; + } + opt->rank_set = true; } else if (!strcmp(arg, "--rdma-device")) { if (i + 1 >= argc) goto missing; opt->rdma_device = argv[++i]; @@ -421,32 +502,38 @@ int ds4_tp_adopt_distributed_options( size_t errlen) { if (!tp || !dist || !tp->requested) return 1; + if (tp->world_size == 0) tp->world_size = 2; + if (tp->world_size != 2 && tp->world_size != 4) { + tp_set_err(err, errlen, "--tensor-parallel-world must be 2 or 4"); + return 0; + } if (tp->role != DS4_TP_NONE) { tp_set_err(err, errlen, - "--tensor-parallel selects its role through --role"); + "network parallelism selects its role through --role"); return 0; } if (dist->role == DS4_DISTRIBUTED_NONE) { tp_set_err(err, errlen, - "--tensor-parallel requires --role coordinator or --role worker"); + "network parallelism requires --role coordinator or --role worker"); return 0; } if (dist->layers.set) { tp_set_err(err, errlen, - "tensor parallelism always uses one 50/50 worker; omit --layers"); + "network parallelism owns model placement; omit --layers"); return 0; } if (dist->prefill_chunk || dist->prefill_window || dist->activation_bits || dist->replay_check || dist->debug) { tp_set_err(err, errlen, - "--dist-* and distributed debug options cannot be used with --tensor-parallel"); + "--dist-* and distributed debug options cannot be used with " + "network parallelism"); return 0; } if (dist->role == DS4_DISTRIBUTED_COORDINATOR) { if (!dist->listen_host || dist->listen_port <= 0) { tp_set_err(err, errlen, - "--role coordinator --tensor-parallel requires --listen HOST PORT"); + "network parallel coordinator requires --listen HOST PORT"); return 0; } if (dist->coordinator_host || dist->coordinator_port) { @@ -455,20 +542,43 @@ int ds4_tp_adopt_distributed_options( return 0; } tp->role = DS4_TP_LEADER; + if (tp->rank_set && tp->rank != 0) { + tp_set_err(err, errlen, + "the coordinator is always --tensor-parallel-rank 0"); + return 0; + } + tp->rank = 0; + tp->rank_set = true; tp->listen_host = dist->listen_host; tp->listen_port = dist->listen_port; } else if (dist->role == DS4_DISTRIBUTED_WORKER) { if (!dist->coordinator_host || dist->coordinator_port <= 0) { tp_set_err(err, errlen, - "--role worker --tensor-parallel requires --coordinator HOST PORT"); + "network parallel worker requires --coordinator HOST PORT"); return 0; } if (dist->listen_host || dist->listen_port) { tp_set_err(err, errlen, - "--role worker --tensor-parallel must not use --listen"); + "network parallel worker must not use --listen"); return 0; } tp->role = DS4_TP_WORKER; + if (!tp->rank_set) { + if (tp->world_size == 2) { + tp->rank = 1; + tp->rank_set = true; + } else { + tp_set_err(err, errlen, + "four-rank workers require --tensor-parallel-rank 1, 2, or 3"); + return 0; + } + } + if (tp->rank == 0 || tp->rank >= tp->world_size) { + tp_set_err(err, errlen, + "worker --tensor-parallel-rank %u is outside 1..%u", + tp->rank, tp->world_size - 1u); + return 0; + } tp->leader_host = dist->coordinator_host; tp->leader_port = dist->coordinator_port; } else { @@ -487,31 +597,82 @@ int ds4_tp_validate_engine_options( { if (!ds4_tp_enabled(&opt->tp)) { if (opt->tp.requested || opt->tp.transport != DS4_TP_TRANSPORT_AUTO || + opt->tp.world_size != 0 || opt->tp.rank_set || opt->tp.rdma_device || opt->tp.rdma_gid_index_set || opt->tp.glm_token_prefill || opt->tp.debug_hash != 0) { tp_set_err(err, errlen, - "tensor-parallel options require --tensor-parallel and --role"); + "network parallel options require --expert-parallel or " + "--tensor-parallel and --role"); return 0; } return 1; } +#ifdef __APPLE__ if (opt->backend != DS4_BACKEND_METAL) { - tp_set_err(err, errlen, "tensor parallelism requires the Metal backend"); + tp_set_err(err, errlen, "network parallelism requires the Metal backend on macOS"); + return 0; + } + if (opt->tp.world_size != 2) { + tp_set_err(err, errlen, "Metal network parallelism currently supports two ranks"); + return 0; + } + if (opt->tp.expert_only) { + tp_set_err(err, errlen, + "--expert-parallel is currently implemented by the CUDA backend"); + return 0; + } + if (opt->tp.transport == DS4_TP_TRANSPORT_NCCL) { + tp_set_err(err, errlen, "--transport nccl requires the CUDA backend"); + return 0; + } +#else +#ifdef DS4_ROCM_BUILD + tp_set_err(err, errlen, + "network parallelism is not implemented for the ROCm backend"); + return 0; +#endif + if (opt->backend != DS4_BACKEND_CUDA) { + tp_set_err(err, errlen, "network parallelism requires the CUDA backend on Linux"); return 0; } + if (opt->tp.transport != DS4_TP_TRANSPORT_AUTO && + opt->tp.transport != DS4_TP_TRANSPORT_NCCL) { + tp_set_err(err, errlen, + "CUDA network parallelism requires --transport nccl (or auto)"); + return 0; + } + if (opt->glm_mtp || opt->glm_mtp_timing) { + tp_set_err(err, errlen, + "CUDA network parallelism does not yet support --glm-mtp"); + return 0; + } +#endif if (opt->ssd_streaming) { - tp_set_err(err, errlen, "tensor parallelism requires resident weights (no --ssd-streaming)"); + tp_set_err(err, errlen, + "network parallelism requires resident weights (no --ssd-streaming)"); + return 0; + } + if ((opt->mtp_path && opt->mtp_path[0]) || opt->dspark) { + tp_set_err(err, errlen, + "network parallelism does not yet support MTP/DSpark speculative drafting"); + return 0; + } + if ((opt->directional_steering_file && + opt->directional_steering_file[0]) || + opt->directional_steering_attn != 0.0f || + opt->directional_steering_ffn != 0.0f) { + tp_set_err(err, errlen, + "network parallelism does not yet support directional steering"); return 0; } if (opt->distributed.role != DS4_DISTRIBUTED_NONE) { - tp_set_err(err, errlen, "tensor parallelism and --role distributed modes are exclusive"); + tp_set_err(err, errlen, + "network parallelism and distributed layer roles are exclusive"); return 0; } - /* Speculative drafting (DSpark/MTP) is allowed on the leader: the - * verify block is mirrored to the worker via DS4_TP_FRAME_VERIFY and - * the legacy MTP path falls back to per-token decode under TP. */ if (opt->load_slice) { - tp_set_err(err, errlen, "tensor parallelism does not use distributed layer slices"); + tp_set_err(err, errlen, + "network parallelism does not use distributed layer slices"); return 0; } return 1; @@ -1238,12 +1399,18 @@ static void tp_rdma_close(ds4_tp *tp) { * Bring-up. * --------------------------------------------------------------------- */ -static int tp_hello_exchange(ds4_tp *tp, const ds4_tp_identity *id, int rdma_ok, - char *err, size_t errlen) { +static int tp_hello_exchange_fd(ds4_tp *tp, int fd, + const ds4_tp_identity *id, int rdma_ok, + uint32_t *peer_rank, int *peer_rdma_ok, + char *err, size_t errlen) { ds4_tp_hello_fixed mine = { .magic = DS4_TP_MAGIC, .version = DS4_TP_PROTOCOL_VERSION, .role = (uint32_t)tp->opt.role, + .rank = (uint32_t)tp->rank, + .world = tp->world, + .collective = tp->collective_active ? 1u : 0u, + .expert_only = tp->opt.expert_only ? 1u : 0u, .rdma_ok = (uint32_t)rdma_ok, .gguf_bytes = id->gguf_bytes, .model_id = id->model_id, @@ -1257,8 +1424,8 @@ static int tp_hello_exchange(ds4_tp *tp, const ds4_tp_identity *id, int rdma_ok, .gates_per_token = id->gates_per_token, }; ds4_tp_hello_fixed theirs; - if (!tp_write_full(tp->control_fd, &mine, sizeof(mine)) || - !tp_read_full(tp->control_fd, &theirs, sizeof(theirs))) { + if (!tp_write_full(fd, &mine, sizeof(mine)) || + !tp_read_full(fd, &theirs, sizeof(theirs))) { tp_set_err(err, errlen, "tp hello exchange failed"); return 0; } @@ -1275,6 +1442,34 @@ static int tp_hello_exchange(ds4_tp *tp, const ds4_tp_identity *id, int rdma_ok, tp_set_err(err, errlen, "tp hello: both sides claim role %u", mine.role); return 0; } + if (theirs.world != mine.world || theirs.rank >= mine.world || + theirs.rank == mine.rank || + (theirs.role == DS4_TP_LEADER && theirs.rank != 0) || + (theirs.role == DS4_TP_WORKER && theirs.rank == 0)) { + tp_set_err(err, errlen, + "tp hello: invalid peer topology rank=%u world=%u role=%u " + "(local rank=%u world=%u)", + theirs.rank, theirs.world, theirs.role, + mine.rank, mine.world); + return 0; + } + if (theirs.collective != mine.collective) { + tp_set_err(err, errlen, + "tp hello: transport mismatch (local collective=%u, peer=%u)", + mine.collective, theirs.collective); + return 0; + } + if (theirs.expert_only != mine.expert_only) { + tp_set_err(err, errlen, + "tp hello: parallel mode mismatch (local expert_only=%u, peer=%u)", + mine.expert_only, theirs.expert_only); + return 0; + } + if (mine.ctx_size == 0 || theirs.ctx_size == 0) { + tp_set_err(err, errlen, + "tp hello: every rank must advertise a nonzero context capacity"); + return 0; + } if (theirs.gguf_bytes != mine.gguf_bytes || theirs.model_id != mine.model_id || theirs.n_layer != mine.n_layer || theirs.n_embd != mine.n_embd || theirs.n_vocab != mine.n_vocab || theirs.quant_bits != mine.quant_bits || @@ -1288,7 +1483,7 @@ static int tp_hello_exchange(ds4_tp *tp, const ds4_tp_identity *id, int rdma_ok, theirs.n_layer, theirs.n_embd, theirs.n_vocab, theirs.quant_bits); return 0; } - tp->peer_ctx = theirs.ctx_size; + if (theirs.ctx_size < tp->peer_ctx) tp->peer_ctx = theirs.ctx_size; tp->n_layer = id->n_layer; tp->n_embd = id->n_embd; tp->vec_bytes = (uint64_t)id->n_embd * sizeof(float); @@ -1297,14 +1492,8 @@ static int tp_hello_exchange(ds4_tp *tp, const ds4_tp_identity *id, int rdma_ok, tp->gate_slot_step = id->gate_slot_step; tp->gates_per_token = id->gates_per_token; tp_slab_layout(tp); - /* Transport decision: RDMA only when both sides can. */ - int want_rdma = tp->opt.transport != DS4_TP_TRANSPORT_TCP; - tp->rdma_active = want_rdma && rdma_ok && theirs.rdma_ok; - if (tp->opt.transport == DS4_TP_TRANSPORT_RDMA && !tp->rdma_active) { - tp_set_err(err, errlen, "tp: --transport rdma but %s side has no active device", - rdma_ok ? "the peer" : "this"); - return 0; - } + if (peer_rank) *peer_rank = theirs.rank; + if (peer_rdma_ok) *peer_rdma_ok = (int)theirs.rdma_ok; return 1; } @@ -1315,16 +1504,43 @@ int ds4_tp_create( char *err, size_t errlen) { + if (!out || !opt || !id) { + tp_set_err(err, errlen, "tp: invalid initialization request"); + return 0; + } *out = NULL; ds4_tp *tp = calloc(1, sizeof(*tp)); if (!tp) { tp_set_err(err, errlen, "tp: out of memory"); return 0; } - tp->opt = *opt; - tp->rank = opt->role == DS4_TP_LEADER ? 0 : 1; + int listener = -1; tp->control_fd = -1; tp->data_fd = -1; + for (uint32_t r = 0; r < DS4_TP_MAX_WORLD; r++) { + tp->control_fds[r] = -1; + } + tp->opt = *opt; + tp->world = opt->world_size ? opt->world_size : 2u; + tp->rank = opt->role == DS4_TP_LEADER ? 0 : (int)opt->rank; + if ((tp->world != 2u && tp->world != 4u) || + tp->world > DS4_TP_MAX_WORLD || tp->rank < 0 || + (uint32_t)tp->rank >= tp->world) { + tp_set_err(err, errlen, "tp: invalid rank %d of %u", tp->rank, tp->world); + goto fail; + } +#ifdef __APPLE__ + tp->collective_active = opt->transport == DS4_TP_TRANSPORT_NCCL; +#else + tp->collective_active = + opt->transport == DS4_TP_TRANSPORT_AUTO || + opt->transport == DS4_TP_TRANSPORT_NCCL; +#endif + if (tp->world > 2u && !tp->collective_active) { + tp_set_err(err, errlen, "tp: four ranks require a collective transport"); + goto fail; + } + tp->peer_ctx = UINT32_MAX; tp->timeout_sec = DS4_TP_DEFAULT_TIMEOUT_SEC; const char *tmo = getenv("DS4_TP_TIMEOUT_SEC"); if (tmo) tp->timeout_sec = (uint64_t)atoi(tmo); @@ -1336,32 +1552,80 @@ int ds4_tp_create( rdma_ok = tp_rdma_probe(&tp->rdma.api); #endif - int listener = -1; if (tp->rank == 0) { listener = tp_listen(opt->listen_host, opt->listen_port, err, errlen); if (listener < 0) goto fail; - fprintf(stderr, "ds4-tp: waiting for worker on %s:%d ...\n", + fprintf(stderr, "ds4-tp: waiting for %u worker%s on %s:%d ...\n", + tp->world - 1u, tp->world == 2u ? "" : "s", opt->listen_host ? opt->listen_host : "0.0.0.0", opt->listen_port); - tp->control_fd = accept(listener, NULL, NULL); - if (tp->control_fd < 0) { - tp_set_err(err, errlen, "tp accept: %s", strerror(errno)); - goto fail; + for (uint32_t connected = 0; connected < tp->world - 1u; connected++) { + int fd = accept(listener, NULL, NULL); + if (fd < 0) { + tp_set_err(err, errlen, "tp accept: %s", strerror(errno)); + goto fail; + } + tp_socket_tune(fd); + uint32_t peer_rank = UINT32_MAX; + int peer_rdma = 0; + if (!tp_hello_exchange_fd(tp, fd, id, rdma_ok, + &peer_rank, &peer_rdma, err, errlen)) { + close(fd); + goto fail; + } + if (peer_rank == 0 || peer_rank >= tp->world || + tp->control_fds[peer_rank] >= 0) { + close(fd); + tp_set_err(err, errlen, + "tp: duplicate or invalid worker rank %u", peer_rank); + goto fail; + } + tp->control_fds[peer_rank] = fd; + if (tp->world == 2u && !tp->collective_active) { + int want_rdma = tp->opt.transport != DS4_TP_TRANSPORT_TCP; + tp->rdma_active = want_rdma && rdma_ok && peer_rdma; + } + fprintf(stderr, "ds4-tp: worker rank %u connected (%u/%u)\n", + peer_rank, connected + 1u, tp->world - 1u); } + tp->control_fd = tp->control_fds[1]; } else { - tp->control_fd = tp_dial(opt->leader_host, opt->leader_port, - (double)tp->timeout_sec, err, errlen); - if (tp->control_fd < 0) goto fail; + int fd = tp_dial(opt->leader_host, opt->leader_port, + (double)tp->timeout_sec, err, errlen); + if (fd < 0) goto fail; + tp_socket_tune(fd); + uint32_t peer_rank = UINT32_MAX; + int peer_rdma = 0; + if (!tp_hello_exchange_fd(tp, fd, id, rdma_ok, + &peer_rank, &peer_rdma, err, errlen)) { + close(fd); + goto fail; + } + if (peer_rank != 0) { + close(fd); + tp_set_err(err, errlen, "tp: worker connected to non-leader rank %u", + peer_rank); + goto fail; + } + tp->control_fds[0] = fd; + tp->control_fd = fd; + if (tp->world == 2u && !tp->collective_active) { + int want_rdma = tp->opt.transport != DS4_TP_TRANSPORT_TCP; + tp->rdma_active = want_rdma && rdma_ok && peer_rdma; + } + } + if (tp->peer_ctx == UINT32_MAX) tp->peer_ctx = id->ctx_size; + if (opt->transport == DS4_TP_TRANSPORT_RDMA && !tp->rdma_active) { + tp_set_err(err, errlen, + "tp: --transport rdma but one side has no active device"); + goto fail; } - tp_socket_tune(tp->control_fd); - - if (!tp_hello_exchange(tp, id, rdma_ok, err, errlen)) goto fail; #ifdef DS4_TP_HAVE_VERBS if (tp->rdma_active) { if (!tp_rdma_open(tp, err, errlen)) goto fail; } #endif - { + if (!tp->collective_active) { /* Second socket dedicated to gate traffic so control frames never * interleave with gate payloads. Created under RDMA too for * headers, verify-block gates, and transport fallback. */ @@ -1379,8 +1643,9 @@ int ds4_tp_create( tp_socket_tune(tp->data_fd); } if (listener >= 0) close(listener); - fprintf(stderr, "ds4-tp: %s connected, transport=%s\n", - tp->rank == 0 ? "worker" : "leader", + fprintf(stderr, "ds4-tp: rank %d/%u ready, transport=%s\n", + tp->rank, tp->world, + tp->collective_active ? "nccl" : tp->rdma_active ? "rdma" : "tcp"); *out = tp; return 1; @@ -1406,14 +1671,25 @@ void ds4_tp_free(ds4_tp *tp) { #ifdef DS4_TP_HAVE_VERBS tp_rdma_close(tp); #endif - if (tp->control_fd >= 0) close(tp->control_fd); + for (uint32_t r = 0; r < DS4_TP_MAX_WORLD; r++) { + if (tp->control_fds[r] >= 0) close(tp->control_fds[r]); + } if (tp->data_fd >= 0) close(tp->data_fd); free(tp); } -int ds4_tp_rank(const ds4_tp *tp) { return tp->rank; } -bool ds4_tp_is_rdma(const ds4_tp *tp) { return tp->rdma_active; } -uint32_t ds4_tp_peer_ctx(const ds4_tp *tp) { return tp->peer_ctx; } +int ds4_tp_rank(const ds4_tp *tp) { return tp ? tp->rank : -1; } +uint32_t ds4_tp_world(const ds4_tp *tp) { return tp ? tp->world : 0u; } +bool ds4_tp_is_rdma(const ds4_tp *tp) { return tp && tp->rdma_active; } +bool ds4_tp_is_collective(const ds4_tp *tp) { + return tp && tp->collective_active; +} +bool ds4_tp_is_expert_only(const ds4_tp *tp) { + return tp && tp->opt.expert_only; +} +uint32_t ds4_tp_peer_ctx(const ds4_tp *tp) { + return tp ? tp->peer_ctx : 0u; +} bool ds4_tp_failed(const ds4_tp *tp) { return tp && atomic_load_explicit(&tp->failed, memory_order_acquire); } @@ -1421,6 +1697,89 @@ void ds4_tp_mark_failed(ds4_tp *tp) { if (tp) atomic_store_explicit(&tp->failed, true, memory_order_release); } +int ds4_tp_collective_preflight(ds4_tp *tp, int local_ready, + char *err, size_t errlen) { + if (!tp || !tp->collective_active) { + tp_set_err(err, errlen, "tp: invalid collective preflight request"); + return 0; + } + uint32_t ready = local_ready ? 1u : 0u; + if (tp->rank == 0) { + for (uint32_t rank = 1; rank < tp->world; rank++) { + uint32_t type = 0, bytes = 0, peer_ready = 0; + const int fd = tp->control_fds[rank]; + if (fd < 0 || !tp_read_frame_header(fd, &type, &bytes) || + type != DS4_TP_FRAME_COLLECTIVE_PREFLIGHT || + bytes != sizeof(peer_ready) || + !tp_read_full(fd, &peer_ready, sizeof(peer_ready))) { + ds4_tp_mark_failed(tp); + tp_set_err(err, errlen, + "tp: collective preflight failed for rank %u", rank); + return 0; + } + ready &= peer_ready == 1u; + } + if (!tp_send_frame_workers(tp, DS4_TP_FRAME_COLLECTIVE_PREFLIGHT, + &ready, sizeof(ready))) { + ds4_tp_mark_failed(tp); + tp_set_err(err, errlen, + "tp: collective preflight result broadcast failed"); + return 0; + } + } else { + if (!tp_send_frame(tp->control_fd, + DS4_TP_FRAME_COLLECTIVE_PREFLIGHT, + &ready, sizeof(ready))) { + ds4_tp_mark_failed(tp); + tp_set_err(err, errlen, + "tp: rank %d collective preflight send failed", tp->rank); + return 0; + } + uint32_t type = 0, bytes = 0; + if (!tp_read_frame_header(tp->control_fd, &type, &bytes) || + type != DS4_TP_FRAME_COLLECTIVE_PREFLIGHT || + bytes != sizeof(ready) || + !tp_read_full(tp->control_fd, &ready, sizeof(ready))) { + ds4_tp_mark_failed(tp); + tp_set_err(err, errlen, + "tp: rank %d collective preflight receive failed", tp->rank); + return 0; + } + } + if (ready != 1u) { + tp_set_err(err, errlen, + "tp: NCCL runtime is unavailable on one or more ranks"); + return 0; + } + return 1; +} + +int ds4_tp_broadcast_blob(ds4_tp *tp, void *data, uint32_t bytes, + char *err, size_t errlen) { + if (!tp || !data || bytes == 0 || !tp->collective_active) { + tp_set_err(err, errlen, "tp: invalid collective bootstrap request"); + return 0; + } + if (tp->rank == 0) { + if (!tp_send_frame_workers(tp, DS4_TP_FRAME_BOOTSTRAP, data, bytes)) { + ds4_tp_mark_failed(tp); + tp_set_err(err, errlen, "tp: collective bootstrap broadcast failed"); + return 0; + } + return 1; + } + uint32_t type = 0, received = 0; + if (!tp_read_frame_header(tp->control_fd, &type, &received) || + type != DS4_TP_FRAME_BOOTSTRAP || received != bytes || + !tp_read_full(tp->control_fd, data, bytes)) { + ds4_tp_mark_failed(tp); + tp_set_err(err, errlen, + "tp: rank %d collective bootstrap receive failed", tp->rank); + return 0; + } + return 1; +} + /* ------------------------------------------------------------------------ * Gate exchange. * --------------------------------------------------------------------- */ @@ -1634,20 +1993,20 @@ static int tp_send_token_command(ds4_tp *tp, uint32_t type, memcpy(payload, &h, sizeof(h)); int32_t *wire_tokens = (int32_t *)(payload + sizeof(h)); for (uint32_t i = 0; i < count; i++) wire_tokens[i] = (int32_t)tokens[i]; - const int ok = tp_send_frame(tp->control_fd, type, payload, bytes); + const int ok = tp_send_frame_workers(tp, type, payload, bytes); free(payload); return ok; } int ds4_tp_send_session_create(ds4_tp *tp, uint64_t session_id, int ctx_size) { ds4_tp_value_command msg = { session_id, (int32_t)ctx_size, 0 }; - return tp_send_frame(tp->control_fd, DS4_TP_FRAME_SESSION_CREATE, - &msg, sizeof(msg)); + return tp_send_frame_workers(tp, DS4_TP_FRAME_SESSION_CREATE, + &msg, sizeof(msg)); } int ds4_tp_send_session_destroy(ds4_tp *tp, uint64_t session_id) { - return tp_send_frame(tp->control_fd, DS4_TP_FRAME_SESSION_DESTROY, - &session_id, sizeof(session_id)); + return tp_send_frame_workers(tp, DS4_TP_FRAME_SESSION_DESTROY, + &session_id, sizeof(session_id)); } int ds4_tp_send_sync(ds4_tp *tp, uint64_t session_id, @@ -1659,18 +2018,18 @@ int ds4_tp_send_sync(ds4_tp *tp, uint64_t session_id, int ds4_tp_send_eval(ds4_tp *tp, uint64_t session_id, uint64_t seq, int token) { ds4_tp_eval_command msg = { session_id, seq, (int32_t)token, 0 }; - return tp_send_frame(tp->control_fd, DS4_TP_FRAME_EVAL, &msg, sizeof(msg)); + return tp_send_frame_workers(tp, DS4_TP_FRAME_EVAL, &msg, sizeof(msg)); } int ds4_tp_send_rewind(ds4_tp *tp, uint64_t session_id, int pos) { ds4_tp_value_command msg = { session_id, (int32_t)pos, 0 }; - return tp_send_frame(tp->control_fd, DS4_TP_FRAME_REWIND, - &msg, sizeof(msg)); + return tp_send_frame_workers(tp, DS4_TP_FRAME_REWIND, + &msg, sizeof(msg)); } int ds4_tp_send_invalidate(ds4_tp *tp, uint64_t session_id) { - return tp_send_frame(tp->control_fd, DS4_TP_FRAME_INVALIDATE, - &session_id, sizeof(session_id)); + return tp_send_frame_workers(tp, DS4_TP_FRAME_INVALIDATE, + &session_id, sizeof(session_id)); } int ds4_tp_send_eval_batch(ds4_tp *tp, const ds4_tp_batch_item *items, @@ -1684,8 +2043,8 @@ int ds4_tp_send_eval_batch(ds4_tp *tp, const ds4_tp_batch_item *items, ds4_tp_batch_command_header h = { count, 0 }; memcpy(payload, &h, sizeof(h)); memcpy(payload + sizeof(h), items, (size_t)count * sizeof(*items)); - const int ok = tp_send_frame(tp->control_fd, DS4_TP_FRAME_EVAL_BATCH, - payload, bytes); + const int ok = tp_send_frame_workers(tp, DS4_TP_FRAME_EVAL_BATCH, + payload, bytes); free(payload); return ok; } @@ -1712,8 +2071,8 @@ int ds4_tp_send_mixed_batch(ds4_tp *tp, uint64_t prefill_session_id, wire_tokens[i] = (int32_t)prompt[i]; } memcpy(payload + sizeof(h) + prompt_bytes, items, (size_t)item_bytes); - const int ok = tp_send_frame(tp->control_fd, DS4_TP_FRAME_MIXED_BATCH, - payload, bytes); + const int ok = tp_send_frame_workers(tp, DS4_TP_FRAME_MIXED_BATCH, + payload, bytes); free(payload); return ok; } @@ -1726,28 +2085,33 @@ int ds4_tp_send_command_ack(ds4_tp *tp, uint64_t session_id, int status) { int ds4_tp_wait_command_ack(ds4_tp *tp, uint64_t session_id, const char *operation, char *err, size_t errlen) { - uint32_t type = 0, bytes = 0; - ds4_tp_command_ack ack; - if (!tp_read_frame_header(tp->control_fd, &type, &bytes) || - type != DS4_TP_FRAME_COMMAND_ACK || bytes != sizeof(ack) || - !tp_read_full(tp->control_fd, &ack, sizeof(ack))) { - ds4_tp_mark_failed(tp); - tp_set_err(err, errlen, "tp: worker failed during %s", - operation ? operation : "command"); - return 0; - } - if (ack.session_id != session_id || ack.status != 0) { - tp_set_err(err, errlen, - "tp: worker %s failed (session %llu, status %d)", - operation ? operation : "command", - (unsigned long long)ack.session_id, (int)ack.status); - return 0; + if (!tp || tp->rank != 0) return 0; + for (uint32_t rank = 1; rank < tp->world; rank++) { + uint32_t type = 0, bytes = 0; + ds4_tp_command_ack ack; + const int fd = tp->control_fds[rank]; + if (fd < 0 || !tp_read_frame_header(fd, &type, &bytes) || + type != DS4_TP_FRAME_COMMAND_ACK || bytes != sizeof(ack) || + !tp_read_full(fd, &ack, sizeof(ack))) { + ds4_tp_mark_failed(tp); + tp_set_err(err, errlen, "tp: worker rank %u failed during %s", + rank, operation ? operation : "command"); + return 0; + } + if (ack.session_id != session_id || ack.status != 0) { + tp_set_err(err, errlen, + "tp: worker rank %u %s failed " + "(session %llu, status %d)", + rank, operation ? operation : "command", + (unsigned long long)ack.session_id, (int)ack.status); + return 0; + } } return 1; } int ds4_tp_send_stop(ds4_tp *tp) { - return tp_send_frame(tp->control_fd, DS4_TP_FRAME_STOP, NULL, 0); + return tp_send_frame_workers(tp, DS4_TP_FRAME_STOP, NULL, 0); } void ds4_tp_command_free(ds4_tp_command *command) { @@ -1894,11 +2258,13 @@ int ds4_tp_recv_command(ds4_tp *tp, ds4_tp_command *command, } int ds4_tp_send_logits_half(ds4_tp *tp, const float *half, uint32_t count) { + if (!tp || tp->world != 2u) return 0; return tp_send_frame(tp->control_fd, DS4_TP_FRAME_LOGITS, half, count * sizeof(float)); } int ds4_tp_recv_logits_half(ds4_tp *tp, float *half, uint32_t count) { + if (!tp || tp->world != 2u) return 0; uint32_t type = 0, bytes = 0; if (!tp_read_frame_header(tp->control_fd, &type, &bytes) || type != DS4_TP_FRAME_LOGITS || bytes != count * sizeof(float)) { @@ -1916,8 +2282,8 @@ int ds4_tp_send_verify(ds4_tp *tp, uint64_t session_id, int ds4_tp_send_verify_commit(ds4_tp *tp, int32_t full_accept, int32_t replay_n) { struct { int32_t full; int32_t replay; } msg = { full_accept, replay_n }; - return tp_send_frame(tp->control_fd, DS4_TP_FRAME_VERIFY_COMMIT, - &msg, sizeof(msg)); + return tp_send_frame_workers(tp, DS4_TP_FRAME_VERIFY_COMMIT, + &msg, sizeof(msg)); } int ds4_tp_recv_verify_commit(ds4_tp *tp, int32_t *full_accept, int32_t *replay_n) { @@ -1937,23 +2303,53 @@ int ds4_tp_recv_verify_commit(ds4_tp *tp, int32_t *full_accept, int32_t *replay_ int ds4_tp_hash_check(ds4_tp *tp, uint64_t seq, uint64_t hash, char *err, size_t errlen) { struct { uint64_t seq; uint64_t hash; } mine = { seq, hash }, theirs; - if (!tp_send_frame(tp->control_fd, DS4_TP_FRAME_HASH, &mine, sizeof(mine))) { - tp_set_err(err, errlen, "tp: hash send failed"); - return 0; - } - uint32_t type = 0, bytes = 0; - if (!tp_read_frame_header(tp->control_fd, &type, &bytes) || - type != DS4_TP_FRAME_HASH || bytes != sizeof(theirs) || - !tp_read_full(tp->control_fd, &theirs, sizeof(theirs))) { - tp_set_err(err, errlen, "tp: hash recv failed"); - return 0; - } - if (theirs.seq != seq || theirs.hash != hash) { - tp_set_err(err, errlen, - "tp: LOCKSTEP DIVERGENCE at seq %llu: local %016llx peer %016llx", - (unsigned long long)seq, - (unsigned long long)hash, (unsigned long long)theirs.hash); - return -1; + if (!tp) return 0; + if (tp->rank == 0) { + if (!tp_send_frame_workers(tp, DS4_TP_FRAME_HASH, &mine, sizeof(mine))) { + tp_set_err(err, errlen, "tp: hash broadcast failed"); + return 0; + } + for (uint32_t rank = 1; rank < tp->world; rank++) { + uint32_t type = 0, bytes = 0; + const int fd = tp->control_fds[rank]; + if (fd < 0 || !tp_read_frame_header(fd, &type, &bytes) || + type != DS4_TP_FRAME_HASH || bytes != sizeof(theirs) || + !tp_read_full(fd, &theirs, sizeof(theirs))) { + tp_set_err(err, errlen, "tp: hash recv from rank %u failed", rank); + return 0; + } + if (theirs.seq != seq || theirs.hash != hash) { + tp_set_err(err, errlen, + "tp: LOCKSTEP DIVERGENCE at seq %llu: " + "rank 0 %016llx rank %u %016llx", + (unsigned long long)seq, + (unsigned long long)hash, rank, + (unsigned long long)theirs.hash); + return -1; + } + } + } else { + if (!tp_send_frame(tp->control_fd, DS4_TP_FRAME_HASH, + &mine, sizeof(mine))) { + tp_set_err(err, errlen, "tp: hash send failed"); + return 0; + } + uint32_t type = 0, bytes = 0; + if (!tp_read_frame_header(tp->control_fd, &type, &bytes) || + type != DS4_TP_FRAME_HASH || bytes != sizeof(theirs) || + !tp_read_full(tp->control_fd, &theirs, sizeof(theirs))) { + tp_set_err(err, errlen, "tp: hash recv failed"); + return 0; + } + if (theirs.seq != seq || theirs.hash != hash) { + tp_set_err(err, errlen, + "tp: LOCKSTEP DIVERGENCE at seq %llu: " + "rank %d %016llx rank 0 %016llx", + (unsigned long long)seq, tp->rank, + (unsigned long long)hash, + (unsigned long long)theirs.hash); + return -1; + } } return 1; } @@ -2023,8 +2419,8 @@ static int tp_worker_send_logits(ds4_tp *tp, ds4_session *session, ds4_tp_send_logits_half(tp, logits + vhalf, vhalf); } -int ds4_tp_worker_run(ds4_engine *engine, const ds4_tp_options *opt) { - char err[256] = ""; +static ds4_tp_identity tp_engine_identity(ds4_engine *engine, + uint32_t ctx_size) { ds4_tp_identity id = { .gguf_bytes = ds4_engine_model_bytes(engine), .model_id = (uint32_t)ds4_engine_model_id(engine), @@ -2032,12 +2428,50 @@ int ds4_tp_worker_run(ds4_engine *engine, const ds4_tp_options *opt) { .n_embd = (uint32_t)ds4_engine_embd_dim(engine), .n_vocab = (uint32_t)ds4_engine_vocab_size(engine), .quant_bits = (uint32_t)ds4_engine_routed_quant_bits(engine), - .ctx_size = 0, /* adopt the leader's */ + .ctx_size = ctx_size, }; ds4_engine_tp_gate_schedule(engine, &id.gate_slot_start, &id.gate_slot_step, &id.gates_per_token); + return id; +} + +int ds4_tp_leader_bind(ds4_engine *engine, const ds4_tp_options *opt, + uint32_t ctx_size, ds4_tp **out, + char *err, size_t errlen) { + if (!engine || !opt || !out || opt->role != DS4_TP_LEADER || + ctx_size == 0) { + tp_set_err(err, errlen, "tp: invalid leader bind request"); + return 0; + } + *out = NULL; + ds4_tp_identity id = tp_engine_identity(engine, ctx_size); + ds4_tp *tp = NULL; + if (!ds4_tp_create(&tp, opt, &id, err, errlen)) return 0; + if (!ds4_engine_tp_bind(engine, tp, err, errlen)) { + (void)ds4_tp_send_stop(tp); + ds4_engine_tp_unbind(engine); + ds4_tp_free(tp); + return 0; + } + *out = tp; + return 1; +} + +void ds4_tp_leader_shutdown(ds4_engine *engine, ds4_tp **tp_ptr) { + if (!tp_ptr || !*tp_ptr) return; + ds4_tp *tp = *tp_ptr; + *tp_ptr = NULL; + (void)ds4_tp_send_stop(tp); + if (engine) ds4_engine_tp_unbind(engine); + ds4_tp_free(tp); +} + +int ds4_tp_worker_run(ds4_engine *engine, const ds4_tp_options *opt, + uint32_t ctx_size) { + char err[256] = ""; + ds4_tp_identity id = tp_engine_identity(engine, ctx_size); ds4_tp *tp = NULL; if (!ds4_tp_create(&tp, opt, &id, err, sizeof(err))) { @@ -2046,6 +2480,7 @@ int ds4_tp_worker_run(ds4_engine *engine, const ds4_tp_options *opt) { } if (!ds4_engine_tp_bind(engine, tp, err, sizeof(err))) { ds4_log(stderr, DS4_LOG_ERROR, "tp worker: %s", err); + ds4_engine_tp_unbind(engine); ds4_tp_free(tp); return 1; } @@ -2055,10 +2490,11 @@ int ds4_tp_worker_run(ds4_engine *engine, const ds4_tp_options *opt) { malloc((size_t)vocab * sizeof(*logits)) : NULL; if (ds4_engine_tp_vocab_split(engine) && !logits) { ds4_log(stderr, DS4_LOG_ERROR, "tp worker: logits buffer allocation failed"); + ds4_engine_tp_unbind(engine); ds4_tp_free(tp); return 1; } - ds4_log(stderr, DS4_LOG_OK, "tp worker ready for mirrored sessions"); + ds4_log(stderr, DS4_LOG_OK, "tp worker ready for mirrored sessions\n"); int rc = 0; ds4_tokens prompt = {0}; @@ -2070,7 +2506,7 @@ int ds4_tp_worker_run(ds4_engine *engine, const ds4_tp_options *opt) { break; } if (command.type == DS4_TP_FRAME_STOP) { - ds4_log(stderr, DS4_LOG_DEFAULT, "tp worker: leader finished"); + ds4_log(stderr, DS4_LOG_DEFAULT, "tp worker: leader finished\n"); ds4_tp_command_free(&command); break; } @@ -2214,6 +2650,7 @@ int ds4_tp_worker_run(ds4_engine *engine, const ds4_tp_options *opt) { } free(sessions.v); free(logits); + ds4_engine_tp_unbind(engine); ds4_tp_free(tp); return rc; } diff --git a/ds4_tp.h b/ds4_tp.h index ddc2e34758..7df46d3e5c 100644 --- a/ds4_tp.h +++ b/ds4_tp.h @@ -8,16 +8,16 @@ #include "ds4.h" -/* Tensor-parallel transport and lockstep protocol. +/* Expert/tensor-parallel transport and lockstep protocol. * - * Two ranks run the same logical model, each with one contiguous half of the - * routed experts resident. Rank 0 (leader) is a normal frontend session that - * mirrors every ds4_session_sync()/ds4_session_eval() call to rank 1 (worker) - * over a TCP control socket, so both engines execute the identical graph - * sequence. - * Inside each decoded token, partial block outputs are exchanged through a - * registered memory slab: two-sided RDMA SEND/RECV when RDMA over - * Thunderbolt is available, or a full-duplex TCP exchange as fallback. + * Two or four ranks run the same logical model, each with one contiguous + * routed-expert range resident. Rank 0 (leader) is a normal frontend session + * that mirrors every session operation to all workers over TCP, so every + * engine executes the identical graph sequence. + * Metal's two-rank path exchanges partial block outputs through a registered + * memory slab using RDMA SEND/RECV or a full-duplex TCP fallback. CUDA's + * two/four-rank path bootstraps NCCL over the same control connection and + * combines f32 partials in place with collectives. * * Layering: ds4.c calls the session-mirroring and slab entry points; * ds4_metal.m only ever sees ds4_tp_gate_exchange() through a callback @@ -26,6 +26,8 @@ typedef struct ds4_tp ds4_tp; +#define DS4_TP_MAX_WORLD 4u + enum { DS4_TP_GATE_ATTN = 0, DS4_TP_GATE_FFN = 1, @@ -57,6 +59,11 @@ typedef struct { bool ds4_tp_enabled(const ds4_tp_options *opt); +/* Contiguous floor partition used by model mapping and ownership-aware + * kernels: [floor(total*rank/world), floor(total*(rank+1)/world)). */ +bool ds4_tp_partition(uint32_t total, uint32_t rank, uint32_t world, + uint32_t *base, uint32_t *count); + typedef enum { DS4_TP_CLI_ERROR = -1, DS4_TP_CLI_NOT_MATCHED = 0, @@ -87,8 +94,8 @@ int ds4_tp_validate_engine_options( char *err, size_t errlen); -/* Connection bring-up. The leader listens and accepts one worker; the - * worker dials with retry. Both then exchange and validate identities. +/* Connection bring-up. The leader listens and accepts world_size-1 workers; + * workers dial with retry. All ranks exchange and validate identities. * Blocking; call after the engine is loaded (identity needs the shape). */ int ds4_tp_create( ds4_tp **out, @@ -99,11 +106,24 @@ int ds4_tp_create( void ds4_tp_free(ds4_tp *tp); int ds4_tp_rank(const ds4_tp *tp); +uint32_t ds4_tp_world(const ds4_tp *tp); bool ds4_tp_is_rdma(const ds4_tp *tp); +bool ds4_tp_is_collective(const ds4_tp *tp); +bool ds4_tp_is_expert_only(const ds4_tp *tp); uint32_t ds4_tp_peer_ctx(const ds4_tp *tp); bool ds4_tp_failed(const ds4_tp *tp); void ds4_tp_mark_failed(ds4_tp *tp); +/* All ranks report whether their collective runtime loaded successfully. + * No rank enters the blocking communicator constructor unless all agree. */ +int ds4_tp_collective_preflight(ds4_tp *tp, int local_ready, + char *err, size_t errlen); + +/* Small startup broadcast used to distribute an opaque collective bootstrap + * payload. Rank 0 supplies initialized bytes; workers receive them. */ +int ds4_tp_broadcast_blob(ds4_tp *tp, void *data, uint32_t bytes, + char *err, size_t errlen); + /* Gate slab. The engine allocates one shared GPU-visible block and hands * its base VA here; ds4_tp registers it with the NIC (RDMA) and exchanges * remote keys. Layout, all offsets from base, S = n_layer * 2 slots: @@ -188,6 +208,8 @@ typedef enum { DS4_TP_FRAME_EVAL_BATCH = 15, DS4_TP_FRAME_MIXED_BATCH = 16, DS4_TP_FRAME_COMMAND_ACK = 17, + DS4_TP_FRAME_BOOTSTRAP = 18, + DS4_TP_FRAME_COLLECTIVE_PREFLIGHT = 19, } ds4_tp_frame_type; typedef struct { @@ -208,8 +230,8 @@ int ds4_tp_recv_command( size_t errlen); void ds4_tp_command_free(ds4_tp_command *command); -/* Debug lockstep check: both sides send their hidden-state hash for a token - * and compare. Returns 0 on transport failure, -1 on hash mismatch. */ +/* Debug lockstep check: every rank sends its hidden-state hash for a token + * and rank 0 compares them. Returns 0 on transport failure, -1 on mismatch. */ int ds4_tp_hash_check(ds4_tp *tp, uint64_t seq, uint64_t hash, char *err, size_t errlen); /* Vocab-split output head: the worker ships its logits half to the leader @@ -217,18 +239,28 @@ int ds4_tp_hash_check(ds4_tp *tp, uint64_t seq, uint64_t hash, char *err, size_t int ds4_tp_send_logits_half(ds4_tp *tp, const float *half, uint32_t count); int ds4_tp_recv_logits_half(ds4_tp *tp, float *half, uint32_t count); -/* Speculative verify mirroring. The leader announces a draft block right - * before both ranks run the expert-split batch verify; the worker then blocks - * on the commit frame, which carries the leader's decision: full_accept keeps - * the pushed rows, otherwise both sides roll back and replay replay_n tokens - * through the gated single-token decode in lockstep. */ +/* Speculative verify mirroring. The leader announces a draft block before all + * ranks run the expert-split batch verify; workers then block on the commit + * frame, which carries the leader's decision: full_accept keeps the pushed + * rows, otherwise every rank rolls back and replays replay_n tokens through + * the gated single-token decode in lockstep. */ int ds4_tp_send_verify(ds4_tp *tp, uint64_t session_id, const int *drafts, uint32_t n); int ds4_tp_send_verify_commit(ds4_tp *tp, int32_t full_accept, int32_t replay_n); int ds4_tp_recv_verify_commit(ds4_tp *tp, int32_t *full_accept, int32_t *replay_n); +/* Shared frontend lifecycle. The leader helper builds the engine identity, + * accepts every worker, and binds the collective before any session is + * created. Shutdown must run after all frontend sessions are destroyed and + * before the engine is closed. */ +int ds4_tp_leader_bind(ds4_engine *engine, const ds4_tp_options *opt, + uint32_t ctx_size, ds4_tp **out, + char *err, size_t errlen); +void ds4_tp_leader_shutdown(ds4_engine *engine, ds4_tp **tp); + /* Standalone worker mode entry. Loads nothing itself: the engine is already * open. */ -int ds4_tp_worker_run(ds4_engine *engine, const ds4_tp_options *opt); +int ds4_tp_worker_run(ds4_engine *engine, const ds4_tp_options *opt, + uint32_t ctx_size); #endif diff --git a/gguf-tools/quality-testing/README.md b/gguf-tools/quality-testing/README.md index 8c1338e2fa..0bd04b56ff 100644 --- a/gguf-tools/quality-testing/README.md +++ b/gguf-tools/quality-testing/README.md @@ -125,6 +125,34 @@ Use `data/flash/manifest.tsv` for Flash GGUFs, which model produced the manifest; the manifest path selects the continuation set. +The DS4 scorer also accepts the normal distributed and network-parallel +options. Start the worker ranks with `./ds4` using the same model, context, +mode, world size, transport, and relevant environment variables, then run the +scorer as rank 0. For example, after starting ranks 1 through 3 as described in +the main README: + +```sh +gguf-tools/quality-testing/score_official \ + /path/to/deepseek-v4-flash-0731.gguf \ + gguf-tools/quality-testing/data/flash/manifest.tsv \ + /tmp/flash-ep4.tsv \ + 4096 \ + --role coordinator --expert-parallel \ + --tensor-parallel-world 4 --listen 192.168.2.30 9911 \ + --transport nccl +``` + +Use `--tensor-parallel` in place of `--expert-parallel` to score network TP. +The scorer waits for the complete rank set before creating its session and +stops the workers after closing the scored session. Set experimental-path +environment variables identically on every rank. `--cuda-tensor-parallel` is +the separate, single-host multi-GPU mode and should not be substituted for +network `--tensor-parallel`. + +Pipeline-distributed scoring is also available on rank 0 through +`--role coordinator --layers A:B --listen HOST PORT`; start its layer workers +with `./ds4` in the usual way. + Add `--quality` to disable DS4's speed-oriented numerical shortcuts. For an independent llama.cpp comparison of a DeepSeek V4 GGUF, use the same manifest and the token-identical DS4 prompt renderer: diff --git a/gguf-tools/quality-testing/score_official.c b/gguf-tools/quality-testing/score_official.c index 9fefb3ba21..b0fa066b31 100644 --- a/gguf-tools/quality-testing/score_official.c +++ b/gguf-tools/quality-testing/score_official.c @@ -1,6 +1,12 @@ +#ifndef _POSIX_C_SOURCE +#define _POSIX_C_SOURCE 200809L +#endif + #include "ds4.h" +#include "ds4_distributed.h" #include "ds4_gpu_args.h" #include "ds4_ssd.h" +#include "ds4_tp.h" #include #include @@ -11,6 +17,7 @@ #include #include #include +#include static void die(const char *msg) { fprintf(stderr, "%s\n", msg); @@ -25,8 +32,11 @@ static void usage(const char *prog) { "[--cuda-tensor-parallel] " "[--ssd-streaming] [--ssd-streaming-cold] " "[--ssd-streaming-cache-experts N|NGB] " - "[--ssd-streaming-preload-experts N]\n", + "[--ssd-streaming-preload-experts N] " + "[distributed/network-parallel options]\n", prog); + ds4_dist_usage(stderr); + ds4_tp_usage(stderr); } static const char *need_arg(int *i, int argc, char **argv, const char *opt) { @@ -510,6 +520,34 @@ static double safe_ratio(long num, long den) { return den ? (double)num / (double)den : 0.0; } +static int wait_distributed_route(ds4_session *session) { + char err[256] = {0}; + char last[256] = {0}; + unsigned ticks = 0; + const struct timespec delay = {0, 250000000L}; + + for (;;) { + int ready = ds4_session_distributed_route_ready(session, err, sizeof(err)); + if (ready > 0) { + if (ticks) fprintf(stderr, "score_official: distributed route ready\n"); + return 0; + } + if (ready < 0) { + fprintf(stderr, + "score_official: distributed route readiness failed: %s\n", + err[0] ? err : "unknown error"); + return 1; + } + const char *why = err[0] ? err : "waiting for workers"; + if (!ticks || strcmp(last, why) != 0 || ticks % 20u == 0u) { + fprintf(stderr, "score_official: waiting for distributed route: %s\n", why); + snprintf(last, sizeof(last), "%s", why); + } + ticks++; + nanosleep(&delay, NULL); + } +} + int main(int argc, char **argv) { if (argc < 4) { usage(argv[0]); @@ -530,9 +568,35 @@ int main(int argc, char **argv) { uint32_t ssd_streaming_cache_experts = 0; uint64_t ssd_streaming_cache_bytes = 0; uint32_t ssd_streaming_preload_experts = 0; + ds4_dist_options dist = {0}; + ds4_tp_options tp = {0}; for (int i = 4; i < argc; i++) { const char *arg = argv[i]; + char dist_parse_err[256] = {0}; + ds4_dist_cli_parse_result dist_parse = + ds4_dist_parse_cli_arg(arg, &i, argc, argv, &dist, + dist_parse_err, sizeof(dist_parse_err)); + if (dist_parse == DS4_DIST_CLI_ERROR) { + fprintf(stderr, "score_official: %s\n", + dist_parse_err[0] ? dist_parse_err : + "invalid distributed option"); + return 2; + } + if (dist_parse == DS4_DIST_CLI_MATCHED) continue; + + char tp_parse_err[256] = {0}; + ds4_tp_cli_parse_result tp_parse = + ds4_tp_parse_cli_arg(arg, &i, argc, argv, &tp, + tp_parse_err, sizeof(tp_parse_err)); + if (tp_parse == DS4_TP_CLI_ERROR) { + fprintf(stderr, "score_official: %s\n", + tp_parse_err[0] ? tp_parse_err : + "invalid network parallel option"); + return 2; + } + if (tp_parse == DS4_TP_CLI_MATCHED) continue; + if (!strcmp(arg, "--quality")) { quality = true; } else if (!strcmp(arg, "--gpu-vram")) { @@ -568,6 +632,26 @@ int main(int argc, char **argv) { } if (ctx_size < 1024) ctx_size = 1024; + char tp_err[256] = {0}; + if (!ds4_tp_adopt_distributed_options(&tp, &dist, + tp_err, sizeof(tp_err))) { + fprintf(stderr, "score_official: %s\n", + tp_err[0] ? tp_err : "invalid network parallel options"); + return 2; + } + char dist_err[256] = {0}; + if (ds4_dist_prepare_engine_options(&dist, NULL, + dist_err, sizeof(dist_err)) != 0) { + fprintf(stderr, "score_official: %s\n", dist_err); + return 2; + } + if (dist.role == DS4_DISTRIBUTED_WORKER) { + fprintf(stderr, + "score_official: --role worker is a serving mode; " + "start pipeline workers with ./ds4\n"); + return 2; + } + ds4_engine_options opt = { .model_path = model_path, #ifdef __APPLE__ @@ -586,7 +670,19 @@ int main(int argc, char **argv) { .cuda_tensor_parallel = cuda_tensor_parallel, .ssd_streaming = ssd_streaming, .ssd_streaming_cold = ssd_streaming_cold, + .distributed = dist, + .tp = tp, }; + if (ds4_dist_prepare_engine_options(&dist, &opt, + dist_err, sizeof(dist_err)) != 0) { + fprintf(stderr, "score_official: %s\n", dist_err); + return 2; + } + if (!ds4_tp_validate_engine_options(&opt, tp_err, sizeof(tp_err))) { + fprintf(stderr, "score_official: %s\n", + tp_err[0] ? tp_err : "invalid network parallel options"); + return 2; + } ds4_engine *engine = NULL; if (gpu_vram_arg || gpu_devices_arg) { @@ -616,21 +712,68 @@ int main(int argc, char **argv) { if (ds4_engine_open(&engine, &opt) != 0) die("failed to open model"); } - ds4_session *session = NULL; - if (ds4_session_create(&session, engine, ctx_size) != 0) die("failed to create session"); - - const int n_vocab = ds4_engine_vocab_size(engine); - float *logits = malloc((size_t)n_vocab * sizeof(logits[0])); - if (!logits) die("out of memory"); + if (tp.role == DS4_TP_WORKER) { + int rc = ds4_tp_worker_run(engine, &tp, (uint32_t)ctx_size); + ds4_engine_close(engine); + return rc; + } FILE *mf = fopen(manifest_path, "rb"); if (!mf) { fprintf(stderr, "open %s: %s\n", manifest_path, strerror(errno)); + ds4_engine_close(engine); return 1; } FILE *out = fopen(out_path, "wb"); if (!out) { fprintf(stderr, "open %s: %s\n", out_path, strerror(errno)); + fclose(mf); + ds4_engine_close(engine); + return 1; + } + + ds4_tp *tp_leader = NULL; + if (tp.role == DS4_TP_LEADER) { + if (!ds4_tp_leader_bind(engine, &tp, (uint32_t)ctx_size, + &tp_leader, tp_err, sizeof(tp_err))) { + fprintf(stderr, "score_official: %s\n", + tp_err[0] ? tp_err : + "network parallel initialization failed"); + fclose(out); + fclose(mf); + ds4_engine_close(engine); + return 1; + } + } + + ds4_session *session = NULL; + if (ds4_session_create(&session, engine, ctx_size) != 0) { + fprintf(stderr, "score_official: failed to create session\n"); + ds4_tp_leader_shutdown(engine, &tp_leader); + fclose(out); + fclose(mf); + ds4_engine_close(engine); + return 1; + } + if (dist.role == DS4_DISTRIBUTED_COORDINATOR && + wait_distributed_route(session) != 0) { + ds4_session_free(session); + ds4_tp_leader_shutdown(engine, &tp_leader); + fclose(out); + fclose(mf); + ds4_engine_close(engine); + return 1; + } + + const int n_vocab = ds4_engine_vocab_size(engine); + float *logits = malloc((size_t)n_vocab * sizeof(logits[0])); + if (!logits) { + fprintf(stderr, "score_official: out of memory allocating logits\n"); + ds4_session_free(session); + ds4_tp_leader_shutdown(engine, &tp_leader); + fclose(out); + fclose(mf); + ds4_engine_close(engine); return 1; } fprintf(out, @@ -651,6 +794,7 @@ int main(int argc, char **argv) { long total_api_ref_tokens = 0; api_metrics total_api = {0}; char err[256]; + int result = 0; while (fgets(line, sizeof(line), mf)) { strip_newline(line); @@ -660,7 +804,11 @@ int main(int argc, char **argv) { char *prompt_path = strtok(NULL, "\t"); char *cont_path = strtok(NULL, "\t"); char *resp_path = strtok(NULL, "\t"); - if (!id || !prompt_path || !cont_path) die("bad manifest row"); + if (!id || !prompt_path || !cont_path) { + fprintf(stderr, "score_official: bad manifest row\n"); + result = 1; + goto cleanup; + } char *prompt_text = read_file(prompt_path); char *cont_text = read_file(cont_path); @@ -676,7 +824,8 @@ int main(int argc, char **argv) { if (prompt.len + target.len + 1 >= ctx_size) { fprintf(stderr, "%s exceeds ctx=%d\n", id, ctx_size); - return 1; + result = 1; + goto case_cleanup; } if (resp_path && resp_path[0]) { have_api = api_ref_load(resp_path, &ref); @@ -695,7 +844,8 @@ int main(int argc, char **argv) { if (ds4_session_sync(session, &prompt, err, sizeof(err)) != 0) { fprintf(stderr, "%s sync failed: %s\n", id, err); - return 1; + result = 1; + goto case_cleanup; } double nll = 0.0; @@ -707,7 +857,8 @@ int main(int argc, char **argv) { int greedy = -1; if (!local_logits(session, logits, n_vocab, &logsum, &greedy)) { fprintf(stderr, "%s logits failed at target token %d\n", id, i); - return 1; + result = 1; + goto case_cleanup; } if (i == 0) first_match = (greedy == target.v[i]); @@ -717,7 +868,8 @@ int main(int argc, char **argv) { const double target_lp = local_logprob(logits, n_vocab, target.v[i], logsum); if (!isfinite(target_lp)) { fprintf(stderr, "%s logprob failed at target token %d\n", id, i); - return 1; + result = 1; + goto case_cleanup; } nll += -target_lp; @@ -782,7 +934,8 @@ int main(int argc, char **argv) { if (ds4_session_eval(session, target.v[i], err, sizeof(err)) != 0) { fprintf(stderr, "%s eval failed at target token %d: %s\n", id, i, err); - return 1; + result = 1; + goto case_cleanup; } } @@ -835,13 +988,17 @@ int main(int argc, char **argv) { safe_ratio(cm.topn_hit, cm.topn_ref), safe_ratio(cm.pair_agree, cm.pair_total)); +case_cleanup: api_ref_free(&ref); ds4_tokens_free(&prompt); ds4_tokens_free(&target); free(prompt_text); free(cont_text); + if (result != 0) break; } + if (result != 0) goto cleanup; + fprintf(stderr, "summary cases=%d tokens=%ld avg_nll=%.9f first_match=%ld avg_lcp=%.3f\n", case_n, @@ -875,10 +1032,12 @@ int main(int argc, char **argv) { total_api.pair_total, safe_ratio(total_api.pair_agree, total_api.pair_total)); +cleanup: fclose(out); fclose(mf); free(logits); ds4_session_free(session); + ds4_tp_leader_shutdown(engine, &tp_leader); ds4_engine_close(engine); - return 0; + return result; } diff --git a/tests/test_gpu_args_cli.sh b/tests/test_gpu_args_cli.sh index 05ef2213db..a4cf96484f 100755 --- a/tests/test_gpu_args_cli.sh +++ b/tests/test_gpu_args_cli.sh @@ -30,10 +30,26 @@ assert_not_grep() { fi } -# Binaries to check +# Binaries exposing the shared CUDA placement parser. BINS=(./ds4 ./ds4-server ./ds4-bench ./ds4-agent) NAMES=(ds4 ds4-server ds4-bench ds4-agent) +# Every executable that can coordinate network expert/tensor parallelism. +TP_BINS=(./ds4 ./ds4-server ./ds4-bench ./ds4-agent ./ds4-eval) +TP_NAMES=(ds4 ds4-server ds4-bench ds4-agent ds4-eval) + +# Network-parallel validation is backend-specific. CUDA currently accepts two- +# or four-rank NCCL topologies, while Metal accepts a two-rank TCP topology. +if [ "$(uname -s)" = "Darwin" ]; then + NETWORK_BACKEND=--metal + NETWORK_WORLD=2 + NETWORK_TRANSPORT=tcp +else + NETWORK_BACKEND=--cuda + NETWORK_WORLD=4 + NETWORK_TRANSPORT=nccl +fi + # 1: each binary's --help mentions both flags. for i in "${!BINS[@]}"; do name=${NAMES[$i]}; bin=${BINS[$i]} @@ -45,12 +61,26 @@ for i in "${!BINS[@]}"; do assert_grep "$name --help mentions --gpu-vram" "gpu-vram" "$LOG" assert_grep "$name --help mentions --gpu-devices" "gpu-devices" "$LOG" assert_grep "$name --help mentions --cuda-tensor-parallel" "cuda-tensor-parallel" "$LOG" - if [ "$name" = "ds4" ]; then - "$bin" --help distributed > "$LOG" 2>&1 || true - assert_grep "$name --help distributed mentions --tensor-parallel-token-prefill" \ - "tensor-parallel-token-prefill" "$LOG" - assert_not_grep "$name --help distributed omits old --tp spellings" "--tp-" "$LOG" +done + +for i in "${!TP_BINS[@]}"; do + name=${TP_NAMES[$i]}; bin=${TP_BINS[$i]} + if [ ! -x "$bin" ]; then + fail "$name not built — skipping network help check" + continue fi + "$bin" --help distributed > "$LOG" 2>&1 || true + assert_grep "$name --help distributed mentions --expert-parallel" \ + "--expert-parallel" "$LOG" + assert_grep "$name --help distributed mentions --tensor-parallel" \ + "--tensor-parallel" "$LOG" + assert_grep "$name --help distributed mentions --tensor-parallel-world" \ + "--tensor-parallel-world" "$LOG" + assert_grep "$name --help distributed mentions --tensor-parallel-rank" \ + "--tensor-parallel-rank" "$LOG" + assert_grep "$name --help distributed mentions --tensor-parallel-token-prefill" \ + "tensor-parallel-token-prefill" "$LOG" + assert_not_grep "$name --help distributed omits old --tp spellings" "--tp-" "$LOG" done # 2: parser error on syntactically invalid value. For ds4-bench, we @@ -141,7 +171,7 @@ if [ -x ./ds4 ]; then ./ds4 --metal --tensor-parallel --role coordinator --listen 127.0.0.1 9911 \ --layers 0:1 -m /dev/null > "$LOG" 2>&1 rc=$? - if [ $rc -ne 0 ] && grep -q "always uses one 50/50 worker" "$LOG"; then + if [ $rc -ne 0 ] && grep -q "owns model placement; omit --layers" "$LOG"; then ok "tensor parallel rejects explicit layer slices" else fail "tensor parallel accepted --layers or returned the wrong error" @@ -161,13 +191,101 @@ if [ -x ./ds4 ]; then --inspect -m /dev/null > "$LOG" 2>&1 rc=$? if [ $rc -ne 0 ] && - grep -qE "model file is too small|another ds4 process is already running" "$LOG" && + grep -qE "model file is too small|another ds4 process is already running|requires the CUDA backend on Linux" "$LOG" && ! grep -q "requires --layers" "$LOG"; then ok "tensor-parallel common options reach model loading" else fail "tensor-parallel common options did not reach model loading" fi + ./ds4 --expert-parallel --tensor-parallel -m /dev/null > "$LOG" 2>&1 + rc=$? + if [ $rc -ne 0 ] && + grep -q -- "--expert-parallel and --tensor-parallel are exclusive" "$LOG"; then + ok "expert and tensor modes are mutually exclusive" + else + fail "expert/tensor mode conflict was not rejected" + fi + + if [ "$(uname -s)" = "Linux" ]; then + ./ds4 --cuda --expert-parallel --glm-mtp \ + --role coordinator --listen 127.0.0.1 9911 \ + -m /dev/null > "$LOG" 2>&1 + rc=$? + if [ $rc -ne 0 ] && + grep -q "CUDA network parallelism does not yet support --glm-mtp" "$LOG"; then + ok "CUDA network parallel rejects GLM MTP" + else + fail "CUDA network parallel accepted GLM MTP" + fi + + ./ds4 --cuda --expert-parallel --mtp /dev/null --dspark \ + --role coordinator --listen 127.0.0.1 9911 \ + -m /dev/null > "$LOG" 2>&1 + rc=$? + if [ $rc -ne 0 ] && + grep -q "does not yet support MTP/DSpark speculative drafting" "$LOG"; then + ok "CUDA network parallel rejects external MTP/DSpark" + else + fail "CUDA network parallel accepted external MTP/DSpark" + fi + + ./ds4 --cuda --expert-parallel \ + --dir-steering-file /dev/null \ + --role coordinator --listen 127.0.0.1 9911 \ + -m /dev/null > "$LOG" 2>&1 + rc=$? + if [ $rc -ne 0 ] && + grep -q "does not yet support directional steering" "$LOG"; then + ok "CUDA network parallel rejects directional steering" + else + fail "CUDA network parallel accepted directional steering" + fi + + ./ds4-bench --cuda --expert-parallel --transport tcp \ + --role coordinator --listen 127.0.0.1 9911 \ + --prompt-file /dev/null -m /dev/null > "$LOG" 2>&1 + rc=$? + if [ $rc -ne 0 ] && + grep -q "CUDA network parallelism requires --transport nccl" "$LOG"; then + ok "ds4-bench validates CUDA network transport before model loading" + else + fail "ds4-bench skipped CUDA network transport validation" + fi + fi + + ./ds4 --expert-parallel --role worker \ + --tensor-parallel-world 3 --coordinator 127.0.0.1 9911 \ + -m /dev/null > "$LOG" 2>&1 + rc=$? + if [ $rc -ne 0 ] && grep -q -- "--tensor-parallel-world must be 2 or 4" "$LOG"; then + ok "network parallel rejects a three-rank world" + else + fail "network parallel accepted a three-rank world" + fi + + ./ds4 --expert-parallel --role worker \ + --tensor-parallel-world 4 --coordinator 127.0.0.1 9911 \ + -m /dev/null > "$LOG" 2>&1 + rc=$? + if [ $rc -ne 0 ] && + grep -q "four-rank workers require --tensor-parallel-rank 1, 2, or 3" "$LOG"; then + ok "four-rank workers require an explicit rank" + else + fail "four-rank worker accepted an implicit rank" + fi + + ./ds4 --expert-parallel --role worker \ + --tensor-parallel-world 2 --tensor-parallel-rank 2 \ + --coordinator 127.0.0.1 9911 -m /dev/null > "$LOG" 2>&1 + rc=$? + if [ $rc -ne 0 ] && + grep -q "worker --tensor-parallel-rank 2 is outside 1..1" "$LOG"; then + ok "two-rank worker rejects a rank outside its world" + else + fail "two-rank worker accepted rank 2" + fi + for old_arg in \ "--tp-coordinator 9911" \ "--tp-lead 9911" \ @@ -176,7 +294,9 @@ if [ -x ./ds4 ]; then "--tp-worker 127.0.0.1 9911" \ "--tp-transport tcp" \ "--tp-debug-hash 2" \ - "--tp-token-prefill" + "--tp-token-prefill" \ + "--tp-world 2" \ + "--tp-rank 1" do # Word splitting is intentional: each item contains one old option # and its former arguments. @@ -190,6 +310,158 @@ if [ -x ./ds4 ]; then done fi +# 6b: every long-lived frontend accepts network leader/worker options and +# reaches engine loading. These use /dev/null, so no socket or GPU work starts. +FRONTEND_BINS=(./ds4-server ./ds4-agent ./ds4-eval) +FRONTEND_NAMES=(ds4-server ds4-agent ds4-eval) +for i in "${!FRONTEND_BINS[@]}"; do + name=${FRONTEND_NAMES[$i]}; bin=${FRONTEND_BINS[$i]} + [ -x "$bin" ] || continue + + "$bin" "$NETWORK_BACKEND" --ctx 128 --tensor-parallel --role worker \ + -m /dev/null > "$LOG" 2>&1 + rc=$? + if [ $rc -ne 0 ] && + grep -q "requires --coordinator HOST PORT" "$LOG"; then + ok "$name network worker requires coordinator address" + else + fail "$name returned the wrong missing worker-address error" + fi + + "$bin" "$NETWORK_BACKEND" --ctx 128 --tensor-parallel \ + --tensor-parallel-world "$NETWORK_WORLD" --tensor-parallel-rank 1 \ + --role worker --coordinator 127.0.0.1 9911 \ + --transport "$NETWORK_TRANSPORT" \ + -m /dev/null > "$LOG" 2>&1 + rc=$? + if [ $rc -ne 0 ] && + grep -q "model file is too small to be GGUF" "$LOG" && + ! grep -qE "unknown option|start workers with ./ds4" "$LOG"; then + ok "$name network worker reaches engine loading" + else + fail "$name network worker did not reach engine loading" + head -10 "$LOG" | sed 's/^/ /' + fi + + "$bin" "$NETWORK_BACKEND" --ctx 128 \ + --expert-parallel --tensor-parallel \ + -m /dev/null > "$LOG" 2>&1 + rc=$? + if [ $rc -ne 0 ] && + grep -q -- "--expert-parallel and --tensor-parallel are exclusive" "$LOG"; then + ok "$name rejects conflicting network modes" + else + fail "$name accepted conflicting network modes" + fi + + if [ "$(uname -s)" = "Linux" ]; then + "$bin" --cuda --ctx 128 --expert-parallel --transport tcp \ + --role coordinator --listen 127.0.0.1 9911 \ + -m /dev/null > "$LOG" 2>&1 + rc=$? + if [ $rc -ne 0 ] && + grep -q "CUDA network parallelism requires --transport nccl" "$LOG"; then + ok "$name validates CUDA network transport before model loading" + else + fail "$name skipped CUDA network transport validation" + fi + fi +done + +if [ -x ./ds4-eval ]; then + ./ds4-eval "$NETWORK_BACKEND" --expert-parallel --role coordinator \ + --listen 127.0.0.1 9911 -m /dev/null > "$LOG" 2>&1 + rc=$? + if [ $rc -ne 0 ] && + grep -q "network parallelism requires an explicit --ctx" "$LOG"; then + ok "ds4-eval network mode requires explicit context capacity" + else + fail "ds4-eval accepted automatic context sizing in network mode" + fi +fi + +# 6c: the official-continuation scorer uses the same network parser and +# validates the rank topology before loading a model. It keeps its historical +# positional MODEL/manifest/output interface, so /dev/null fills those slots. +QUALITY_SCORER=./gguf-tools/quality-testing/score_official +if [ ! -x "$QUALITY_SCORER" ]; then + fail "score_official not built — skipping network parser checks" +else + "$QUALITY_SCORER" > "$LOG" 2>&1 + rc=$? + if [ $rc -ne 0 ] && + grep -q -- "--expert-parallel" "$LOG" && + grep -q -- "--tensor-parallel" "$LOG"; then + ok "score_official usage advertises network EP/TP" + else + fail "score_official usage omits network EP/TP" + fi + + "$QUALITY_SCORER" /dev/null /dev/null /dev/null 1024 \ + --tensor-parallel --role worker > "$LOG" 2>&1 + rc=$? + if [ $rc -ne 0 ] && + grep -q "requires --coordinator HOST PORT" "$LOG"; then + ok "score_official network worker requires coordinator address" + else + fail "score_official returned the wrong missing worker-address error" + fi + + "$QUALITY_SCORER" /dev/null /dev/null /dev/null 1024 \ + --expert-parallel --tensor-parallel > "$LOG" 2>&1 + rc=$? + if [ $rc -ne 0 ] && + grep -q -- "--expert-parallel and --tensor-parallel are exclusive" "$LOG"; then + ok "score_official rejects conflicting network modes" + else + fail "score_official accepted conflicting network modes" + fi + + if [ "$(uname -s)" = "Linux" ]; then + "$QUALITY_SCORER" /dev/null /dev/null /dev/null 1024 \ + --expert-parallel --transport tcp \ + --role coordinator --listen 127.0.0.1 9911 > "$LOG" 2>&1 + rc=$? + if [ $rc -ne 0 ] && + grep -q "CUDA network parallelism requires --transport nccl" "$LOG"; then + ok "score_official validates CUDA network transport before model loading" + else + fail "score_official skipped CUDA network transport validation" + fi + fi + + "$QUALITY_SCORER" /dev/null /dev/null /dev/null 1024 \ + --tensor-parallel --tensor-parallel-world "$NETWORK_WORLD" \ + --tensor-parallel-rank 1 --role worker \ + --coordinator 127.0.0.1 9911 \ + --transport "$NETWORK_TRANSPORT" > "$LOG" 2>&1 + rc=$? + if [ $rc -ne 0 ] && + grep -q "model file is too small to be GGUF" "$LOG" && + ! grep -qE "unknown option|start pipeline workers with ./ds4" "$LOG"; then + ok "score_official network worker reaches engine loading" + else + fail "score_official network worker did not reach engine loading" + head -10 "$LOG" | sed 's/^/ /' + fi +fi + +if [ -x ./ds4-server ]; then + ./ds4-server "$NETWORK_BACKEND" --ctx 128 --batched-session 2 \ + --tensor-parallel --tensor-parallel-world "$NETWORK_WORLD" \ + --tensor-parallel-rank 1 --role worker \ + --coordinator 127.0.0.1 9911 --transport "$NETWORK_TRANSPORT" \ + -m /dev/null > "$LOG" 2>&1 + rc=$? + if [ $rc -ne 0 ] && + grep -q "model file is too small to be GGUF" "$LOG" && + ! grep -q "batched-session" "$LOG"; then + ok "ds4-server accepts native batching in network worker mode" + else + fail "ds4-server rejected native batching in network mode" + fi +fi + # 7: --gpu-vram 40,12 layout line. if [ -x ./ds4 ]; then ./ds4 --gpu-vram 40,12 -m /dev/null > "$LOG" 2>&1 diff --git a/tests/test_gpu_xdev.c b/tests/test_gpu_xdev.c index 0e6edbaba6..db1fd45f2b 100644 --- a/tests/test_gpu_xdev.c +++ b/tests/test_gpu_xdev.c @@ -16,6 +16,7 @@ #include "ds4_gpu_mgpu.h" #include +#include #include #include #include @@ -982,79 +983,266 @@ static int run_q8_kslice(void) { const uint64_t in_dim = 128; const uint64_t out_dim = 96; - const uint64_t split = 64; + const uint64_t world = 4; + const uint64_t split = in_dim / world; + const uint64_t n_tok = 16; const uint64_t blocks = (in_dim + 31u) / 32u; const uint64_t model_size = out_dim * blocks * 34u; unsigned char *model = (unsigned char *)malloc((size_t)model_size); - float *host_x = (float *)malloc((size_t)in_dim * sizeof(float)); - float *host_full = (float *)malloc((size_t)out_dim * sizeof(float)); - float *host_sum = (float *)malloc((size_t)out_dim * sizeof(float)); - CHECK(model && host_x && host_full && host_sum, "q8_kslice host alloc"); + float *host_x = (float *)malloc((size_t)(n_tok * in_dim) * sizeof(float)); + float *host_local = (float *)malloc((size_t)(n_tok * split) * sizeof(float)); + float *host_full = (float *)malloc((size_t)(n_tok * out_dim) * sizeof(float)); + float *host_part = (float *)malloc((size_t)(n_tok * out_dim) * sizeof(float)); + float *host_sum = (float *)calloc((size_t)(n_tok * out_dim), sizeof(float)); + CHECK(model && host_x && host_local && host_full && host_part && host_sum, + "q8_kslice host alloc"); pack_q8_identity_scale(model, in_dim, out_dim); - for (uint64_t i = 0; i < in_dim; i++) { - host_x[i] = (float)((int)(i % 37u) - 18) * 0.03125f; + for (uint64_t t = 0; t < n_tok; t++) { + for (uint64_t i = 0; i < in_dim; i++) { + host_x[t * in_dim + i] = + (float)((int)((i * 5u + t * 11u) % 53u) - 26) * 0.0234375f; + } } CHECK(ds4_gpu_set_model_map(model, model_size), "q8_kslice set model map"); ds4_gpu_tensor x; memset(&x, 0, sizeof(x)); ds4_gpu_tensor full; memset(&full, 0, sizeof(full)); - ds4_gpu_tensor p0; memset(&p0, 0, sizeof(p0)); - ds4_gpu_tensor p1; memset(&p1, 0, sizeof(p1)); - ds4_gpu_tensor sum; memset(&sum, 0, sizeof(sum)); - CHECK(ds4_gpu_tensor_alloc_on(&x, 0, in_dim * sizeof(float)) == 0, + CHECK(ds4_gpu_tensor_alloc_on(&x, 0, n_tok * in_dim * sizeof(float)) == 0, "q8_kslice alloc x"); - CHECK(ds4_gpu_tensor_alloc_on(&full, 0, out_dim * sizeof(float)) == 0, + CHECK(ds4_gpu_tensor_alloc_on(&full, 0, n_tok * out_dim * sizeof(float)) == 0, "q8_kslice alloc full"); - CHECK(ds4_gpu_tensor_alloc_on(&p0, 0, out_dim * sizeof(float)) == 0, - "q8_kslice alloc p0"); - CHECK(ds4_gpu_tensor_alloc_on(&p1, 0, out_dim * sizeof(float)) == 0, - "q8_kslice alloc p1"); - CHECK(ds4_gpu_tensor_alloc_on(&sum, 0, out_dim * sizeof(float)) == 0, - "q8_kslice alloc sum"); - CHECK(ds4_gpu_tensor_write(&x, 0, host_x, in_dim * sizeof(float)), + CHECK(ds4_gpu_tensor_write(&x, 0, host_x, + n_tok * in_dim * sizeof(float)), "q8_kslice write x"); - ds4_gpu_tensor *x0 = ds4_gpu_tensor_view(&x, 0, split * sizeof(float)); - ds4_gpu_tensor *x1 = ds4_gpu_tensor_view(&x, split * sizeof(float), - (in_dim - split) * sizeof(float)); - CHECK(x0 && x1, "q8_kslice views"); int ok = ds4_gpu_matmul_q8_0_tensor(&full, model, model_size, 0, - in_dim, out_dim, &x, 1) && - ds4_gpu_matmul_q8_0_kslice_rows_tensor(&p0, model, model_size, 0, - in_dim, out_dim, 0, split, x0, 1) && - ds4_gpu_matmul_q8_0_kslice_rows_tensor(&p1, model, model_size, 0, - in_dim, out_dim, split, - in_dim - split, x1, 1) && - ds4_gpu_add_tensor(&sum, &p0, &p1, (uint32_t)out_dim) && - ds4_gpu_tensor_read(&full, 0, host_full, out_dim * sizeof(float)) && - ds4_gpu_tensor_read(&sum, 0, host_sum, out_dim * sizeof(float)); + in_dim, out_dim, &x, n_tok); + for (uint64_t rank = 0; ok && rank < world; rank++) { + for (uint64_t t = 0; t < n_tok; t++) { + memcpy(host_local + t * split, + host_x + t * in_dim + rank * split, + (size_t)split * sizeof(float)); + } + ds4_gpu_tensor local; memset(&local, 0, sizeof(local)); + ds4_gpu_tensor part; memset(&part, 0, sizeof(part)); + ok = ds4_gpu_tensor_alloc_on(&local, 0, + n_tok * split * sizeof(float)) == 0 && + ds4_gpu_tensor_alloc_on(&part, 0, + n_tok * out_dim * sizeof(float)) == 0 && + ds4_gpu_tensor_write(&local, 0, host_local, + n_tok * split * sizeof(float)) && + ds4_gpu_matmul_q8_0_kslice_rows_tensor( + &part, model, model_size, 0, + in_dim, out_dim, rank * split, split, &local, n_tok) && + ds4_gpu_tensor_read(&part, 0, host_part, + n_tok * out_dim * sizeof(float)); + if (ok) { + for (uint64_t i = 0; i < n_tok * out_dim; i++) { + host_sum[i] += host_part[i]; + } + } + ds4_gpu_tensor_free_in_place(&part); + ds4_gpu_tensor_free_in_place(&local); + } + ok = ok && ds4_gpu_tensor_read(&full, 0, host_full, + n_tok * out_dim * sizeof(float)); CHECK(ok, "q8_kslice compute"); - for (uint64_t i = 0; i < out_dim; i++) { + for (uint64_t i = 0; i < n_tok * out_dim; i++) { float diff = fabsf(host_full[i] - host_sum[i]); - if (diff > 1.0e-4f) { + if (diff > 2.0e-4f) { fprintf(stderr, - "FAIL: q8_kslice mismatch row=%llu full=%f sum=%f diff=%g\n", - (unsigned long long)i, host_full[i], host_sum[i], diff); + "FAIL: q8_kslice mismatch elem=%llu full=%f sum=%f diff=%g\n", + (unsigned long long)i, + host_full[i], host_sum[i], diff); return 1; } } - ds4_gpu_tensor_free(x0); - ds4_gpu_tensor_free(x1); ds4_gpu_tensor_free_in_place(&x); ds4_gpu_tensor_free_in_place(&full); - ds4_gpu_tensor_free_in_place(&p0); - ds4_gpu_tensor_free_in_place(&p1); - ds4_gpu_tensor_free_in_place(&sum); free(model); free(host_x); + free(host_local); free(host_full); + free(host_part); free(host_sum); ds4_gpu_cleanup(); fprintf(stderr, " q8_kslice OK\n"); return 0; } +static int run_q8_ffn_tp4(void) { + int dev_count = 0; + (void)cudaGetDeviceCount(&dev_count); + if (dev_count < 1) { + fprintf(stderr, " skipping q8_ffn_tp4 (need CUDA device)\n"); + return 0; + } + + ds4_gpu_config cfg; memset(&cfg, 0, sizeof(cfg)); + cfg.n_gpus = 1; + cfg.device_indices[0] = 0; + CHECK(ds4_gpu_init_multi(&cfg), "q8_ffn_tp4 init_multi"); + + const uint64_t in_dim = 128; + const uint64_t hidden = 256; + const uint64_t out_dim = 96; + const uint64_t world = 4; + const uint64_t split = hidden / world; + const uint64_t n_tok = 16; + const uint64_t gu_row_bytes = ((in_dim + 31u) / 32u) * 34u; + const uint64_t down_row_bytes = ((hidden + 31u) / 32u) * 34u; + const uint64_t gate_off = 0; + const uint64_t up_off = gate_off + hidden * gu_row_bytes; + const uint64_t down_off = up_off + hidden * gu_row_bytes; + const uint64_t model_size = down_off + out_dim * down_row_bytes; + unsigned char *model = (unsigned char *)malloc((size_t)model_size); + float *host_x = (float *)malloc((size_t)(n_tok * in_dim) * sizeof(float)); + float *host_mid = (float *)malloc((size_t)(n_tok * hidden) * sizeof(float)); + float *host_local_mid = (float *)malloc((size_t)(n_tok * split) * sizeof(float)); + float *host_full = (float *)malloc((size_t)(n_tok * out_dim) * sizeof(float)); + float *host_part = (float *)malloc((size_t)(n_tok * out_dim) * sizeof(float)); + float *host_sum = (float *)calloc((size_t)(n_tok * out_dim), sizeof(float)); + CHECK(model && host_x && host_mid && host_local_mid && host_full && + host_part && host_sum, "q8_ffn_tp4 host alloc"); + pack_q8_identity_scale(model + gate_off, in_dim, hidden); + pack_q8_identity_scale(model + up_off, in_dim, hidden); + pack_q8_identity_scale(model + down_off, hidden, out_dim); + for (uint64_t t = 0; t < n_tok; t++) { + for (uint64_t i = 0; i < in_dim; i++) { + host_x[t * in_dim + i] = + (float)((int)((i * 7u + t * 13u) % 61u) - 30) * 0.01953125f; + } + } + CHECK(ds4_gpu_set_model_map(model, model_size), + "q8_ffn_tp4 set model map"); + + ds4_gpu_tensor x; memset(&x, 0, sizeof(x)); + ds4_gpu_tensor full_gate; memset(&full_gate, 0, sizeof(full_gate)); + ds4_gpu_tensor full_up; memset(&full_up, 0, sizeof(full_up)); + ds4_gpu_tensor full_mid; memset(&full_mid, 0, sizeof(full_mid)); + ds4_gpu_tensor full_out; memset(&full_out, 0, sizeof(full_out)); + CHECK(ds4_gpu_tensor_alloc_on(&x, 0, n_tok * in_dim * sizeof(float)) == 0 && + ds4_gpu_tensor_alloc_on(&full_gate, 0, + n_tok * hidden * sizeof(float)) == 0 && + ds4_gpu_tensor_alloc_on(&full_up, 0, + n_tok * hidden * sizeof(float)) == 0 && + ds4_gpu_tensor_alloc_on(&full_mid, 0, + n_tok * hidden * sizeof(float)) == 0 && + ds4_gpu_tensor_alloc_on(&full_out, 0, + n_tok * out_dim * sizeof(float)) == 0, + "q8_ffn_tp4 alloc full tensors"); + CHECK(ds4_gpu_tensor_write(&x, 0, host_x, + n_tok * in_dim * sizeof(float)), + "q8_ffn_tp4 write x"); + + int ok = ds4_gpu_shared_gate_up_swiglu_q8_0_rows_tensor( + &full_gate, &full_up, &full_mid, + model, model_size, gate_off, up_off, + in_dim, hidden, &x, n_tok, 0.0f) && + ds4_gpu_matmul_q8_0_tensor(&full_out, + model, model_size, down_off, + hidden, out_dim, &full_mid, n_tok) && + ds4_gpu_tensor_read(&full_mid, 0, host_mid, + n_tok * hidden * sizeof(float)); + for (uint64_t rank = 0; ok && rank < world; rank++) { + ds4_gpu_tensor gate; memset(&gate, 0, sizeof(gate)); + ds4_gpu_tensor up; memset(&up, 0, sizeof(up)); + ds4_gpu_tensor mid; memset(&mid, 0, sizeof(mid)); + ds4_gpu_tensor part; memset(&part, 0, sizeof(part)); + ok = ds4_gpu_tensor_alloc_on(&gate, 0, + n_tok * split * sizeof(float)) == 0 && + ds4_gpu_tensor_alloc_on(&up, 0, + n_tok * split * sizeof(float)) == 0 && + ds4_gpu_tensor_alloc_on(&mid, 0, + n_tok * split * sizeof(float)) == 0 && + ds4_gpu_tensor_alloc_on(&part, 0, + n_tok * out_dim * sizeof(float)) == 0 && + ds4_gpu_shared_gate_up_swiglu_q8_0_rows_tensor( + &gate, &up, &mid, + model, model_size, + gate_off + rank * split * gu_row_bytes, + up_off + rank * split * gu_row_bytes, + in_dim, split, &x, n_tok, 0.0f) && + ds4_gpu_tensor_read(&mid, 0, host_local_mid, + n_tok * split * sizeof(float)); + for (uint64_t t = 0; ok && t < n_tok; t++) { + for (uint64_t i = 0; i < split; i++) { + const float want = host_mid[t * hidden + rank * split + i]; + const float got = host_local_mid[t * split + i]; + if (want != got) { + fprintf(stderr, + "FAIL: q8_ffn_tp4 mid mismatch rank=%llu token=%llu " + "row=%llu full=%f local=%f diff=%g\n", + (unsigned long long)rank, + (unsigned long long)t, + (unsigned long long)i, + want, got, fabsf(want - got)); + ok = 0; + break; + } + } + } + if (ok) { + ok = ds4_gpu_matmul_q8_0_kslice_rows_tensor( + &part, model, model_size, down_off, + hidden, out_dim, rank * split, split, + &mid, n_tok) && + ds4_gpu_tensor_read(&part, 0, host_part, + n_tok * out_dim * sizeof(float)); + } + if (ok) { + for (uint64_t i = 0; i < n_tok * out_dim; i++) { + host_sum[i] += host_part[i]; + } + } + ds4_gpu_tensor_free_in_place(&part); + ds4_gpu_tensor_free_in_place(&mid); + ds4_gpu_tensor_free_in_place(&up); + ds4_gpu_tensor_free_in_place(&gate); + } + ok = ok && ds4_gpu_tensor_read(&full_out, 0, host_full, + n_tok * out_dim * sizeof(float)); + CHECK(ok, "q8_ffn_tp4 compute"); + float max_abs = 0.0f; + float max_scaled = 0.0f; + for (uint64_t i = 0; i < n_tok * out_dim; i++) { + const float diff = fabsf(host_full[i] - host_sum[i]); + const float scale = fmaxf(1.0f, fabsf(host_full[i])); + const float scaled = diff / scale; + if (diff > max_abs) max_abs = diff; + if (scaled > max_scaled) max_scaled = scaled; + /* Four partial dot products change only the floating-point + * association. Allow a small number of ulps plus an absolute + * floor for values around zero. */ + if (diff > 2.0e-4f + 4.0f * FLT_EPSILON * scale) { + fprintf(stderr, + "FAIL: q8_ffn_tp4 output mismatch elem=%llu " + "full=%f sum=%f diff=%g scaled=%g\n", + (unsigned long long)i, + host_full[i], host_sum[i], diff, scaled); + return 1; + } + } + + ds4_gpu_tensor_free_in_place(&full_out); + ds4_gpu_tensor_free_in_place(&full_mid); + ds4_gpu_tensor_free_in_place(&full_up); + ds4_gpu_tensor_free_in_place(&full_gate); + ds4_gpu_tensor_free_in_place(&x); + free(model); + free(host_x); + free(host_mid); + free(host_local_mid); + free(host_full); + free(host_part); + free(host_sum); + ds4_gpu_cleanup(); + fprintf(stderr, " q8_ffn_tp4 OK (max_abs=%g max_scaled=%g)\n", + max_abs, max_scaled); + return 0; +} + static int run_q8_matmul_top1_fused(void) { int dev_count = 0; (void)cudaGetDeviceCount(&dev_count); @@ -1854,6 +2042,7 @@ int main(void) { if (run_glm_decode_attention_staged()) return 1; if (run_moe_handoff_pack()) return 1; if (run_q8_kslice()) return 1; + if (run_q8_ffn_tp4()) return 1; if (run_q8_matmul_top1_fused()) return 1; if (run_f16_small_out()) return 1; if (run_f16_small_batch()) return 1; diff --git a/tests/test_metal_session_batch.c b/tests/test_metal_session_batch.c index edd6285c5c..e215c3c71d 100644 --- a/tests/test_metal_session_batch.c +++ b/tests/test_metal_session_batch.c @@ -152,7 +152,7 @@ int main(void) { if (ds4_engine_open(&engine, &opt) != 0) fail("engine open", -1, -1); if (tp_worker) { - const int worker_rc = ds4_tp_worker_run(engine, &opt.tp); + const int worker_rc = ds4_tp_worker_run(engine, &opt.tp, TEST_CTX); ds4_engine_close(engine); return worker_rc; } diff --git a/tests/test_tp_protocol.c b/tests/test_tp_protocol.c new file mode 100644 index 0000000000..e2c8d59084 --- /dev/null +++ b/tests/test_tp_protocol.c @@ -0,0 +1,492 @@ +#include "ds4_tp.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +static int reserve_loopback_port(void) { + int fd = socket(AF_INET, SOCK_STREAM, 0); + if (fd < 0) return -1; + struct sockaddr_in addr = { + .sin_family = AF_INET, + .sin_port = 0, + .sin_addr.s_addr = htonl(INADDR_LOOPBACK), + }; + socklen_t len = sizeof(addr); + if (bind(fd, (const struct sockaddr *)&addr, sizeof(addr)) != 0 || + getsockname(fd, (struct sockaddr *)&addr, &len) != 0) { + close(fd); + return -1; + } + const int port = ntohs(addr.sin_port); + close(fd); + return port; +} + +static ds4_tp_identity test_identity(uint32_t ctx_size) { + return (ds4_tp_identity) { + .gguf_bytes = UINT64_C(211075856448), + .model_id = 52, + .n_layer = 78, + .n_embd = 6144, + .n_vocab = 128000, + .quant_bits = 2, + .ctx_size = ctx_size, + .gate_slot_start = 7, + .gate_slot_step = 2, + .gates_per_token = 75, + }; +} + +static void child_error(uint32_t rank, const char *what, const char *detail, + int code) { + fprintf(stderr, "test_tp_protocol: rank %u %s%s%s\n", + rank, what, detail && detail[0] ? ": " : "", + detail && detail[0] ? detail : ""); + _exit(code); +} + +static void run_success_worker(uint32_t rank, uint32_t world, int port, + bool expert_only) { + alarm(30); + /* Make rank 3 connect first, then rank 2, then rank 1. The leader must + * index peers by their advertised rank rather than accept order. */ + if (world == 4u) usleep((useconds_t)(world - 1u - rank) * 20000u); + ds4_tp_options opt = { + .role = DS4_TP_WORKER, + .requested = true, + .expert_only = expert_only, + .leader_host = "127.0.0.1", + .leader_port = port, + .world_size = world, + .rank = rank, + .rank_set = true, + .transport = DS4_TP_TRANSPORT_NCCL, + }; + const uint32_t ctx_size = 4096u - rank * 512u; + const ds4_tp_identity id = test_identity(ctx_size); + char err[256] = ""; + ds4_tp *tp = NULL; + if (!ds4_tp_create(&tp, &opt, &id, err, sizeof(err))) { + child_error(rank, "create failed", err, 10); + } + if (!ds4_tp_collective_preflight(tp, 1, err, sizeof(err))) { + child_error(rank, "preflight failed", err, 11); + } + uint8_t blob[32] = {0}; + if (!ds4_tp_broadcast_blob(tp, blob, sizeof(blob), err, sizeof(err))) { + child_error(rank, "bootstrap receive failed", err, 12); + } + for (uint32_t i = 0; i < sizeof(blob); i++) { + const uint8_t expected = (uint8_t)(i ^ (world << 4u)); + if (blob[i] != expected) { + child_error(rank, "bootstrap payload mismatch", NULL, 13); + } + } + ds4_tp_command command; + if (!ds4_tp_recv_command(tp, &command, err, sizeof(err))) { + child_error(rank, "command receive failed", err, 14); + } + const uint64_t session_id = UINT64_C(0x12340000) + world; + const uint64_t session_id_2 = UINT64_C(0x56780000) + world; + if (command.type != DS4_TP_FRAME_SESSION_CREATE || + command.session_id != session_id || command.value != 1024) { + ds4_tp_command_free(&command); + child_error(rank, "session-create payload mismatch", NULL, 15); + } + ds4_tp_command_free(&command); + if (!ds4_tp_send_command_ack(tp, session_id, 0)) { + child_error(rank, "ACK send failed", NULL, 16); + } + + if (!ds4_tp_recv_command(tp, &command, err, sizeof(err)) || + command.type != DS4_TP_FRAME_SESSION_CREATE || + command.session_id != session_id_2 || command.value != 768) { + ds4_tp_command_free(&command); + child_error(rank, "second session-create mismatch", err, 17); + } + ds4_tp_command_free(&command); + if (!ds4_tp_send_command_ack(tp, session_id_2, 0)) { + child_error(rank, "second create ACK failed", NULL, 18); + } + + if (!ds4_tp_recv_command(tp, &command, err, sizeof(err)) || + command.type != DS4_TP_FRAME_SYNC || + command.session_id != session_id || command.n_tokens != 3 || + command.tokens[0] != 11 || command.tokens[1] != 22 || + command.tokens[2] != 33) { + ds4_tp_command_free(&command); + child_error(rank, "sync payload mismatch", err, 19); + } + ds4_tp_command_free(&command); + if (!ds4_tp_send_command_ack(tp, session_id, 0)) { + child_error(rank, "sync ACK failed", NULL, 20); + } + + if (!ds4_tp_recv_command(tp, &command, err, sizeof(err)) || + command.type != DS4_TP_FRAME_EVAL || + command.session_id != session_id || command.seq != 77 || + command.value != 44) { + ds4_tp_command_free(&command); + child_error(rank, "eval payload mismatch", err, 21); + } + ds4_tp_command_free(&command); + + if (!ds4_tp_recv_command(tp, &command, err, sizeof(err)) || + command.type != DS4_TP_FRAME_REWIND || + command.session_id != session_id || command.value != 2) { + ds4_tp_command_free(&command); + child_error(rank, "rewind payload mismatch", err, 22); + } + ds4_tp_command_free(&command); + + if (!ds4_tp_recv_command(tp, &command, err, sizeof(err)) || + command.type != DS4_TP_FRAME_INVALIDATE || + command.session_id != session_id_2) { + ds4_tp_command_free(&command); + child_error(rank, "invalidate payload mismatch", err, 23); + } + ds4_tp_command_free(&command); + + if (!ds4_tp_recv_command(tp, &command, err, sizeof(err)) || + command.type != DS4_TP_FRAME_EVAL_BATCH || command.n_items != 2 || + command.items[0].session_id != session_id || + command.items[0].token != 55 || + command.items[1].session_id != session_id_2 || + command.items[1].token != 66) { + ds4_tp_command_free(&command); + child_error(rank, "eval-batch payload mismatch", err, 24); + } + ds4_tp_command_free(&command); + if (!ds4_tp_send_command_ack(tp, 0, 0)) { + child_error(rank, "eval-batch ACK failed", NULL, 25); + } + + if (!ds4_tp_recv_command(tp, &command, err, sizeof(err)) || + command.type != DS4_TP_FRAME_MIXED_BATCH || + command.session_id != session_id_2 || command.n_tokens != 2 || + command.tokens[0] != 7 || command.tokens[1] != 8 || + command.n_items != 1 || + command.items[0].session_id != session_id || + command.items[0].token != 9) { + ds4_tp_command_free(&command); + child_error(rank, "mixed-batch payload mismatch", err, 26); + } + ds4_tp_command_free(&command); + if (!ds4_tp_send_command_ack(tp, session_id_2, 0)) { + child_error(rank, "mixed-batch ACK failed", NULL, 27); + } + + if (!ds4_tp_recv_command(tp, &command, err, sizeof(err)) || + command.type != DS4_TP_FRAME_SESSION_DESTROY || + command.session_id != session_id) { + ds4_tp_command_free(&command); + child_error(rank, "first session-destroy mismatch", err, 28); + } + ds4_tp_command_free(&command); + if (!ds4_tp_send_command_ack(tp, session_id, 0)) { + child_error(rank, "first destroy ACK failed", NULL, 29); + } + + if (!ds4_tp_recv_command(tp, &command, err, sizeof(err)) || + command.type != DS4_TP_FRAME_SESSION_DESTROY || + command.session_id != session_id_2) { + ds4_tp_command_free(&command); + child_error(rank, "second session-destroy mismatch", err, 30); + } + ds4_tp_command_free(&command); + if (!ds4_tp_send_command_ack(tp, session_id_2, 0)) { + child_error(rank, "second destroy ACK failed", NULL, 31); + } + + if (!ds4_tp_recv_command(tp, &command, err, sizeof(err))) { + child_error(rank, "stop receive failed", err, 32); + } + if (command.type != DS4_TP_FRAME_STOP) { + ds4_tp_command_free(&command); + child_error(rank, "expected stop", NULL, 33); + } + ds4_tp_command_free(&command); + ds4_tp_free(tp); + _exit(0); +} + +static int wait_workers(pid_t *pids, uint32_t count) { + int ok = 1; + for (uint32_t i = 0; i < count; i++) { + int status = 0; + if (waitpid(pids[i], &status, 0) != pids[i] || + !WIFEXITED(status) || WEXITSTATUS(status) != 0) { + fprintf(stderr, + "test_tp_protocol: worker pid %ld failed (status=0x%x)\n", + (long)pids[i], status); + ok = 0; + } + } + return ok; +} + +static void terminate_workers(pid_t *pids, uint32_t count) { + for (uint32_t i = 0; i < count; i++) { + if (pids[i] > 0) (void)kill(pids[i], SIGTERM); + } + (void)wait_workers(pids, count); +} + +static int run_success_case(uint32_t world, bool expert_only) { + const int port = reserve_loopback_port(); + if (port <= 0) return 0; + pid_t pids[DS4_TP_MAX_WORLD - 1u] = {0}; + for (uint32_t rank = 1; rank < world; rank++) { + const pid_t pid = fork(); + if (pid == 0) run_success_worker(rank, world, port, expert_only); + if (pid < 0) { + terminate_workers(pids, rank - 1u); + return 0; + } + pids[rank - 1u] = pid; + } + + alarm(30); + ds4_tp_options opt = { + .role = DS4_TP_LEADER, + .requested = true, + .expert_only = expert_only, + .listen_host = "127.0.0.1", + .listen_port = port, + .world_size = world, + .rank = 0, + .rank_set = true, + .transport = DS4_TP_TRANSPORT_NCCL, + }; + const ds4_tp_identity id = test_identity(4096u); + char err[256] = ""; + ds4_tp *tp = NULL; + if (!ds4_tp_create(&tp, &opt, &id, err, sizeof(err))) { + fprintf(stderr, "test_tp_protocol: leader create: %s\n", err); + terminate_workers(pids, world - 1u); + return 0; + } + const uint32_t expected_ctx = 4096u - (world - 1u) * 512u; + int ok = ds4_tp_rank(tp) == 0 && ds4_tp_world(tp) == world && + ds4_tp_is_collective(tp) && + ds4_tp_is_expert_only(tp) == expert_only && + ds4_tp_peer_ctx(tp) == expected_ctx; + if (!ok) fprintf(stderr, "test_tp_protocol: leader topology mismatch\n"); + if (ok && !ds4_tp_collective_preflight(tp, 1, err, sizeof(err))) { + fprintf(stderr, "test_tp_protocol: leader preflight: %s\n", err); + ok = 0; + } + uint8_t blob[32]; + for (uint32_t i = 0; i < sizeof(blob); i++) { + blob[i] = (uint8_t)(i ^ (world << 4u)); + } + if (ok && !ds4_tp_broadcast_blob(tp, blob, sizeof(blob), err, + sizeof(err))) { + fprintf(stderr, "test_tp_protocol: leader bootstrap: %s\n", err); + ok = 0; + } + const uint64_t session_id = UINT64_C(0x12340000) + world; + const uint64_t session_id_2 = UINT64_C(0x56780000) + world; + if (ok && (!ds4_tp_send_session_create(tp, session_id, 1024) || + !ds4_tp_wait_command_ack(tp, session_id, "test create", + err, sizeof(err)))) { + fprintf(stderr, "test_tp_protocol: command fan-out: %s\n", err); + ok = 0; + } + if (ok && (!ds4_tp_send_session_create(tp, session_id_2, 768) || + !ds4_tp_wait_command_ack(tp, session_id_2, "test create 2", + err, sizeof(err)))) { + fprintf(stderr, "test_tp_protocol: second create fan-out: %s\n", err); + ok = 0; + } + const int sync_tokens[] = {11, 22, 33}; + if (ok && (!ds4_tp_send_sync(tp, session_id, sync_tokens, 3) || + !ds4_tp_wait_command_ack(tp, session_id, "test sync", + err, sizeof(err)))) { + fprintf(stderr, "test_tp_protocol: sync fan-out: %s\n", err); + ok = 0; + } + if (ok && (!ds4_tp_send_eval(tp, session_id, 77, 44) || + !ds4_tp_send_rewind(tp, session_id, 2) || + !ds4_tp_send_invalidate(tp, session_id_2))) { + fprintf(stderr, "test_tp_protocol: eval/rewind/invalidate fan-out failed\n"); + ok = 0; + } + const ds4_tp_batch_item batch[] = { + {.session_id = session_id, .token = 55}, + {.session_id = session_id_2, .token = 66}, + }; + if (ok && (!ds4_tp_send_eval_batch(tp, batch, 2) || + !ds4_tp_wait_command_ack(tp, 0, "test eval batch", + err, sizeof(err)))) { + fprintf(stderr, "test_tp_protocol: eval-batch fan-out: %s\n", err); + ok = 0; + } + const int mixed_prompt[] = {7, 8}; + const ds4_tp_batch_item mixed[] = { + {.session_id = session_id, .token = 9}, + }; + if (ok && (!ds4_tp_send_mixed_batch(tp, session_id_2, + mixed_prompt, 2, mixed, 1) || + !ds4_tp_wait_command_ack(tp, session_id_2, + "test mixed batch", + err, sizeof(err)))) { + fprintf(stderr, "test_tp_protocol: mixed-batch fan-out: %s\n", err); + ok = 0; + } + if (ok && (!ds4_tp_send_session_destroy(tp, session_id) || + !ds4_tp_wait_command_ack(tp, session_id, "test destroy", + err, sizeof(err)) || + !ds4_tp_send_session_destroy(tp, session_id_2) || + !ds4_tp_wait_command_ack(tp, session_id_2, "test destroy 2", + err, sizeof(err)))) { + fprintf(stderr, "test_tp_protocol: destroy fan-out: %s\n", err); + ok = 0; + } + if (!ds4_tp_send_stop(tp)) ok = 0; + ds4_tp_free(tp); + if (!wait_workers(pids, world - 1u)) ok = 0; + alarm(0); + if (ok) { + fprintf(stderr, "test_tp_protocol: %u-rank %s path OK\n", + world, expert_only ? "expert" : "tensor"); + } + return ok; +} + +static void run_failed_preflight_worker(int port) { + alarm(30); + ds4_tp_options opt = { + .role = DS4_TP_WORKER, + .requested = true, + .expert_only = true, + .leader_host = "127.0.0.1", + .leader_port = port, + .world_size = 2, + .rank = 1, + .rank_set = true, + .transport = DS4_TP_TRANSPORT_NCCL, + }; + const ds4_tp_identity id = test_identity(4096u); + char err[256] = ""; + ds4_tp *tp = NULL; + if (!ds4_tp_create(&tp, &opt, &id, err, sizeof(err))) { + child_error(1, "failure-case create failed", err, 30); + } + if (ds4_tp_collective_preflight(tp, 0, err, sizeof(err)) || + strstr(err, "unavailable") == NULL) { + child_error(1, "failure-case preflight was not rejected", err, 31); + } + ds4_tp_command command; + if (!ds4_tp_recv_command(tp, &command, err, sizeof(err)) || + command.type != DS4_TP_FRAME_STOP) { + child_error(1, "failure-case stop receive failed", err, 32); + } + ds4_tp_command_free(&command); + ds4_tp_free(tp); + _exit(0); +} + +static int run_failed_preflight_case(void) { + const int port = reserve_loopback_port(); + if (port <= 0) return 0; + const pid_t pid = fork(); + if (pid == 0) run_failed_preflight_worker(port); + if (pid < 0) return 0; + pid_t pids[1] = {pid}; + + alarm(30); + ds4_tp_options opt = { + .role = DS4_TP_LEADER, + .requested = true, + .expert_only = true, + .listen_host = "127.0.0.1", + .listen_port = port, + .world_size = 2, + .rank = 0, + .rank_set = true, + .transport = DS4_TP_TRANSPORT_NCCL, + }; + const ds4_tp_identity id = test_identity(4096u); + char err[256] = ""; + ds4_tp *tp = NULL; + if (!ds4_tp_create(&tp, &opt, &id, err, sizeof(err))) { + fprintf(stderr, "test_tp_protocol: failure-case leader create: %s\n", + err); + terminate_workers(pids, 1); + return 0; + } + int ok = !ds4_tp_collective_preflight(tp, 1, err, sizeof(err)) && + strstr(err, "unavailable") != NULL; + if (!ok) { + fprintf(stderr, + "test_tp_protocol: unavailable-rank preflight was not rejected: %s\n", + err); + } + if (!ds4_tp_send_stop(tp)) ok = 0; + ds4_tp_free(tp); + if (!wait_workers(pids, 1)) ok = 0; + alarm(0); + if (ok) fprintf(stderr, "test_tp_protocol: failed preflight path OK\n"); + return ok; +} + +static int run_partition_cases(void) { + for (uint32_t world = 2u; world <= 4u; world += 2u) { + uint32_t next = 0u; + uint32_t sum = 0u; + for (uint32_t rank = 0u; rank < world; rank++) { + uint32_t base = UINT32_MAX; + uint32_t count = 0u; + if (!ds4_tp_partition(257u, rank, world, &base, &count) || + base != next || count == 0u) { + fprintf(stderr, + "test_tp_protocol: %u-rank partition failed at rank %u\n", + world, rank); + return 0; + } + next = base + count; + sum += count; + } + if (next != 257u || sum != 257u) { + fprintf(stderr, + "test_tp_protocol: %u-rank partition has a gap/overlap\n", + world); + return 0; + } + } + uint32_t base = 0u; + uint32_t count = 0u; + if (ds4_tp_partition(256u, 4u, 4u, &base, &count) || + ds4_tp_partition(0u, 0u, 2u, &base, &count) || + ds4_tp_partition(1u, 0u, 4u, &base, &count)) { + fprintf(stderr, "test_tp_protocol: invalid partition accepted\n"); + return 0; + } + fprintf(stderr, "test_tp_protocol: 2/4-rank partitions OK\n"); + return 1; +} + +int main(void) { + signal(SIGPIPE, SIG_IGN); + setenv("DS4_TP_TIMEOUT_SEC", "10", 1); + if (!run_partition_cases() || + !run_success_case(2, true) || + !run_success_case(2, false) || + !run_success_case(4, true) || + !run_success_case(4, false) || + !run_failed_preflight_case()) { + fprintf(stderr, "test_tp_protocol: FAIL\n"); + return 1; + } + fprintf(stderr, "test_tp_protocol: PASS\n"); + return 0; +}