GitHub - PJHkorea/branchless-nccl-router: A Branchless, Zero-Jitter Ingress Router for 32-GPU Distributed Mesh Networks utilizing JAX/XLA and NCCL. · GitHub

/" data-turbo-transient="true" />

Skip to content

Search or jump to...

Search code, repositories, users, issues, pull requests...

-->

Search

Clear

Search syntax tips

Provide feedback

--><br>We read every piece of feedback, and take your input very seriously.

Include my email address so I can be contacted

Cancel

Submit feedback

Saved searches

Use saved searches to filter your results more quickly

-->

Name

Query

To see all available qualifiers, see our documentation.

Cancel

Create saved search

Sign in

/;ref_cta:Sign up;ref_loc:header logged out"}"<br>Sign up

Appearance settings

Resetting focus

You signed in with another tab or window. Reload to refresh your session.<br>You signed out in another tab or window. Reload to refresh your session.<br>You switched accounts on another tab or window. Reload to refresh your session.

Dismiss alert

{{ message }}

PJHkorea

branchless-nccl-router

Public

Notifications<br>You must be signed in to change notification settings

Fork

Star

main

BranchesTags

Go to file

CodeOpen more actions menu

Folders and files<br>NameNameLast commit message<br>Last commit date<br>Latest commit

History<br>11 Commits<br>11 Commits

LICENSE

LICENSE

README.md

README.md

nccl_ingress_router_jax.py

nccl_ingress_router_jax.py

View all files

Repository files navigation

🌌 branchless-nccl-router (v1.0)

[KR] "분산 통신망의 규칙을 다른 방식으로 생각해 봤습니다. 제어 흐름(Control Flow) 분기를 없애고, 대규모 클러스터의 집합 통신과 패킷 정화를 단일 융합 기계어 수식 레일로 해결하려 합니다. 이것은 그것에 대한 아키텍처 방향성을 정립하기 위한 독창적인 청사진(Blueprint) 개념 실증 모델입니다."

[EN] "We re-engineered the principles of distributed networks from the ground up. By obliterating control flow branches, we channel multi-node collective communication and stream auditing through a single, fused hardware-aligned algebraic rail. This project serves as an original blueprint concept designed explicitly to establish and validate this architectural direction."

🛰️ Overview

[KR] 본 레포지토리는 다중 노드 분산 가속기 환경(Multi-Node Distributed Infrastructure)의 입구 경계면(Ingress Gate)에서 발생하는 제어 분기문(if/else)을 제거하고, 네트워크 동기화 지터(Synchronous Jitter) 비용을 물리적으로 0ns로 수축 시키기 위해 설계된 하드웨어 밀착형 분산 라우터 커널 모듈입니다.

[EN] This repository delivers a hardware-aligned distributed router kernel module, engineered to eliminate runtime control flow branches (if/else) at the ingress boundary of multi-node accelerator clusters, effectively shrinking network synchronous jitter overheads to exactly 0ns .

1. jax.lax.psum과 제어문(if)의 거부 (Zero-Jitter Collective Op)

[KR] 다중 노드 분산 학습/추론 환경에서 엔지니어들을 가장 괴롭히는 난제는 "특정 노드가 패킷을 다 받았는지 확인하기 위해 멈춰 서는 하드웨어 동기화 펜스(Fence Stall) 병목"입니다. 본 커널은 네트워크 패킷 유실 및 오염 여부를 판정하는 if (is_packet_corrupted) 같은 제어문을 제거했습니다. 대신 JAX/XLA 백엔드에 상주하는 NCCL All-Reduce 코어 프리미티브인 jax.lax.psum 연산 결과를 불린이 아닌 단정밀도 부동소수점 마스크 수식 (global_sync_mask == 0.0).astype(target_dtype) 내부의 피연산자로 직결했습니다. 32개 디바이스는 네트워크 전송 상태나 지터 발생 여부와 무관하게, 물리 NCCL 링(Ring) 레일 위를 동일한 속도로 움직이게 됩니다.

[EN] The most critical bottleneck in large-scale distributed scaling is the hardware synchronous fence (Fence Stall), where accelerators halt execution waiting for straggler nodes to complete packet tracking. This kernel eradicates runtime control flow branches like if (is_packet_corrupted). Instead, it directly channels the output of jax.lax.psum—the underlying core primitive for NCCL All-Reduce—as a hardware arithmetic operand using the float-mask equation (global_sync_mask == 0.0).astype(target_dtype). Consequently, all 32 distributed devices stream through the physical NCCL Ring topology at a perfectly uniform velocity, thoroughly decoupled from dynamic network fluctuations or transport jitters.

# Eradicate control branches and stream directly into the algebraic mask circuit<br>global_sync_mask = jax.lax.psum(is_packet_corrupted, axis_name=cluster_axis_name)<br>is_mesh_clean = (global_sync_mask == literal_zero).astype(target_dtype)<br>next_jitter_flag = network_jitter_flag * is_mesh_clean

2. 다중 노드 SRAM Spill Over 방어형 클로저 스캔 (SRAM Register Locking)

[KR] 32개 디바이스의 대규모 패킷 스트림 데이터를 루프 단계마다 통째로 들고 이동(Carry)하면 가속기는 고속 온칩 메모리 용량 부족으로 데이터를 글로벌 VRAM으로 빼두는 레지스터 스필오버(Register Spill) 병목이 생깁니다. 본 엔진은 입력 스트림 행렬을 중첩 함수의 클로저(Closure) 영역에 박아두고 컴파일러 단에 정적 상수로 락을 걸었습니다. 그리고 오직 jnp.take(..., dev_idx, axis=0)라는 고속 가상 포인터 인덱싱 주소 변환만으로 데이터를 수집하도록 강제하여, XLA 컴파일러가 32개 분산 영토의 Stride Map 뼈대를 SRAM 레지스터 맵에 정적으로 고정 캐싱하도록 유도했습니다. 가변 패킷 유입 속에서도 메모리 버스 대역폭 낭비를 제거하는 것이 목표입니다.

[EN] Passing massive distributed packet streams as dynamic parameters (Carry) across loop boundaries triggers catastrophic SRAM register spilling, forcing the accelerator to offload intermediate activations to high-latency global VRAM. This engine isolates the input stream matrix entirely within a nested function closure, locking it as a compile-time static constant layout. By strictly enforcing target collections via jnp.take(..., dev_idx, axis=0)—a high-speed virtual pointer lookup—the...