diff --git a/benchmarks/comm/bench_ulysses_head_chunk.py b/benchmarks/comm/bench_ulysses_head_chunk.py new file mode 100644 index 00000000000..230b7bcd0b0 --- /dev/null +++ b/benchmarks/comm/bench_ulysses_head_chunk.py @@ -0,0 +1,445 @@ +"""Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +Reference and benchmark for head-chunked Ulysses primitives. + +This is intentionally a benchmark-local scheduler, not a FlashInfer runtime +API. It demonstrates how a framework can use two process groups, +communicators, side streams, and the head-chunk transport primitives. + +Examples (MiniMax H3 T2V physical attention shape): + + torchrun --standalone --nproc_per_node=4 \ + benchmarks/comm/bench_ulysses_head_chunk.py \ + --global-seq 37760 --heads 56 --head-dim 128 --schedule 2,10,2 + + torchrun --standalone --nproc_per_node=2 \ + benchmarks/comm/bench_ulysses_head_chunk.py \ + --global-seq 37760 --heads 56 --head-dim 128 --schedule 7,14,7 + +The reported time is the maximum rank latency for each iteration. ``ordinary`` +uses three input scatter collectives and one output gather. ``whole_fused`` +uses one fused-QKV input collective and one output collective. The two chunked +variants use the same schedule either serially or with the reference three- +stream pipeline. Attention is PyTorch SDPA and treats the full physical +sequence as dense; model integrations remain responsible for varlen metadata. +""" + +import argparse +import json +import math +import os +import statistics +import time + +import torch +import torch.distributed as dist +import torch.nn.functional as F + +from flashinfer.comm import UlyssesCommunicator + + +def _parse_schedule(text: str, local_heads: int): + if text: + schedule = tuple(int(item) for item in text.split(",")) + elif local_heads == 14: + schedule = (2, 10, 2) + elif local_heads == 28: + schedule = (7, 14, 7) + else: + left = max(1, local_heads // 4) + middle = local_heads - 2 * left + schedule = (left, middle, left) if middle > 0 else (local_heads,) + if not schedule or any(item <= 0 for item in schedule): + raise ValueError(f"schedule must contain positive head counts: {schedule}") + if sum(schedule) != local_heads: + raise ValueError( + f"schedule {schedule} sums to {sum(schedule)}, expected " + f"local_heads={local_heads}" + ) + return schedule + + +def _attention(q, k, v): + # BSHD -> BHSD -> BSHD. The explicit contiguous output gives every + # transport variant the same attention-output layout. + result = F.scaled_dot_product_attention( + q.permute(0, 2, 1, 3), + k.permute(0, 2, 1, 3), + v.permute(0, 2, 1, 3), + ) + return result.permute(0, 2, 1, 3).contiguous() + + +def _attention_from_fused(payload, head_dim): + q, k, v = payload.split(head_dim, dim=-1) + return _attention(q, k, v) + + +class ReferenceHeadChunkPipeline: + """Benchmark-only input-A2A/attention/output-A2A stream pipeline.""" + + def __init__( + self, + *, + input_comm, + output_comm, + batch, + local_seq, + heads, + head_dim, + schedule, + dtype, + device, + ): + self.input_comm = input_comm + self.output_comm = output_comm + self.batch = batch + self.local_seq = local_seq + self.heads = heads + self.local_heads = heads // input_comm.world_size + self.head_dim = head_dim + self.schedule = schedule + self.device = device + self.input_stream = torch.cuda.Stream(device=device) + self.output_stream = torch.cuda.Stream(device=device) + self.input_ready = [torch.cuda.Event() for _ in schedule] + self.compute_ready = [torch.cuda.Event() for _ in schedule] + + max_chunk = max(schedule) + input_capacity = ( + 3 * batch * local_seq * input_comm.world_size * max_chunk * head_dim + ) + output_capacity = ( + batch * local_seq * output_comm.world_size * max_chunk * head_dim + ) + self.input_workspace = input_comm.create_workspace(max_elems=input_capacity) + self.output_workspace = output_comm.create_workspace(max_elems=output_capacity) + global_seq = local_seq * input_comm.world_size + self.qkv_chunks = [ + torch.empty( + batch, + global_seq, + head_count, + 3 * head_dim, + dtype=dtype, + device=device, + ) + for head_count in schedule + ] + self.output = torch.empty( + batch, local_seq, heads, head_dim, dtype=dtype, device=device + ) + + def sequential(self, q, k, v): + offset = 0 + for index, head_count in enumerate(self.schedule): + payload = self.input_comm.scatter_qkv_head_chunk( + q, + k, + v, + head_offset=offset, + head_count=head_count, + out=self.qkv_chunks[index], + workspace=self.input_workspace, + ) + attention_out = _attention_from_fused(payload, self.head_dim) + self.output_comm.gather_output_head_chunk( + attention_out, + local_heads=self.local_heads, + head_offset=offset, + out=self.output, + workspace=self.output_workspace, + ) + offset += head_count + return self.output + + def overlap(self, q, k, v): + compute_stream = torch.cuda.current_stream(self.device) + self.input_stream.wait_stream(compute_stream) + outputs = [] + offset = 0 + for index, head_count in enumerate(self.schedule): + with torch.cuda.stream(self.input_stream): + self.input_comm.scatter_qkv_head_chunk( + q, + k, + v, + head_offset=offset, + head_count=head_count, + out=self.qkv_chunks[index], + workspace=self.input_workspace, + ) + self.input_ready[index].record(self.input_stream) + + compute_stream.wait_event(self.input_ready[index]) + attention_out = _attention_from_fused(self.qkv_chunks[index], self.head_dim) + self.compute_ready[index].record(compute_stream) + outputs.append(attention_out) + + with torch.cuda.stream(self.output_stream): + self.output_stream.wait_event(self.compute_ready[index]) + self.output_comm.gather_output_head_chunk( + attention_out, + local_heads=self.local_heads, + head_offset=offset, + out=self.output, + workspace=self.output_workspace, + ) + attention_out.record_stream(self.output_stream) + offset += head_count + + compute_stream.wait_stream(self.output_stream) + # Keep SDPA outputs alive until their cross-stream consumers have been + # enqueued and joined to the caller stream. + del outputs + return self.output + + +def _rank_max_elapsed_ms(function, warmup, iterations, device): + for _ in range(warmup): + function() + torch.cuda.synchronize(device) + dist.barrier() + samples = [] + for _ in range(iterations): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + function() + end.record() + end.synchronize() + elapsed = torch.tensor( + [start.elapsed_time(end)], dtype=torch.float64, device=device + ) + dist.all_reduce(elapsed, op=dist.ReduceOp.MAX) + samples.append(float(elapsed.item())) + return samples + + +def _summarize(samples): + ordered = sorted(samples) + p90_index = min(len(ordered) - 1, max(0, math.ceil(0.9 * len(ordered)) - 1)) + return { + "median_ms": statistics.median(samples), + "p90_ms": ordered[p90_index], + "min_ms": min(samples), + "max_ms": max(samples), + "samples_ms": samples, + } + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--global-seq", type=int, default=37760) + parser.add_argument("--heads", type=int, default=56) + parser.add_argument("--head-dim", type=int, default=128) + parser.add_argument("--batch", type=int, default=1) + parser.add_argument("--schedule", type=str, default="") + parser.add_argument("--backend", choices=("auto", "nccl", "nvlink"), default="nccl") + parser.add_argument("--warmup", type=int, default=3) + parser.add_argument("--iters", type=int, default=10) + parser.add_argument("--json", type=str, default="") + args = parser.parse_args() + if args.iters <= 0: + parser.error("--iters must be greater than zero") + + dist.init_process_group(backend="nccl") + rank = dist.get_rank() + world_size = dist.get_world_size() + local_rank = int(os.environ.get("LOCAL_RANK", rank)) + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + if args.global_seq % world_size != 0: + raise ValueError("global sequence length must be divisible by world size") + if args.heads % world_size != 0: + raise ValueError("head count must be divisible by world size") + local_seq = args.global_seq // world_size + local_heads = args.heads // world_size + schedule = _parse_schedule(args.schedule, local_heads) + + # Both groups contain the same ranks but own independent NCCL + # communicators, allowing input and output directions to overlap. + ranks = list(range(world_size)) + input_group = dist.new_group(ranks=ranks, backend="nccl") + output_group = dist.new_group(ranks=ranks, backend="nccl") + dtype = torch.bfloat16 + whole_input_elems = 3 * args.batch * local_seq * args.heads * args.head_dim + whole_output_elems = args.batch * local_seq * args.heads * args.head_dim + input_comm = UlyssesCommunicator( + input_group, + max_elems=whole_input_elems, + dtype=dtype, + backend=args.backend, + device=device, + ) + output_comm = UlyssesCommunicator( + output_group, + max_elems=whole_output_elems, + dtype=dtype, + backend=args.backend, + device=device, + ) + + generator = torch.Generator(device=device).manual_seed(1234 + rank) + q = torch.randn( + args.batch, + local_seq, + args.heads, + args.head_dim, + dtype=dtype, + device=device, + generator=generator, + ) + k = torch.randn(q.shape, dtype=dtype, device=device, generator=generator) + v = torch.randn(q.shape, dtype=dtype, device=device, generator=generator) + + input_workspace = input_comm.create_workspace() + output_workspace = output_comm.create_workspace() + q_global = torch.empty( + args.batch, + args.global_seq, + local_heads, + args.head_dim, + dtype=dtype, + device=device, + ) + k_global = torch.empty_like(q_global) + v_global = torch.empty_like(q_global) + ordinary_output = torch.empty_like(q) + whole_payload = torch.empty( + args.batch, + args.global_seq, + local_heads, + 3 * args.head_dim, + dtype=dtype, + device=device, + ) + whole_output = torch.empty_like(q) + + def ordinary(): + input_comm.scatter_heads(q, out=q_global, workspace=input_workspace) + input_comm.scatter_heads(k, out=k_global, workspace=input_workspace) + input_comm.scatter_heads(v, out=v_global, workspace=input_workspace) + attention_out = _attention(q_global, k_global, v_global) + output_comm.gather_heads( + attention_out, out=ordinary_output, workspace=output_workspace + ) + return ordinary_output + + def whole_fused(): + input_comm.scatter_qkv_head_chunk( + q, + k, + v, + head_offset=0, + head_count=local_heads, + out=whole_payload, + workspace=input_workspace, + ) + attention_out = _attention_from_fused(whole_payload, args.head_dim) + output_comm.gather_output_head_chunk( + attention_out, + local_heads=local_heads, + head_offset=0, + out=whole_output, + workspace=output_workspace, + ) + return whole_output + + pipeline = ReferenceHeadChunkPipeline( + input_comm=input_comm, + output_comm=output_comm, + batch=args.batch, + local_seq=local_seq, + heads=args.heads, + head_dim=args.head_dim, + schedule=schedule, + dtype=dtype, + device=device, + ) + + # Compile kernels and verify the complete mapping before timing. + reference = ordinary() + fused_result = whole_fused() + # The reference pipeline reuses one output buffer for both variants, so + # preserve the serial result before overlap() overwrites that buffer. + # This clone is correctness-only and remains outside the timed paths. + sequential_result = pipeline.sequential(q, k, v).clone() + overlap_result = pipeline.overlap(q, k, v) + torch.cuda.synchronize(device) + for name, result in ( + ("whole_fused", fused_result), + ("chunk_sequential", sequential_result), + ("chunk_overlap", overlap_result), + ): + torch.testing.assert_close(result, reference, rtol=2e-2, atol=2e-2) + if rank == 0: + max_abs = float((result - reference).abs().max().item()) + print(f"correctness {name}: max_abs={max_abs:.6g}") + + variants = ( + ("ordinary", ordinary), + ("whole_fused", whole_fused), + ("chunk_sequential", lambda: pipeline.sequential(q, k, v)), + ("chunk_overlap", lambda: pipeline.overlap(q, k, v)), + ) + results = {} + for name, function in variants: + dist.barrier() + samples = _rank_max_elapsed_ms(function, args.warmup, args.iters, device) + results[name] = _summarize(samples) + if rank == 0: + print( + f"{name:>18}: median={results[name]['median_ms']:.3f} ms " + f"p90={results[name]['p90_ms']:.3f} ms" + ) + + baseline = results["ordinary"]["median_ms"] + for value in results.values(): + value["speedup_vs_ordinary"] = baseline / value["median_ms"] + record = { + "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S%z"), + "gpu": torch.cuda.get_device_name(device), + "world_size": world_size, + "backend": input_comm.backend, + "shape": { + "batch": args.batch, + "global_seq": args.global_seq, + "local_seq": local_seq, + "heads": args.heads, + "local_heads": local_heads, + "head_dim": args.head_dim, + }, + "schedule": schedule, + "warmup": args.warmup, + "iterations": args.iters, + "results": results, + } + if rank == 0: + print(json.dumps(record, indent=2)) + if args.json: + with open(args.json, "w") as output_file: + json.dump(record, output_file, indent=2) + + input_comm.close() + output_comm.close() + dist.barrier() + dist.destroy_process_group(input_group) + dist.destroy_process_group(output_group) + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/docs/api/comm.rst b/docs/api/comm.rst index 78c90f2a1fa..f1d49858f4d 100644 --- a/docs/api/comm.rst +++ b/docs/api/comm.rst @@ -289,6 +289,14 @@ run on the current CUDA stream; all ranks must issue the same call sequence with consistent shapes, one collective in flight per communicator at a time. +**Destination-passing and workspace mode.** Both collectives accept an +optional preallocated ``out=`` tensor. The NCCL backend additionally accepts +a reusable :class:`UlyssesWorkspace`, which owns the packed send and receive +buffers and removes their per-call allocations. Omitting both keywords keeps +the original allocation path unchanged. The NVLink backend uses ``out=`` but +does not consume the workspace because its IPC staging allocation is owned by +the communicator. + **Known limitations.** - PyTorch builds without ``torch.cuda.get_device_properties(...).uuid`` @@ -318,12 +326,74 @@ for the full integration):: o_ = attention(q_, k_, v_) o = comm.gather_heads(o_) # [B,S_global,H_local,D] -> [B,S_local,H,D] +Preallocated outputs and NCCL staging can be reused across serialized calls:: + + workspace = comm.create_workspace() + q_global = torch.empty(B, S_global, H_local, D, + dtype=q.dtype, device=q.device) + comm.scatter_heads(q, out=q_global, workspace=workspace) + .. autoclass:: UlyssesCommunicator :members: :show-inheritance: .. automethod:: __init__ +.. autoclass:: UlyssesWorkspace + :members: + :show-inheritance: + + .. automethod:: __init__ + +Head-Chunk Layout and Transport Primitives +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Head-chunk APIs are opt-in building blocks for framework-owned +communication/compute pipelines. They do not choose a chunk schedule, create +streams/events or process groups, or alter the ordinary ``scatter_heads`` / +``gather_heads`` path. + +``pack_ulysses_qkv_head_chunk`` accepts independent positive-strided Q/K/V +views, selects the same destination-local head band from every Ulysses rank's +head slice, and fuses Q/K/V into the last dimension. The communicator's +``scatter_qkv_head_chunk`` method combines that transform with one all-to-all. +After attention, ``gather_output_head_chunk`` performs the reverse all-to-all +and writes the compact band directly into the full output. A framework that +overlaps the two communication directions must use separate communicators and +workspaces:: + + schedule = ((0, 3), (3, 8), (11, 3)) # local_heads == 14 + full_output = torch.empty_like(q) + for head_offset, head_count in schedule: + qkv = input_comm.scatter_qkv_head_chunk( + q, k, v, + head_offset=head_offset, + head_count=head_count, + out=qkv_chunk_buffers[head_count], + workspace=input_workspace, + ) + q_chunk, k_chunk, v_chunk = qkv.split(head_dim, dim=-1) + o_chunk = attention(q_chunk, k_chunk, v_chunk) + output_comm.gather_output_head_chunk( + o_chunk, + local_heads=14, + head_offset=head_offset, + out=full_output, + workspace=output_workspace, + ) + +Head chunking is not automatically profitable. In particular, tensor +parallelism reduces the effective local head count to approximately +``H / (TP * Ulysses)``. Frameworks should enable a schedule only after all +ranks agree that the attention backend supports every chunk head count and an +offline benchmark shows a positive result. Small local-head counts, GQA/MLA +with unequal Q/K/V heads, Ring Attention, and Attention2D should use the +ordinary path unless separately implemented and measured. + +.. autofunction:: pack_ulysses_qkv_head_chunk + +.. autofunction:: merge_ulysses_output_head_chunk + Topology Probing and Backend Selection ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/docs/design_docs/ulysses_head_chunk_primitives.md b/docs/design_docs/ulysses_head_chunk_primitives.md new file mode 100644 index 00000000000..91047e27713 --- /dev/null +++ b/docs/design_docs/ulysses_head_chunk_primitives.md @@ -0,0 +1,266 @@ +# Head-Chunk Ulysses Transport Primitives + +## Scope + +This design extends `UlyssesCommunicator` with allocation-stable destination +passing and opt-in head-chunk transport primitives. It does not add a model +scheduler, choose a chunk schedule, create process groups or side streams, or +change the default whole-tensor Ulysses path. + +The intended framework pipeline is: + +```text +input stream: pack+A2A(chunk 0) pack+A2A(chunk 1) pack+A2A(chunk 2) +compute stream: attention(chunk 0) attention(chunk 1) ... +output stream: A2A+merge(chunk 0) ... +``` + +Input and output communication must use separate communicators and reusable +workspaces when they can be in flight at the same time. One communicator +continues to permit at most one in-flight collective. + +## Public building blocks + +| API | Purpose | +| --- | --- | +| `UlyssesWorkspace` | Own fixed NCCL send and receive staging buffers. | +| `create_workspace()` | Size a workspace from a communicator or a smaller chunk bound. | +| `scatter_heads(..., out=, workspace=)` | Use caller-owned output and NCCL staging for ordinary Ulysses. | +| `gather_heads(..., out=, workspace=)` | Reverse ordinary Ulysses without hot-path allocation. | +| `pack_ulysses_qkv_head_chunk(...)` | Pack a destination-local head band from independent positive-strided Q/K/V views. | +| `merge_ulysses_output_head_chunk(...)` | Merge a received compact band into a full output. | +| `scatter_qkv_head_chunk(...)` | Fuse Q/K/V packing with one input all-to-all. | +| `gather_output_head_chunk(...)` | Pack, all-to-all, and merge one attention-output band. | + +All operations enqueue on the caller's current CUDA stream. The standalone +pack operation allocates only if `out` is omitted; the standalone merge +requires `out`. Communicator operations are allocation-stable after warmup +when both `out` and a sufficiently large workspace are supplied. + +## Layout + +For world size `W`, global head count `H`, local head count `HL = H / W`, and +a selected band `[head_offset, head_offset + HC)`: + +```text +Q, K, V input: [B, S_local, H, D] +packed local QKV: [B, S_local, W * HC, 3 * D] +input A2A output: [B, W * S_local, HC, 3 * D] +attention output: [B, W * S_local, HC, D] +output A2A payload: [W, B, S_local, HC, D] +full merged output: [B, S_local, H, D] +``` + +The NCCL input pack writes send-major `[W, B, S_local, HC, 3 * D]` directly +into workspace, avoiding a separate `permute().contiguous()` allocation. For +`B == 1`, the receive-major NCCL buffer is already a valid view of the public +output layout. If no explicit output is supplied, the method returns this +workspace view; its lifetime ends when that workspace is reused. + +## Compatibility and safety boundary + +The primitives are deliberately opt-in. Frameworks own admission and must +make one rank-consistent decision before entering a data collective. + +Initial integrations should require: + +- pure Ulysses context parallelism with `world_size > 1`; +- equal Q/K/V head geometry (no implicit GQA or MLA head remapping); +- `H % world_size == 0` and a non-overlapping schedule summing exactly to + `H / world_size`; +- an attention backend instance valid for every scheduled `head_count`; +- separate input/output communicator and workspace when directions overlap; +- valid device, dtype, positive strides, output shape, and capacity on every + rank; +- model-owned varlen or padding metadata that remains identical to the + ordinary attention path. + +The implementation rejects byte-range overlap between outputs, inputs, and +workspace storage. Payloads are capped at the int32 element range used by the +transport ABI, while Triton source address arithmetic uses int64 so a legal +positive-strided view cannot wrap its source offset. + +Tensor parallelism reduces the effective attention heads available after +Ulysses, approximately `H / (TP * Ulysses)`. Expert parallelism does not +change this layout, but concurrent MoE communication can contend for the same +links. Ring Attention and Attention2D have different collective ordering and +metadata semantics and are outside this implementation. + +Head count alone is not a profitability rule. A long sequence can make even +one-head attention chunks large enough to hide communication. Conversely, +short sequences with many heads can still lose to launch and event overhead. +Production schedulers should use an offline `(GPU, topology, backend, dtype, +sequence, local_heads)` whitelist, not a universal head threshold. Unknown +shapes stay on whole-tensor Ulysses. + +There is no safe rank-local fallback after one rank has entered a collective. +An error at that point aborts the distributed call; fallback can only be +selected jointly before the call or at the next initialization. + +## Path-selection guidance + +The four benchmark paths serve different purposes and are not ordered from +slowest to fastest: + +| Path | Intended use | +| --- | --- | +| Ordinary Ulysses | Allocation-stable production baseline and fallback. | +| Whole-QKV fusion | Reduce three Q/K/V input collectives to one without splitting attention into head chunks. Enable only when the exact backend and shape benchmark faster than ordinary Ulysses. | +| Sequential head chunks | Correctness and attribution control; it isolates chunking overhead and is not intended as a production speed path. | +| Overlapped head chunks | Pipeline input communication, attention, and output communication. Enable only for an offline-whitelisted configuration with a repeatable gain beyond the measurement noise. | + +Whole-QKV fusion is useful when collective-count reduction is worth more than +its fused pack and output-merge work. Head-chunk overlap is useful when enough +communication remains exposed for attention on another chunk to hide it. Fast +links or fast attention can reverse either tradeoff because pack, merge, event, +and stream-scheduling costs remain. Therefore an integration should benchmark +all applicable paths against the allocation-stable ordinary baseline and pick +the lowest measured latency; neither fusion nor overlap is a universal default. + +The whitelist key should include at least `(GPU, topology, transport backend, +attention backend, dtype, sequence geometry, local_heads, world_size)`. A +framework-defined confidence margin should reject gains comparable to run-to-run +noise. The selected path and schedule must be agreed by all ranks before any +collective begins. + +## Backend behavior + +- NCCL uses caller-owned send/receive workspace in the opt-in path. +- The existing NVLink backend already owns IPC staging for ordinary Ulysses; + `out` is honored but the optional ordinary workspace is not consumed. +- Head-chunk NVLink transport reuses caller workspace for local fused input or + compact output around the existing raw all-to-all. +- FP16, BF16, and FP32 follow the communicator's existing dtype contract. + FP8 QKV/output communication is not part of this change. +- No process group is created inside FlashInfer. + +## Validation + +The test suite covers: + +- world sizes 1, 2, 3, and 4 on NCCL; +- FP16 and BF16 head-chunk transport; +- independent non-contiguous fused-projection Q/K/V views; +- uneven schedules `[3, 8, 3]` and `[5, 5, 4]`; +- batch sizes 1 and 2; +- default and non-default current streams; +- destination identity, direct receive-workspace lifetime, capacity, dtype, + shape, and storage-overlap failures; +- stable workspace pointers and no PyTorch CUDA allocator events after + warmup in explicit-buffer communication paths; +- exact reconstruction against ordinary scatter/gather references. + +The custom-NVLink head-chunk path has been functionally validated on a B200 +full-NVLink mesh, including zero-difference reconstruction against ordinary +scatter/gather. Its performance is shape-dependent and is not admitted by +default. The SM120 PCIe runs exercise the NCCL path and the NVLink topology +rejection/fallback behavior. + +## Reference performance + +The benchmark uses BF16 PyTorch SDPA and reports rank-maximum latency. The +ordinary baseline already uses destination passing and reusable workspace, so +the comparison isolates QKV fusion, head chunking, and overlap instead of +allocator noise. These are operator-pipeline measurements, not model E2E. + +The following RTX PRO 6000 Blackwell Server Edition results were recorded with +the pre-review implementation at commit `45052623`, using NCCL over PCIe, 3 +warmups and 7 timed iterations. They explain the motivation for the primitives, +but are not a current-head performance claim: commit `e782390a` changed several +Triton divisors/strides from compile-time to runtime arguments and the current +code has not yet been rerun on SM120 hardware. + +| Physical shape `[B, S, H, D]` | Ulysses | Schedule | Ordinary median | Overlap median | Speedup | +| --- | ---: | --- | ---: | ---: | ---: | +| `[1, 37760, 56, 128]` | 2 | `[7, 14, 7]` | 99.332 ms | 89.486 ms | 1.110x | +| `[1, 37760, 56, 128]` | 4 | `[2, 10, 2]` | 50.881 ms | 47.589 ms | 1.069x | +| `[1, 39808, 56, 128]` | 2 | `[7, 14, 7]` | 107.418 ms | 99.553 ms | 1.079x | +| `[1, 39808, 56, 128]` | 4 | `[2, 10, 2]` | 57.195 ms | 51.524 ms | 1.110x | + +All compared outputs had zero maximum absolute difference in these runs. The +benchmark treats the full physical sequence as dense. Model integrations are +still responsible for preserving true `cu_seqlens` and tail-padding semantics. + +An SM120 NCU run on the U4 10-head QKV pack measured 341.34 us, 1.48 TB/s +DRAM throughput, 92.74% of peak memory throughput, and 11.72% compute +throughput. An aggregate Nsys capture attributed about 69.5% of GPU kernel +time to NCCL SendRecv, 28.5% to attention, and about 0.6% to the three layout +kernels. The next material optimization target is transport exposure and +resource contention, not arithmetic inside pack/merge. + +### B200 full-NVLink reference + +The following current-head (`e782390a`) BF16 measurements used 6 available +NVIDIA B200 GPUs (SM100) on an NV18 full mesh. MiniMax H3 has 56 attention +heads, so its exact shape is valid for Ulysses world sizes 2 and 4 but not 6. +Each row uses 5 warmups and 20 timed iterations in each of three independent +process runs. Latencies are the median of the three run medians; speedups are +the median of three paired ordinary/variant ratios. The NCCL and custom-NVLink +results are separate transport experiments. + +| Transport | Shape | Ulysses | Schedule | Ordinary | Whole QKV | Whole speedup | Chunk overlap | Overlap speedup | +| --- | --- | ---: | --- | ---: | ---: | ---: | ---: | ---: | +| NCCL | T2V `[1,37760,56,128]` | 2 | `[7,14,7]` | 18.184 ms | 18.049 ms | 0.998x | 17.863 ms | 1.020x | +| NCCL | T2V `[1,37760,56,128]` | 4 | `[2,10,2]` | 8.977 ms | 8.981 ms | 0.999x | 9.026 ms | 0.998x | +| NCCL | I2V `[1,39808,56,128]` | 2 | `[7,14,7]` | 20.044 ms | 19.818 ms | 1.011x | 19.854 ms | 1.025x | +| NCCL | I2V `[1,39808,56,128]` | 4 | `[2,10,2]` | 9.954 ms | 9.868 ms | 1.013x | 10.163 ms | 0.979x | +| Custom NVLink | T2V `[1,37760,56,128]` | 2 | `[7,14,7]` | 17.732 ms | 19.120 ms | 0.927x | 18.293 ms | 0.974x | +| Custom NVLink | T2V `[1,37760,56,128]` | 4 | `[2,10,2]` | 8.824 ms | 9.261 ms | 0.953x | 9.173 ms | 0.962x | +| Custom NVLink | I2V `[1,39808,56,128]` | 2 | `[7,14,7]` | 20.173 ms | 20.752 ms | 0.977x | 20.080 ms | 1.004x | +| Custom NVLink | I2V `[1,39808,56,128]` | 4 | `[2,10,2]` | 9.765 ms | 10.381 ms | 0.944x | 10.211 ms | 0.953x | + +All compared outputs had zero maximum absolute difference in these runs. On +this machine, NCCL W2 overlap gains are only about 2%--2.5% and W4 does not +show a repeatable gain. With the custom NVLink transport, ordinary Ulysses is +the correct default for these shapes; its already-low communication exposure +does not amortize the current extra pack, merge, and multi-stream scheduling. +These results are an admission counterexample, not a claim that head chunking +is unhelpful on all SM100 systems. + +The same current head was also tested at physical `S=37888`, matching a +separate MiniMax-H3 distributed-FA4 experiment exactly: + +| Transport | Ulysses | Ordinary | Whole QKV | Whole speedup | Chunk overlap | Overlap speedup | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | +| Custom NVLink | 2 | 17.735 ms | 18.889 ms | 0.934x | 18.294 ms | 0.969x | +| Custom NVLink | 4 | 8.745 ms | 9.310 ms | 0.939x | 9.164 ms | 0.956x | + +### Relationship to SM100 attention-kernel fusion + +Head-chunk transport and distributed attention-kernel fusion are complementary +portfolio backends, not two layers that should always be nested. A separate +local SM100 prototype pushes K/V into symmetric windows, loads remote Q tiles +inside FA4, and owner-scatters O from the FA4 epilogue. That removes bulk Q and +O collective stages instead of scheduling the same stages on side streams. +It is not implemented by this change and depends on an architecture-specific +FA4 patch. + +In the same pinned B200 environment and physical work domain +`[1,37888,56,128]`, a common one-time dispatcher selecting that separate +backend measured: + +| Ulysses | Ordinary dense-prefix | Distributed FA4 | Speedup | Reduction | +| ---: | ---: | ---: | ---: | ---: | +| 2 | 18.379 ms | 16.149 ms | 1.138x | 12.13% | +| 4 | 9.266 ms | 8.182 ms | 1.132x | 11.70% | + +Valid-prefix maximum absolute error was `0` on U2 and `0.0002441` on U4. These +numbers show why SM100 should select an architecture-specific kernel backend +when it is available; they are not performance claims for the primitives in +this PR. A framework can share the request contract, rank-consistent planning, +workspace lifetime, and ordinary fallback while registering different SM90, +SM100, and SM120 execution backends. + +## Non-goals and follow-up work + +- The benchmark-local three-stream executor is a reference, not a runtime API. +- Schedule selection and TP/EP/CP policy belong in an integrating framework. +- CUDA Graph capture needs explicit workspace/stream lifetime tests. +- GQA/MLA needs an explicit Q-to-KV head mapping API. +- FP8 communication needs independent scale format, numerical, and quality + review; FP8 attention compute with BF16 communication can integrate first. +- A custom PCIe P2P or topology-aware transport may reduce the dominant NCCL + cost, but should be developed independently from these layout primitives. +- Remote-Q loading and owner-O scattering inside an attention kernel are not + part of this change; an SM100 distributed-FA4 backend requires a separate + kernel-focused contribution and validation matrix. diff --git a/flashinfer/comm/__init__.py b/flashinfer/comm/__init__.py index e50c2dd6c06..cd85da00de9 100644 --- a/flashinfer/comm/__init__.py +++ b/flashinfer/comm/__init__.py @@ -57,11 +57,18 @@ resolve_pcie_ipc_profile as resolve_pcie_ipc_profile, ) from .ulysses import UlyssesCommunicator as UlyssesCommunicator +from .ulysses import UlyssesWorkspace as UlyssesWorkspace from .ulysses import dispose_ulysses_a2a as dispose_ulysses_a2a from .ulysses import gen_ulysses_a2a_module as gen_ulysses_a2a_module from .ulysses import get_ulysses_a2a_module as get_ulysses_a2a_module from .ulysses import init_ulysses_a2a as init_ulysses_a2a from .ulysses import ulysses_a2a as ulysses_a2a +from .ulysses_head_chunk import ( + merge_ulysses_output_head_chunk as merge_ulysses_output_head_chunk, +) +from .ulysses_head_chunk import ( + pack_ulysses_qkv_head_chunk as pack_ulysses_qkv_head_chunk, +) from .ulysses_topology import ULYSSES_BACKENDS as ULYSSES_BACKENDS from .ulysses_topology import UlyssesBackendDecision as UlyssesBackendDecision from .ulysses_topology import UlyssesBackendError as UlyssesBackendError diff --git a/flashinfer/comm/ulysses.py b/flashinfer/comm/ulysses.py index 7819d63cca7..fb1659f4913 100644 --- a/flashinfer/comm/ulysses.py +++ b/flashinfer/comm/ulysses.py @@ -46,6 +46,97 @@ _OPEN, _CLOSING, _CLOSED = "open", "closing", "closed" +def _storage_ranges_overlap(left: torch.Tensor, right: torch.Tensor) -> bool: + """Conservatively test whether two positive-strided tensors share bytes.""" + if left.device != right.device or left.numel() == 0 or right.numel() == 0: + return False + + def storage_end(tensor: torch.Tensor) -> int: + max_element_offset = sum( + (size - 1) * stride + for size, stride in zip(tensor.shape, tensor.stride(), strict=True) + if size > 0 + ) + return tensor.data_ptr() + (max_element_offset + 1) * tensor.element_size() + + return left.data_ptr() < storage_end(right) and right.data_ptr() < storage_end(left) + + +class UlyssesWorkspace: + r"""Reusable staging storage for NCCL Ulysses collectives. + + A workspace owns two flat CUDA buffers: one for the packed all-to-all + input and one for the receive result. Passing it to + :meth:`UlyssesCommunicator.scatter_heads` or + :meth:`UlyssesCommunicator.gather_heads` removes the per-call NCCL staging + allocations. The NVLink backend already owns IPC staging storage, so it + validates but otherwise does not use this object. + + Workspaces are local allocations: they own no process group and create no + streams. A workspace may be shared by serialized scatter/gather calls, but + must not be used by concurrent collectives. + + Parameters + ---------- + max_elems : int + Capacity of each staging buffer in elements. + dtype : torch.dtype + Element dtype; float16, bfloat16, or float32. + device : torch.device or str or int, optional + CUDA device. ``None`` uses the current CUDA device. + """ + + @flashinfer_api + def __init__( + self, + *, + max_elems: int, + dtype: torch.dtype, + device: Optional[Union[torch.device, str, int]] = None, + ): + r"""Initialize reusable NCCL send and receive staging buffers. + + Parameters + ---------- + max_elems : int + Capacity of each staging buffer in elements. + dtype : torch.dtype + Element dtype; float16, bfloat16, or float32. + device : torch.device or str or int, optional + CUDA device. ``None`` uses the current CUDA device. + """ + if type(max_elems) is not int or max_elems <= 0: + raise ValueError(f"max_elems must be a positive int, got {max_elems!r}") + if max_elems > _INT32_MAX: + raise ValueError( + f"max_elems must be at most {_INT32_MAX} (int32 index range), " + f"got {max_elems}" + ) + if dtype not in _SUPPORTED_DTYPES: + raise ValueError(f"dtype must be one of {_SUPPORTED_DTYPES}, got {dtype!r}") + ordinal, error = UlyssesCommunicator._parse_cuda_ordinal(device) + if error is not None: + raise ValueError(f"invalid workspace device: {error}") + if ordinal is None: + ordinal = torch.cuda.current_device() + self.max_elems = max_elems + self.dtype = dtype + self.device = torch.device("cuda", ordinal) + with torch.cuda.device(self.device): + self._send_buffer = torch.empty(max_elems, dtype=dtype, device=self.device) + self._recv_buffer = torch.empty(max_elems, dtype=dtype, device=self.device) + + @property + def send_buffer(self) -> torch.Tensor: + """Flat send staging buffer (advanced/debugging use).""" + return self._send_buffer + + @property + def recv_buffer(self) -> torch.Tensor: + """Flat receive staging buffer (advanced/debugging use).""" + return self._recv_buffer + + class UlyssesCommunicator: r"""Ulysses context-parallelism all-to-all communicator. @@ -662,10 +753,53 @@ def __enter__(self) -> "UlyssesCommunicator": def __exit__(self, exc_type, exc_val, exc_tb) -> None: self.close() + @flashinfer_api + def create_workspace(self, *, max_elems: Optional[int] = None) -> UlyssesWorkspace: + r"""Allocate reusable NCCL send/receive staging buffers. + + This operation is rank-local and non-collective. ``max_elems`` + defaults to the communicator capacity and may be smaller when a + caller knows the maximum chunk size it will communicate. + + Parameters + ---------- + max_elems : int, optional + Capacity of each send and receive staging buffer in elements. + ``None`` uses the communicator's ``max_elems`` capacity. + + Returns + ------- + UlyssesWorkspace + A workspace on the communicator's device with its dtype. + """ + if self._state != _OPEN: + raise RuntimeError( + "create_workspace called on a " + f"{self._state} UlyssesCommunicator (use-after-close)" + ) + if max_elems is None: + max_elems = self.max_elems + if type(max_elems) is not int or max_elems <= 0: + raise ValueError(f"max_elems must be a positive int, got {max_elems!r}") + if max_elems > self.max_elems: + raise ValueError( + f"workspace max_elems={max_elems} exceeds communicator " + f"capacity max_elems={self.max_elems}" + ) + return UlyssesWorkspace( + max_elems=max_elems, dtype=self.dtype, device=self.device + ) + # ---- collectives ----------------------------------------------------------- @flashinfer_api - def scatter_heads(self, x: torch.Tensor) -> torch.Tensor: + def scatter_heads( + self, + x: torch.Tensor, + *, + out: Optional[torch.Tensor] = None, + workspace: Optional[UlyssesWorkspace] = None, + ) -> torch.Tensor: r"""``[B, S_local, H, D] -> [B, S_global, H_local, D]``. Scatter the global heads across ranks and gather the full sequence: @@ -678,6 +812,15 @@ def scatter_heads(self, x: torch.Tensor) -> torch.Tensor: ---------- x : torch.Tensor Contiguous 4-D CUDA tensor with shape ``[B, S_local, H, D]``. + out : torch.Tensor, optional + Preallocated contiguous output with shape + ``[B, S_local * world_size, H // world_size, D]``. Supplying it + removes the public output allocation. It must not alias ``x`` + when ``world_size > 1``. + workspace : UlyssesWorkspace, optional + Reusable NCCL pack/receive storage. Supplying it removes the two + NCCL staging allocations. It is validated but unused by the + NVLink backend, which owns IPC staging internally. Returns ------- @@ -693,23 +836,220 @@ def scatter_heads(self, x: torch.Tensor) -> torch.Tensor: f"divisible by world size {self.world_size}, got shape " f"{tuple(x.shape)}" ) + output_shape = ( + B, + S_local * self.world_size, + H // self.world_size, + D, + ) + out = self._prepare_out(x, out, output_shape, "scatter_heads") + self._validate_workspace(workspace, x.numel(), "scatter_heads") + self._validate_workspace_out_alias(out, workspace, "scatter_heads") if self.world_size == 1: - return x + if out is None or out is x: + return x + out.copy_(x) + return out if self.backend == "nvlink": - out = torch.empty( + if out is None: + out = torch.empty(output_shape, dtype=x.dtype, device=x.device) + ulysses_a2a(self._fa, x, out, B, S_local, H, D, 0) + return out + return self._nccl_scatter_heads(x, out=out, workspace=workspace) + + @flashinfer_api + def scatter_qkv_head_chunk( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + *, + head_offset: int, + head_count: int, + out: Optional[torch.Tensor] = None, + workspace: Optional[UlyssesWorkspace] = None, + ) -> torch.Tensor: + r"""Pack and scatter one Q/K/V head band with one collective. + + ``query``, ``key`` and ``value`` have shape + ``[B, S_local, H, D]`` and may be independent positive-strided views. + The output is a fused ``[B, S_global, head_count, 3 * D]`` payload; + slicing its final dimension produces the three attention operands. + + The method is a transport primitive, not a pipeline scheduler. A + caller that overlaps input and output communication must use separate + communicators and workspaces. Runs on the current CUDA stream. + + When the NCCL backend, ``B == 1``, ``out is None`` and a workspace is + supplied, the returned tensor aliases the workspace receive buffer and + remains valid only until that workspace is reused. Pass ``out`` for an + explicit lifetime. + + Parameters + ---------- + query : torch.Tensor + Positive-strided CUDA tensor with shape ``[B, S_local, H, D]``. + key : torch.Tensor + Positive-strided CUDA tensor with the same shape, dtype, and + device as ``query``. + value : torch.Tensor + Positive-strided CUDA tensor with the same shape, dtype, and + device as ``query``. + head_offset : int + Start of the selected band within each destination rank's + ``H // world_size`` local heads. + head_count : int + Number of consecutive local heads in the selected band. + out : torch.Tensor, optional + Preallocated contiguous result with shape + ``[B, S_local * world_size, head_count, 3 * D]``. It must not + alias ``query``, ``key``, or ``value``. + workspace : UlyssesWorkspace, optional + Reusable staging storage large enough for the fused Q/K/V + payload. A workspace cannot service concurrent collectives. + + Returns + ------- + torch.Tensor + Fused Q/K/V payload with shape + ``[B, S_global, head_count, 3 * D]``. + """ + from .ulysses_head_chunk import ( + _launch_pack_qkv, + _validate_qkv_geometry, + ) + + self._require_open("scatter_qkv_head_chunk") + B, S_local, local_heads, D, payload_elems = _validate_qkv_geometry( + query, + key, + value, + world_size=self.world_size, + head_offset=head_offset, + head_count=head_count, + ) + for name, tensor in (("query", query), ("key", key), ("value", value)): + if tensor.device != self.device: + raise ValueError( + f"scatter_qkv_head_chunk {name} is on {tensor.device}, but " + f"this communicator is bound to {self.device}" + ) + if tensor.dtype != self.dtype: + raise ValueError( + f"scatter_qkv_head_chunk {name} dtype {tensor.dtype} does " + f"not match communicator dtype {self.dtype}" + ) + self._validate_capacity(payload_elems, "scatter_qkv_head_chunk") + self._validate_workspace(workspace, payload_elems, "scatter_qkv_head_chunk") + output_shape = ( + B, + S_local * self.world_size, + head_count, + 3 * D, + ) + out = self._prepare_out(query, out, output_shape, "scatter_qkv_head_chunk") + self._validate_workspace_out_alias( + out, + workspace, + "scatter_qkv_head_chunk", + allow_recv_alias=(self.backend == "nccl" and B == 1), + ) + if out is not None and any( + _storage_ranges_overlap(out, tensor) for tensor in (query, key, value) + ): + raise ValueError( + "scatter_qkv_head_chunk out must not alias query, key, or value" + ) + + if self.world_size == 1: + if out is None: + out = torch.empty(output_shape, dtype=query.dtype, device=query.device) + _launch_pack_qkv( + out, + query, + key, + value, + world_size=1, + local_heads=local_heads, + head_offset=head_offset, + head_count=head_count, + nccl_layout=False, + ) + return out + + if workspace is None: + send_storage = torch.empty( + payload_elems, dtype=self.dtype, device=self.device + ) + recv_storage = None + else: + send_storage = workspace._send_buffer[:payload_elems] + recv_storage = workspace._recv_buffer[:payload_elems] + + if self.backend == "nvlink": + packed = send_storage.view(B, S_local, self.world_size * head_count, 3 * D) + _launch_pack_qkv( + packed, + query, + key, + value, + world_size=self.world_size, + local_heads=local_heads, + head_offset=head_offset, + head_count=head_count, + nccl_layout=False, + ) + if out is None: + out = torch.empty(output_shape, dtype=self.dtype, device=self.device) + ulysses_a2a( + self._fa, + packed, + out, B, - S_local * self.world_size, - H // self.world_size, - D, - dtype=x.dtype, - device=x.device, + S_local, + self.world_size * head_count, + 3 * D, + 0, ) - ulysses_a2a(self._fa, x, out, B, S_local, H, D, 0) return out - return self._nccl_scatter_heads(x) + + send = send_storage.view(self.world_size, B, S_local, head_count, 3 * D) + if recv_storage is None: + recv = torch.empty_like(send) + else: + recv = recv_storage.view_as(send) + _launch_pack_qkv( + send, + query, + key, + value, + world_size=self.world_size, + local_heads=local_heads, + head_offset=head_offset, + head_count=head_count, + nccl_layout=True, + ) + dist.all_to_all_single(recv, send, group=self.group) + recv_as_output = recv.view(output_shape) if B == 1 else None + if out is None and recv_as_output is not None: + return recv_as_output + if out is None: + out = torch.empty(output_shape, dtype=self.dtype, device=self.device) + if recv_as_output is not None and out.data_ptr() == recv.data_ptr(): + return out + out.view(B, self.world_size, S_local, head_count, 3 * D).copy_( + recv.permute(1, 0, 2, 3, 4) + ) + return out @flashinfer_api - def gather_heads(self, x: torch.Tensor) -> torch.Tensor: + def gather_heads( + self, + x: torch.Tensor, + *, + out: Optional[torch.Tensor] = None, + workspace: Optional[UlyssesWorkspace] = None, + ) -> torch.Tensor: r"""``[B, S_global, H_local, D] -> [B, S_local, H, D]``. Inverse of :meth:`scatter_heads`: gather all head slices for this @@ -720,6 +1060,15 @@ def gather_heads(self, x: torch.Tensor) -> torch.Tensor: ---------- x : torch.Tensor Contiguous 4-D CUDA tensor with shape ``[B, S_global, H_local, D]``. + out : torch.Tensor, optional + Preallocated contiguous output with shape + ``[B, S_global // world_size, H_local * world_size, D]``. + Supplying it removes the public output allocation. It must not + alias ``x`` when ``world_size > 1``. + workspace : UlyssesWorkspace, optional + Reusable NCCL pack/receive storage. Supplying it removes the two + NCCL staging allocations. It is validated but unused by the + NVLink backend. Returns ------- @@ -735,50 +1084,283 @@ def gather_heads(self, x: torch.Tensor) -> torch.Tensor: f"be divisible by world size {self.world_size}, got shape " f"{tuple(x.shape)}" ) - if self.world_size == 1: - return x S_local = S_global // self.world_size H = H_local * self.world_size + output_shape = (B, S_local, H, D) + out = self._prepare_out(x, out, output_shape, "gather_heads") + self._validate_workspace(workspace, x.numel(), "gather_heads") + self._validate_workspace_out_alias(out, workspace, "gather_heads") + if self.world_size == 1: + if out is None or out is x: + return x + out.copy_(x) + return out if self.backend == "nvlink": - out = torch.empty(B, S_local, H, D, dtype=x.dtype, device=x.device) + if out is None: + out = torch.empty(output_shape, dtype=x.dtype, device=x.device) ulysses_a2a(self._fa, x, out, B, S_local, H, D, 1) return out - return self._nccl_gather_heads(x) + return self._nccl_gather_heads(x, out=out, workspace=workspace) + + @flashinfer_api + def gather_output_head_chunk( + self, + x: torch.Tensor, + *, + local_heads: int, + head_offset: int, + out: torch.Tensor, + workspace: Optional[UlyssesWorkspace] = None, + ) -> torch.Tensor: + r"""Gather one attention-output head band into a full destination. + + ``x`` has shape ``[B, S_global, head_count, D]``. ``out`` is the + preallocated full output ``[B, S_local, world_size * local_heads, D]``. + Only the selected head band is written; calling this method for every + non-overlapping band in a complete schedule reconstructs ordinary + whole-head Ulysses output communication. + + The NCCL path packs the sequence split directly into reusable staging + and merges the received head band directly into ``out``. The NVLink + path uses its fused-transpose collective followed by the same merge + primitive. Runs on the current CUDA stream. + + Parameters + ---------- + x : torch.Tensor + Positive-strided CUDA tensor with shape + ``[B, S_global, head_count, D]``. + local_heads : int + Total number of attention-output heads owned by this rank before + the output all-to-all. + head_offset : int + Start of this band within the rank's ``local_heads``. + out : torch.Tensor + Preallocated contiguous destination with shape + ``[B, S_global // world_size, world_size * local_heads, D]``. + Only this head band is modified. + workspace : UlyssesWorkspace, optional + Reusable staging storage large enough for ``x``. A workspace + cannot service concurrent collectives. + + Returns + ------- + torch.Tensor + ``out`` after merging the gathered head band. + """ + from .ulysses_head_chunk import ( + _launch_merge_rank_major, + _launch_pack_output_sequence, + _positive_int, + _nonnegative_int, + _validate_cuda_tensor, + ) + + self._require_open("gather_output_head_chunk") + local_heads = _positive_int(local_heads, "local_heads") + head_offset = _nonnegative_int(head_offset, "head_offset") + x = _validate_cuda_tensor( + x, "gather_output_head_chunk input", ndim=4, contiguous=False + ) + B, S_global, head_count, D = x.shape + if x.device != self.device or x.dtype != self.dtype: + raise ValueError( + f"gather_output_head_chunk input device/dtype " + f"({x.device}, {x.dtype}) must match communicator " + f"({self.device}, {self.dtype})" + ) + if S_global % self.world_size != 0: + raise ValueError( + f"global sequence length {S_global} must be divisible by " + f"world size {self.world_size}" + ) + if head_offset + head_count > local_heads: + raise ValueError( + f"head band [{head_offset}, {head_offset + head_count}) " + f"exceeds local_heads={local_heads}" + ) + S_local = S_global // self.world_size + output_shape = (B, S_local, self.world_size * local_heads, D) + out = self._prepare_out(x, out, output_shape, "gather_output_head_chunk") + if out is None: # ``out`` is required by the public signature. + raise TypeError("gather_output_head_chunk requires out") + if _storage_ranges_overlap(out, x): + raise ValueError("gather_output_head_chunk out must not alias the input") + payload_elems = x.numel() + self._validate_capacity(payload_elems, "gather_output_head_chunk") + self._validate_workspace(workspace, payload_elems, "gather_output_head_chunk") + self._validate_workspace_out_alias(out, workspace, "gather_output_head_chunk") + + if self.world_size == 1: + _launch_merge_rank_major( + x.unsqueeze(0), + out, + world_size=1, + local_heads=local_heads, + head_offset=head_offset, + ) + return out + + if workspace is None: + send_storage = torch.empty( + payload_elems, dtype=self.dtype, device=self.device + ) + recv_storage = torch.empty_like(send_storage) + else: + send_storage = workspace._send_buffer[:payload_elems] + recv_storage = workspace._recv_buffer[:payload_elems] + + if self.backend == "nvlink": + if x.is_contiguous(): + contiguous_input = x + else: + contiguous_input = send_storage.view_as(x) + contiguous_input.copy_(x) + compact = recv_storage.view(B, S_local, self.world_size * head_count, D) + ulysses_a2a( + self._fa, + contiguous_input, + compact, + B, + S_local, + self.world_size * head_count, + D, + 1, + ) + rank_major = compact.view( + B, S_local, self.world_size, head_count, D + ).permute(2, 0, 1, 3, 4) + else: + send = send_storage.view(self.world_size, B, S_local, head_count, D) + recv = recv_storage.view_as(send) + _launch_pack_output_sequence( + send, + x, + world_size=self.world_size, + ) + dist.all_to_all_single(recv, send, group=self.group) + rank_major = recv + + _launch_merge_rank_major( + rank_major, + out, + world_size=self.world_size, + local_heads=local_heads, + head_offset=head_offset, + ) + return out # ---- NCCL fallback --------------------------------------------------------- # The conventional all_to_all_single path with explicit permute/contiguous # glue before and after (exactly the data movement the fused NVLink kernel # folds into its cross-GPU writes). Bit-identical to the NVLink backend. - def _nccl_scatter_heads(self, x: torch.Tensor) -> torch.Tensor: + def _nccl_scatter_heads( + self, + x: torch.Tensor, + *, + out: Optional[torch.Tensor], + workspace: Optional[UlyssesWorkspace], + ) -> torch.Tensor: B, S_local, H, D = x.shape W = self.world_size H_local = H // W - xt = x.reshape(B, S_local, W, H_local, D).permute(2, 0, 1, 3, 4).contiguous() - recv = torch.empty_like(xt) - dist.all_to_all_single(recv, xt, group=self.group) - # chunk j == rank j's contribution to sequence block j - return recv.permute(1, 0, 2, 3, 4).reshape(B, W * S_local, H_local, D) + # Preserve the pre-existing allocation/layout path byte-for-byte when + # destination passing is not requested. This keeps the default API's + # behavior and performance independent of the new opt-in feature. + if out is None and workspace is None: + xt = ( + x.reshape(B, S_local, W, H_local, D).permute(2, 0, 1, 3, 4).contiguous() + ) + recv = torch.empty_like(xt) + dist.all_to_all_single(recv, xt, group=self.group) + # chunk j == rank j's contribution to sequence block j + return recv.permute(1, 0, 2, 3, 4).reshape(B, W * S_local, H_local, D) + + if out is None: + out = torch.empty( + B, + W * S_local, + H_local, + D, + dtype=x.dtype, + device=x.device, + ) + if workspace is None: + send = torch.empty( + W, B, S_local, H_local, D, dtype=x.dtype, device=x.device + ) + recv = torch.empty_like(send) + else: + send = workspace._send_buffer[: x.numel()].view(W, B, S_local, H_local, D) + recv = workspace._recv_buffer[: x.numel()].view(W, B, S_local, H_local, D) + send.copy_(x.reshape(B, S_local, W, H_local, D).permute(2, 0, 1, 3, 4)) + dist.all_to_all_single(recv, send, group=self.group) + out.view(B, W, S_local, H_local, D).copy_(recv.permute(1, 0, 2, 3, 4)) + return out - def _nccl_gather_heads(self, x: torch.Tensor) -> torch.Tensor: + def _nccl_gather_heads( + self, + x: torch.Tensor, + *, + out: Optional[torch.Tensor], + workspace: Optional[UlyssesWorkspace], + ) -> torch.Tensor: B, S_global, H_local, D = x.shape W = self.world_size S_local = S_global // W - xt = x.reshape(B, W, S_local, H_local, D).permute(1, 0, 2, 3, 4).contiguous() - recv = torch.empty_like(xt) - dist.all_to_all_single(recv, xt, group=self.group) - # chunk p == this rank's sequence block, head slice p - return ( - recv.permute(1, 2, 0, 3, 4).reshape(B, S_local, W * H_local, D).contiguous() - ) + if out is None and workspace is None: + xt = ( + x.reshape(B, W, S_local, H_local, D).permute(1, 0, 2, 3, 4).contiguous() + ) + recv = torch.empty_like(xt) + dist.all_to_all_single(recv, xt, group=self.group) + # chunk p == this rank's sequence block, head slice p + return ( + recv.permute(1, 2, 0, 3, 4) + .reshape(B, S_local, W * H_local, D) + .contiguous() + ) + + if out is None: + out = torch.empty( + B, + S_local, + W * H_local, + D, + dtype=x.dtype, + device=x.device, + ) + if workspace is None: + send = torch.empty( + W, B, S_local, H_local, D, dtype=x.dtype, device=x.device + ) + recv = torch.empty_like(send) + else: + send = workspace._send_buffer[: x.numel()].view(W, B, S_local, H_local, D) + recv = workspace._recv_buffer[: x.numel()].view(W, B, S_local, H_local, D) + send.copy_(x.reshape(B, W, S_local, H_local, D).permute(1, 0, 2, 3, 4)) + dist.all_to_all_single(recv, send, group=self.group) + out.view(B, S_local, W, H_local, D).copy_(recv.permute(1, 2, 0, 3, 4)) + return out # ---- validation ------------------------------------------------------------ - def _validate(self, x, op: str) -> None: + def _require_open(self, op: str) -> None: if self._state != _OPEN: raise RuntimeError( f"{op} called on a {self._state} UlyssesCommunicator (use-after-close)" ) + + def _validate_capacity(self, required_elems: int, op: str) -> None: + if required_elems > self.max_elems: + raise ValueError( + f"{op} payload has {required_elems} elements, exceeding the " + f"communicator capacity max_elems={self.max_elems}" + ) + + def _validate(self, x, op: str) -> None: + self._require_open(op) if not isinstance(x, torch.Tensor): raise TypeError(f"{op} expects a torch.Tensor, got {type(x).__name__}") if x.dim() != 4: @@ -809,6 +1391,95 @@ def _validate(self, x, op: str) -> None: f"(which is capped at the int32 index range {_INT32_MAX})" ) + def _prepare_out( + self, + x: torch.Tensor, + out: Optional[torch.Tensor], + expected_shape: Tuple[int, ...], + op: str, + ) -> Optional[torch.Tensor]: + if out is None: + return None + if not isinstance(out, torch.Tensor): + raise TypeError( + f"{op} out expects a torch.Tensor, got {type(out).__name__}" + ) + if tuple(out.shape) != expected_shape: + raise ValueError( + f"{op} out has shape {tuple(out.shape)}, expected {expected_shape}" + ) + if out.device != x.device: + raise ValueError(f"{op} out is on {out.device}, but input is on {x.device}") + if out.dtype != x.dtype: + raise ValueError( + f"{op} out dtype {out.dtype} does not match input dtype {x.dtype}" + ) + if not out.is_contiguous(): + raise ValueError(f"{op} out must be contiguous") + if _storage_ranges_overlap(out, x): + if self.world_size > 1 or out.data_ptr() != x.data_ptr(): + raise ValueError(f"{op} out must not alias the input") + return out + + def _validate_workspace( + self, + workspace: Optional[UlyssesWorkspace], + required_elems: int, + op: str, + ) -> None: + if workspace is None: + return + if not isinstance(workspace, UlyssesWorkspace): + raise TypeError( + f"{op} workspace expects UlyssesWorkspace, got " + f"{type(workspace).__name__}" + ) + if workspace.device != self.device: + raise ValueError( + f"{op} workspace is on {workspace.device}, but communicator " + f"is bound to {self.device}" + ) + if workspace.dtype != self.dtype: + raise ValueError( + f"{op} workspace dtype {workspace.dtype} does not match " + f"communicator dtype {self.dtype}" + ) + if workspace.max_elems < required_elems: + raise ValueError( + f"{op} requires {required_elems} workspace elements, but " + f"workspace capacity is {workspace.max_elems}" + ) + for name, buffer in ( + ("send", workspace._send_buffer), + ("recv", workspace._recv_buffer), + ): + if ( + buffer.device != self.device + or buffer.dtype != self.dtype + or not buffer.is_contiguous() + or buffer.numel() < workspace.max_elems + ): + raise ValueError(f"{op} workspace {name} buffer was modified") + + @staticmethod + def _validate_workspace_out_alias( + out: Optional[torch.Tensor], + workspace: Optional[UlyssesWorkspace], + op: str, + *, + allow_recv_alias: bool = False, + ) -> None: + if out is None or workspace is None: + return + if _storage_ranges_overlap(out, workspace._send_buffer): + raise ValueError(f"{op} out must not alias the workspace send buffer") + if _storage_ranges_overlap(out, workspace._recv_buffer): + exact_recv_alias = out.data_ptr() == workspace._recv_buffer.data_ptr() + if not (allow_recv_alias and exact_recv_alias): + raise ValueError( + f"{op} out must not alias the workspace receive buffer" + ) + # ============================================================================= # Raw (advanced) kernel entry points diff --git a/flashinfer/comm/ulysses_head_chunk.py b/flashinfer/comm/ulysses_head_chunk.py new file mode 100644 index 00000000000..db12b0ec998 --- /dev/null +++ b/flashinfer/comm/ulysses_head_chunk.py @@ -0,0 +1,399 @@ +"""Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +Layout primitives for head-chunked Ulysses attention. + +These functions deliberately do not create process groups, streams, events, +or an attention scheduler. They only transform Q/K/V and output head bands so +framework integrations can build their own communication/compute pipeline. +""" + +from typing import Optional, Tuple + +import torch + +from ..api_logging import flashinfer_api + +_INT32_MAX = 2**31 - 1 +_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32) + + +def _storage_ranges_overlap(left: torch.Tensor, right: torch.Tensor) -> bool: + """Conservatively test whether two positive-strided tensors share bytes.""" + if left.device != right.device or left.numel() == 0 or right.numel() == 0: + return False + + def storage_end(tensor: torch.Tensor) -> int: + max_element_offset = sum( + (size - 1) * stride + for size, stride in zip(tensor.shape, tensor.stride(), strict=True) + if size > 0 + ) + return tensor.data_ptr() + (max_element_offset + 1) * tensor.element_size() + + return left.data_ptr() < storage_end(right) and right.data_ptr() < storage_end(left) + + +def _positive_int(value, name: str) -> int: + if type(value) is not int or value <= 0: + raise ValueError(f"{name} must be a positive int, got {value!r}") + return value + + +def _nonnegative_int(value, name: str) -> int: + if type(value) is not int or value < 0: + raise ValueError(f"{name} must be a nonnegative int, got {value!r}") + return value + + +def _validate_cuda_tensor( + tensor, name: str, *, ndim: int, contiguous: bool +) -> torch.Tensor: + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor, got {type(tensor).__name__}") + if not tensor.is_cuda: + raise ValueError(f"{name} must be a CUDA tensor") + if tensor.dim() != ndim: + raise ValueError(f"{name} must be {ndim}-D, got shape {tuple(tensor.shape)}") + if tensor.dtype not in _SUPPORTED_DTYPES: + raise ValueError( + f"{name} dtype must be float16/bfloat16/float32, got {tensor.dtype}" + ) + if any(size <= 0 for size in tensor.shape): + raise ValueError(f"{name} dims must be positive, got {tuple(tensor.shape)}") + if contiguous and not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous") + if not contiguous and any(stride <= 0 for stride in tensor.stride()): + raise ValueError(f"{name} must have positive strides, got {tensor.stride()}") + return tensor + + +def _validate_qkv_geometry( + query, + key, + value, + *, + world_size, + head_offset, + head_count, +) -> Tuple[int, int, int, int, int]: + world_size = _positive_int(world_size, "world_size") + head_count = _positive_int(head_count, "head_count") + head_offset = _nonnegative_int(head_offset, "head_offset") + query = _validate_cuda_tensor(query, "query", ndim=4, contiguous=False) + key = _validate_cuda_tensor(key, "key", ndim=4, contiguous=False) + value = _validate_cuda_tensor(value, "value", ndim=4, contiguous=False) + for name, tensor in (("key", key), ("value", value)): + if tensor.shape != query.shape: + raise ValueError( + f"{name} shape {tuple(tensor.shape)} does not match query " + f"shape {tuple(query.shape)}" + ) + if tensor.device != query.device: + raise ValueError( + f"{name} is on {tensor.device}, but query is on {query.device}" + ) + if tensor.dtype != query.dtype: + raise ValueError( + f"{name} dtype {tensor.dtype} does not match query dtype {query.dtype}" + ) + batch, seq_len, global_heads, head_dim = query.shape + if global_heads % world_size != 0: + raise ValueError( + f"global head count {global_heads} must be divisible by world_size " + f"{world_size}" + ) + local_heads = global_heads // world_size + if head_offset + head_count > local_heads: + raise ValueError( + f"head band [{head_offset}, {head_offset + head_count}) exceeds " + f"local_heads={local_heads}" + ) + payload_elems = batch * seq_len * world_size * head_count * 3 * head_dim + if payload_elems > _INT32_MAX: + raise ValueError( + f"packed QKV payload has {payload_elems} elements, exceeding the " + f"int32 index range {_INT32_MAX}" + ) + return batch, seq_len, local_heads, head_dim, payload_elems + + +def _prepare_qkv_output( + query, + out, + *, + world_size, + head_count, +) -> torch.Tensor: + batch, seq_len, _, head_dim = query.shape + expected = (batch, seq_len, world_size * head_count, 3 * head_dim) + if out is None: + return torch.empty(expected, dtype=query.dtype, device=query.device) + _validate_cuda_tensor(out, "out", ndim=4, contiguous=True) + if tuple(out.shape) != expected: + raise ValueError(f"out has shape {tuple(out.shape)}, expected {expected}") + if out.device != query.device or out.dtype != query.dtype: + raise ValueError( + f"out device/dtype ({out.device}, {out.dtype}) must match query " + f"({query.device}, {query.dtype})" + ) + return out + + +def _launch_pack_qkv( + out, + query, + key, + value, + *, + world_size, + local_heads, + head_offset, + head_count, + nccl_layout, +) -> None: + from ..triton.ulysses import _pack_ulysses_qkv_head_chunk + + _pack_ulysses_qkv_head_chunk( + out, + query, + key, + value, + world_size=world_size, + local_heads=local_heads, + head_offset=head_offset, + head_count=head_count, + nccl_layout=nccl_layout, + ) + + +@flashinfer_api +def pack_ulysses_qkv_head_chunk( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + *, + world_size: int, + head_offset: int, + head_count: int, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + r"""Pack one destination-local head band of Q/K/V. + + The three inputs use ``[B, S_local, H, D]`` and may be independent + positive-strided views (including views produced by a fused projection). + ``H`` must be divisible by ``world_size``. The selected local band is + applied inside every destination rank's head slice, and the result is a + contiguous fused payload ``[B, S_local, world_size * head_count, 3 * D]``. + + The payload can be passed directly to + :meth:`flashinfer.comm.UlyssesCommunicator.scatter_heads`; its result has + shape ``[B, S_global, head_count, 3 * D]``. Slicing the final dimension + yields Q/K/V views for the attention backend. + + This operation runs on the caller's current CUDA stream and allocates only + when ``out`` is omitted. + + Parameters + ---------- + query : torch.Tensor + Positive-strided CUDA tensor with shape ``[B, S_local, H, D]``. + key : torch.Tensor + Positive-strided CUDA tensor with the same shape, dtype, and device + as ``query``. + value : torch.Tensor + Positive-strided CUDA tensor with the same shape, dtype, and device + as ``query``. + world_size : int + Number of Ulysses ranks. ``H`` must be divisible by ``world_size``. + head_offset : int + Start of the selected band within every destination rank's + ``H // world_size`` local heads. + head_count : int + Number of consecutive local heads in the selected band. + out : torch.Tensor, optional + Preallocated contiguous result with shape + ``[B, S_local, world_size * head_count, 3 * D]``. It must not alias + ``query``, ``key``, or ``value``. + + Returns + ------- + torch.Tensor + Contiguous fused Q/K/V payload with shape + ``[B, S_local, world_size * head_count, 3 * D]``. + """ + _, _, local_heads, _, _ = _validate_qkv_geometry( + query, + key, + value, + world_size=world_size, + head_offset=head_offset, + head_count=head_count, + ) + out = _prepare_qkv_output(query, out, world_size=world_size, head_count=head_count) + if any(_storage_ranges_overlap(out, tensor) for tensor in (query, key, value)): + raise ValueError("out must not alias query, key, or value") + _launch_pack_qkv( + out, + query, + key, + value, + world_size=world_size, + local_heads=local_heads, + head_offset=head_offset, + head_count=head_count, + nccl_layout=False, + ) + return out + + +def _validate_merge_geometry( + received, + out, + *, + world_size, + local_heads, + head_offset, +): + world_size = _positive_int(world_size, "world_size") + local_heads = _positive_int(local_heads, "local_heads") + head_offset = _nonnegative_int(head_offset, "head_offset") + received = _validate_cuda_tensor(received, "received", ndim=4, contiguous=True) + out = _validate_cuda_tensor(out, "out", ndim=4, contiguous=True) + batch, local_seq, compact_heads, head_dim = received.shape + if compact_heads % world_size != 0: + raise ValueError( + f"received compact head count {compact_heads} must be divisible by " + f"world_size {world_size}" + ) + head_count = compact_heads // world_size + if head_offset + head_count > local_heads: + raise ValueError( + f"head band [{head_offset}, {head_offset + head_count}) exceeds " + f"local_heads={local_heads}" + ) + expected = (batch, local_seq, world_size * local_heads, head_dim) + if tuple(out.shape) != expected: + raise ValueError(f"out has shape {tuple(out.shape)}, expected {expected}") + if out.device != received.device or out.dtype != received.dtype: + raise ValueError( + f"out device/dtype ({out.device}, {out.dtype}) must match received " + f"({received.device}, {received.dtype})" + ) + if _storage_ranges_overlap(out, received): + raise ValueError("out must not alias received") + if received.numel() > _INT32_MAX: + raise ValueError( + f"received has {received.numel()} elements, exceeding the int32 " + f"index range {_INT32_MAX}" + ) + return batch, local_seq, head_count, head_dim + + +def _launch_merge_rank_major( + received_rank_major, + out, + *, + world_size, + local_heads, + head_offset, +) -> None: + from ..triton.ulysses import _merge_ulysses_output_head_chunk + + _merge_ulysses_output_head_chunk( + received_rank_major, + out, + world_size=world_size, + local_heads=local_heads, + head_offset=head_offset, + ) + + +@flashinfer_api +def merge_ulysses_output_head_chunk( + received: torch.Tensor, + *, + world_size: int, + local_heads: int, + head_offset: int, + out: torch.Tensor, +) -> torch.Tensor: + r"""Merge one compact Ulysses output head band into a full output. + + ``received`` is the result of gathering a chunked attention output and has + shape ``[B, S_local, world_size * head_count, D]``. ``out`` is the full + ``[B, S_local, world_size * local_heads, D]`` destination. Only the band + ``[rank * local_heads + head_offset, + head_count)`` for each source rank + is written; other output elements remain unchanged. + + Calling this function once for every non-overlapping chunk in a schedule + whose head counts sum to ``local_heads`` reconstructs the ordinary + whole-head Ulysses result. It runs on the caller's current CUDA stream and + performs no allocation. + + Parameters + ---------- + received : torch.Tensor + Contiguous gathered head band with shape + ``[B, S_local, world_size * head_count, D]``. + world_size : int + Number of Ulysses ranks represented in ``received`` and ``out``. + local_heads : int + Number of attention-output heads contributed by each rank to the full + output. + head_offset : int + Start of the selected band within each rank's ``local_heads``. + out : torch.Tensor + Preallocated contiguous destination with shape + ``[B, S_local, world_size * local_heads, D]``. It must not alias + ``received``; elements outside the selected band remain unchanged. + + Returns + ------- + torch.Tensor + ``out`` after merging the selected head band. + """ + batch, local_seq, head_count, head_dim = _validate_merge_geometry( + received, + out, + world_size=world_size, + local_heads=local_heads, + head_offset=head_offset, + ) + rank_major = received.view( + batch, local_seq, world_size, head_count, head_dim + ).permute(2, 0, 1, 3, 4) + _launch_merge_rank_major( + rank_major, + out, + world_size=world_size, + local_heads=local_heads, + head_offset=head_offset, + ) + return out + + +def _launch_pack_output_sequence( + out_rank_major, + source, + *, + world_size, +) -> None: + from ..triton.ulysses import _pack_ulysses_output_sequence_chunk + + _pack_ulysses_output_sequence_chunk( + out_rank_major, + source, + world_size=world_size, + ) diff --git a/flashinfer/triton/ulysses.py b/flashinfer/triton/ulysses.py new file mode 100644 index 00000000000..b6ee9c68349 --- /dev/null +++ b/flashinfer/triton/ulysses.py @@ -0,0 +1,294 @@ +""" +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import triton +import triton.language as tl + +# Model/topology invariants stay specialized so index division can be folded. +# Resolution, head-band schedule, offsets, and tensor strides remain runtime +# arguments to avoid recompiling for each request shape or positive-strided view. + + +@triton.jit +def _pack_ulysses_qkv_head_chunk_kernel( + output, + query, + key, + value, + total, + batch, + seq_len, + world_size: tl.constexpr, + local_heads: tl.constexpr, + chunk_heads, + head_offset, + head_dim: tl.constexpr, + q_batch_stride, + q_seq_stride, + q_head_stride, + q_dim_stride, + k_batch_stride, + k_seq_stride, + k_head_stride, + k_dim_stride, + v_batch_stride, + v_seq_stride, + v_head_stride, + v_dim_stride, + nccl_layout: tl.constexpr, + BLOCK: tl.constexpr, +): + """Pack Q/K/V into fused ``[..., 3 * D]`` Ulysses payloads. + + ``nccl_layout=False`` writes ``[B,S,W*HC,3D]``. ``True`` writes the + send-major ``[W,B,S,HC,3D]`` representation consumed directly by + ``all_to_all_single``. + """ + # The compact payload is int32-sized, but a positive-strided source view + # can address beyond that range. Keep all derived source offsets in i64. + offsets = (tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)).to(tl.int64) + valid = offsets < total + dim = offsets % head_dim + slot = offsets // head_dim + head_in_chunk = slot % chunk_heads + slot = slot // chunk_heads + destination = slot % world_size + slot = slot // world_size + seq = slot % seq_len + batch_idx = slot // seq_len + + global_head = destination * local_heads + head_offset + head_in_chunk + q_offset = ( + batch_idx * q_batch_stride + + seq * q_seq_stride + + global_head * q_head_stride + + dim * q_dim_stride + ) + k_offset = ( + batch_idx * k_batch_stride + + seq * k_seq_stride + + global_head * k_head_stride + + dim * k_dim_stride + ) + v_offset = ( + batch_idx * v_batch_stride + + seq * v_seq_stride + + global_head * v_head_stride + + dim * v_dim_stride + ) + + if nccl_layout: + record = ( + (destination * batch + batch_idx) * seq_len + seq + ) * chunk_heads + head_in_chunk + else: + record = ( + (batch_idx * seq_len + seq) * world_size + destination + ) * chunk_heads + head_in_chunk + output_base = record * (3 * head_dim) + dim + tl.store(output + output_base, tl.load(query + q_offset, mask=valid), mask=valid) + tl.store( + output + output_base + head_dim, + tl.load(key + k_offset, mask=valid), + mask=valid, + ) + tl.store( + output + output_base + 2 * head_dim, + tl.load(value + v_offset, mask=valid), + mask=valid, + ) + + +@triton.jit +def _pack_ulysses_output_sequence_chunk_kernel( + output, + source, + total, + batch, + local_seq, + world_size: tl.constexpr, + chunk_heads, + head_dim: tl.constexpr, + source_batch_stride, + source_seq_stride, + source_head_stride, + source_dim_stride, + BLOCK: tl.constexpr, +): + """Pack ``[B,W*S,HC,D]`` into NCCL send-major ``[W,B,S,HC,D]``.""" + offsets = (tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)).to(tl.int64) + valid = offsets < total + dim = offsets % head_dim + slot = offsets // head_dim + head = slot % chunk_heads + slot = slot // chunk_heads + seq = slot % local_seq + slot = slot // local_seq + batch_idx = slot % batch + destination = slot // batch + global_seq = destination * local_seq + seq + source_offset = ( + batch_idx * source_batch_stride + + global_seq * source_seq_stride + + head * source_head_stride + + dim * source_dim_stride + ) + tl.store(output + offsets, tl.load(source + source_offset, mask=valid), mask=valid) + + +@triton.jit +def _merge_ulysses_output_head_chunk_kernel( + received, + output, + total, + batch, + local_seq, + world_size: tl.constexpr, + local_heads: tl.constexpr, + chunk_heads, + head_offset, + head_dim: tl.constexpr, + recv_rank_stride, + recv_batch_stride, + recv_seq_stride, + recv_head_stride, + recv_dim_stride, + out_batch_stride, + out_seq_stride, + out_head_stride, + out_dim_stride, + BLOCK: tl.constexpr, +): + """Merge logical ``[source_rank,B,S,HC,D]`` bands into full output.""" + offsets = (tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)).to(tl.int64) + valid = offsets < total + dim = offsets % head_dim + slot = offsets // head_dim + head_in_chunk = slot % chunk_heads + slot = slot // chunk_heads + seq = slot % local_seq + slot = slot // local_seq + batch_idx = slot % batch + source_rank = slot // batch + + recv_offset = ( + source_rank * recv_rank_stride + + batch_idx * recv_batch_stride + + seq * recv_seq_stride + + head_in_chunk * recv_head_stride + + dim * recv_dim_stride + ) + global_head = source_rank * local_heads + head_offset + head_in_chunk + output_offset = ( + batch_idx * out_batch_stride + + seq * out_seq_stride + + global_head * out_head_stride + + dim * out_dim_stride + ) + tl.store( + output + output_offset, + tl.load(received + recv_offset, mask=valid), + mask=valid, + ) + + +def _pack_ulysses_qkv_head_chunk( + output, + query, + key, + value, + *, + world_size, + local_heads, + head_offset, + head_count, + nccl_layout=False, +): + batch, seq_len, _, head_dim = query.shape + total = batch * seq_len * world_size * head_count * head_dim + grid = ((total + 255) // 256,) + _pack_ulysses_qkv_head_chunk_kernel[grid]( + output, + query, + key, + value, + total, + batch, + seq_len, + world_size, + local_heads, + head_count, + head_offset, + head_dim, + *query.stride(), + *key.stride(), + *value.stride(), + nccl_layout, + BLOCK=256, + ) + + +def _pack_ulysses_output_sequence_chunk( + output, + source, + *, + world_size, +): + batch, global_seq, chunk_heads, head_dim = source.shape + local_seq = global_seq // world_size + total = source.numel() + grid = ((total + 255) // 256,) + _pack_ulysses_output_sequence_chunk_kernel[grid]( + output, + source, + total, + batch, + local_seq, + world_size, + chunk_heads, + head_dim, + *source.stride(), + BLOCK=256, + ) + + +def _merge_ulysses_output_head_chunk( + received_rank_major, + output, + *, + world_size, + local_heads, + head_offset, +): + world, batch, local_seq, chunk_heads, head_dim = received_rank_major.shape + assert world == world_size + total = received_rank_major.numel() + grid = ((total + 255) // 256,) + _merge_ulysses_output_head_chunk_kernel[grid]( + received_rank_major, + output, + total, + batch, + local_seq, + world_size, + local_heads, + chunk_heads, + head_offset, + head_dim, + *received_rank_major.stride(), + *output.stride(), + BLOCK=256, + ) diff --git a/tests/comm/test_ulysses_communicator.py b/tests/comm/test_ulysses_communicator.py index 15d4b9631df..f6ebbea059e 100644 --- a/tests/comm/test_ulysses_communicator.py +++ b/tests/comm/test_ulysses_communicator.py @@ -16,7 +16,7 @@ import torch import torch.distributed as dist -from flashinfer.comm import UlyssesCommunicator +from flashinfer.comm import UlyssesCommunicator, UlyssesWorkspace from flashinfer.comm.ulysses_topology import UlyssesBackendError, UlyssesRankTopology @@ -228,6 +228,44 @@ def test_w1_passthrough_no_copy(gloo_pg, monkeypatch): comm.close() +@requires_cuda +def test_w1_destination_passing_and_workspace(gloo_pg, monkeypatch): + _forbid_ipc_and_jit(monkeypatch) + comm = _make_w1(gloo_pg, monkeypatch, max_elems=4096) + x = torch.randn(2, 8, 4, 16, dtype=torch.float16, device="cuda") + out = torch.empty_like(x) + workspace = comm.create_workspace(max_elems=x.numel()) + + assert isinstance(workspace, UlyssesWorkspace) + assert comm.scatter_heads(x, out=out, workspace=workspace) is out + assert torch.equal(out, x) + out.zero_() + assert comm.gather_heads(x, out=out, workspace=workspace) is out + assert torch.equal(out, x) + + with pytest.raises(ValueError, match="expected"): + comm.scatter_heads(x, out=torch.empty(1, device="cuda", dtype=x.dtype)) + with pytest.raises(TypeError, match="UlyssesWorkspace"): + comm.scatter_heads(x, workspace=object()) + with pytest.raises(ValueError, match="capacity"): + comm.scatter_heads(x, workspace=comm.create_workspace(max_elems=x.numel() - 1)) + with pytest.raises(ValueError, match="exceeds communicator"): + comm.create_workspace(max_elems=comm.max_elems + 1) + with pytest.raises(ValueError, match="workspace send buffer"): + comm.scatter_heads( + x, + out=workspace.send_buffer[: x.numel()].view_as(x), + workspace=workspace, + ) + with pytest.raises(ValueError, match="workspace receive buffer"): + comm.gather_heads( + x, + out=workspace.recv_buffer[: x.numel()].view_as(x), + workspace=workspace, + ) + comm.close() + + # ---- operand validation -------------------------------------------------------- @@ -508,6 +546,242 @@ def _correctness_body(rank, world_size, group, backend): return ("ok", "correct") +def _destination_passing_body(rank, world_size, group, backend): + B, S_local, H, D = 2, 7, 24, 32 + max_elems = B * S_local * H * D + comm = UlyssesCommunicator( + group, max_elems=max_elems, dtype=torch.bfloat16, backend=backend + ) + assert comm.backend == backend + workspace = comm.create_workspace() + send_ptr = workspace.send_buffer.data_ptr() + recv_ptr = workspace.recv_buffer.data_ptr() + + for iteration in range(3): + torch.manual_seed(1000 + iteration + rank) + x = torch.randn(B, S_local, H, D, dtype=torch.bfloat16, device="cuda") + scatter_out = torch.empty( + B, + S_local * world_size, + H // world_size, + D, + dtype=x.dtype, + device=x.device, + ) + returned = comm.scatter_heads(x, out=scatter_out, workspace=workspace) + ref = _ref_scatter_heads(x, world_size, rank, group) + torch.cuda.synchronize() + assert returned is scatter_out + assert torch.equal(scatter_out, ref) + + torch.manual_seed(2000 + iteration + rank) + y = torch.randn_like(scatter_out) + gather_out = torch.empty_like(x) + returned = comm.gather_heads(y, out=gather_out, workspace=workspace) + ref = _ref_gather_heads(y, world_size, rank, group) + torch.cuda.synchronize() + assert returned is gather_out + assert torch.equal(gather_out, ref) + assert workspace.send_buffer.data_ptr() == send_ptr + assert workspace.recv_buffer.data_ptr() == recv_ptr + + # After the communicator and NCCL path have warmed, explicit destinations + # plus workspace must not call Python-level tensor allocation APIs. + torch.manual_seed(4000 + rank) + x = torch.randn(B, S_local, H, D, dtype=torch.bfloat16, device="cuda") + scatter_out = torch.empty( + B, + S_local * world_size, + H // world_size, + D, + dtype=x.dtype, + device=x.device, + ) + y = torch.randn_like(scatter_out) + gather_out = torch.empty_like(x) + scatter_ref = _ref_scatter_heads(x, world_size, rank, group) + gather_ref = _ref_gather_heads(y, world_size, rank, group) + original_empty = torch.empty + original_empty_like = torch.empty_like + + def reject_allocation(*args, **kwargs): + raise AssertionError("stable workspace path attempted a tensor allocation") + + torch.empty = reject_allocation + torch.empty_like = reject_allocation + allocations_before = torch.cuda.memory_stats()["allocation.all.allocated"] + try: + comm.scatter_heads(x, out=scatter_out, workspace=workspace) + comm.gather_heads(y, out=gather_out, workspace=workspace) + finally: + torch.empty = original_empty + torch.empty_like = original_empty_like + torch.cuda.synchronize() + allocations_after = torch.cuda.memory_stats()["allocation.all.allocated"] + assert allocations_after == allocations_before + assert torch.equal(scatter_out, scatter_ref) + assert torch.equal(gather_out, gather_ref) + + comm.close() + return ("ok", backend) + + +def _head_chunk_body(rank, world_size, group, arg): + if isinstance(arg, tuple): + backend, dtype_name = arg + else: + backend, dtype_name = arg, "bfloat16" + dtype = getattr(torch, dtype_name) + B, S_local, H, D = 2, 7, 14 * world_size, 32 + schedule = (3, 8, 3) + max_payload_elems = 3 * B * S_local * world_size * max(schedule) * D + comm = UlyssesCommunicator( + group, + max_elems=max_payload_elems, + dtype=dtype, + backend=backend, + ) + assert comm.backend == backend + input_workspace = comm.create_workspace() + output_workspace = comm.create_workspace() + + # Q/K/V are independent non-contiguous views of one fused projection. + torch.manual_seed(3000 + rank) + projection = torch.randn(B, S_local, H, 3, D, dtype=dtype, device="cuda") + query, key, value = projection.unbind(dim=3) + assert not query.is_contiguous() + expected_q = _ref_scatter_heads(query, world_size, rank, group) + expected_k = _ref_scatter_heads(key, world_size, rank, group) + expected_v = _ref_scatter_heads(value, world_size, rank, group) + full_output = torch.full_like(query.contiguous(), float("nan")) + + offset = 0 + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for head_count in schedule: + fused_out = torch.empty( + B, + S_local * world_size, + head_count, + 3 * D, + dtype=query.dtype, + device=query.device, + ) + returned = comm.scatter_qkv_head_chunk( + query, + key, + value, + head_offset=offset, + head_count=head_count, + out=fused_out, + workspace=input_workspace, + ) + assert returned is fused_out + expected = torch.cat( + ( + expected_q[:, :, offset : offset + head_count], + expected_k[:, :, offset : offset + head_count], + expected_v[:, :, offset : offset + head_count], + ), + dim=-1, + ) + assert torch.equal(fused_out, expected) + + # Identity attention makes the complete reverse path reconstruct + # the original Q tensor after all bands have been merged. + attention_out = fused_out[..., :D] + comm.gather_output_head_chunk( + attention_out, + local_heads=14, + head_offset=offset, + out=full_output, + workspace=output_workspace, + ) + offset += head_count + stream.synchronize() + assert torch.equal(full_output, query) + + # The B=1 NCCL fast lifetime path may return a direct view of the + # workspace receive buffer, avoiding both an output allocation and an + # unpack copy. Its contents are valid until this workspace is reused. + if backend == "nccl": + projection_b1 = projection[:1] + q1, k1, v1 = projection_b1.unbind(dim=3) + direct = comm.scatter_qkv_head_chunk( + q1, + k1, + v1, + head_offset=3, + head_count=8, + workspace=input_workspace, + ) + expected = torch.cat( + ( + _ref_scatter_heads(q1, world_size, rank, group)[:, :, 3:11], + _ref_scatter_heads(k1, world_size, rank, group)[:, :, 3:11], + _ref_scatter_heads(v1, world_size, rank, group)[:, :, 3:11], + ), + dim=-1, + ) + torch.cuda.synchronize() + assert direct.data_ptr() == input_workspace.recv_buffer.data_ptr() + assert torch.equal(direct, expected) + + # The explicit-buffer head-chunk transport itself stays allocation-free + # after warmup (attention allocation is outside this primitive's scope). + head_count = 3 + fused_out = torch.empty( + B, + S_local * world_size, + head_count, + 3 * D, + dtype=dtype, + device="cuda", + ) + noalloc_output = torch.full_like(query.contiguous(), float("nan")) + original_empty = torch.empty + original_empty_like = torch.empty_like + + def reject_allocation(*args, **kwargs): + raise AssertionError("head-chunk transport attempted a tensor allocation") + + torch.empty = reject_allocation + torch.empty_like = reject_allocation + allocations_before = torch.cuda.memory_stats()["allocation.all.allocated"] + try: + comm.scatter_qkv_head_chunk( + query, + key, + value, + head_offset=0, + head_count=head_count, + out=fused_out, + workspace=input_workspace, + ) + comm.gather_output_head_chunk( + fused_out[..., :D], + local_heads=14, + head_offset=0, + out=noalloc_output, + workspace=output_workspace, + ) + finally: + torch.empty = original_empty + torch.empty_like = original_empty_like + torch.cuda.synchronize() + allocations_after = torch.cuda.memory_stats()["allocation.all.allocated"] + assert allocations_after == allocations_before + for source_rank in range(world_size): + start = source_rank * 14 + assert torch.equal( + noalloc_output[:, :, start : start + head_count], + query[:, :, start : start + head_count], + ) + comm.close() + return ("ok", backend) + + def _api_body(rank, world_size, group, backend): comm = UlyssesCommunicator( group, max_elems=1 << 16, dtype=torch.bfloat16, backend=backend @@ -1211,6 +1485,24 @@ def test_correctness_forced_nccl(world_size): _run_multi_rank("_correctness_body", world_size, "nccl") +def test_destination_passing_forced_nccl(): + _run_multi_rank("_destination_passing_body", 2, "nccl") + + +def test_destination_passing_forced_nvlink(): + _run_multi_rank("_destination_passing_body", 2, "nvlink", allow_skip=True) + + +@pytest.mark.parametrize("dtype_name", ["float16", "bfloat16"]) +@pytest.mark.parametrize("world_size", [2, 4]) +def test_head_chunk_forced_nccl(world_size, dtype_name): + _run_multi_rank("_head_chunk_body", world_size, ("nccl", dtype_name)) + + +def test_head_chunk_forced_nvlink(): + _run_multi_rank("_head_chunk_body", 2, ("nvlink", "bfloat16"), allow_skip=True) + + def test_api_auto_ws3_falls_back_to_nccl(): # 3 is not a fused-kernel world size: auto must fall back and say why. results = _run_multi_rank("_api_body", 3, "auto") diff --git a/tests/comm/test_ulysses_head_chunk.py b/tests/comm/test_ulysses_head_chunk.py new file mode 100644 index 00000000000..371b5cb92ad --- /dev/null +++ b/tests/comm/test_ulysses_head_chunk.py @@ -0,0 +1,178 @@ +""" +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import pytest +import torch + +from flashinfer.comm import ( + merge_ulysses_output_head_chunk, + pack_ulysses_qkv_head_chunk, +) + + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="needs a CUDA device" +) + + +def _reference_pack(q, k, v, world_size, head_offset, head_count): + local_heads = q.shape[2] // world_size + planes = [] + for tensor in (q, k, v): + bands = [ + tensor[ + :, + :, + rank * local_heads + head_offset : rank * local_heads + + head_offset + + head_count, + :, + ] + for rank in range(world_size) + ] + planes.append(torch.cat(bands, dim=2)) + return torch.cat(planes, dim=-1).contiguous() + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("schedule", [(3, 8, 3), (5, 5, 4)]) +def test_pack_qkv_uneven_schedule_noncontiguous_views(dtype, schedule): + B, S, W, H_LOCAL, D = 2, 11, 4, 14, 32 + projection = torch.randn( + B, + S, + W * H_LOCAL, + 3, + D, + device="cuda", + dtype=dtype, + ) + query, key, value = projection.unbind(dim=3) + assert not query.is_contiguous() + + head_offset = 0 + for head_count in schedule: + out = torch.empty( + B, + S, + W * head_count, + 3 * D, + device=query.device, + dtype=query.dtype, + ) + returned = pack_ulysses_qkv_head_chunk( + query, + key, + value, + world_size=W, + head_offset=head_offset, + head_count=head_count, + out=out, + ) + expected = _reference_pack(query, key, value, W, head_offset, head_count) + assert returned is out + assert torch.equal(out, expected) + head_offset += head_count + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("schedule", [(3, 8, 3), (5, 5, 4)]) +def test_merge_output_reconstructs_uneven_schedule_on_nondefault_stream( + dtype, schedule +): + B, S, W, H_LOCAL, D = 2, 13, 4, 14, 32 + full = torch.randn(B, S, W * H_LOCAL, D, device="cuda", dtype=dtype) + reconstructed = torch.full_like(full, float("nan")) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + offset = 0 + with torch.cuda.stream(stream): + for head_count in schedule: + compact = torch.cat( + [ + full[ + :, + :, + rank * H_LOCAL + offset : rank * H_LOCAL + offset + head_count, + :, + ] + for rank in range(W) + ], + dim=2, + ).contiguous() + returned = merge_ulysses_output_head_chunk( + compact, + world_size=W, + local_heads=H_LOCAL, + head_offset=offset, + out=reconstructed, + ) + assert returned is reconstructed + offset += head_count + stream.synchronize() + assert torch.equal(reconstructed, full) + + +def test_head_chunk_validation(): + qkv = torch.randn(1, 4, 8, 3, 16, device="cuda", dtype=torch.float16) + q, k, v = qkv.unbind(dim=3) + with pytest.raises(ValueError, match="divisible"): + pack_ulysses_qkv_head_chunk( + q, + k, + v, + world_size=3, + head_offset=0, + head_count=1, + ) + with pytest.raises(ValueError, match="exceeds local_heads"): + pack_ulysses_qkv_head_chunk( + q, + k, + v, + world_size=2, + head_offset=3, + head_count=2, + ) + with pytest.raises(ValueError, match="does not match query"): + pack_ulysses_qkv_head_chunk( + q, + k[:, :, :-1], + v, + world_size=2, + head_offset=0, + head_count=1, + ) + with pytest.raises(ValueError, match="must not alias"): + pack_ulysses_qkv_head_chunk( + q, + k, + v, + world_size=2, + head_offset=0, + head_count=4, + out=qkv.view(1, 4, 8, 3 * 16), + ) + + compact = torch.randn(1, 4, 4, 16, device="cuda", dtype=torch.float16) + with pytest.raises(ValueError, match="expected"): + merge_ulysses_output_head_chunk( + compact, + world_size=2, + local_heads=4, + head_offset=0, + out=torch.empty(1, 4, 7, 16, device="cuda", dtype=torch.float16), + )