Skip to content

Local decode with distributed prefill #343

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
lobanov wants to merge 15 commits into
base: main
Choose a base branch
from
+24,421 −337
Draft

Local decode with distributed prefill #343

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
9952733
create plan
lobanov Jun 3, 2026
116e358
phase 0 and phase 1 completed
lobanov Jun 3, 2026
9c83dac
phase 2 complete
lobanov Jun 3, 2026
f7a0baa
refine phase 3 plan
lobanov Jun 4, 2026
85c13af
phase 3.5, route-specificity
lobanov Jun 4, 2026
c05eed4
refine phase 4 scope
lobanov Jun 4, 2026
2ad3829
phase 4 closure
lobanov Jun 5, 2026
e9e5a48
phases 5+ scope refinement
lobanov Jun 5, 2026
5c621a0
remove compiled binaries
lobanov Jun 5, 2026
30c9832
phase 5 plan
lobanov Jun 5, 2026
e2bd47a
phase 5 first implementation with initial tests
lobanov Jun 5, 2026
b7ed1ec
phase 5 local decode with sampling
lobanov Jun 5, 2026
a9d673b
fix: stale coordinator logits
lobanov Jun 5, 2026
42f6021
phase 5 frontend testing and fixes
lobanov Jun 5, 2026
97d0688
phase 5 performance eval [WIP]
lobanov Jun 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
97 changes: 95 additions & 2 deletions Makefile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ CUDA_LDLIBS ?= -lm -Xcompiler -pthread -L$(CUDA_HOME)/targets/sbsa-linux/lib -L$
METAL_LDLIBS := $(LDLIBS)
endif

.PHONY: all help clean test cpu cuda cuda-spark cuda-generic cuda-regression
.PHONY: all help clean test cpu cuda cuda-spark cuda-generic cuda-regression issue304-phase0-local

ifeq ($(UNAME_S),Darwin)
all: ds4 ds4-server ds4-bench ds4-eval ds4-agent
Expand Down Expand Up @@ -190,6 +190,99 @@ ds4_cuda.o: ds4_cuda.cu ds4_gpu.h ds4_iq2_tables_cuda.inc
tests/cuda_long_context_smoke: tests/cuda_long_context_smoke.o ds4_cuda.o
$(NVCC) $(NVCCFLAGS) -o $@ $^ $(CUDA_LDLIBS)

tests/issue304_phase0_local.o: tests/issue304_phase0_local.c ds4.h ds4_distributed.h
$(CC) $(CFLAGS) -c -o $@ tests/issue304_phase0_local.c

tests/issue304_phase0_local: tests/issue304_phase0_local.o $(CORE_OBJS)
ifeq ($(UNAME_S),Darwin)
$(CC) $(CFLAGS) -o $@ tests/issue304_phase0_local.o $(CORE_OBJS) $(METAL_LDLIBS)
else
$(NVCC) $(NVCCFLAGS) -o $@ tests/issue304_phase0_local.o $(CORE_OBJS) $(CUDA_LDLIBS)
endif

tests/issue304_phase0_dgx.o: tests/issue304_phase0_dgx.c ds4.h ds4_distributed.h
$(CC) $(CFLAGS) -c -o $@ tests/issue304_phase0_dgx.c

tests/issue304_phase0_dgx: tests/issue304_phase0_dgx.o $(CORE_OBJS)
ifeq ($(UNAME_S),Darwin)
$(CC) $(CFLAGS) -o $@ tests/issue304_phase0_dgx.o $(CORE_OBJS) $(METAL_LDLIBS)
else
$(NVCC) $(NVCCFLAGS) -o $@ tests/issue304_phase0_dgx.o $(CORE_OBJS) $(CUDA_LDLIBS)
endif

tests/issue304_phase1_matrix.o: tests/issue304_phase1_matrix.c ds4.h
$(CC) $(CFLAGS) -c -o $@ tests/issue304_phase1_matrix.c

tests/issue304_phase1_matrix: tests/issue304_phase1_matrix.o $(CORE_OBJS)
ifeq ($(UNAME_S),Darwin)
$(CC) $(CFLAGS) -o $@ tests/issue304_phase1_matrix.o $(CORE_OBJS) $(METAL_LDLIBS)
else
$(NVCC) $(NVCCFLAGS) -o $@ tests/issue304_phase1_matrix.o $(CORE_OBJS) $(CUDA_LDLIBS)
endif

tests/issue304_phase2_handoff.o: tests/issue304_phase2_handoff.c ds4.h ds4_distributed.h
$(CC) $(CFLAGS) -c -o $@ tests/issue304_phase2_handoff.c

tests/issue304_phase2_handoff: tests/issue304_phase2_handoff.o $(CORE_OBJS)
ifeq ($(UNAME_S),Darwin)
$(CC) $(CFLAGS) -o $@ tests/issue304_phase2_handoff.o $(CORE_OBJS) $(METAL_LDLIBS)
else
$(NVCC) $(NVCCFLAGS) -o $@ tests/issue304_phase2_handoff.o $(CORE_OBJS) $(CUDA_LDLIBS)
endif

