| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Deadlock-free fused RS + short-conv + AG + residual + RMSNorm. |
| 4 | |
| 5 | The public integration surface is ``LamportRSConv.rs_sconv_ag_add_norm``. |
| 6 | |
| 7 | Liveness |
| 8 | -------- |
| 9 | Large grids are deliberately split at the two communication dependencies: |
| 10 | |
| 11 | 1. ``_publish_input_kernel`` only publishes rank partials. |
| 12 | 2. ``_reduce_insert_kernel`` waits, reduces, and inserts the local shard. |
| 13 | 3. ``_sconv_publish_kernel`` only computes and publishes the local result. |
| 14 | 4. ``_gather_norm_kernel`` waits, gathers, and normalizes. |
| 15 | |
| 16 | Without PDL, CUDA stream order completes a producer before its consumer. With |
| 17 | PDL, every producer CTA posts all peer stores before triggering its dependent, |
| 18 | and the consumer executes ``gdc_wait`` before polling. A consumer therefore |
| 19 | waits only for stores from a producer that is already running (or complete) on |
| 20 | another GPU. It never waits for another block in its own grid. Consequently |
| 21 | spinning consumers cannot occupy resources needed by any producer, and the |
| 22 | wait-for graph has no cycle. The proof is independent of grid size, block |
| 23 | dispatch order, and occupancy. |
| 24 | |
| 25 | For one token, the first three phases use eight independent channel slices to |
| 26 | expose enough CTA parallelism for decode latency. Each slice has exclusive |
| 27 | ownership of its cache and Lamport columns; the gather/RMSNorm phase retains |
| 28 | one CTA per token so no cross-CTA reduction or completion counter is needed. |
| 29 | |
| 30 | Immediate buffer reuse (including replay of a captured CUDA graph) is also |
| 31 | safe. Rank R cannot republish input for call n+1 until its gather for n has |
| 32 | finished; that gather waited for owner O's output, which O publishes only |
| 33 | after consuming R's input for n. Likewise, R cannot republish output for n+1 |
| 34 | until its reduction for n+1 has observed destination D's input; D publishes |
| 35 | that input only after its gather consumed R's output for n. Thus every prior |
| 36 | read happens-before a same-slot rewrite. Three generations reduce incidental |
| 37 | coupling for ordinary launches, but correctness does not depend on rotation. |
| 38 | |
| 39 | The payload itself is the Lamport flag. A 32-bit store publishes two bf16s |
| 40 | atomically; 0x80008000 (two negative zeroes) denotes an empty pair. Real |
| 41 | negative zeroes are changed to positive zero before publication. Consumers |
| 42 | use volatile 32-bit loads and restore the sentinel after consuming a slot. |
| 43 | """ |
| 44 | |
| 45 | from __future__ import annotations |
| 46 | |
| 47 | import os |
| 48 | |
| 49 | import torch |
| 50 | |
| 51 | from vllm.distributed import get_tp_group |
| 52 | from vllm.distributed.parallel_state import in_the_same_node_as |
| 53 | from vllm.logger import init_logger |
| 54 | from vllm.triton_utils import HAS_TRITON, tl, triton |
| 55 | |
| 56 | logger = init_logger(__name__) |
| 57 | |
| 58 | _MAX_TOKENS = 16384 |
| 59 | _EMPTY_PAIR = tl.constexpr(0x80008000) if HAS_TRITON else 0x80008000 |
| 60 | |
| 61 | |
| 62 | @triton.jit |
| 63 | def _pack_bf16_pairs(values): |
| 64 | """Pack bf16 values into atomic u32 pairs and reserve negative zero.""" |
| 65 | lo, hi = tl.split(values.reshape([values.shape[0] // 2, 2])) |
| 66 | lo = lo.to(tl.uint16, bitcast=True) |
| 67 | hi = hi.to(tl.uint16, bitcast=True) |
| 68 | lo = tl.where(lo == 0x8000, 0, lo).to(tl.uint32) |
| 69 | hi = tl.where(hi == 0x8000, 0, hi).to(tl.uint32) |
| 70 | return lo | (hi << 16) |
| 71 | |
| 72 | |
| 73 | @triton.jit |
| 74 | def _unpack_bf16_pairs(values): |
| 75 | lo = (values & 0xFFFF).to(tl.uint16).to(tl.bfloat16, bitcast=True) |
| 76 | hi = (values >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True) |
| 77 | return tl.interleave(lo, hi) |
| 78 | |
| 79 | |
| 80 | @triton.jit |
| 81 | def _wait_pairs(ptr, offsets, mask): |
| 82 | values = tl.load(ptr + offsets, mask=mask, other=0, volatile=True) |
| 83 | while tl.max(tl.where(mask & (values == _EMPTY_PAIR), 1, 0)) != 0: |
| 84 | values = tl.load(ptr + offsets, mask=mask, other=0, volatile=True) |
| 85 | return values |
| 86 | |
| 87 | |
| 88 | @triton.jit |
| 89 | def _publish_input_kernel( |
| 90 | stage_ptr, |
| 91 | shared_ptr, |
| 92 | peer_ptrs, |
| 93 | peer_offset_u32, |
| 94 | stride_stage_t, |
| 95 | stride_shared_t, |
| 96 | C: tl.constexpr, |
| 97 | CS: tl.constexpr, |
| 98 | CS_P2: tl.constexpr, |
| 99 | SPLITS: tl.constexpr, |
| 100 | RANK: tl.constexpr, |
| 101 | WORLD: tl.constexpr, |
| 102 | HAS_SHARED: tl.constexpr, |
| 103 | USE_PDL: tl.constexpr, |
| 104 | launch_pdl: tl.constexpr, |
| 105 | ): |
| 106 | """Publish this rank's full partial row into every shard owner's slots.""" |
| 107 | token = tl.program_id(0).to(tl.int64) |
| 108 | split = tl.program_id(1).to(tl.int64) |
| 109 | CSS: tl.constexpr = CS // SPLITS |
| 110 | pair = tl.arange(0, CS_P2 // 2) |
| 111 | pair_mask = pair < CSS // 2 |
| 112 | elem = tl.arange(0, CS_P2) |
| 113 | elem_mask = elem < CSS |
| 114 | ptrs = peer_ptrs.to(tl.pointer_type(tl.uint64)) |
| 115 | # Address generation is independent of the preceding stream kernel. The |
| 116 | # acquire remains before stage/peer loads, which may consume its writes. |
| 117 | if USE_PDL: |
| 118 | tl.extra.cuda.gdc_wait() |
| 119 | |
| 120 | for owner in tl.static_range(WORLD): |
| 121 | values = tl.load( |
| 122 | stage_ptr + token * stride_stage_t + owner * CS + split * CSS + elem, |
| 123 | mask=elem_mask, |
| 124 | other=0.0, |
| 125 | ) |
| 126 | if HAS_SHARED: |
| 127 | shared = tl.load( |
| 128 | shared_ptr + token * stride_shared_t + owner * CS + split * CSS + elem, |
| 129 | mask=elem_mask, |
| 130 | other=0.0, |
| 131 | ) |
| 132 | values = (values.to(tl.float32) + shared.to(tl.float32)).to(tl.bfloat16) |
| 133 | packed = _pack_bf16_pairs(values) |
| 134 | base = tl.load(ptrs + owner).to(tl.pointer_type(tl.uint32)) |
| 135 | dst = (token * WORLD + RANK) * (CS // 2) + split * (CSS // 2) + pair |
| 136 | tl.store(base + peer_offset_u32 + dst, packed, mask=pair_mask) |
| 137 | if USE_PDL: |
| 138 | tl.extra.cuda.gdc_launch_dependents() |
| 139 | |
| 140 | |
| 141 | @triton.jit |
| 142 | def _reduce_insert_kernel( |
| 143 | input_peer_ptrs, |
| 144 | input_peer_offset_u32, |
| 145 | cache_ptr, |
| 146 | slot_ptr, |
| 147 | stride_cache_block, |
| 148 | stride_cache_head, |
| 149 | stride_cache_token, |
| 150 | stride_cache_dim, |
| 151 | block_size, |
| 152 | C: tl.constexpr, |
| 153 | CS: tl.constexpr, |
| 154 | CS_P2: tl.constexpr, |
| 155 | SPLITS: tl.constexpr, |
| 156 | HEAD_SIZE: tl.constexpr, |
| 157 | CACHE_OFFSET: tl.constexpr, |
| 158 | RANK: tl.constexpr, |
| 159 | WORLD: tl.constexpr, |
| 160 | USE_PDL: tl.constexpr, |
| 161 | launch_pdl: tl.constexpr, |
| 162 | ): |
| 163 | """Consume all partials for this rank and insert the reduced cache row.""" |
| 164 | token = tl.program_id(0).to(tl.int64) |
| 165 | split = tl.program_id(1).to(tl.int64) |
| 166 | CSS: tl.constexpr = CS // SPLITS |
| 167 | source = tl.arange(0, WORLD) |
| 168 | pair = tl.arange(0, CS_P2 // 2) |
| 169 | pair_mask = pair < CSS // 2 |
| 170 | offsets = ( |
| 171 | (token * WORLD + source)[:, None] * (CS // 2) |
| 172 | + split * (CSS // 2) |
| 173 | + pair[None, :] |
| 174 | ) |
| 175 | mask = tl.full([WORLD], True, tl.int1)[:, None] & pair_mask[None, :] |
| 176 | input_ptrs = input_peer_ptrs.to(tl.pointer_type(tl.uint64)) |
| 177 | input_u32 = tl.load(input_ptrs + RANK).to(tl.pointer_type(tl.uint32)) |
| 178 | input_u32 += input_peer_offset_u32 |
| 179 | # Slot metadata and the cache destination do not depend on input publish. |
| 180 | slot = tl.load(slot_ptr + token) |
| 181 | valid = slot >= 0 |
| 182 | channel = tl.arange(0, CS_P2) |
| 183 | channel_mask = channel < CSS |
| 184 | global_channel = split * CSS + channel |
| 185 | head = tl.minimum(global_channel // HEAD_SIZE, CS // HEAD_SIZE - 1) |
| 186 | dim = CACHE_OFFSET + global_channel % HEAD_SIZE |
| 187 | safe_slot = tl.maximum(slot, 0).to(tl.int64) |
| 188 | dst = ( |
| 189 | cache_ptr |
| 190 | + (safe_slot // block_size) * stride_cache_block |
| 191 | + head * stride_cache_head |
| 192 | + (safe_slot % block_size) * stride_cache_token |
| 193 | + dim * stride_cache_dim |
| 194 | ) |
| 195 | if USE_PDL: |
| 196 | tl.extra.cuda.gdc_wait() |
| 197 | packed = _wait_pairs(input_u32, offsets, mask) |
| 198 | |
| 199 | lo = (packed & 0xFFFF).to(tl.uint16).to(tl.bfloat16, bitcast=True) |
| 200 | hi = (packed >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True) |
| 201 | reduced = tl.interleave( |
| 202 | tl.sum(lo.to(tl.float32), axis=0).to(tl.bfloat16), |
| 203 | tl.sum(hi.to(tl.float32), axis=0).to(tl.bfloat16), |
| 204 | ) |
| 205 | tl.store( |
| 206 | input_u32 + offsets, |
| 207 | tl.full([WORLD, CS_P2 // 2], _EMPTY_PAIR, tl.uint32), |
| 208 | mask=mask, |
| 209 | ) |
| 210 | |
| 211 | tl.store(dst, reduced, mask=valid & channel_mask) |
| 212 | if USE_PDL: |
| 213 | tl.extra.cuda.gdc_launch_dependents() |
| 214 | |
| 215 | |
| 216 | @triton.jit |
| 217 | def _sconv_publish_kernel( |
| 218 | peer_ptrs, |
| 219 | peer_offset_u32, |
| 220 | residual_ptr, |
| 221 | weight_ptr, |
| 222 | cache_ptr, |
| 223 | position_ptr, |
| 224 | sequence_ptr, |
| 225 | slot_ptr, |
| 226 | block_table_ptr, |
| 227 | stride_residual_t, |
| 228 | stride_cache_block, |
| 229 | stride_cache_head, |
| 230 | stride_cache_token, |
| 231 | stride_cache_dim, |
| 232 | stride_block_table_r, |
| 233 | max_blocks, |
| 234 | block_size, |
| 235 | C: tl.constexpr, |
| 236 | CS: tl.constexpr, |
| 237 | CS_P2: tl.constexpr, |
| 238 | SPLITS: tl.constexpr, |
| 239 | HEAD_SIZE: tl.constexpr, |
| 240 | CACHE_OFFSET: tl.constexpr, |
| 241 | RANK: tl.constexpr, |
| 242 | WORLD: tl.constexpr, |
| 243 | WINDOW: tl.constexpr, |
| 244 | USE_PDL: tl.constexpr, |
| 245 | launch_pdl: tl.constexpr, |
| 246 | ): |
| 247 | """Compute this rank's shard, then publish it to every rank.""" |
| 248 | tl.static_assert(WINDOW == 4) |
| 249 | token = tl.program_id(0).to(tl.int64) |
| 250 | split = tl.program_id(1).to(tl.int64) |
| 251 | CSS: tl.constexpr = CS // SPLITS |
| 252 | channel = tl.arange(0, CS_P2) |
| 253 | channel_mask = channel < CSS |
| 254 | global_channel = split * CSS + channel |
| 255 | head = tl.minimum(global_channel // HEAD_SIZE, CS // HEAD_SIZE - 1) |
| 256 | dim = CACHE_OFFSET + global_channel % HEAD_SIZE |
| 257 | slot = tl.load(slot_ptr + token) |
| 258 | valid = slot >= 0 |
| 259 | position = tl.load(position_ptr + token) |
| 260 | sequence = tl.load(sequence_ptr + token) |
| 261 | safe_slot = tl.maximum(slot, 0).to(tl.int64) |
| 262 | own_ptr = ( |
| 263 | cache_ptr |
| 264 | + (safe_slot // block_size) * stride_cache_block |
| 265 | + head * stride_cache_head |
| 266 | + (safe_slot % block_size) * stride_cache_token |
| 267 | + dim * stride_cache_dim |
| 268 | ) |
| 269 | # These loads were made visible before reduce/insert could signal us and |
| 270 | # are independent of its cache store. Hoisting them gives PDL useful work |
| 271 | # to overlap while retaining the acquire before every cache read. |
| 272 | residual = tl.load( |
| 273 | residual_ptr + token * stride_residual_t + RANK * CS + global_channel, |
| 274 | mask=channel_mask, |
| 275 | other=0.0, |
| 276 | ) |
| 277 | weight0 = tl.load( |
| 278 | weight_ptr + global_channel * WINDOW, |
| 279 | mask=channel_mask, |
| 280 | other=0.0, |
| 281 | ).to(tl.float32) |
| 282 | weight1 = tl.load( |
| 283 | weight_ptr + global_channel * WINDOW + 1, |
| 284 | mask=channel_mask, |
| 285 | other=0.0, |
| 286 | ).to(tl.float32) |
| 287 | weight2 = tl.load( |
| 288 | weight_ptr + global_channel * WINDOW + 2, |
| 289 | mask=channel_mask, |
| 290 | other=0.0, |
| 291 | ).to(tl.float32) |
| 292 | weight3 = tl.load( |
| 293 | weight_ptr + global_channel * WINDOW + 3, |
| 294 | mask=channel_mask, |
| 295 | other=0.0, |
| 296 | ).to(tl.float32) |
| 297 | ptrs = peer_ptrs.to(tl.pointer_type(tl.uint64)) |
| 298 | if USE_PDL: |
| 299 | tl.extra.cuda.gdc_wait() |
| 300 | current = tl.load(own_ptr, mask=valid & channel_mask, other=0.0) |
| 301 | |
| 302 | conv = tl.zeros([CS_P2], tl.float32) |
| 303 | for tap_idx in tl.static_range(WINDOW): |
| 304 | source_position = position - (WINDOW - 1) + tap_idx |
| 305 | take = valid & (source_position >= 0) |
| 306 | if tap_idx == WINDOW - 1: |
| 307 | value = tl.where(take, current.to(tl.float32), 0.0) |
| 308 | else: |
| 309 | safe_position = tl.maximum(source_position, 0) |
| 310 | logical_block = tl.minimum(safe_position // block_size, max_blocks - 1) |
| 311 | physical_block = tl.load( |
| 312 | block_table_ptr + sequence * stride_block_table_r + logical_block, |
| 313 | mask=take, |
| 314 | other=0, |
| 315 | ).to(tl.int64) |
| 316 | source_ptr = ( |
| 317 | cache_ptr |
| 318 | + physical_block * stride_cache_block |
| 319 | + head * stride_cache_head |
| 320 | + (safe_position % block_size) * stride_cache_token |
| 321 | + dim * stride_cache_dim |
| 322 | ) |
| 323 | cached = tl.load(source_ptr, mask=take & channel_mask, other=0.0) |
| 324 | value = tl.where(take, cached.to(tl.float32), 0.0) |
| 325 | if tap_idx == 0: |
| 326 | weight = weight0 |
| 327 | elif tap_idx == 1: |
| 328 | weight = weight1 |
| 329 | elif tap_idx == 2: |
| 330 | weight = weight2 |
| 331 | else: |
| 332 | weight = weight3 |
| 333 | conv += value * weight |
| 334 | |
| 335 | # Preserve both bf16 rounding points of the original sublayer. |
| 336 | short_conv_with_skip = (conv + current.to(tl.float32)).to(tl.bfloat16) |
| 337 | output = (residual.to(tl.float32) + short_conv_with_skip.to(tl.float32)).to( |
| 338 | tl.bfloat16 |
| 339 | ) |
| 340 | |
| 341 | pair = tl.arange(0, CS_P2 // 2) |
| 342 | pair_mask = pair < CSS // 2 |
| 343 | packed = _pack_bf16_pairs(output) |
| 344 | row_offset = token * (C // 2) + RANK * (CS // 2) + split * (CSS // 2) + pair |
| 345 | |
| 346 | for destination in tl.static_range(WORLD): |
| 347 | base = tl.load(ptrs + destination).to(tl.pointer_type(tl.uint32)) |
| 348 | tl.store(base + peer_offset_u32 + row_offset, packed, mask=pair_mask) |
| 349 | if USE_PDL: |
| 350 | tl.extra.cuda.gdc_launch_dependents() |
| 351 | |
| 352 | |
| 353 | @triton.jit |
| 354 | def _gather_norm_kernel( |
| 355 | output_peer_ptrs, |
| 356 | output_peer_offset_u32, |
| 357 | norm_weight_ptr, |
| 358 | normed_ptr, |
| 359 | residual_out_ptr, |
| 360 | eps, |
| 361 | stride_output_t, |
| 362 | C: tl.constexpr, |
| 363 | C_P2: tl.constexpr, |
| 364 | RANK: tl.constexpr, |
| 365 | HAS_NORM: tl.constexpr, |
| 366 | USE_PDL: tl.constexpr, |
| 367 | launch_pdl: tl.constexpr, |
| 368 | ): |
| 369 | """Consume a complete row; one CTA owns all outputs for one token.""" |
| 370 | token = tl.program_id(0).to(tl.int64) |
| 371 | pair = tl.arange(0, C_P2 // 2) |
| 372 | pair_mask = pair < C // 2 |
| 373 | output_ptrs = output_peer_ptrs.to(tl.pointer_type(tl.uint64)) |
| 374 | output_u32 = tl.load(output_ptrs + RANK).to(tl.pointer_type(tl.uint32)) |
| 375 | output_u32 += output_peer_offset_u32 |
| 376 | offsets = token * (C // 2) + pair |
| 377 | channel = tl.arange(0, C_P2) |
| 378 | channel_mask = channel < C |
| 379 | # Gamma is independent of the preceding sconv publication. Keep the |
| 380 | # acquire immediately before polling the Lamport output slots. |
| 381 | if HAS_NORM: |
| 382 | weight = tl.load(norm_weight_ptr + channel, mask=channel_mask, other=0.0) |
| 383 | if USE_PDL: |
| 384 | tl.extra.cuda.gdc_wait() |
| 385 | packed = _wait_pairs(output_u32, offsets, pair_mask) |
| 386 | row = _unpack_bf16_pairs(packed) |
| 387 | tl.store( |
| 388 | residual_out_ptr + token * stride_output_t + channel, row, mask=channel_mask |
| 389 | ) |
| 390 | if HAS_NORM: |
| 391 | row_f32 = tl.where(channel_mask, row.to(tl.float32), 0.0) |
| 392 | inv_rms = tl.rsqrt(tl.sum(row_f32 * row_f32, axis=0) / C + eps) |
| 393 | tl.store( |
| 394 | normed_ptr + token * stride_output_t + channel, |
| 395 | (row_f32 * inv_rms * weight.to(tl.float32)).to(tl.bfloat16), |
| 396 | mask=channel_mask, |
| 397 | ) |
| 398 | tl.store( |
| 399 | output_u32 + offsets, |
| 400 | tl.full([C_P2 // 2], _EMPTY_PAIR, tl.uint32), |
| 401 | mask=pair_mask, |
| 402 | ) |
| 403 | if USE_PDL: |
| 404 | tl.extra.cuda.gdc_launch_dependents() |
| 405 | |
| 406 | |
| 407 | @triton.jit |
| 408 | def _validate_lamport_init_kernel( |
| 409 | input_peer_ptrs, |
| 410 | output_peer_ptrs, |
| 411 | bad_ptr, |
| 412 | num_pairs, |
| 413 | RANK: tl.constexpr, |
| 414 | BLOCK: tl.constexpr, |
| 415 | ): |
| 416 | """Validate both complete local allocations through their fabric pointers.""" |
| 417 | offsets = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) |
| 418 | mask = offsets < num_pairs |
| 419 | ptrs_in = input_peer_ptrs.to(tl.pointer_type(tl.uint64)) |
| 420 | ptrs_out = output_peer_ptrs.to(tl.pointer_type(tl.uint64)) |
| 421 | local_in = tl.load(ptrs_in + RANK).to(tl.pointer_type(tl.uint32)) |
| 422 | local_out = tl.load(ptrs_out + RANK).to(tl.pointer_type(tl.uint32)) |
| 423 | value_in = tl.load(local_in + offsets, mask=mask, other=_EMPTY_PAIR) |
| 424 | value_out = tl.load(local_out + offsets, mask=mask, other=_EMPTY_PAIR) |
| 425 | bad = tl.max( |
| 426 | tl.where(mask & ((value_in != _EMPTY_PAIR) | (value_out != _EMPTY_PAIR)), 1, 0) |
| 427 | ) |
| 428 | if bad != 0: |
| 429 | tl.atomic_max(bad_ptr, 1) |
| 430 | |
| 431 | |
| 432 | class LamportRSConv: |
| 433 | """Persistent symmetric buffers for one TP group.""" |
| 434 | |
| 435 | def _initialize_lamport_buffers(self) -> None: |
| 436 | """Arm every slot and collectively verify the exact sentinel bits. |
| 437 | |
| 438 | Do not replace this with a zero-fill: Lamport readiness distinguishes |
| 439 | the bf16 bit pattern for -0.0 (0x8000) from every published payload. |
| 440 | The validation is deliberately collective so one rank cannot enter the |
| 441 | first polling kernel while another rank still has an unarmed buffer. |
| 442 | """ |
| 443 | if not self.buf_in.is_contiguous() or not self.buf_out.is_contiguous(): |
| 444 | raise RuntimeError("Lamport symmetric buffers must be contiguous") |
| 445 | if self.buf_in.numel() % 2 or self.buf_out.numel() % 2: |
| 446 | raise RuntimeError("Lamport buffers must contain whole uint32 pairs") |
| 447 | |
| 448 | # Fill through int16 so each bf16 lane receives the exact -0.0 bits. |
| 449 | # This covers all three generations and all max-token slots, including |
| 450 | # slots that a smaller first invocation does not touch. |
| 451 | self.buf_in.view(torch.int16).fill_(-0x8000) |
| 452 | self.buf_out.view(torch.int16).fill_(-0x8000) |
| 453 | torch.accelerator.synchronize(self.device) |
| 454 | self.tp.barrier() |
| 455 | |
| 456 | # Scan the complete local allocations. Since every rank performs the |
| 457 | # scan and participates in MAX, success proves every symmetric backing |
| 458 | # allocation was armed before any rank is allowed to use generation 0. |
| 459 | bad = torch.logical_or( |
| 460 | self.buf_in.view(torch.int16).ne(-0x8000).any(), |
| 461 | self.buf_out.view(torch.int16).ne(-0x8000).any(), |
| 462 | ).to(dtype=torch.float32) |
| 463 | bad = self.tp.all_reduce(bad) |
| 464 | if int(bad.item()) != 0: |
| 465 | raise RuntimeError("Lamport sentinel initialization failed on a TP rank") |
| 466 | self.tp.barrier() |
| 467 | |
| 468 | def _initialize_mnnvl_buffers(self) -> None: |
| 469 | """Initialize and validate FlashInfer fabric-mapped Lamport storage.""" |
| 470 | self._mnnvl_input_handle.lamport_initialize(self.rank, torch.bfloat16) |
| 471 | self._mnnvl_output_handle.lamport_initialize(self.rank, torch.bfloat16) |
| 472 | torch.accelerator.synchronize(self.device) |
| 473 | self.tp.barrier() |
| 474 | |
| 475 | bad = torch.zeros((), dtype=torch.int32, device=self.device) |
| 476 | num_pairs = self.num_buffers * self.max_tokens * self.hidden_size // 2 |
| 477 | _validate_lamport_init_kernel[(triton.cdiv(num_pairs, 256),)]( |
| 478 | self.input_peer_ptrs, |
| 479 | self.output_peer_ptrs, |
| 480 | bad, |
| 481 | num_pairs, |
| 482 | RANK=self.rank, |
| 483 | BLOCK=256, |
| 484 | num_warps=4, |
| 485 | ) |
| 486 | bad = self.tp.all_reduce(bad.to(torch.float32)) |
| 487 | if int(bad.item()) != 0: |
| 488 | raise RuntimeError("MNNVL Lamport sentinel initialization failed") |
| 489 | self.tp.barrier() |
| 490 | |
| 491 | def __init__( |
| 492 | self, hidden_size: int, window_size: int, max_tokens: int = _MAX_TOKENS |
| 493 | ) -> None: |
| 494 | import torch.distributed._symmetric_memory as symm_mem |
| 495 | |
| 496 | tp = get_tp_group() |
| 497 | self.tp = tp |
| 498 | self.group = tp.device_group |
| 499 | self.world_size = tp.world_size |
| 500 | self.rank = tp.rank_in_group |
| 501 | self.device = torch.device(tp.device) |
| 502 | if self.world_size not in (2, 4, 8): |
| 503 | raise ValueError(f"TP world size must be 2, 4, or 8, got {self.world_size}") |
| 504 | if hidden_size % (2 * self.world_size) != 0: |
| 505 | raise ValueError("hidden size must produce an even shard on every rank") |
| 506 | if window_size != 4: |
| 507 | raise ValueError(f"short-conv window size must be 4, got {window_size}") |
| 508 | if max_tokens < 1 or max_tokens > _MAX_TOKENS: |
| 509 | raise ValueError(f"max_tokens must be in [1, {_MAX_TOKENS}]") |
| 510 | |
| 511 | is_cross_node = not all(in_the_same_node_as(tp.cpu_group)) |
| 512 | self.hidden_size = hidden_size |
| 513 | self.window_size = window_size |
| 514 | self.max_tokens = max_tokens |
| 515 | self.shard_size = hidden_size // self.world_size |
| 516 | # Three generations follow FlashInfer's Lamport layout. A generation |
| 517 | # is reused only after two intervening collective calls. |
| 518 | self.num_buffers = 3 |
| 519 | self.input_generation_bytes = max_tokens * hidden_size * 2 |
| 520 | self.output_generation_bytes = max_tokens * hidden_size * 2 |
| 521 | self.use_pdl = torch.cuda.get_device_capability(self.device)[0] >= 9 |
| 522 | if is_cross_node: |
| 523 | try: |
| 524 | from flashinfer.comm.mnnvl import ( |
| 525 | McastGPUBuffer, |
| 526 | TorchDistBackend, |
| 527 | is_mnnvl_fabric_supported, |
| 528 | ) |
| 529 | except ImportError as error: |
| 530 | raise RuntimeError( |
| 531 | "cross-node TP requires FlashInfer MNNVL support" |
| 532 | ) from error |
| 533 | |
| 534 | local_supported = int( |
| 535 | is_mnnvl_fabric_supported(torch.accelerator.current_device_index()) |
| 536 | ) |
| 537 | unsupported = torch.tensor( |
| 538 | 1 - local_supported, dtype=torch.float32, device=self.device |
| 539 | ) |
| 540 | unsupported = tp.all_reduce(unsupported) |
| 541 | if int(unsupported.item()) != 0: |
| 542 | raise RuntimeError("cross-node TP is supported only on MNNVL fabric") |
| 543 | |
| 544 | comm_backend = TorchDistBackend(self.group) |
| 545 | allocation_bytes = self.num_buffers * max_tokens * hidden_size * 2 |
| 546 | self._mnnvl_input_handle = McastGPUBuffer( |
| 547 | allocation_bytes, |
| 548 | self.world_size, |
| 549 | self.rank, |
| 550 | self.device, |
| 551 | comm_backend, |
| 552 | ) |
| 553 | self._mnnvl_output_handle = McastGPUBuffer( |
| 554 | allocation_bytes, |
| 555 | self.world_size, |
| 556 | self.rank, |
| 557 | self.device, |
| 558 | comm_backend, |
| 559 | ) |
| 560 | self.input_peer_ptrs = self._mnnvl_input_handle.get_buffer_ptrs_dev() |
| 561 | self.output_peer_ptrs = self._mnnvl_output_handle.get_buffer_ptrs_dev() |
| 562 | self._initialize_mnnvl_buffers() |
| 563 | logger.info("using FlashInfer fabric-mapped MNNVL Lamport buffers") |
| 564 | else: |
| 565 | self.buf_in = symm_mem.empty( |
| 566 | self.num_buffers, |
| 567 | max_tokens, |
| 568 | self.world_size, |
| 569 | self.shard_size, |
| 570 | dtype=torch.bfloat16, |
| 571 | device=self.device, |
| 572 | ) |
| 573 | self.buf_out = symm_mem.empty( |
| 574 | self.num_buffers, |
| 575 | max_tokens, |
| 576 | hidden_size, |
| 577 | dtype=torch.bfloat16, |
| 578 | device=self.device, |
| 579 | ) |
| 580 | group_name = self.group.group_name |
| 581 | input_handle = symm_mem.rendezvous(self.buf_in, group_name) |
| 582 | output_handle = symm_mem.rendezvous(self.buf_out, group_name) |
| 583 | self.input_peer_ptrs = input_handle.buffer_ptrs_dev |
| 584 | self.output_peer_ptrs = output_handle.buffer_ptrs_dev |
| 585 | self._input_handle = input_handle |
| 586 | self._output_handle = output_handle |
| 587 | self._initialize_lamport_buffers() |
| 588 | self.generation = 0 |
| 589 | |
| 590 | def usable(self, num_tokens: int) -> bool: |
| 591 | return 0 < num_tokens <= self.max_tokens |
| 592 | |
| 593 | def rs_sconv_ag_add_norm( |
| 594 | self, |
| 595 | input_tensor: torch.Tensor, |
| 596 | residual: torch.Tensor, |
| 597 | conv_weight: torch.Tensor, |
| 598 | norm_weight: torch.Tensor | None, |
| 599 | eps: float, |
| 600 | cache: torch.Tensor, |
| 601 | positions: torch.Tensor, |
| 602 | block_table: torch.Tensor, |
| 603 | seq_idx: torch.Tensor, |
| 604 | slot_mapping: torch.Tensor, |
| 605 | off_s: int, |
| 606 | ws: int, |
| 607 | block_size: int, |
| 608 | shared_tensor: torch.Tensor | None = None, |
| 609 | ) -> tuple[torch.Tensor | None, torch.Tensor]: |
| 610 | """Return ``(normed | None, new_residual)``, both shaped ``[T, 6144]``.""" |
| 611 | tokens, hidden_size = residual.shape |
| 612 | if not self.usable(tokens): |
| 613 | raise ValueError(f"num_tokens must be in [1, {self.max_tokens}]") |
| 614 | if hidden_size != self.hidden_size or residual.dtype != torch.bfloat16: |
| 615 | raise ValueError("residual must be bf16 [T, 6144]") |
| 616 | if ( |
| 617 | input_tensor.shape != residual.shape |
| 618 | or input_tensor.dtype != torch.bfloat16 |
| 619 | or input_tensor.stride(1) != 1 |
| 620 | ): |
| 621 | raise ValueError("input_tensor must be channel-contiguous bf16 [T, 6144]") |
| 622 | if shared_tensor is not None and ( |
| 623 | shared_tensor.shape != input_tensor.shape |
| 624 | or shared_tensor.dtype != input_tensor.dtype |
| 625 | or shared_tensor.device != input_tensor.device |
| 626 | or shared_tensor.stride(1) != 1 |
| 627 | ): |
| 628 | raise ValueError("shared_tensor must match input_tensor") |
| 629 | shard_size = hidden_size // self.world_size |
| 630 | if conv_weight.shape != (shard_size, self.window_size): |
| 631 | raise ValueError( |
| 632 | f"conv_weight must have shape [{shard_size}, {self.window_size}]" |
| 633 | ) |
| 634 | if ( |
| 635 | conv_weight.dtype != torch.bfloat16 |
| 636 | or conv_weight.stride(0) != self.window_size |
| 637 | ): |
| 638 | raise ValueError("conv_weight must be contiguous bf16") |
| 639 | if norm_weight is not None and ( |
| 640 | norm_weight.shape != (hidden_size,) or norm_weight.dtype != torch.bfloat16 |
| 641 | ): |
| 642 | raise ValueError("norm_weight must be bf16 [6144] or None") |
| 643 | if cache.dtype != torch.bfloat16 or cache.ndim != 4: |
| 644 | raise ValueError("cache must be a 4-D bf16 tensor") |
| 645 | if shard_size % ws != 0 or cache.shape[1] != shard_size // ws: |
| 646 | raise ValueError("cache head layout is inconsistent with ws") |
| 647 | if off_s < 0 or off_s + ws > cache.shape[3]: |
| 648 | raise ValueError("cache channel offset is out of bounds") |
| 649 | |
| 650 | index = self.generation |
| 651 | input_offset = index * self.input_generation_bytes // 4 |
| 652 | output_offset = index * self.output_generation_bytes // 4 |
| 653 | normed = torch.empty_like(residual) if norm_weight is not None else None |
| 654 | residual_out = torch.empty_like(residual) |
| 655 | phase_splits = 8 if tokens == 1 and shard_size % 8 == 0 else 1 |
| 656 | phase_tile_p2 = triton.next_power_of_2(shard_size // phase_splits) |
| 657 | phase_grid = (tokens, phase_splits) |
| 658 | # Wide CTAs win before the grid saturates; smaller CTAs reduce register |
| 659 | # pressure once high-throughput batches provide enough parallelism. |
| 660 | if 128 <= tokens <= 2048: |
| 661 | phase_warps = 16 |
| 662 | elif tokens > 2048: |
| 663 | phase_warps = 8 |
| 664 | else: |
| 665 | phase_warps = 4 |
| 666 | gather_warps = 4 if tokens >= 256 else 8 |
| 667 | _publish_input_kernel[phase_grid]( |
| 668 | input_tensor, |
| 669 | shared_tensor if shared_tensor is not None else input_tensor, |
| 670 | self.input_peer_ptrs, |
| 671 | input_offset, |
| 672 | input_tensor.stride(0), |
| 673 | ( |
| 674 | shared_tensor.stride(0) |
| 675 | if shared_tensor is not None |
| 676 | else input_tensor.stride(0) |
| 677 | ), |
| 678 | C=hidden_size, |
| 679 | CS=shard_size, |
| 680 | CS_P2=phase_tile_p2, |
| 681 | SPLITS=phase_splits, |
| 682 | RANK=self.rank, |
| 683 | WORLD=self.world_size, |
| 684 | HAS_SHARED=shared_tensor is not None, |
| 685 | USE_PDL=self.use_pdl, |
| 686 | launch_pdl=self.use_pdl, |
| 687 | num_warps=phase_warps, |
| 688 | ) |
| 689 | _reduce_insert_kernel[phase_grid]( |
| 690 | self.input_peer_ptrs, |
| 691 | input_offset, |
| 692 | cache, |
| 693 | slot_mapping, |
| 694 | cache.stride(0), |
| 695 | cache.stride(1), |
| 696 | cache.stride(2), |
| 697 | cache.stride(3), |
| 698 | block_size, |
| 699 | C=hidden_size, |
| 700 | CS=shard_size, |
| 701 | CS_P2=phase_tile_p2, |
| 702 | SPLITS=phase_splits, |
| 703 | HEAD_SIZE=ws, |
| 704 | CACHE_OFFSET=off_s, |
| 705 | RANK=self.rank, |
| 706 | WORLD=self.world_size, |
| 707 | USE_PDL=self.use_pdl, |
| 708 | launch_pdl=self.use_pdl, |
| 709 | num_warps=phase_warps, |
| 710 | ) |
| 711 | _sconv_publish_kernel[phase_grid]( |
| 712 | self.output_peer_ptrs, |
| 713 | output_offset, |
| 714 | residual, |
| 715 | conv_weight, |
| 716 | cache, |
| 717 | positions, |
| 718 | seq_idx, |
| 719 | slot_mapping, |
| 720 | block_table, |
| 721 | residual.stride(0), |
| 722 | cache.stride(0), |
| 723 | cache.stride(1), |
| 724 | cache.stride(2), |
| 725 | cache.stride(3), |
| 726 | block_table.stride(0), |
| 727 | block_table.shape[1], |
| 728 | block_size, |
| 729 | C=hidden_size, |
| 730 | CS=shard_size, |
| 731 | CS_P2=phase_tile_p2, |
| 732 | SPLITS=phase_splits, |
| 733 | HEAD_SIZE=ws, |
| 734 | CACHE_OFFSET=off_s, |
| 735 | RANK=self.rank, |
| 736 | WORLD=self.world_size, |
| 737 | WINDOW=self.window_size, |
| 738 | USE_PDL=self.use_pdl, |
| 739 | launch_pdl=self.use_pdl, |
| 740 | num_warps=phase_warps, |
| 741 | ) |
| 742 | _gather_norm_kernel[(tokens,)]( |
| 743 | self.output_peer_ptrs, |
| 744 | output_offset, |
| 745 | norm_weight if norm_weight is not None else residual, |
| 746 | normed if normed is not None else residual_out, |
| 747 | residual_out, |
| 748 | eps, |
| 749 | residual.stride(0), |
| 750 | C=hidden_size, |
| 751 | C_P2=triton.next_power_of_2(hidden_size), |
| 752 | RANK=self.rank, |
| 753 | HAS_NORM=norm_weight is not None, |
| 754 | USE_PDL=self.use_pdl, |
| 755 | launch_pdl=self.use_pdl, |
| 756 | num_warps=gather_warps, |
| 757 | ) |
| 758 | self.generation = (index + 1) % self.num_buffers |
| 759 | return normed, residual_out |
| 760 | |
| 761 | |
| 762 | _STATE: LamportRSConv | None = None |
| 763 | _STATE_FAILED = False |
| 764 | |
| 765 | |
| 766 | def initialize_lamport_rs_conv( |
| 767 | hidden_size: int, window_size: int, max_num_batched_tokens: int |
| 768 | ) -> None: |
| 769 | """Collectively initialize the TP-group state during model construction.""" |
| 770 | global _STATE, _STATE_FAILED |
| 771 | if _STATE is not None: |
| 772 | if _STATE.hidden_size != hidden_size or _STATE.window_size != window_size: |
| 773 | raise RuntimeError("all Lamport users must share hidden and window sizes") |
| 774 | return |
| 775 | if _STATE_FAILED or os.environ.get("LAMPORT_RS_SCONV", "1") == "0": |
| 776 | return |
| 777 | try: |
| 778 | max_tokens = min(_MAX_TOKENS, max_num_batched_tokens) |
| 779 | _STATE = LamportRSConv(hidden_size, window_size, max_tokens=max_tokens) |
| 780 | except Exception: |
| 781 | _STATE_FAILED = True |
| 782 | logger.exception("fused collective unavailable; use the NCCL fallback") |
| 783 | |
| 784 | |
| 785 | def get_lamport_rs_conv(hidden_size: int, window_size: int) -> LamportRSConv | None: |
| 786 | """Return the state initialized with the model, or ``None`` for fallback.""" |
| 787 | if _STATE is not None and ( |
| 788 | _STATE.hidden_size != hidden_size or _STATE.window_size != window_size |
| 789 | ): |
| 790 | raise RuntimeError("all Lamport users must share hidden and window sizes") |
| 791 | return _STATE |