Skip to content

Fused ring RS: per-tile bitmask + drain-ring with owner-boundary flus… #1

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

Open
OpenDarrenlu wants to merge 1 commit into
base: feat/sm100-bf16-gemm-reduce-scatter
Choose a base branch
from
Open

Fused ring RS: per-tile bitmask + drain-ring with owner-boundary flus… #1

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
13 changes: 9 additions & 4 deletions csrc/jit_kernels/impls/sm100_bf16_gemm_reduce_scatter_ring.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,7 +90,8 @@ static void __instantiate_kernel() {{
//
// `out_sym_buffer` symmetric layout (identical on every rank):
// [ barrier (32 B) ][ BF16 output: m_per_rank*n ][ BF16 ring recv: R * m_per_rank * n ]
// [ int flags: num_m_blocks * num_n_blocks ]
// [ mask: R * ceil(seg_tiles/64) * uint64_t ] (per-tile ready bitmask; R = num_ranks,
// seg_tiles = (m_per_rank/BLOCK_M)*(n/BLOCK_N))
// `sym_buffer_ptrs`: all peers' base pointers into that symmetric allocation.
static void sm100_bf16_gemm_reduce_scatter_ring(const torch::Tensor& a,
const torch::Tensor& b,
Expand Down Expand Up @@ -126,7 +127,7 @@ static void sm100_bf16_gemm_reduce_scatter_ring(const torch::Tensor& a,

// Single-CTA NON-swap layout: token(M)=UMMA_M=128, hidden(N)=UMMA_N=BLOCK_N=128. BLOCK_M LOCKED
// to 128 (owner-aligned token tile; m_per_rank padded to 128 by the wrapper). Standard (non-
// transposed) epilogue -> no STSM_T. Whole [128,128] tile per store + one segment flag.
// transposed) epilogue -> no STSM_T. Whole [128,128] tile per store + per-tile bitmask signaling.
const int chosen_block_m = 128;
const auto all_candidates = SM100ArchSpec::get_layout_candidates(desc);
std::vector<Layout> candidates;
Expand Down Expand Up @@ -165,8 +166,12 @@ static void sm100_bf16_gemm_reduce_scatter_ring(const torch::Tensor& a,
// [ partial_buf (part_stages of [128,128]) ][ A: ns ][ B: ns ]
// [ barriers: (2*ns + tmem*2 + part*3)*8 ][ tmem_ptr: 4 ]. NO epi_smem.
// Stage counts — MUST match the kernel (.cuh): kNumEpilogueStages=4 (TMEM), kNumPartialStages=2.
// WG1 adds upstream in-registers and writes the running-sum into partial_buf; partial_buf has 3
// barrier groups (part_full w3->WG1, epi_full WG1->WG2, part_empty WG2->w3).
// part_stages=2 measured best: it keeps the A/B pipeline at ns=4 (part_stages=3 drops ns to
// 3, part_stages=4 to 2 — both net losses on bench), while WG2's drain-ring still hides one
// store's NVLink RTT per tile.
// NOTE: this value is compiled INTO _C.so — after changing it you MUST rebuild the extension
// (`touch csrc/python_api.cpp && python setup.py build_ext --inplace`), otherwise the JIT
// allocates smem for the OLD stage count while the kernel uses the NEW layout → IMA.
const int tmem_stages = 4;
const int part_stages = 2;
{
Expand Down
21 changes: 14 additions & 7 deletions deep_gemm/comm/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -353,13 +353,16 @@ class ReduceScatterRingFusedBuffer:

Layout (single symmetric allocation, shared across ranks over NVLink):
[ barrier region (32 B) ]
[ output : m_per_rank_pad x n bf16 ] (END writes here)
[ ring recv: num_ranks x m_per_rank_pad x n bf16 ] (upstream running-sum, per owner)
[ flags : num_m_blocks x num_n_blocks int32 ] (per-tile ready flag)
[ output : m_per_rank_pad x n bf16 ] (END writes here)
[ ring recv: num_ranks x m_per_rank_pad x n bf16 ] (upstream running-sum, per owner)
[ mask : num_ranks x ceil(seg_tiles/64) x uint64_t ] (per-tile ready bitmask)

swap-AB padding mirrors `ReduceScatterBuffer`: token(M) is the UMMA N-dim tiled by
BLOCK_M=`rs_block_m(m_per_rank)`; each owner segment is padded to a multiple of BLOCK_M so
every tile is single-owner. `num_m_blocks = m_pad / BLOCK_M`, `num_n_blocks = n / 128`.
`seg_tiles = num_m_blocks / num_ranks * num_n_blocks` — each tile owns one bit in its
segment's uint64_t bitmask; upstream (WG2 on rank i+1) sets bits per-tile via
`red.release.sys.or`, downstream (rank i, w3) polls with `ld.acquire.sys`.
"""

BARRIER_BYTES = 32
Expand All @@ -383,9 +386,13 @@ def __init__(self, group: 'dist.ProcessGroup', m: int, n: int):
shard_elems = self.m_per_rank_pad * n
self._out_bytes = shard_elems * torch.bfloat16.itemsize
ring_bytes = self.world_size * shard_elems * torch.bfloat16.itemsize
# Per-OWNER-SEGMENT flag protocol: R flags + R device-scope counters (int32 each).
flag_bytes = 2 * self.world_size * 4
num_bytes = self.BARRIER_BYTES + self._out_bytes + ring_bytes + flag_bytes
# Per-TILE bitmask protocol: one bit per tile in the owner segment, uint64_t words.
# Kernel does `red.release.sys.or` per-tile; downstream polls `ld.acquire.sys`.
bpo = self.m_per_rank_pad // self.block_m # blocks per owner
seg_tiles = bpo * self.num_n_blocks # tiles per owner segment
words_per_seg = (seg_tiles + 63) // 64
mask_bytes = self.world_size * words_per_seg * 8 # uint64_t per word
num_bytes = self.BARRIER_BYTES + self._out_bytes + ring_bytes + mask_bytes

self.buffer = symm_mem.empty(num_bytes, dtype=torch.int8, device='cuda')
self.handle = symm_mem.rendezvous(self.buffer, group=group)
Expand All @@ -404,7 +411,7 @@ def output(self) -> torch.Tensor:
return padded[:self.m_per_rank]

def zero_(self):
# Ring + flags must be cleared every call (the kernel also self-resets flags at entry,
# Ring + mask must be cleared every call (the kernel also self-resets the mask at entry,
# but a full zero keeps the ring recv region clean for repeated launches).
self.buffer.zero_()
self.group.barrier()
Expand Down
Loading