tests/issue304_phase3_vectors.o: tests/issue304_phase3_vectors.c ds4.h ds4_distributed.h
$(CC) $(CFLAGS) -c -o $@ tests/issue304_phase3_vectors.c

tests/issue304_phase3_vectors: tests/issue304_phase3_vectors.o $(CORE_OBJS)
ifeq ($(UNAME_S),Darwin)
$(CC) $(CFLAGS) -o $@ tests/issue304_phase3_vectors.o $(CORE_OBJS) $(METAL_LDLIBS)
else
$(NVCC) $(NVCCFLAGS) -o $@ tests/issue304_phase3_vectors.o $(CORE_OBJS) $(CUDA_LDLIBS)
endif

tests/issue304_phase35_vectors.o: tests/issue304_phase35_vectors.c ds4.h ds4_distributed.h
$(CC) $(CFLAGS) -c -o $@ tests/issue304_phase35_vectors.c

tests/issue304_phase35_vectors: tests/issue304_phase35_vectors.o $(CORE_OBJS)
ifeq ($(UNAME_S),Darwin)
$(CC) $(CFLAGS) -o $@ tests/issue304_phase35_vectors.o $(CORE_OBJS) $(METAL_LDLIBS)
else
$(NVCC) $(NVCCFLAGS) -o $@ tests/issue304_phase35_vectors.o $(CORE_OBJS) $(CUDA_LDLIBS)
endif

tests/issue304_phase4_handoff.o: tests/issue304_phase4_handoff.c ds4.h ds4_distributed.h
$(CC) $(CFLAGS) -c -o $@ tests/issue304_phase4_handoff.c

tests/issue304_phase4_handoff: tests/issue304_phase4_handoff.o $(CORE_OBJS)
ifeq ($(UNAME_S),Darwin)
$(CC) $(CFLAGS) -o $@ tests/issue304_phase4_handoff.o $(CORE_OBJS) $(METAL_LDLIBS)
else
$(NVCC) $(NVCCFLAGS) -o $@ tests/issue304_phase4_handoff.o $(CORE_OBJS) $(CUDA_LDLIBS)
endif

tests/issue304_phase5_multiturn.o: tests/issue304_phase5_multiturn.c ds4.h ds4_distributed.h
$(CC) $(CFLAGS) -c -o $@ tests/issue304_phase5_multiturn.c

tests/issue304_phase5_multiturn: tests/issue304_phase5_multiturn.o $(CORE_OBJS)
ifeq ($(UNAME_S),Darwin)
$(CC) $(CFLAGS) -o $@ tests/issue304_phase5_multiturn.o $(CORE_OBJS) $(METAL_LDLIBS)
else
$(NVCC) $(NVCCFLAGS) -o $@ tests/issue304_phase5_multiturn.o $(CORE_OBJS) $(CUDA_LDLIBS)
endif

tests/issue304_phase4_diagnose.o: tests/issue304_phase4_diagnose.c ds4.h ds4_distributed.h
$(CC) $(CFLAGS) -c -o $@ tests/issue304_phase4_diagnose.c

tests/issue304_phase4_diagnose: tests/issue304_phase4_diagnose.o $(CORE_OBJS)
ifeq ($(UNAME_S),Darwin)
$(CC) $(CFLAGS) -o $@ tests/issue304_phase4_diagnose.o $(CORE_OBJS) $(METAL_LDLIBS)
else
$(NVCC) $(NVCCFLAGS) -o $@ tests/issue304_phase4_diagnose.o $(CORE_OBJS) $(CUDA_LDLIBS)
endif

issue304-phase0-local: ds4 tests/issue304_phase0_local
./tests/issue304_phase0_local

ds4_test: ds4_test.o ds4_help.o ds4_kvstore.o rax.o $(CORE_OBJS)
ifeq ($(UNAME_S),Darwin)
$(CC) $(CFLAGS) -o $@ ds4_test.o ds4_help.o ds4_kvstore.o rax.o $(CORE_OBJS) $(METAL_LDLIBS)
Expand All @@ -206,4 +299,4 @@ q4k-dot-test: tests/test_q4k_dot.c
./tests/test_q4k_dot

clean:
rm -f ds4 ds4-server ds4-bench ds4-eval ds4-agent ds4_cpu ds4_native ds4_server_test ds4_test tests/test_q4k_dot *.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 tests/test_q4k_dot *.o tests/cuda_long_context_smoke tests/cuda_long_context_smoke.o tests/issue304_phase0_local tests/issue304_phase0_local.o tests/issue304_phase0_dgx tests/issue304_phase0_dgx.o tests/issue304_phase1_matrix tests/issue304_phase1_matrix.o tests/issue304_phase2_handoff tests/issue304_phase2_handoff.o tests/issue304_phase3_vectors tests/issue304_phase3_vectors.o tests/issue304_phase35_vectors tests/issue304_phase35_vectors.o tests/issue304_phase4_handoff tests/issue304_phase4_handoff.o tests/issue304_phase5_multiturn tests/issue304_phase5_multiturn.o
Loading