feat: support Apple RDMA in layer parallel distributed mode - #715
Closed
ryan5rdx wants to merge 13 commits into
Closed
feat: support Apple RDMA in layer parallel distributed mode#715ryan5rdx wants to merge 13 commits into
ryan5rdx wants to merge 13 commits into
Conversation
ryan5rdx
force-pushed
the
apple-rdma-pp
branch
2 times, most recently
from
August 6, 2026 07:22
3ad3b54 to
962cdf9
Compare
ryan5rdx
marked this pull request as ready for review
August 6, 2026 21:45
ryan5rdx
changed the title
ryan5rdx
changed the title
kernel_dsv4_indexer_score_one_direct accumulated 64 head contributions through a four-entry threadgroup scratch that every loop iteration reused. Guarding that reuse cost two threadgroup barriers per iteration, so a single scored row paid 33 barriers to add 64 floats - and three of the four simdgroups idled at each one while thread 0 did four adds. Give each head its own slot instead. The writes never collide, so both per-iteration barriers go and one barrier before the final sum is enough. Threadgroup memory grows from (128+4) to (128+64) floats, 528 to 768 B, which does not change occupancy at 128 threads. Bit-identical: the final accumulation walks heads in the same ascending order with the same max(s,0)*(w[head]*scale) grouping, so every score is unchanged bit for bit and the selected top-512 cannot move. This kernel is the dominant context-scaling term at decode. Measured at 155k context (DS4_METAL_INDEXER_STAGE_PROFILE=1): score 0.925 ms/layer over 21 ratio-4 layers, which deflates to roughly 12 ms of a 47.6 ms token once the profiler's own serialization is removed. It runs at 690 GFLOP/s, 2.6% of peak, and the barriers are the reason.
A symbolicated time profile of the decode thread at 190k (2104 ms over 631 tokens, 3.33 ms/token) showed most of its walkable time going to work that recomputes constants. ds4_gpu_get_pipeline formatted an NSString on every call to probe an NSMutableDictionary, to recover a pointer that never changes. That was the largest single item: 0.49 ms/token across __CFStringAppendFormatCore, __CFStringChangeSizeMultiple and __CFStringCreateImmutableFunnel3, plus a share of the objc_msgSend/retain/release traffic behind it. Probe a content-hashed C-string table first. The dictionary still owns the reference so entries are unretained and never need invalidating, and the strcmp means a recycled name pointer cannot alias another kernel; a full bucket just falls through to the dictionary. glm_graph_env_present called getenv from the per-layer graph path, which is a locked linear scan of the environment - 0.15 ms/token in __findenv_locked. Memoized on the name pointers, which are literals. ds4_gpu_wrap_model_range rescanned g_model_views[] from zero on every lookup - 0.12 ms/token. Consecutive lookups almost always hit the same view, so start from the last hit. Ordering hint only; the range check still decides. Together ~0.87 ms/token of the 4.8 ms/token that is not GPU time. Whether that converts to throughput depends on how much of the encode thread overlaps GPU execution, which commit 42d099f just changed - so measure the split flush first and treat this as hygiene rather than a headline. No numerical effect: tests/bench_indexer_score reports byte-identical scores, and the layer_pack, gpu_args, mgpu_placement and mxfp4_metal suites all pass.
DS4_METAL_Q8_MV_ROWS=4 widens nr0 for every Q8_0 matvec, but three consumers dispatch a kernel whose row count is a compile-time constant (kernel_mul_mv_q8_0_f32_pair and the two fused HC-expand variants all bind NR0 = N_R0_Q8_0 and ignore args.nr0) while still sizing their grid as out_dim / nr0. The grid halved while each threadgroup kept writing two rows, so the upper half of every output block went unwritten and decode produced garbage. Pin those three paths back to the 2-row shape. The other seven callers of ds4_gpu_make_q8_0_mv_dispatch either take the kernel name from the dispatch or branch on nr0 >= 4, so they keep the wide path. Standalone kernel benchmarks on an M3 Ultra put attn q_b [1024x32768] at 336 GB/s with the shipped nsg4/nr2 and 695 GB/s at nsg2/nr4, worth about 2.4 ms of a 37.2 ms decode token on two-node pipeline-parallel Flash MXFP4. Note the wide path changes matvec reduction order, so greedy output is not bit-identical to the nr0=2 default; run the greedy-identity QA before relying on it.
ds4_gpu_decode_attn_rope_fuse_available() and ds4_gpu_kv_rope_fp8_fuse_available() tested the device name for "M3" or "M5", as did the two sites selecting the in-place pair and compact affine RoPE specialisation that those predicates mirror. An M3 Ultra satisfies both that name test and ds4_gpu_device_is_pre_m5_apple_silicon(), so e6a6ab5's ported decode fusions were exercised there while every other pre-M5 part fell through to the unfused path: the gate was a device-name test where the rest of that port used a family predicate. Route all four sites through one predicate so the fuse gates cannot drift from the RoPE path they replicate. Decode only - callers pair it with n_tok == 1, so long prefill keeps the proven shared4 schedule and the YaRN fast-math caveat recorded at the shared-coefficient site cannot apply. Unblocks the qkv-norm + KV-RoPE + FP8-store triple, 3 dispatches to 1 over 43 layers, and the compressor emit finalize, 7-11 to 1 over 21 layers. Measured on a pair of M2 Ultra Mac Studios over RDMA, DeepSeek V4 Flash MXFP4, 0 context: 37.00 -> 38.66 t/s alone, +4.5%, 27.03 -> 25.87 ms/token. Not bit-identical by construction: these are separate kernels, so exactness under Apple8 fast-math codegen is empirical. The full-vocabulary logit harness cannot run here - metal_decode_schedule_bench is single-node and this model does not fit one 128 GB machine - so the available check is a greedy-output comparison at --temp 0 over several hundred tokens. DS4_METAL_DISABLE_PRE_M5_DECODE_ROPE_PAIR_AFFINE=1 rolls it back.
Two more device-name gates of the kind 01a56db fixed. Both tested
ds4_gpu_device_name_contains("M3") or ("M5") where the surrounding port
used ds4_gpu_device_is_pre_m5_apple_silicon(), so an M3 Ultra exercised
them and no other pre-M5 part ever did.
The gathered KV stage collapses the raw f32->f16 ring copy, the compressed
f16 copy and the pad into one dispatch; the persistent zero mask replaces
a per-token fill_f16 of literal zero with a buffer memset once at startup.
Three dispatches per gathered layer, and past position 127 every layer is
gathered, so roughly 123 dispatches per token.
Measured on the 2x M2 Ultra pair over RDMA, DeepSeek V4 Flash MXFP4,
0 context, ~1000 generated tokens, four interleaved arms:
both rollbacks set 39.37 t/s 25.400 ms 8.187 GPU ms/cb
stage off, mask on 39.65 25.221 8.131
stage on, mask off 40.77 24.528 7.892
both on 41.21 24.266 7.795
Stage alone -0.872 ms, mask alone -0.179 ms, together -1.134 ms against a
sum of singles of 1.052, so additive with a small positive interaction.
The GPU busy counter says what kind of win this is. Busy time per token
falls 23.58 -> 22.45 ms, a 1.13 ms drop against a 1.134 ms wall-clock
drop: the entire gain comes out of GPU busy, one for one. Busy fraction is
92.8/92.8/92.7/92.5% across the four arms and the non-busy remainder is a
flat 1.80-1.82 ms in every one. This pair is GPU-bound, and removing
dispatches still shortens the token because a dispatch costs command
processor time that is counted as busy -- it is not bubble recovery, which
is why 32ef898's host-side changes measured nothing.
Implied marginal cost, given 41 gathered layers past position 127: about
4.4 us for the one-threadgroup mask fill and about 10.6 us for each staged
copy, the latter higher because the fusion also removes duplicate reads
and writes of the KV rows rather than only a launch.
Also fixes latent test failures. ds4_test --metal-kernels on Apple8 goes
29 failures to 23: sites 2728, 2778 and 2779 disappear. 2728 is
test_metal_gathered_kv_stage_exact asserting a nonzero return under
DS4_METAL_REQUIRE_GATHERED_KV_STAGE, which could not succeed on this
family because eligible was false; 2778 and 2779 are the byte-for-byte
comparisons against the unfused reference that the require assert was
gating, now live over four raw_start values including two ring wraps. No
new assertion site appears. A control with all four pre-M5 knobs set
reproduces 29 with identical sites, so neither 01a56db nor 65b4cf2
contributes to the remainder. The 23 that remain are pre-existing on
Apple8, all inside test_metal_f16_compressor_pair_state_store_exact_case
(sites 1098 and 1174-1178), and want a separate look.
Deliberately not touching use_shared_kvpad (ds4_metal.m:27858). Leaving it
M3-only keeps shared_pad at 0 so the fused stage writes both the K and V
pad copies, matching the standalone pad kernel. It removes no dispatch and
carries a layout hazard.
DS4_METAL_DISABLE_PRE_M5_GATHERED_KV_STAGE=1 and
DS4_METAL_DISABLE_PRE_M5_PERSISTENT_ZERO_ATTN_MASK=1 roll it back; set
both on both ranks.
ds4_test --metal-kernels reported 23 failures on Apple8. None were numerical. ds4_gpu_matmul_f16_pair_compressor_store_tensor rejected any device whose name contains neither "M3" nor "M5" (ds4_metal.m:20038) and returned 0, while test_metal_f16_compressor_pair_state_store_exact_case asserts it returns 1 with no hardware guard (tests/ds4_test.c:1093-1098). The first assert therefore failed, the fused destinations were never written, and the five comparisons below it read uninitialised memory - which is why the log showed a total mismatch (proj=256/256) and absurd ULP values in the billions rather than a small numeric drift. Same device-name-versus-family mismatch as 01a56db and ba132ba: the surrounding port used ds4_gpu_device_is_pre_m5_apple_silicon(), which matches M1 through M4, but these gates tested the marketing string. An M3 Ultra satisfies both, so the path was exercised there and nowhere else. Widen to the family predicate. The test then runs its byte-for-byte comparison for the first time on this hardware and passes with zero mismatches and zero ULP across all four cases: width 256, 512 and 1024, ratio 4 and 128, emit on and off, APE f32 and f16, decode_pack on and off, exact_reduce on and off. metal-kernels goes from 23 failures to OK. Expected to be throughput-neutral: the decode graph reaches this entry point only when quad_store == 0 (ds4.c:22701), and the pre-M5 quad store fuse already sets it above zero. If some path does reach it, the kernel is byte-exact by the test above and collapses dispatches, so the change is safe in either direction. DS4_METAL_DISABLE_PRE_M5_COMPRESSOR_PAIR_STORE=1 rolls it back.
Last of the device-name gates from this family. use_weights4 required
ds4_gpu_device_name_contains("M3") or the DS4_GPU_TEST_OUTPUT_HC_WEIGHTS4
test flag (ds4_metal.m:42759), so on Apple8 production took the multi-
dispatch fallback while the test took the fused kernel.
That split is what makes this one safe. Because the test flag arm is
family-independent, test_metal_output_hc_weights4_exact
(tests/ds4_test.c:3660) has been driving this exact kernel on Apple8 all
along and passing - the kernel's exactness on this hardware is measured,
not argued. Production and test reach the same code with the same buffers
and the same n_hc == 4, n_tokens == 1 shape; only the predicate that sets
use_weights4 differs.
Worth about 3 dispatches per token, so roughly 0.02 ms against a 24.27 ms
token at the 4.4-6 us per dispatch ba132ba measured. This is hygiene, not
a win; it is here because it is provable locally, which is the property
that separated the two changes that held this week from the two that did
not.
ds4_test --metal-kernels stays OK.
DS4_METAL_DISABLE_PRE_M5_OUTPUT_HC_WEIGHTS4=1 rolls it back.
…irez#770 Cherry-picks the hunks of antirez#770 (kk1987) that this branch had not already covered. That PR fixes the same defect found here independently - literal ds4_gpu_device_name_contains("M3")/("M5") whitelists where the surrounding port used ds4_gpu_device_is_pre_m5_apple_silicon() - and reports the same 29 forced-fusion exactness assertions failing, which is the count this branch saw go 29 -> 23 -> 0 across ba132ba and bb576ed. Of its fifteen sites, eight were already converted here (01a56db x4, ba132ba x2, bb576ed, 36d2758). Every one of the remaining seven is a prefill or batch path except shared kvpad, which is deliberately left out below. Six are taken here: 2012 zero-prefix prefill mask cache +4.5% prefill on M4 Max 20623 HC rms-scale-project, n_rows > 8 21869 compressor score + APE add 22677 concat f32 dim1 / ratio-4 pack 22723 ratio-4 direct pool 32096 router select batch weights, n_tokens > 1 +0.8% prefill This branch's decode audit had already classified all six as prefill-only and skipped them for that reason, so no decode change is expected; they are here for the prefill side, which nobody had swept. Taken as upstream wrote them, replacing the name test outright rather than adding an or-arm as the earlier commits here did, so the tree converges with antirez#770 if it merges. The eight sites already converted are left in their existing form: they are measured and working, and rewriting them for cosmetic convergence would risk a regression for no gain. Not applied: 27563, shared kvpad. It removes no dispatch, and ba132ba records the reason for leaving it M3-only - shared_pad stays 0 so the fused stage writes both the K and V pad copies, matching the standalone pad kernel. Upstream measured it in its neutral bucket, so there is nothing to gain against a layout hazard nobody has tested on Apple8. ds4_test --metal-kernels stays OK. Prefill effect unmeasured on this hardware; every site keeps its own DS4_METAL_DISABLE_* switch.
Distributed pipeline inference over a UC dual-QP-pair channel (ds4_dist_dchan): a full-duplex reliable byte stream that replaces the TCP data plane so pipelined prefill runs over RDMA. Opt in with --dist-transport rdma; --dist-rdma-adj-devices selects the local device per neighbour (a bare name for a routable fabric, or a peerhost=device map for point-to-point Thunderbolt links). librdma is loaded at runtime and the link bring-up (device/GID discovery, UC INIT->RTR->RTS) is shared with the tensor-parallel path in ds4_rdma_verbs; off Apple the RDMA files compile to TCP-only stubs. See Apple TN3205.
ryan5rdx
force-pushed
the
apple-rdma-pp
branch
from
August 17, 2026 04:57
962cdf9 to
f7d3837
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Supporting RDMA layer parallel mode:
The changeset is large here - but a lot of that is extracting shared RDMA logic from the tp module to a shared one.
sample worker command:
benchmarks:
2 nodes, TCP:
2 nodes, RDMA:
but the real benefit is at >2 nodes, where the TCP hops causes decode to collapse <<10t/s on M3 Ultra(benchmark takes forever). With RDMA 3 nodes:
3 nodes, RDMA: