diff --git a/README.md b/README.md index 332bdd2a..e1d4c545 100644 --- a/README.md +++ b/README.md @@ -146,6 +146,9 @@ print(f'Final state shape: {final_state.shape}') # [2, 32, 128, 128] See [USAGE.md](USAGE.md) for detailed usage examples and notes. See [the standalone kernel inventory](docs/KERNEL_VALUE.md) for the public kernel boundary and the overlap assessment against FlashInfer and cuDNN. +The optional [cuDNN hybrid-attention runtime](docs/HYBRID_CUDNN_RUNTIME.md) +documents packed index pooling, memory-bounded DSA selection, and CP layout +ownership without precision downgrades. ## Benchmarks diff --git a/REPO_LAYOUT.md b/REPO_LAYOUT.md index d9df540e..d3c53d74 100644 --- a/REPO_LAYOUT.md +++ b/REPO_LAYOUT.md @@ -21,6 +21,10 @@ cuLA/ │ ├── distributed/ # Optional cuDNN KDA Kimi CP runtime │ │ ├── plan.py # CPU-only packed-THD boundary planning │ │ └── kda.py # Collectives and distributed autograd +│ ├── hybrid/ # Layout-aware cuDNN hybrid-attention runtime +│ │ ├── index_pool.py # Differentiable CSA index pooling +│ │ ├── dsa.py # Bounded packed-THD DSA block selection +│ │ └── plan.py # CP block ownership and seam-halo plans │ │ │ ├── lightning/ # [non-KDA] Lightning Attention public wrappers │ │ @@ -48,6 +52,13 @@ cuLA/ │ └── experimental/ # [non-KDA] unwired prototypes │ ├── benchmarks/ # Accuracy and performance runners +│ └── large_cp_emulation/ # Logical CP-on-fewer-GPUs simulation layer +│ ├── topology.py # Logical/physical rank separation +│ ├── transport.py # Native-CP protocol cost emulation +│ ├── workload.py # Packed-THD workload descriptions +│ ├── bench_components.py # DSA and ring microbenchmarks +│ ├── bench_kda.py # KDA logical-CP E2E +│ └── bench_hybrid.py # KDA + DSA serial/overlap E2E ├── tests/ # CPU registry tests and architecture-marked GPU tests ├── docs/ # Pipeline and backend-dispatch design notes ├── scripts/build_wheel.sh @@ -61,6 +72,8 @@ cuLA/ |-----------|----------|-------------| | `cula/kernels/` | Python | Framework-independent public surface and static kernel inventory. Imports are lazy and do not compile kernels. | | `cula/kda/` | Python / CuTeDSL | KDA prefill/decode wrappers and the pure-CuTeDSL SM100 modular composition. `kda_prefill` calls FlashKDA directly without a native fallback. See [`cula/kda/README.md`](cula/kda/README.md). | +| `cula/hybrid/` | Python | Optional cuDNN-backed composition layer for packed index pooling, bounded DSA selection, and CP layout planning. See [`docs/HYBRID_CUDNN_RUNTIME.md`](docs/HYBRID_CUDNN_RUNTIME.md). | +| `benchmarks/large_cp_emulation/` | Python | Benchmark-only logical-CP layer built on the native `cula.distributed` and `cula.hybrid` APIs; it never participates in production imports. | | `cula/ops/kda/` | Python (CuTeDSL) | CuTeDSL KDA kernels organized into SM100 modular kernels, SM90 FlashKDA K1+K2 and intracard CP, decode, and experimental code. | | `cula/ops/lightning/` · `cula/ops/experimental/` | Python (CuTeDSL) | `[non-KDA]` Lightning/linear-attention kernels and prototypes. Lightning prefill dispatches to the SM90 or SM100 backend from Q's device capability. | | `cula/ops/{inv,ptx}.py` · `cula/ops/sm100/ptx.py` | Python | Shared low-level helpers used across operators. | diff --git a/USAGE.md b/USAGE.md index e269c64d..8719afb9 100644 --- a/USAGE.md +++ b/USAGE.md @@ -4,6 +4,21 @@ This document provides usage notes and examples for cuLA kernels. --- +## cuDNN hybrid index pooling and DSA selection + +The optional hybrid runtime preserves packed THD layouts between cuDNN index +pooling and DSA selection. It owns static metadata and bounded scratch, while +cuDNN owns the fused kernels and precision contract. Install it with: + +```bash +pip install 'cuda-linear-attention[cudnn-hybrid]' +``` + +See [`docs/HYBRID_CUDNN_RUNTIME.md`](docs/HYBRID_CUDNN_RUNTIME.md) for the API, +128K heterogeneous/CP4 measurements, and current implementation boundary. + +--- + ## cuDNN KDA Kimi context parallelism The optional distributed runtime executes cuDNN KDA over contiguous packed-THD diff --git a/benchmarks/bench_hybrid_cp_emulation.py b/benchmarks/bench_hybrid_cp_emulation.py new file mode 100644 index 00000000..078b009d --- /dev/null +++ b/benchmarks/bench_hybrid_cp_emulation.py @@ -0,0 +1,18 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 +# ruff: noqa: E402, I001 + +"""Compatibility entry point for large-CP component emulation.""" + +from __future__ import annotations + +import pathlib +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) + +from benchmarks.large_cp_emulation.bench_components import main + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_hybrid_kda_dsa_e2e.py b/benchmarks/bench_hybrid_kda_dsa_e2e.py new file mode 100644 index 00000000..63d26bbb --- /dev/null +++ b/benchmarks/bench_hybrid_kda_dsa_e2e.py @@ -0,0 +1,18 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 +# ruff: noqa: E402, I001 + +"""Compatibility entry point for the large-CP hybrid emulator.""" + +from __future__ import annotations + +import pathlib +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) + +from benchmarks.large_cp_emulation.bench_hybrid import main + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_index_pool_cp.py b/benchmarks/bench_index_pool_cp.py new file mode 100644 index 00000000..fb240cd4 --- /dev/null +++ b/benchmarks/bench_index_pool_cp.py @@ -0,0 +1,350 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Compare gather-then-pool with cuLA's distributed pool-before-gather path.""" + +from __future__ import annotations + +import argparse +import gc +import json +import os +import pathlib +import statistics +import sys + +import torch +import torch.distributed as dist + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1])) + +from cula.hybrid import ( + IndexPoolPlan, + IndexPoolShardPlan, + cudnn_index_pool, + cudnn_index_pool_cp, +) + + +class _GatherReferenceRows(torch.autograd.Function): + """Reference raw-row all-gather with FP32 gradient reduction.""" + + @staticmethod + def forward(ctx, local, group, offsets, rank): + counts = tuple(right - left for left, right in zip(offsets, offsets[1:])) + if any(count != counts[0] for count in counts[1:]): + raise ValueError("the benchmark reference requires equal token shards") + output = local.new_empty((sum(counts), *local.shape[1:])) + if hasattr(dist, "all_gather_single"): + dist.all_gather_single(output, local.contiguous(), group=group) + else: + dist.all_gather_into_tensor(output, local.contiguous(), group=group) + ctx.group = group + ctx.local_count = counts[rank] + return output + + @staticmethod + def backward(ctx, grad_output): + local_grad = torch.empty( + (ctx.local_count, *grad_output.shape[1:]), + dtype=torch.float32, + device=grad_output.device, + ) + source = grad_output.float().contiguous() + if hasattr(dist, "reduce_scatter_single"): + dist.reduce_scatter_single(local_grad, source, group=ctx.group) + else: + dist.reduce_scatter_tensor(local_grad, source, group=ctx.group) + return local_grad.to(grad_output.dtype), None, None, None + + +def _boundaries(total: int, short_long_mix: bool) -> tuple[int, ...]: + if not short_long_mix: + return (0, total) + short_total = 20 * 1024 + remaining = total - short_total + if remaining <= 0 or remaining % 2: + raise ValueError("20-short+2-long requires an even positive long-sequence remainder") + lengths = [1024] * 20 + [remaining // 2] * 2 + result = [0] + for length in lengths: + result.append(result[-1] + length) + return tuple(result) + + +def _measure_once(callable_) -> float: + dist.barrier() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + result = callable_() + end.record() + end.synchronize() + elapsed = start.elapsed_time(end) + del result + return elapsed + + +def _measure_pair(left, right, repeats: int) -> tuple[list[float], list[float]]: + """Interleave paths and reverse order to limit drift and ordering bias.""" + + left_samples = [] + right_samples = [] + for repeat in range(repeats): + if repeat % 2: + right_samples.append(_measure_once(right)) + left_samples.append(_measure_once(left)) + else: + left_samples.append(_measure_once(left)) + right_samples.append(_measure_once(right)) + return left_samples, right_samples + + +def _measure_incremental_peak(callable_) -> int: + gc.collect() + torch.cuda.empty_cache() + torch.cuda.synchronize() + baseline = torch.cuda.memory_allocated() + torch.cuda.reset_peak_memory_stats() + result = callable_() + torch.cuda.synchronize() + peak = torch.cuda.max_memory_allocated() - baseline + del result + return peak + + +def _local_backward_oracle(kv, score, ape, grad_blocks, plan): + members = [] + for start, end in zip(plan.cu_seqlens, plan.cu_seqlens[1:]): + complete_end = start + (end - start) // plan.ratio * plan.ratio + if complete_end > start: + members.append(torch.arange(start, complete_end, device=kv.device).view(-1, plan.ratio)) + members = torch.cat(members) + grouped_kv = kv.index_select(0, members.flatten()).view(-1, plan.ratio, plan.head_dim).float() + grouped_score = score.index_select(0, members.flatten()).view_as(grouped_kv).float() + weights = torch.softmax(grouped_score + ape.float().unsqueeze(0), dim=1) + dy = grad_blocks.float().unsqueeze(1) + dweight = grouped_kv * dy + dscore_grouped = weights * (dweight - (dweight * weights).sum(dim=1, keepdim=True)) + dkv = torch.zeros_like(kv, dtype=torch.float32) + dscore = torch.zeros_like(score, dtype=torch.float32) + dkv.index_copy_(0, members.flatten(), (weights * dy).flatten(0, 1)) + dscore.index_copy_(0, members.flatten(), dscore_grouped.flatten(0, 1)) + return dkv, dscore, dscore_grouped.sum(dim=0) + + +def _errors(actual, oracle): + error = actual.float() - oracle.float() + return { + "max_abs": float(error.abs().max()), + "rmse": float(error.square().mean().sqrt()), + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--global-tokens", type=int, default=1 << 17) + parser.add_argument("--ratio", type=int, default=4) + parser.add_argument("--head-dim", type=int, default=128) + parser.add_argument("--repeats", type=int, default=7) + parser.add_argument("--thd-20s2l", action="store_true") + args = parser.parse_args() + + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group("nccl") + rank = dist.get_rank() + world_size = dist.get_world_size() + device = torch.device("cuda", local_rank) + boundaries = _boundaries(args.global_tokens, args.thd_20s2l) + shard_plan = IndexPoolShardPlan.from_global_cu_seqlens( + boundaries, + rank=rank, + world_size=world_size, + head_dim=args.head_dim, + ratio=args.ratio, + ) + if not shard_plan.all_shards_direct or shard_plan.local_plan is None: + raise ValueError("benchmark geometry must use the aligned pool-before-gather path") + + torch.manual_seed(1104 + rank) + kv = torch.randn( + shard_plan.local_tokens, + args.head_dim, + dtype=torch.bfloat16, + device=device, + ) + score = torch.randn_like(kv) + torch.manual_seed(47) + ape = torch.randn(args.ratio, args.head_dim, dtype=torch.float32, device=device) + local_cu, local_cb = shard_plan.local_plan.materialize(device) + global_plan = IndexPoolPlan.from_cu_seqlens( + boundaries, + head_dim=args.head_dim, + ratio=args.ratio, + ) + global_cu, global_cb = global_plan.materialize(device) + token_offsets = tuple(shard * args.global_tokens // world_size for shard in range(world_size + 1)) + + def reference(kv_, score_, ape_): + global_kv = _GatherReferenceRows.apply(kv_, dist.group.WORLD, token_offsets, rank) + global_score = _GatherReferenceRows.apply(score_, dist.group.WORLD, token_offsets, rank) + return cudnn_index_pool( + global_kv, + global_score, + ape_, + global_cu, + global_cb, + plan=global_plan, + deterministic=False, + ) + + def candidate(kv_, score_, ape_): + return cudnn_index_pool_cp( + kv_, + score_, + ape_, + local_cu, + local_cb, + plan=shard_plan, + group=dist.group.WORLD, + deterministic=False, + ) + + ref_inputs = tuple(t.detach().clone().requires_grad_(True) for t in (kv, score, ape)) + opt_inputs = tuple(t.detach().clone().requires_grad_(True) for t in (kv, score, ape)) + ref_out = reference(*ref_inputs) + opt_out = candidate(*opt_inputs) + torch.manual_seed(9000 + rank) + grad_out = torch.randn_like(ref_out) + ref_grads = list(torch.autograd.grad(ref_out, ref_inputs, grad_out)) + opt_grads = list(torch.autograd.grad(opt_out, opt_inputs, grad_out)) + total_grad_out = grad_out.float() + dist.all_reduce(total_grad_out) + local_begin = shard_plan.block_offsets[rank] + local_end = shard_plan.block_offsets[rank + 1] + oracle_grads = list( + _local_backward_oracle( + kv, + score, + ape, + total_grad_out[local_begin:local_end], + shard_plan.local_plan, + ) + ) + # APE is replicated, so compare the CP-global parameter gradient. + dist.all_reduce(ref_grads[2]) + dist.all_reduce(opt_grads[2]) + dist.all_reduce(oracle_grads[2]) + correctness = { + "out_bitwise": bool(torch.equal(ref_out, opt_out)), + "dk_bitwise": bool(torch.equal(ref_grads[0], opt_grads[0])), + "dk_max_abs": float((ref_grads[0].float() - opt_grads[0].float()).abs().max()), + "dscore_bitwise": bool(torch.equal(ref_grads[1], opt_grads[1])), + "dscore_max_abs": float((ref_grads[1].float() - opt_grads[1].float()).abs().max()), + "dape_max_abs": float((ref_grads[2] - opt_grads[2]).abs().max()), + "dape_max_rel": float(((ref_grads[2] - opt_grads[2]).abs() / ref_grads[2].abs().clamp_min(1e-12)).max()), + "reference_vs_fp32_oracle": { + "dk": _errors(ref_grads[0], oracle_grads[0]), + "dscore": _errors(ref_grads[1], oracle_grads[1]), + "dape": _errors(ref_grads[2], oracle_grads[2]), + }, + "candidate_vs_fp32_oracle": { + "dk": _errors(opt_grads[0], oracle_grads[0]), + "dscore": _errors(opt_grads[1], oracle_grads[1]), + "dape": _errors(opt_grads[2], oracle_grads[2]), + }, + } + + with torch.no_grad(): + reference(kv, score, ape) + candidate(kv, score, ape) + torch.cuda.synchronize() + ref_samples, opt_samples = _measure_pair( + lambda: reference(kv, score, ape), + lambda: candidate(kv, score, ape), + args.repeats, + ) + + def reference_forward_backward(): + inputs = tuple(t.detach().requires_grad_(True) for t in (kv, score, ape)) + output = reference(*inputs) + return torch.autograd.grad(output, inputs, grad_out) + + def candidate_forward_backward(): + inputs = tuple(t.detach().requires_grad_(True) for t in (kv, score, ape)) + output = candidate(*inputs) + return torch.autograd.grad(output, inputs, grad_out) + + reference_forward_backward() + candidate_forward_backward() + torch.cuda.synchronize() + ref_peak = _measure_incremental_peak(reference_forward_backward) + opt_peak = _measure_incremental_peak(candidate_forward_backward) + ref_fb_samples, opt_fb_samples = _measure_pair( + reference_forward_backward, + candidate_forward_backward, + args.repeats, + ) + + local = torch.tensor( + [ + ref_samples, + opt_samples, + ref_fb_samples, + opt_fb_samples, + [ref_peak] * args.repeats, + [opt_peak] * args.repeats, + ], + dtype=torch.float64, + device=device, + ) + gathered = torch.empty( + world_size * 6 * args.repeats, + dtype=torch.float64, + device=device, + ) + if hasattr(dist, "all_gather_single"): + dist.all_gather_single(gathered, local.flatten()) + else: + dist.all_gather_into_tensor(gathered, local.flatten()) + if rank == 0: + samples = gathered.cpu().view(world_size, 6, args.repeats) + ref_critical = samples[:, 0].max(dim=0).values.tolist() + opt_critical = samples[:, 1].max(dim=0).values.tolist() + ref_fb_critical = samples[:, 2].max(dim=0).values.tolist() + opt_fb_critical = samples[:, 3].max(dim=0).values.tolist() + ref_peak_bytes = int(samples[:, 4, 0].max()) + opt_peak_bytes = int(samples[:, 5, 0].max()) + ref_median = statistics.median(ref_critical) + opt_median = statistics.median(opt_critical) + ref_fb_median = statistics.median(ref_fb_critical) + opt_fb_median = statistics.median(opt_fb_critical) + print( + json.dumps( + { + "global_tokens": args.global_tokens, + "world_size": world_size, + "ratio": args.ratio, + "sequence_lengths": [right - left for left, right in zip(boundaries, boundaries[1:])], + "raw_collective_bytes_per_rank": 2 * kv.numel() * kv.element_size(), + "pooled_collective_bytes_per_rank": (shard_plan.local_blocks * args.head_dim * kv.element_size()), + "correctness": correctness, + "gather_then_pool_median_ms": ref_median, + "pool_then_gather_median_ms": opt_median, + "speedup": ref_median / opt_median, + "gather_then_pool_forward_backward_median_ms": ref_fb_median, + "pool_then_gather_forward_backward_median_ms": opt_fb_median, + "forward_backward_speedup": ref_fb_median / opt_fb_median, + "gather_then_pool_incremental_peak_bytes": ref_peak_bytes, + "pool_then_gather_incremental_peak_bytes": opt_peak_bytes, + "incremental_peak_reduction": ref_peak_bytes / opt_peak_bytes, + } + ), + flush=True, + ) + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_kimi_cp_e2e.py b/benchmarks/bench_kimi_cp_e2e.py new file mode 100644 index 00000000..47962e8c --- /dev/null +++ b/benchmarks/bench_kimi_cp_e2e.py @@ -0,0 +1,18 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 +# ruff: noqa: E402, I001 + +"""Compatibility entry point for the large-CP KDA emulator.""" + +from __future__ import annotations + +import pathlib +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) + +from benchmarks.large_cp_emulation.bench_kda import main + + +if __name__ == "__main__": + main() diff --git a/benchmarks/large_cp_emulation/README.md b/benchmarks/large_cp_emulation/README.md new file mode 100644 index 00000000..6f4d4c8d --- /dev/null +++ b/benchmarks/large_cp_emulation/README.md @@ -0,0 +1,33 @@ +# Large logical-CP emulation + +This directory is a benchmark-only layer built on cuLA's native CP public +APIs. Production code under `cula/distributed` and `cula/hybrid` never imports +this package. + +The emulators preserve real per-rank tensor shapes and execute real local +cuDNN kernels plus real NCCL hops on the available physical GPUs. They model a +logical dependency length; they do not claim to reproduce a 64-GPU topology, +cross-node congestion, or aggregate bandwidth. + +## Structure + +- `topology.py`: logical-to-physical rank mapping and transport-span plans. +- `transport.py`: benchmark-only implementation of the native KDA boundary + transport protocol and shared dependent-ring cost emulation. +- `workload.py`: reusable packed-THD workload descriptions. +- `bench_components.py`: isolated DSA and ring microbenchmarks. +- `bench_kda.py`: complete native KDA forward/backward with logical CP cost. +- `bench_hybrid.py`: native KDA plus DSA-indexer serial/overlap benchmark. + +`full-ring` emulates `logical_cp - 1` hops. `sample-chain` uses the longest +packed sample's actual CP span. Neither mode changes cuDNN's KDA summary +composition arithmetic or DSA arithmetic. + +Run from the repository root, for example: + +```bash +torchrun --standalone --nproc-per-node=4 \ + benchmarks/large_cp_emulation/bench_hybrid.py logical \ + --global-tokens 1048576 --logical-cp 64 --logical-transport sample-chain \ + --thd-20s2l --schedule both +``` diff --git a/benchmarks/large_cp_emulation/__init__.py b/benchmarks/large_cp_emulation/__init__.py new file mode 100644 index 00000000..a7cd240c --- /dev/null +++ b/benchmarks/large_cp_emulation/__init__.py @@ -0,0 +1,17 @@ +"""Large logical-CP emulation built exclusively on cuLA's native CP APIs.""" + +from benchmarks.large_cp_emulation.topology import ( + LogicalCPEmulationPlan, + representative_logical_rank, + sample_chain_hops, +) +from benchmarks.large_cp_emulation.transport import LogicalCPRingTransport +from benchmarks.large_cp_emulation.workload import PackedTHDWorkload + +__all__ = [ + "LogicalCPEmulationPlan", + "LogicalCPRingTransport", + "PackedTHDWorkload", + "representative_logical_rank", + "sample_chain_hops", +] diff --git a/benchmarks/large_cp_emulation/bench_components.py b/benchmarks/large_cp_emulation/bench_components.py new file mode 100644 index 00000000..9671b2a2 --- /dev/null +++ b/benchmarks/large_cp_emulation/bench_components.py @@ -0,0 +1,222 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Emulate large logical CP selection and ring communication on fewer GPUs. + +DSA mode executes one logical rank's real query/block geometry. Ring mode must +be launched with torchrun and sends the real per-rank payload for ``CP - 1`` +hops. It measures the serial dependency path, not a claim about 64-GPU fabric +contention or topology. +""" + +from __future__ import annotations + +import argparse +import json +import os +import pathlib +import statistics +import sys +from functools import partial + +import torch +import torch.distributed as dist + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[2])) + +from benchmarks.large_cp_emulation.transport import emulate_ring_hops # noqa: E402 +from cula.hybrid import DsaSelectionPlan, DsaSelectionWorkspace, cudnn_dsa_select # noqa: E402 + + +def _csv_ints(value: str) -> tuple[int, ...]: + try: + result = tuple(int(part) for part in value.split(",")) + except ValueError as error: + raise argparse.ArgumentTypeError("expected comma-separated integers") from error + if not result or any(item <= 0 for item in result): + raise argparse.ArgumentTypeError("all values must be positive") + return result + + +def _positive_int(value: str) -> int: + result = int(value) + if result < 1: + raise argparse.ArgumentTypeError("value must be positive") + return result + + +def _measure(callable_, repeats: int) -> list[float]: + values = [] + for _ in range(repeats): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + callable_() + end.record() + end.synchronize() + values.append(start.elapsed_time(end)) + return values + + +def _candidate_prefix(position: int, block_size: int) -> int: + quotient, remainder = divmod(position, block_size) + return block_size * quotient * (quotient - 1) // 2 + quotient * (remainder + 1) + + +def run_dsa(args) -> None: + if args.global_tokens % args.logical_cp: + raise ValueError("global_tokens must be divisible by logical_cp") + if args.global_tokens % args.block_size: + raise ValueError("global_tokens must be divisible by block_size") + local_tokens = args.global_tokens // args.logical_cp + logical_rank = args.logical_rank if args.logical_rank is not None else args.logical_cp - 1 + if not 0 <= logical_rank < args.logical_cp: + raise ValueError("logical_rank must be in [0, logical_cp)") + blocks = args.global_tokens // args.block_size + query_offset = logical_rank * local_tokens + device = torch.device("cuda", args.device) + torch.cuda.set_device(device) + torch.manual_seed(args.seed) + qi = torch.randn(local_tokens, args.index_heads, args.head_dim, dtype=torch.bfloat16, device=device) + kbar = torch.randn(blocks, 1, args.head_dim, dtype=torch.bfloat16, device=device) + weights = torch.randn(local_tokens, args.index_heads, dtype=torch.float32, device=device) + full_floats = _candidate_prefix(query_offset + local_tokens, args.block_size) - _candidate_prefix( + query_offset, args.block_size + ) + print( + json.dumps( + { + "mode": "dsa", + "global_tokens": args.global_tokens, + "logical_cp": args.logical_cp, + "logical_rank": logical_rank, + "local_tokens": local_tokens, + "index_heads": args.index_heads, + "compressed_blocks": blocks, + "full_candidate_gib": full_floats * 4 / 2**30, + } + ), + flush=True, + ) + + for budget_mib in args.budgets_mib: + plan = DsaSelectionPlan.from_cu_seqlens( + [0, local_tokens], + [0, blocks], + top_k=args.top_k, + block_size=args.block_size, + q_causal_offsets=[query_offset], + candidate_budget_bytes=budget_mib << 20, + ) + workspace = DsaSelectionWorkspace.from_plan(plan, device) + + execute = partial(cudnn_dsa_select, qi, kbar, weights, plan=plan, workspace=workspace) + + execute() + torch.cuda.synchronize(device) + elapsed = _measure(execute, args.repeats) + print( + json.dumps( + { + "index_heads": args.index_heads, + "budget_mib": budget_mib, + "windows": len(plan.windows), + "workspace_mib": plan.candidate_workspace_floats * 4 / 2**20, + "samples_ms": elapsed, + "median_ms": statistics.median(elapsed), + } + ), + flush=True, + ) + del execute, workspace + torch.cuda.empty_cache() + + +def run_ring(args) -> None: + if args.logical_cp < 2: + raise ValueError("ring emulation requires logical_cp >= 2") + if not dist.is_initialized(): + dist.init_process_group("nccl") + rank = dist.get_rank() + world_size = dist.get_world_size() + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + steps = args.logical_cp - 1 + elements = (args.payload_bytes + 3) // 4 + send = torch.randn(elements, dtype=torch.float32, device=device) + receive = torch.empty_like(send) + + def execute(): + emulate_ring_hops( + send, + receive, + physical_rank=rank, + physical_world_size=world_size, + hops=steps, + ) + + execute() + torch.cuda.synchronize(device) + dist.barrier(device_ids=[local_rank]) + elapsed = _measure(execute, args.repeats) + local = torch.tensor(elapsed, dtype=torch.float64, device=device) + gathered = torch.empty(world_size * args.repeats, dtype=torch.float64, device=device) + all_gather_single = dist.all_gather_single if hasattr(dist, "all_gather_single") else dist.all_gather_into_tensor + all_gather_single(gathered, local) + if rank == 0: + samples = gathered.cpu().reshape(world_size, args.repeats) + rank_medians = [statistics.median(row.tolist()) for row in samples] + critical_samples = samples.max(dim=0).values.tolist() + critical = statistics.median(critical_samples) + print( + json.dumps( + { + "mode": "ring", + "physical_gpus": world_size, + "logical_cp": args.logical_cp, + "steps": steps, + "payload_bytes": args.payload_bytes, + "rank_median_ms": rank_medians, + "critical_samples_ms": critical_samples, + "critical_ms": critical, + "per_hop_us": critical * 1000 / steps, + } + ), + flush=True, + ) + dist.destroy_process_group() + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + subparsers = parser.add_subparsers(dest="mode", required=True) + dsa = subparsers.add_parser("dsa") + dsa.add_argument("--global-tokens", type=_positive_int, default=1 << 20) + dsa.add_argument("--logical-cp", type=_positive_int, default=64) + dsa.add_argument("--logical-rank", type=int) + dsa.add_argument("--budgets-mib", type=_csv_ints, default=(512, 1024, 2048, 4096)) + dsa.add_argument("--block-size", type=_positive_int, default=4) + dsa.add_argument("--top-k", type=_positive_int, default=512) + dsa.add_argument("--index-heads", type=_positive_int, default=64) + dsa.add_argument("--head-dim", type=_positive_int, default=128) + dsa.add_argument("--device", type=int, default=0) + dsa.add_argument("--repeats", type=_positive_int, default=5) + dsa.add_argument("--seed", type=int, default=53) + ring = subparsers.add_parser("ring") + ring.add_argument("--logical-cp", type=_positive_int, default=64) + ring.add_argument("--payload-bytes", type=_positive_int, default=4 << 20) + ring.add_argument("--repeats", type=_positive_int, default=5) + return parser.parse_args() + + +def main(): + args = parse_args() + if args.mode == "dsa": + run_dsa(args) + else: + run_ring(args) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/large_cp_emulation/bench_dsa_layout.py b/benchmarks/large_cp_emulation/bench_dsa_layout.py new file mode 100644 index 00000000..d9f8b72a --- /dev/null +++ b/benchmarks/large_cp_emulation/bench_dsa_layout.py @@ -0,0 +1,325 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Measure DSA indexer cost under contiguous and sample-alternating CP layouts.""" + +from __future__ import annotations + +import argparse +import json +import os +import pathlib +import statistics +import sys + +import torch +import torch.distributed as dist + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[2])) + +from benchmarks.large_cp_emulation.query_layout import ( # noqa: E402 + QueryLayout, + build_alternating_sequence_query_layout, + build_contiguous_query_layout, + build_planned_query_layout, +) +from benchmarks.large_cp_emulation.workload import PackedTHDWorkload # noqa: E402 +from cula.hybrid import ( # noqa: E402 + DsaSelectionPlan, + DsaSelectionWorkspace, + IndexPoolPlan, + cudnn_dsa_select, + cudnn_index_pool, +) + + +def _positive_int(value: str) -> int: + result = int(value) + if result < 1: + raise argparse.ArgumentTypeError("value must be positive") + return result + + +def _sequence_lengths(value: str) -> tuple[int, ...]: + try: + result = tuple(int(item) for item in value.split(",")) + except ValueError as error: + raise argparse.ArgumentTypeError("sequence lengths must be comma-separated integers") from error + if not result or any(length < 1 for length in result): + raise argparse.ArgumentTypeError("sequence lengths must be positive") + return result + + +def _global_block_prefix(boundaries: tuple[int, ...], block_size: int) -> tuple[int, ...]: + result = [0] + for left, right in zip(boundaries, boundaries[1:]): + result.append(result[-1] + (right - left) // block_size) + return tuple(result) + + +def _layout_for_rank( + name: str, + boundaries: tuple[int, ...], + *, + rank: int, + world_size: int, + block_size: int, +) -> QueryLayout: + builder = { + "contiguous": build_contiguous_query_layout, + "alternating": build_alternating_sequence_query_layout, + "planned": build_planned_query_layout, + }[name] + return builder( + boundaries, + rank=rank, + world_size=world_size, + block_size=block_size, + ) + + +def _block_ids(layout: QueryLayout, block_size: int) -> tuple[int, ...]: + return tuple( + block for span in layout.spans for block in range(span.token_start // block_size, span.token_end // block_size) + ) + + +def _selection_key_ids( + layout: QueryLayout, + block_prefix: tuple[int, ...], +) -> tuple[int, ...]: + return tuple( + block for span in layout.spans for block in range(block_prefix[span.sequence], block_prefix[span.sequence + 1]) + ) + + +def _validate_block_aligned_workload(workload: PackedTHDWorkload, block_size: int) -> None: + if any(length % block_size for length in workload.sequence_lengths): + raise ValueError("every packed sequence length must be divisible by the DSA block size") + + +def _run(args: argparse.Namespace) -> None: + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + dist.init_process_group("nccl") + rank = dist.get_rank() + world_size = dist.get_world_size() + + if args.sequence_lengths is None: + workload = PackedTHDWorkload.short_long_mix( + args.global_tokens, + short_sequences=args.short_sequences, + short_tokens=args.short_tokens, + long_sequences=args.long_sequences, + ) + else: + boundaries = [0] + for length in args.sequence_lengths: + boundaries.append(boundaries[-1] + length) + workload = PackedTHDWorkload(tuple(boundaries)) + boundaries = workload.global_cu_seqlens + _validate_block_aligned_workload(workload, args.block_size) + layouts = tuple( + _layout_for_rank( + args.layout, + boundaries, + rank=owner, + world_size=world_size, + block_size=args.block_size, + ) + for owner in range(world_size) + ) + layout = layouts[rank] + local_tokens = layout.local_tokens + local_blocks = local_tokens // args.block_size + if any(candidate.local_tokens != local_tokens for candidate in layouts): + raise ValueError("this benchmark requires equal tokens per physical rank") + if any(len(_block_ids(candidate, args.block_size)) != local_blocks for candidate in layouts): + raise ValueError("this benchmark requires equal pooled blocks per physical rank") + + block_prefix = _global_block_prefix(boundaries, args.block_size) + gathered_block_ids = torch.tensor( + [block for candidate in layouts for block in _block_ids(candidate, args.block_size)], + dtype=torch.int64, + device=device, + ) + selection_key_ids = torch.tensor( + _selection_key_ids(layout, block_prefix), + dtype=torch.int64, + device=device, + ) + select_cu_k = [0] + for span in layout.spans: + select_cu_k.append(select_cu_k[-1] + block_prefix[span.sequence + 1] - block_prefix[span.sequence]) + + pool_plan = IndexPoolPlan.from_cu_seqlens( + layout.local_cu_seqlens, + head_dim=args.index_dim, + ratio=args.block_size, + ) + pool_cu_tokens, pool_cu_blocks = pool_plan.materialize(device) + budget = None if args.dsa_budget_mib == 0 else args.dsa_budget_mib << 20 + select_plan = DsaSelectionPlan.from_cu_seqlens( + layout.local_cu_seqlens, + select_cu_k, + top_k=args.top_k, + block_size=args.block_size, + q_causal_offsets=layout.causal_offsets, + candidate_budget_bytes=budget, + ) + workspace = DsaSelectionWorkspace.from_plan(select_plan, device) + + torch.manual_seed(args.seed + rank) + qi = torch.randn(local_tokens, args.index_heads, args.index_dim, dtype=torch.bfloat16, device=device) + local_kv = torch.randn(local_tokens, args.index_dim, dtype=torch.bfloat16, device=device) + pool_score = torch.randn_like(local_kv) + ape = torch.randn(args.block_size, args.index_dim, dtype=torch.float32, device=device) + weights = torch.randn(local_tokens, args.index_heads, dtype=torch.float32, device=device) + gathered_kbar = torch.empty( + world_size * local_blocks, + args.index_dim, + dtype=torch.bfloat16, + device=device, + ) + global_kbar = torch.empty( + boundaries[-1] // args.block_size, + args.index_dim, + dtype=torch.bfloat16, + device=device, + ) + selection_kbar = torch.empty( + selection_key_ids.numel(), + 1, + args.index_dim, + dtype=torch.bfloat16, + device=device, + ) + + def execute(measure: bool) -> dict[str, object]: + events = [torch.cuda.Event(enable_timing=True) for _ in range(4)] + events[0].record() + local_kbar = cudnn_index_pool( + local_kv, + pool_score, + ape, + pool_cu_tokens, + pool_cu_blocks, + plan=pool_plan, + deterministic=True, + ) + events[1].record() + if hasattr(dist, "all_gather_single"): + dist.all_gather_single(gathered_kbar, local_kbar) + else: # pragma: no cover - compatibility with older PyTorch releases + dist.all_gather_into_tensor(gathered_kbar, local_kbar) + global_kbar.index_copy_(0, gathered_block_ids, gathered_kbar) + torch.index_select(global_kbar, 0, selection_key_ids, out=selection_kbar[:, 0]) + events[2].record() + result = cudnn_dsa_select( + qi, + selection_kbar, + weights, + plan=select_plan, + workspace=workspace, + deterministic=True, + ) + events[3].record() + events[3].synchronize() + return { + "total_ms": events[0].elapsed_time(events[3]) if measure else None, + "pool_ms": events[0].elapsed_time(events[1]) if measure else None, + "remap_comm_ms": events[1].elapsed_time(events[2]) if measure else None, + "select_ms": events[2].elapsed_time(events[3]) if measure else None, + "id_sum": int(result.block_ids.to(torch.int64).sum().item()), + "score_sum": float(result.scores[torch.isfinite(result.scores)].double().sum().item()), + } + + for _ in range(args.warmups): + execute(False) + torch.cuda.reset_peak_memory_stats(device) + samples = [execute(True) for _ in range(args.repeats)] + record = { + "rank": rank, + "tokens": local_tokens, + "spans": len(layout.spans), + "candidate_floats": layout.candidate_floats, + "windows": len(select_plan.windows), + "candidate_workspace_gib": select_plan.candidate_workspace_floats * 4 / 2**30, + "peak_allocated_gib": torch.cuda.max_memory_allocated(device) / 2**30, + "samples": samples, + } + records = [None] * world_size + dist.all_gather_object(records, record) + if rank == 0: + critical = [ + max(records[owner]["samples"][sample]["total_ms"] for owner in range(world_size)) for sample in range(args.repeats) + ] + print( + json.dumps( + { + "layout": args.layout, + "global_tokens": workload.global_tokens, + "index_heads": args.index_heads, + "physical_gpus": world_size, + "sequence_lengths": workload.sequence_lengths, + "critical_samples_ms": critical, + "critical_median_ms": statistics.median(critical), + "ranks": [ + { + key: record[key] + for key in ( + "rank", + "tokens", + "spans", + "candidate_floats", + "windows", + "candidate_workspace_gib", + "peak_allocated_gib", + ) + } + | { + f"{stage}_median_ms": statistics.median(sample[stage] for sample in record["samples"]) + for stage in ("pool_ms", "remap_comm_ms", "select_ms", "total_ms") + } + | { + "id_sum_stable": all( + sample["id_sum"] == record["samples"][0]["id_sum"] for sample in record["samples"] + ), + "score_sum_span": max(sample["score_sum"] for sample in record["samples"]) + - min(sample["score_sum"] for sample in record["samples"]), + } + for record in records + ], + } + ), + flush=True, + ) + dist.destroy_process_group() + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--layout", choices=("contiguous", "alternating", "planned"), required=True) + parser.add_argument("--global-tokens", type=_positive_int, default=1 << 20) + parser.add_argument( + "--sequence-lengths", + type=_sequence_lengths, + help="explicit comma-separated THD lengths; overrides the generated short/long mix", + ) + parser.add_argument("--short-sequences", type=_positive_int, default=20) + parser.add_argument("--short-tokens", type=_positive_int, default=1024) + parser.add_argument("--long-sequences", type=_positive_int, default=2) + parser.add_argument("--block-size", type=_positive_int, default=4) + parser.add_argument("--top-k", type=_positive_int, default=512) + parser.add_argument("--index-heads", type=_positive_int, default=64) + parser.add_argument("--index-dim", type=_positive_int, default=128) + parser.add_argument("--dsa-budget-mib", type=int, default=8192) + parser.add_argument("--warmups", type=_positive_int, default=2) + parser.add_argument("--repeats", type=_positive_int, default=5) + parser.add_argument("--seed", type=int, default=20261104) + return parser.parse_args() + + +if __name__ == "__main__": + _run(_parse_args()) diff --git a/benchmarks/large_cp_emulation/bench_hybrid.py b/benchmarks/large_cp_emulation/bench_hybrid.py new file mode 100644 index 00000000..a9fa45dc --- /dev/null +++ b/benchmarks/large_cp_emulation/bench_hybrid.py @@ -0,0 +1,728 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Benchmark one cuDNN KDA + DSA-indexer layer on context-parallel shards. + +Two modes deliberately answer different questions: + +* ``real`` runs a real CP group. KDA uses cuLA's prepared all-gather + transport, while the DSA branch pools local index keys and all-gathers the + compressed rows before deterministic selection. +* ``logical`` maps representative ranks from a larger logical CP group onto + the physical processes. It keeps the real per-rank tensor shapes and runs + real NCCL ring hops for both KDA boundary summaries and compressed DSA + blocks. ``full-ring`` uses ``logical_cp - 1`` hops as a conservative + baseline. ``sample-chain`` uses the longest packed sample's CP span, which + models a sample-aware THD communicator. Both are critical-path emulators, + not claims about a real scale-out topology. + +The benchmark compares a conservative serial schedule with an overlap schedule +that runs the independent DSA indexer branch on a dedicated CUDA stream while +KDA executes on the caller stream. Arithmetic precision and deterministic +selection are identical in both schedules. + +The DSA timing covers index pooling, compressed-block communication and Top-K +selection. It does not claim to include the block-sparse attention kernel, +which is not yet part of cuLA's public hybrid runtime. +""" + +from __future__ import annotations + +import argparse +import json +import os +import pathlib +import statistics +import sys +import time +from dataclasses import dataclass + +import torch +import torch.distributed as dist + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[2])) + +from benchmarks.large_cp_emulation.topology import LogicalCPEmulationPlan # noqa: E402 +from benchmarks.large_cp_emulation.transport import ( # noqa: E402 + LogicalCPRingTransport, + emulate_ring_hops, +) +from benchmarks.large_cp_emulation.workload import PackedTHDWorkload # noqa: E402 +from cula.distributed import ( # noqa: E402 + AllGatherKimiCPTransport, + KimiCPPlan, + build_kimi_cp_plan, + cudnn_kda_cp, +) +from cula.hybrid import ( # noqa: E402 + DsaSelectionPlan, + DsaSelectionWorkspace, + IndexPoolPlan, + cudnn_dsa_select, + cudnn_index_pool, +) + + +def _positive_int(value: str) -> int: + result = int(value) + if result < 1: + raise argparse.ArgumentTypeError("value must be positive") + return result + + +@dataclass(slots=True) +class _DsaEvents: + start: torch.cuda.Event + after_pool: torch.cuda.Event + after_comm: torch.cuda.Event + end: torch.cuda.Event + + def timings(self) -> dict[str, float]: + return { + "pool_ms": self.start.elapsed_time(self.after_pool), + "comm_ms": self.after_pool.elapsed_time(self.after_comm), + "select_ms": self.after_comm.elapsed_time(self.end), + "total_ms": self.start.elapsed_time(self.end), + } + + +@dataclass(slots=True) +class _TimedKimiCPTransport: + """Benchmark-only CUDA timing wrapper for KDA boundary propagation.""" + + transport: object + events: list[tuple[bool, torch.cuda.Event, torch.cuda.Event]] + + @classmethod + def wrap(cls, transport: object) -> _TimedKimiCPTransport: + return cls(transport=transport, events=[]) + + @property + def workspace_bytes(self) -> int: + return self.transport.workspace_bytes + + def clear_events(self) -> None: + self.events.clear() + + def boundary_summaries( + self, + h: torch.Tensor, + transition: torch.Tensor, + *, + plan: KimiCPPlan, + group: dist.ProcessGroup | None, + reverse: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + begin = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + begin.record() + result = self.transport.boundary_summaries( + h, + transition, + plan=plan, + group=group, + reverse=reverse, + ) + end.record() + self.events.append((reverse, begin, end)) + return result + + def timings(self) -> dict[str, float]: + totals = {False: 0.0, True: 0.0} + for reverse, begin, end in self.events: + totals[reverse] += begin.elapsed_time(end) + return { + "forward_ms": totals[False], + "reverse_ms": totals[True], + "total_ms": totals[False] + totals[True], + } + + +@dataclass(slots=True) +class _DsaIndexerBranch: + mode: str + logical_cp: int + logical_rank: int + physical_rank: int + physical_world_size: int + transport_hops: int + local_tokens: int + local_blocks: int + global_block_start: int + key_start: int + key_end: int + qi: torch.Tensor + local_kv: torch.Tensor + pool_score: torch.Tensor + ape: torch.Tensor + weights: torch.Tensor + pool_plan: IndexPoolPlan + pool_cu_tokens: torch.Tensor + pool_cu_blocks: torch.Tensor + select_plan: DsaSelectionPlan + select_workspace: DsaSelectionWorkspace + global_kbar: torch.Tensor + ring_receive: torch.Tensor + last_result: object | None = None + + @classmethod + def create( + cls, + *, + mode: str, + logical_cp: int, + logical_rank: int, + physical_rank: int, + physical_world_size: int, + transport_hops: int, + global_cu_seqlens: tuple[int, ...], + shard_plan: KimiCPPlan, + head_dim: int, + index_heads: int, + top_k: int, + block_size: int, + candidate_budget_bytes: int | None, + device: torch.device, + ) -> _DsaIndexerBranch: + global_tokens = global_cu_seqlens[-1] + local_tokens = shard_plan.local_tokens + if local_tokens % block_size: + raise ValueError("local_tokens must be divisible by the DSA block size") + sequence_lengths = [right - left for left, right in zip(global_cu_seqlens, global_cu_seqlens[1:])] + if any(length % block_size for length in sequence_lengths): + raise ValueError("every packed sequence length must be divisible by the DSA block size") + local_blocks = local_tokens // block_size + global_blocks = global_tokens // block_size + pool_plan = IndexPoolPlan.from_cu_seqlens( + shard_plan.local_cu_seqlens, + head_dim=head_dim, + ratio=block_size, + ) + if pool_plan.total_comp != local_blocks: + raise ValueError("the CP shard must cut packed samples on DSA block boundaries") + pool_cu_tokens, pool_cu_blocks = pool_plan.materialize(device) + global_block_prefix = [0] + for length in sequence_lengths: + global_block_prefix.append(global_block_prefix[-1] + length // block_size) + first_sequence = shard_plan.first_sequence + last_sequence = shard_plan.last_sequence + select_cu_k = [0] + select_offsets = [] + for sequence in range(first_sequence, last_sequence + 1): + select_cu_k.append(select_cu_k[-1] + global_block_prefix[sequence + 1] - global_block_prefix[sequence]) + local_sequence_start = max(shard_plan.token_start, global_cu_seqlens[sequence]) + select_offsets.append(local_sequence_start - global_cu_seqlens[sequence]) + select_plan = DsaSelectionPlan.from_cu_seqlens( + shard_plan.local_cu_seqlens, + select_cu_k, + top_k=top_k, + block_size=block_size, + q_causal_offsets=select_offsets, + candidate_budget_bytes=candidate_budget_bytes, + ) + return cls( + mode=mode, + logical_cp=logical_cp, + logical_rank=logical_rank, + physical_rank=physical_rank, + physical_world_size=physical_world_size, + transport_hops=transport_hops, + local_tokens=local_tokens, + local_blocks=local_blocks, + global_block_start=shard_plan.token_start // block_size, + key_start=global_block_prefix[first_sequence], + key_end=global_block_prefix[last_sequence + 1], + qi=torch.randn(local_tokens, index_heads, head_dim, dtype=torch.bfloat16, device=device), + local_kv=torch.randn(local_tokens, head_dim, dtype=torch.bfloat16, device=device), + pool_score=torch.randn(local_tokens, head_dim, dtype=torch.bfloat16, device=device), + ape=torch.randn(block_size, head_dim, dtype=torch.float32, device=device), + weights=torch.randn(local_tokens, index_heads, dtype=torch.float32, device=device), + pool_plan=pool_plan, + pool_cu_tokens=pool_cu_tokens, + pool_cu_blocks=pool_cu_blocks, + select_plan=select_plan, + select_workspace=DsaSelectionWorkspace.from_plan(select_plan, device), + global_kbar=torch.empty(global_blocks, head_dim, dtype=torch.bfloat16, device=device), + ring_receive=torch.empty(local_blocks, head_dim, dtype=torch.bfloat16, device=device), + ) + + @property + def workspace_bytes(self) -> int: + tensors = ( + self.select_workspace.candidates, + self.select_workspace.block_ids, + self.select_workspace.scores, + self.global_kbar, + self.ring_receive, + ) + return sum(tensor.numel() * tensor.element_size() for tensor in tensors) + + def _communicate(self, local_kbar: torch.Tensor) -> None: + if self.mode == "real": + if hasattr(dist, "all_gather_single"): + dist.all_gather_single(self.global_kbar, local_kbar.contiguous()) + else: # pragma: no cover - compatibility with older PyTorch releases + dist.all_gather_into_tensor(self.global_kbar, local_kbar.contiguous()) + return + + local_slice = self.global_kbar[self.global_block_start : self.global_block_start + self.local_blocks] + local_slice.copy_(local_kbar) + # The ring alternates buffers; local_kbar is scratch from here on. + emulate_ring_hops( + local_kbar, + self.ring_receive, + physical_rank=self.physical_rank, + physical_world_size=self.physical_world_size, + hops=self.transport_hops, + ) + + def execute(self) -> _DsaEvents: + events = _DsaEvents(*(torch.cuda.Event(enable_timing=True) for _ in range(4))) + events.start.record() + with torch.cuda.nvtx.range("dsa.index_pool"): + local_kbar = cudnn_index_pool( + self.local_kv, + self.pool_score, + self.ape, + self.pool_cu_tokens, + self.pool_cu_blocks, + plan=self.pool_plan, + deterministic=True, + ) + events.after_pool.record() + with torch.cuda.nvtx.range("dsa.compressed_block_transport"): + self._communicate(local_kbar) + events.after_comm.record() + with torch.cuda.nvtx.range("dsa.select"): + self.last_result = cudnn_dsa_select( + self.qi, + self.global_kbar[self.key_start : self.key_end, None, :], + self.weights, + plan=self.select_plan, + workspace=self.select_workspace, + deterministic=True, + ) + events.end.record() + return events + + +def _run(args: argparse.Namespace) -> None: + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + dist.init_process_group("nccl") + rank = dist.get_rank() + world_size = dist.get_world_size() + if world_size < 2: + raise ValueError("the hybrid CP benchmark requires at least two physical processes") + + workload = ( + PackedTHDWorkload.short_long_mix( + args.global_tokens, + short_sequences=args.short_sequences, + short_tokens=args.short_tokens, + long_sequences=args.long_sequences, + ) + if args.thd_20s2l + else PackedTHDWorkload.single_sequence(args.global_tokens) + ) + global_cu_seqlens = workload.global_cu_seqlens + emulation = None + if args.mode == "real": + logical_cp = world_size + if args.global_tokens % logical_cp: + raise ValueError("global_tokens must be divisible by the CP size") + local_tokens = args.global_tokens // logical_cp + logical_rank = rank + transport_hops = logical_cp - 1 + else: + emulation = LogicalCPEmulationPlan.build( + global_cu_seqlens, + logical_cp=args.logical_cp, + physical_rank=rank, + physical_world_size=world_size, + transport_mode=args.logical_transport, + ) + logical_cp = emulation.logical_cp + logical_rank = emulation.logical_rank + local_tokens = emulation.logical_plan.local_tokens + transport_hops = emulation.transport_hops + torch.manual_seed(args.seed + logical_rank) + + def randn(*shape, dtype=torch.bfloat16): + return torch.randn(*shape, dtype=dtype, device=device) + + q = torch.nn.functional.normalize(randn(local_tokens, args.heads, args.key_dim).float(), dim=-1).to(torch.bfloat16) + k = torch.nn.functional.normalize(randn(local_tokens, args.heads, args.key_dim).float(), dim=-1).to(torch.bfloat16) + v = randn(local_tokens, args.heads, args.value_dim) + g = -(0.01 + 0.04 * torch.rand(local_tokens, args.heads, args.key_dim, dtype=torch.float32, device=device)) + beta = torch.sigmoid(randn(local_tokens, args.heads, dtype=torch.float32)) + grad_output = (randn(local_tokens, args.heads, args.value_dim) / float(args.global_tokens)).contiguous() + + if args.mode == "real": + kda_plan = build_kimi_cp_plan(global_cu_seqlens, group=dist.group.WORLD) + logical_shard_plan = kda_plan + else: + assert emulation is not None + logical_shard_plan = emulation.logical_plan + kda_plan = emulation.runtime_plan + kda_cu_seqlens = torch.tensor(kda_plan.local_cu_seqlens, dtype=torch.int32, device=device) + if args.mode == "real": + kda_transport = AllGatherKimiCPTransport.prepare( + kda_plan, + heads=args.heads, + key_dim=args.key_dim, + value_dim=args.value_dim, + device=device, + ) + if args.profile_kda_transport: + kda_transport = _TimedKimiCPTransport.wrap(kda_transport) + else: + kda_transport = LogicalCPRingTransport( + logical_cp, + rank, + world_size, + ring_hops=transport_hops, + ) + + candidate_budget_bytes = None if args.dsa_budget_mib == 0 else args.dsa_budget_mib << 20 + dsa = _DsaIndexerBranch.create( + mode=args.mode, + logical_cp=logical_cp, + logical_rank=logical_rank, + physical_rank=rank, + physical_world_size=world_size, + transport_hops=transport_hops, + global_cu_seqlens=global_cu_seqlens, + shard_plan=logical_shard_plan, + head_dim=args.index_dim, + index_heads=args.index_heads, + top_k=args.top_k, + block_size=args.block_size, + candidate_budget_bytes=candidate_budget_bytes, + device=device, + ) + if args.mode == "logical": + # Populate remote logical block ranges once. Timed iterations update + # the local logical slice after emulating the selected transport span. + dsa.global_kbar.normal_() + + dsa_stream = torch.cuda.Stream(device=device) + profile_marker = torch.zeros(1, dtype=torch.int32, device=device) + + def execute(schedule: str, measure: bool) -> dict[str, object]: + leaves = [tensor.detach().requires_grad_(True) for tensor in (q, k, v, g, beta)] + if isinstance(kda_transport, LogicalCPRingTransport | _TimedKimiCPTransport): + kda_transport.clear_events() + start = torch.cuda.Event(enable_timing=True) + kda_start = torch.cuda.Event(enable_timing=True) + kda_forward_end = torch.cuda.Event(enable_timing=True) + kda_end = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + if args.cuda_profiler_capture: + # Tiny kernels delimit the exact GPU interval in nsys. An NVTX + # range alone is projected onto child kernels and cannot represent + # stream waits, which made the apparent idle fraction ambiguous. + with torch.cuda.nvtx.range("hybrid.hotpath.begin_marker"): + profile_marker.add_(1) + torch.cuda.synchronize(device) + hotpath_wall_start = time.perf_counter() + torch.cuda.nvtx.range_push(f"hybrid.hotpath.{schedule}") + start.record() + + if schedule == "overlap": + dsa_stream.wait_event(start) + with torch.cuda.stream(dsa_stream): + dsa_events = dsa.execute() + kda_start.record() + else: + dsa_events = dsa.execute() + kda_start.record() + + with torch.cuda.nvtx.range("kda.forward_backward"): + output = cudnn_kda_cp( + *leaves, + kda_cu_seqlens, + plan=kda_plan, + group=dist.group.WORLD, + boundary_transport=kda_transport, + checkpoint_every_n_tokens=args.kda_checkpoint_tokens, + ) + kda_forward_end.record() + output.backward(grad_output) + kda_end.record() + if schedule == "overlap": + torch.cuda.current_stream(device).wait_stream(dsa_stream) + end.record() + torch.cuda.nvtx.range_pop() + if args.cuda_profiler_capture: + with torch.cuda.nvtx.range("hybrid.hotpath.end_marker"): + profile_marker.add_(1) + end.synchronize() + hotpath_wall_ms = (time.perf_counter() - hotpath_wall_start) * 1000 + + result = dsa.last_result + assert result is not None + sample: dict[str, object] = { + "cuda_total_ms": start.elapsed_time(end) if measure else None, + "wall_ms": hotpath_wall_ms if measure else None, + "kda_ms": kda_start.elapsed_time(kda_end) if measure else None, + "kda_forward_ms": kda_start.elapsed_time(kda_forward_end) if measure else None, + "kda_backward_ms": kda_forward_end.elapsed_time(kda_end) if measure else None, + "dsa_ms": dsa_events.timings() if measure else None, + "proxy_loss": (output.float() * grad_output.float()).sum().item(), + "output_norm": output.float().norm().item(), + "grad_norms": [tensor.grad.float().norm().item() for tensor in leaves], + "kda_finite": bool(torch.isfinite(output).all()) + and all(bool(torch.isfinite(tensor.grad).all()) for tensor in leaves), + "dsa_has_nan": bool(torch.isnan(result.scores).any()), + "dsa_id_sum": int(result.block_ids.to(torch.int64).sum().item()), + "dsa_finite_score_sum": float(result.scores[torch.isfinite(result.scores)].double().sum().item()), + } + if measure and isinstance(kda_transport, LogicalCPRingTransport): + sample["kda_transport_ms"] = { + f"{direction}_{stage}_ms": ( + begin.elapsed_time(after_comm) if stage == "comm" else after_comm.elapsed_time(after_prepare) + ) + for direction, begin, after_comm, after_prepare in kda_transport.events + for stage in ("comm", "prepare") + } + if measure and isinstance(kda_transport, _TimedKimiCPTransport): + sample["kda_boundary_ms"] = kda_transport.timings() + return sample + + schedules = ("serial", "overlap") if args.schedule == "both" else (args.schedule,) + if args.cuda_profiler_capture and len(schedules) != 1: + raise ValueError("--cuda-profiler-capture requires --schedule serial or --schedule overlap") + all_samples: dict[str, list[dict[str, object]]] = {} + memory_by_schedule = {} + for schedule in schedules: + dist.barrier(device_ids=[local_rank]) + for _ in range(args.warmups): + execute(schedule, False) + dist.barrier(device_ids=[local_rank]) + torch.cuda.reset_peak_memory_stats(device) + if args.cuda_profiler_capture: + torch.cuda.cudart().cudaProfilerStart() + samples = [] + for _ in range(args.repeats): + dist.barrier(device_ids=[local_rank]) + with torch.cuda.nvtx.range(f"hybrid.{schedule}"): + sample = execute(schedule, True) + samples.append(sample) + if args.cuda_profiler_capture: + torch.cuda.cudart().cudaProfilerStop() + all_samples[schedule] = samples + memory_by_schedule[schedule] = { + "peak_allocated_gib": torch.cuda.max_memory_allocated(device) / 2**30, + "peak_reserved_gib": torch.cuda.max_memory_reserved(device) / 2**30, + } + + record = { + "physical_rank": rank, + "logical_rank": logical_rank, + "samples": all_samples, + "memory_by_schedule": memory_by_schedule, + "kda_transport_workspace_gib": kda_transport.workspace_bytes / 2**30, + "dsa_workspace_gib": dsa.workspace_bytes / 2**30, + "dsa_windows": len(dsa.select_plan.windows), + "dsa_candidate_workspace_gib": dsa.select_plan.candidate_workspace_floats * 4 / 2**30, + } + records = [None] * world_size + dist.all_gather_object(records, record) + if rank == 0: + summary = {} + for schedule in schedules: + critical_cuda = [ + max(records[r]["samples"][schedule][i]["cuda_total_ms"] for r in range(world_size)) + for i in range(args.repeats) + ] + critical_wall = [ + max(records[r]["samples"][schedule][i]["wall_ms"] for r in range(world_size)) for i in range(args.repeats) + ] + summary[schedule] = { + "critical_cuda_samples_ms": critical_cuda, + "critical_cuda_median_ms": statistics.median(critical_cuda), + "critical_wall_samples_ms": critical_wall, + "critical_wall_median_ms": statistics.median(critical_wall), + "rank_stage_medians_ms": [ + { + "physical_rank": r, + "logical_rank": records[r]["logical_rank"], + "kda": statistics.median(sample["kda_ms"] for sample in records[r]["samples"][schedule]), + "kda_forward": statistics.median( + sample["kda_forward_ms"] for sample in records[r]["samples"][schedule] + ), + "kda_backward": statistics.median( + sample["kda_backward_ms"] for sample in records[r]["samples"][schedule] + ), + **( + { + f"kda_boundary_{stage}": statistics.median( + sample["kda_boundary_ms"][stage] for sample in records[r]["samples"][schedule] + ) + for stage in ("forward_ms", "reverse_ms", "total_ms") + } + if args.profile_kda_transport + else {} + ), + **{ + f"dsa_{stage}": statistics.median( + sample["dsa_ms"][stage] for sample in records[r]["samples"][schedule] + ) + for stage in ("pool_ms", "comm_ms", "select_ms", "total_ms") + }, + } + for r in range(world_size) + ], + } + if set(schedules) == {"serial", "overlap"}: + serial = summary["serial"]["critical_cuda_median_ms"] + overlap = summary["overlap"]["critical_cuda_median_ms"] + summary["overlap"]["speedup_vs_serial"] = serial / overlap + print( + json.dumps( + { + "mode": f"hybrid_kda_dsa_{args.mode}", + "global_tokens": args.global_tokens, + "logical_cp": logical_cp, + "logical_transport": args.logical_transport if args.mode == "logical" else "physical", + "kda_transport": "allgather" if args.mode == "real" else "logical-ring", + "transport_hops": transport_hops, + "physical_gpus": world_size, + "local_tokens": local_tokens, + "index_heads": args.index_heads, + "kda_checkpoint_tokens": args.kda_checkpoint_tokens, + "sequence_lengths": workload.sequence_lengths, + "schedules": summary, + "rank_resources": [ + { + "physical_rank": r, + "logical_rank": records[r]["logical_rank"], + "dsa_windows": records[r]["dsa_windows"], + "dsa_candidate_workspace_gib": records[r]["dsa_candidate_workspace_gib"], + "dsa_workspace_gib": records[r]["dsa_workspace_gib"], + "kda_transport_workspace_gib": records[r]["kda_transport_workspace_gib"], + "memory_by_schedule": records[r]["memory_by_schedule"], + } + for r in range(world_size) + ], + "rank_checks": [ + { + "physical_rank": r, + "logical_rank": records[r]["logical_rank"], + "kda_finite": all( + sample["kda_finite"] for schedule in schedules for sample in records[r]["samples"][schedule] + ), + "dsa_has_nan": any( + sample["dsa_has_nan"] for schedule in schedules for sample in records[r]["samples"][schedule] + ), + "repeat_stable": all( + all( + sample[key] == records[r]["samples"][schedules[0]][0][key] + for key in ( + "proxy_loss", + "output_norm", + "grad_norms", + "dsa_id_sum", + "dsa_finite_score_sum", + ) + ) + for schedule in schedules + for sample in records[r]["samples"][schedule] + ), + **{ + key: records[r]["samples"][schedules[0]][0][key] + for key in ( + "proxy_loss", + "output_norm", + "grad_norms", + "dsa_id_sum", + "dsa_finite_score_sum", + ) + }, + } + for r in range(world_size) + ], + "records": None if args.summary_only else records, + } + ), + flush=True, + ) + dist.destroy_process_group() + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + subparsers = parser.add_subparsers(dest="mode", required=True) + real = subparsers.add_parser("real", help="run the physical processes as the real CP group") + real.add_argument("--global-tokens", type=_positive_int, default=1 << 17) + real.add_argument( + "--profile-kda-transport", + action="store_true", + help="record CUDA time around real KDA boundary transport calls", + ) + logical = subparsers.add_parser("logical", help="emulate a larger logical CP group") + logical.add_argument("--global-tokens", type=_positive_int, default=1 << 20) + logical.add_argument("--logical-cp", type=_positive_int, default=64) + logical.add_argument( + "--logical-transport", + choices=("full-ring", "sample-chain"), + default="full-ring", + help="emulate all CP edges or only the longest packed sample's CP chain", + ) + real.set_defaults(logical_transport="physical") + logical.set_defaults(profile_kda_transport=False) + for subparser in (real, logical): + subparser.add_argument("--heads", type=_positive_int, default=32) + subparser.add_argument("--key-dim", type=_positive_int, default=128) + subparser.add_argument("--value-dim", type=_positive_int, default=128) + subparser.add_argument("--index-heads", type=_positive_int, default=64) + subparser.add_argument("--index-dim", type=_positive_int, default=128) + subparser.add_argument("--top-k", type=_positive_int, default=512) + subparser.add_argument("--block-size", type=_positive_int, default=4) + subparser.add_argument( + "--kda-checkpoint-tokens", + type=int, + choices=(0, 16), + default=0, + help="16 retains every native KDA chunk state; 0 recomputes them in backward", + ) + subparser.add_argument( + "--dsa-budget-mib", + type=int, + default=0, + help="0 selects cuLA's adaptive exact-workspace policy", + ) + subparser.add_argument("--schedule", choices=("serial", "overlap", "both"), default="both") + subparser.add_argument( + "--warmups", + type=_positive_int, + default=3, + help="untimed iterations used to amortize cuDNN JIT and communicator setup", + ) + subparser.add_argument("--repeats", type=_positive_int, default=5) + subparser.add_argument("--seed", type=int, default=20261102) + subparser.add_argument( + "--thd-20s2l", + action="store_true", + help="pack 20 short samples plus two long samples into global_tokens", + ) + subparser.add_argument("--short-sequences", type=_positive_int, default=20) + subparser.add_argument("--short-tokens", type=_positive_int, default=1024) + subparser.add_argument("--long-sequences", type=_positive_int, default=2) + subparser.add_argument("--cuda-profiler-capture", action="store_true") + subparser.add_argument( + "--summary-only", + action="store_true", + help="omit per-iteration records from the JSON output", + ) + return parser.parse_args() + + +def main() -> None: + _run(_parse_args()) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/large_cp_emulation/bench_kda.py b/benchmarks/large_cp_emulation/bench_kda.py new file mode 100644 index 00000000..35bfe7c3 --- /dev/null +++ b/benchmarks/large_cp_emulation/bench_kda.py @@ -0,0 +1,185 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Emulate a large logical Kimi CP group on fewer physical GPUs. + +The benchmark executes the complete cuLA KDA path: local summary, boundary +transport and cuDNN-composed apply, local forward, and distributed backward. It maps one +representative logical rank to each physical process and runs ``CP - 1`` real +NCCL ring hops per direction. This is a critical-path emulator, not a claim +about real CP64 topology or contention. + +Example: + torchrun --standalone --nproc-per-node=4 benchmarks/large_cp_emulation/bench_kda.py +""" + +from __future__ import annotations + +import argparse +import json +import os +import pathlib +import statistics +import sys +import time + +import torch +import torch.distributed as dist + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[2])) + +from benchmarks.large_cp_emulation.topology import LogicalCPEmulationPlan # noqa: E402 +from benchmarks.large_cp_emulation.transport import LogicalCPRingTransport # noqa: E402 +from cula.distributed import KimiCPBoundaryTransport, cudnn_kda_cp # noqa: E402 + + +def _positive_int(value: str) -> int: + result = int(value) + if result < 1: + raise argparse.ArgumentTypeError("value must be positive") + return result + + +def _run(args) -> None: + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + dist.init_process_group("nccl") + rank = dist.get_rank() + world_size = dist.get_world_size() + if world_size < 2: + raise ValueError("logical CP emulation requires at least two physical processes") + emulation = LogicalCPEmulationPlan.build( + [0, args.global_tokens], + logical_cp=args.logical_cp, + physical_rank=rank, + physical_world_size=world_size, + ) + local_tokens = emulation.logical_plan.local_tokens + logical_rank = emulation.logical_rank + torch.manual_seed(args.seed + logical_rank) + + def randn(*shape, dtype=torch.bfloat16): + return torch.randn(*shape, dtype=dtype, device=device) + + q = torch.nn.functional.normalize(randn(local_tokens, args.heads, args.key_dim).float(), dim=-1).to(torch.bfloat16) + k = torch.nn.functional.normalize(randn(local_tokens, args.heads, args.key_dim).float(), dim=-1).to(torch.bfloat16) + v = randn(local_tokens, args.heads, args.value_dim) + g = -(0.01 + 0.04 * torch.rand(local_tokens, args.heads, args.key_dim, dtype=torch.float32, device=device)) + beta = torch.sigmoid(randn(local_tokens, args.heads, dtype=torch.float32)) + grad_output = (randn(local_tokens, args.heads, args.value_dim) / float(args.global_tokens)).contiguous() + + plan = emulation.runtime_plan + cu_seqlens = torch.tensor(plan.local_cu_seqlens, dtype=torch.int32, device=device) + transport = LogicalCPRingTransport( + args.logical_cp, + rank, + world_size, + ring_hops=emulation.transport_hops, + ) + assert isinstance(transport, KimiCPBoundaryTransport) + + def execute(measure: bool): + inputs = [tensor.detach().requires_grad_(True) for tensor in (q, k, v, g, beta)] + transport.clear_events() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + if measure: + start.record() + output = cudnn_kda_cp( + *inputs, + cu_seqlens, + plan=plan, + group=dist.group.WORLD, + boundary_transport=transport, + ) + output.backward(grad_output) + if measure: + end.record() + end.synchronize() + else: + torch.cuda.synchronize(device) + transport_times = {} + if measure: + for direction, begin, after_comm, after_prepare in transport.events: + transport_times[f"comm_{direction}_ms"] = begin.elapsed_time(after_comm) + transport_times[f"prepare_{direction}_ms"] = after_comm.elapsed_time(after_prepare) + return { + "cuda_total_ms": start.elapsed_time(end) if measure else None, + "transport_ms": transport_times, + "proxy_loss": (output.float() * grad_output.float()).sum().item(), + "output_norm": output.float().norm().item(), + "grad_norms": [tensor.grad.float().norm().item() for tensor in inputs], + "finite": bool(torch.isfinite(output).all()) and all(bool(torch.isfinite(tensor.grad).all()) for tensor in inputs), + } + + dist.barrier(device_ids=[local_rank]) + execute(False) + dist.barrier(device_ids=[local_rank]) + torch.cuda.reset_peak_memory_stats(device) + samples = [] + for _ in range(args.repeats): + dist.barrier(device_ids=[local_rank]) + wall_start = time.perf_counter() + sample = execute(True) + sample["wall_ms"] = (time.perf_counter() - wall_start) * 1000 + samples.append(sample) + + record = { + "physical_rank": rank, + "logical_rank": logical_rank, + "samples": samples, + "peak_allocated_gib": torch.cuda.max_memory_allocated(device) / 2**30, + "peak_reserved_gib": torch.cuda.max_memory_reserved(device) / 2**30, + "transport_workspace_gib": transport.workspace_bytes / 2**30, + } + records = [None] * world_size + dist.all_gather_object(records, record) + if rank == 0: + critical_cuda = [ + max(records[r]["samples"][i]["cuda_total_ms"] for r in range(world_size)) for i in range(args.repeats) + ] + critical_wall = [max(records[r]["samples"][i]["wall_ms"] for r in range(world_size)) for i in range(args.repeats)] + print( + json.dumps( + { + "mode": "kda_logical_cp_e2e", + "logical_global_tokens": args.global_tokens, + "logical_cp": args.logical_cp, + "local_tokens": local_tokens, + "physical_gpus": world_size, + "shape": { + "heads": args.heads, + "key_dim": args.key_dim, + "value_dim": args.value_dim, + }, + "critical_cuda_samples_ms": critical_cuda, + "critical_cuda_median_ms": statistics.median(critical_cuda), + "critical_wall_samples_ms": critical_wall, + "critical_wall_median_ms": statistics.median(critical_wall), + "records": records, + } + ), + flush=True, + ) + dist.destroy_process_group() + + +def _parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--global-tokens", type=_positive_int, default=1 << 20) + parser.add_argument("--logical-cp", type=_positive_int, default=64) + parser.add_argument("--heads", type=_positive_int, default=32) + parser.add_argument("--key-dim", type=_positive_int, default=128) + parser.add_argument("--value-dim", type=_positive_int, default=128) + parser.add_argument("--repeats", type=_positive_int, default=3) + parser.add_argument("--seed", type=int, default=20261102) + return parser.parse_args() + + +def main() -> None: + _run(_parse_args()) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/large_cp_emulation/bench_thd_planner.py b/benchmarks/large_cp_emulation/bench_thd_planner.py new file mode 100644 index 00000000..e254bfa2 --- /dev/null +++ b/benchmarks/large_cp_emulation/bench_thd_planner.py @@ -0,0 +1,41 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Report modeled CP load balance over representative packed THD batches.""" + +from __future__ import annotations + +from cula.distributed.thd_planner import plan_thd_layout + +CASES = { + "all-short": [1024] * 1024, + "single-long": [1 << 20], + "short+two-long": [1024] * 20 + [514048] * 2, + "long-tail": [1024] * 256 + [786432], + "multi-long": [4096] * 32 + [131072, 262144, 524288], + "power-law": [1024, 4096, 16384, 65536, 131072, 307200, 523264], +} + + +def _boundaries(lengths: list[int]) -> list[int]: + result = [0] + for length in lengths: + result.append(result[-1] + length) + return result + + +def main() -> None: + print("| THD distribution | selected | token max/avg | causal max/avg | score | contiguous score | remote boundaries |") + print("|---|---:|---:|---:|---:|---:|---:|") + for name, lengths in CASES.items(): + plan = plan_thd_layout(_boundaries(lengths), world_size=4) + baseline = dict(plan.evaluated_scores)["contiguous"] + print( + f"| {name} | {plan.policy} | {plan.metrics.token_imbalance:.3f} | " + f"{plan.metrics.causal_imbalance:.3f} | {plan.metrics.score:.3f} | " + f"{baseline:.3f} | {sum(plan.metrics.boundary_loads) // 2} |" + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/large_cp_emulation/query_layout.py b/benchmarks/large_cp_emulation/query_layout.py new file mode 100644 index 00000000..92292d45 --- /dev/null +++ b/benchmarks/large_cp_emulation/query_layout.py @@ -0,0 +1,203 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Query-span layouts used to study causal-work balance independently of kernels.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from cula._metadata import as_boundaries, balanced_shards +from cula.distributed.thd_planner import THDPlannerConfig, plan_thd_layout + + +@dataclass(frozen=True, slots=True) +class QuerySpan: + sequence: int + token_start: int + token_end: int + sequence_start: int + sequence_end: int + + @property + def tokens(self) -> int: + return self.token_end - self.token_start + + @property + def causal_offset(self) -> int: + return self.token_start - self.sequence_start + + +@dataclass(frozen=True, slots=True) +class QueryLayout: + rank: int + world_size: int + block_size: int + spans: tuple[QuerySpan, ...] + + @property + def local_tokens(self) -> int: + return sum(span.tokens for span in self.spans) + + @property + def local_cu_seqlens(self) -> tuple[int, ...]: + result = [0] + for span in self.spans: + result.append(result[-1] + span.tokens) + return tuple(result) + + @property + def causal_offsets(self) -> tuple[int, ...]: + return tuple(span.causal_offset for span in self.spans) + + @property + def candidate_floats(self) -> int: + return sum(_candidate_floats(span, self.block_size) for span in self.spans) + + +def _candidate_prefix(position: int, block_size: int) -> int: + quotient, remainder = divmod(position, block_size) + return block_size * quotient * (quotient - 1) // 2 + quotient * (remainder + 1) + + +def _candidate_floats(span: QuerySpan, block_size: int) -> int: + start = span.causal_offset + return _candidate_prefix(start + span.tokens, block_size) - _candidate_prefix(start, block_size) + + +def _validate_layout_inputs( + global_cu_seqlens: tuple[int, ...] | list[int], + *, + rank: int, + world_size: int, + block_size: int, +) -> tuple[int, ...]: + boundaries = as_boundaries(global_cu_seqlens, name="global_cu_seqlens") + if not isinstance(world_size, int) or world_size < 1: + raise ValueError("world_size must be a positive integer") + if not isinstance(rank, int) or not 0 <= rank < world_size: + raise ValueError("rank must be in [0, world_size)") + if not isinstance(block_size, int) or isinstance(block_size, bool) or block_size < 1: + raise ValueError("block_size must be a positive integer") + if any((right - left) % block_size for left, right in zip(boundaries, boundaries[1:])): + raise ValueError("every sequence length must be block aligned") + return boundaries + + +def build_contiguous_query_layout( + global_cu_seqlens: tuple[int, ...] | list[int], + *, + rank: int, + world_size: int, + block_size: int = 4, +) -> QueryLayout: + boundaries = _validate_layout_inputs( + global_cu_seqlens, + rank=rank, + world_size=world_size, + block_size=block_size, + ) + block_offsets = balanced_shards(boundaries[-1] // block_size, world_size) + shard_start, shard_end = (offset * block_size for offset in block_offsets[rank : rank + 2]) + spans = [] + for sequence, (sequence_start, sequence_end) in enumerate(zip(boundaries, boundaries[1:])): + start = max(shard_start, sequence_start) + end = min(shard_end, sequence_end) + if start < end: + spans.append(QuerySpan(sequence, start, end, sequence_start, sequence_end)) + return QueryLayout(rank, world_size, block_size, tuple(spans)) + + +def build_alternating_sequence_query_layout( + global_cu_seqlens: tuple[int, ...] | list[int], + *, + rank: int, + world_size: int, + block_size: int = 4, + min_shard_tokens: int = 16 << 10, +) -> QueryLayout: + """Keep short samples whole and alternate long-sequence shard direction. + + Long sequence 0 maps its ordered shards to ascending ranks, long sequence + 1 to descending ranks, and so on. Causal work from early and late regions + therefore cancels across adjacent long samples. Whole short samples are + greedily assigned to the rank with the fewest tokens, so they never pay a + max-CP zigzag tax. + """ + + boundaries = _validate_layout_inputs( + global_cu_seqlens, + rank=rank, + world_size=world_size, + block_size=block_size, + ) + if ( + not isinstance(min_shard_tokens, int) + or isinstance(min_shard_tokens, bool) + or min_shard_tokens < block_size + or min_shard_tokens % block_size + ): + raise ValueError("min_shard_tokens must be a positive block-aligned integer") + + assignments: list[list[QuerySpan]] = [[] for _ in range(world_size)] + assigned_tokens = [0] * world_size + long_index = 0 + short_sequences = [] + for sequence, (sequence_start, sequence_end) in enumerate(zip(boundaries, boundaries[1:])): + length = sequence_end - sequence_start + if length // world_size < min_shard_tokens: + short_sequences.append((sequence, sequence_start, sequence_end)) + continue + block_offsets = balanced_shards(length // block_size, world_size) + reverse = bool(long_index % 2) + for owner in range(world_size): + shard = world_size - 1 - owner if reverse else owner + start = sequence_start + block_offsets[shard] * block_size + end = sequence_start + block_offsets[shard + 1] * block_size + assignments[owner].append(QuerySpan(sequence, start, end, sequence_start, sequence_end)) + assigned_tokens[owner] += end - start + long_index += 1 + + for sequence, sequence_start, sequence_end in short_sequences: + owner = min(range(world_size), key=lambda candidate: (assigned_tokens[candidate], candidate)) + assignments[owner].append(QuerySpan(sequence, sequence_start, sequence_end, sequence_start, sequence_end)) + assigned_tokens[owner] += sequence_end - sequence_start + + spans = tuple(sorted(assignments[rank], key=lambda span: (span.sequence, span.token_start))) + return QueryLayout(rank, world_size, block_size, spans) + + +def build_planned_query_layout( + global_cu_seqlens: tuple[int, ...] | list[int], + *, + rank: int, + world_size: int, + block_size: int = 4, +) -> QueryLayout: + """Adapt the production THD planner output to the DSA benchmark layout.""" + + plan = plan_thd_layout( + global_cu_seqlens, + world_size=world_size, + config=THDPlannerConfig(block_size=block_size), + ) + spans = tuple( + QuerySpan( + span.sequence, + span.token_start, + span.token_end, + span.sequence_start, + span.sequence_end, + ) + for span in plan.for_rank(rank).spans + ) + return QueryLayout(rank, world_size, block_size, spans) + + +__all__ = [ + "QueryLayout", + "QuerySpan", + "build_alternating_sequence_query_layout", + "build_contiguous_query_layout", + "build_planned_query_layout", +] diff --git a/benchmarks/large_cp_emulation/topology.py b/benchmarks/large_cp_emulation/topology.py new file mode 100644 index 00000000..255e0367 --- /dev/null +++ b/benchmarks/large_cp_emulation/topology.py @@ -0,0 +1,109 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Logical-to-physical topology mapping for large-CP emulation.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass, replace + +from cula.distributed import KimiCPPlan + + +def representative_logical_rank(physical_rank: int, physical_world_size: int, logical_cp: int) -> int: + """Spread physical processes over the logical CP rank interval.""" + if physical_world_size < 1: + raise ValueError("physical_world_size must be positive") + if logical_cp < physical_world_size: + raise ValueError("logical_cp must be at least physical_world_size") + if not 0 <= physical_rank < physical_world_size: + raise ValueError("physical_rank must be in [0, physical_world_size)") + if physical_world_size == 1: + return logical_cp - 1 + return physical_rank * (logical_cp - 1) // (physical_world_size - 1) + + +def sample_chain_hops(global_cu_seqlens: Sequence[int], local_tokens: int) -> int: + """Return the longest packed sample's number of rank-to-rank edges.""" + boundaries = tuple(global_cu_seqlens) + if local_tokens < 1: + raise ValueError("local_tokens must be positive") + if len(boundaries) < 2 or boundaries[0] != 0: + raise ValueError("global_cu_seqlens must contain at least two entries and start at zero") + if any(left >= right for left, right in zip(boundaries, boundaries[1:])): + raise ValueError("global_cu_seqlens must be strictly increasing") + max_ranks = 1 + for left, right in zip(boundaries, boundaries[1:]): + first_rank = left // local_tokens + last_rank = (right - 1) // local_tokens + max_ranks = max(max_ranks, last_rank - first_rank + 1) + return max_ranks - 1 + + +@dataclass(frozen=True, slots=True) +class LogicalCPEmulationPlan: + """Keep logical shard semantics separate from the physical process group. + + ``logical_plan`` describes the representative CP64 shard. ``runtime_plan`` + preserves that shard's THD boundary metadata while carrying the physical + rank/world size required by cuLA's native process-group validation. + """ + + logical_plan: KimiCPPlan + runtime_plan: KimiCPPlan + logical_cp: int + logical_rank: int + physical_rank: int + physical_world_size: int + transport_hops: int + transport_mode: str + + @classmethod + def build( + cls, + global_cu_seqlens: Sequence[int], + *, + logical_cp: int, + physical_rank: int, + physical_world_size: int, + transport_mode: str = "full-ring", + ) -> LogicalCPEmulationPlan: + boundaries = tuple(global_cu_seqlens) + if logical_cp < 1: + raise ValueError("logical_cp must be positive") + if len(boundaries) < 2 or boundaries[0] != 0: + raise ValueError("global_cu_seqlens must contain at least two entries and start at zero") + if any(left >= right for left, right in zip(boundaries, boundaries[1:])): + raise ValueError("global_cu_seqlens must be strictly increasing") + if boundaries[-1] % logical_cp: + raise ValueError("global token count must be divisible by logical_cp") + logical_rank = representative_logical_rank(physical_rank, physical_world_size, logical_cp) + logical_plan = KimiCPPlan.from_global_cu_seqlens( + boundaries, + rank=logical_rank, + world_size=logical_cp, + ) + if transport_mode == "full-ring": + hops = logical_cp - 1 + elif transport_mode == "sample-chain": + hops = sample_chain_hops(boundaries, logical_plan.local_tokens) + else: + raise ValueError("transport_mode must be 'full-ring' or 'sample-chain'") + runtime_plan = replace( + logical_plan, + rank=physical_rank, + world_size=physical_world_size, + token_start=physical_rank * logical_plan.local_tokens, + token_end=(physical_rank + 1) * logical_plan.local_tokens, + ) + return cls( + logical_plan=logical_plan, + runtime_plan=runtime_plan, + logical_cp=logical_cp, + logical_rank=logical_rank, + physical_rank=physical_rank, + physical_world_size=physical_world_size, + transport_hops=hops, + transport_mode=transport_mode, + ) diff --git a/benchmarks/large_cp_emulation/transport.py b/benchmarks/large_cp_emulation/transport.py new file mode 100644 index 00000000..225bcb7d --- /dev/null +++ b/benchmarks/large_cp_emulation/transport.py @@ -0,0 +1,132 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Benchmark-only transports that emulate logical CP dependency lengths.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import torch +import torch.distributed as dist + + +def emulate_ring_hops( + source: torch.Tensor, + receive: torch.Tensor, + *, + physical_rank: int, + physical_world_size: int, + hops: int, +) -> None: + """Execute dependent NCCL hops without pretending to exchange CP64 data.""" + if source.shape != receive.shape or source.dtype != receive.dtype or source.device != receive.device: + raise ValueError("source and receive buffers must have identical shape, dtype, and device") + if not 0 <= physical_rank < physical_world_size: + raise ValueError("physical_rank must be in [0, physical_world_size)") + if hops < 0: + raise ValueError("hops must be non-negative") + send, target = source, receive + next_rank = (physical_rank + 1) % physical_world_size + previous_rank = (physical_rank - 1 + physical_world_size) % physical_world_size + for _ in range(hops): + requests = dist.batch_isend_irecv( + [ + dist.P2POp(dist.isend, send, next_rank), + dist.P2POp(dist.irecv, target, previous_rank), + ] + ) + for request in requests: + request.wait() + send, target = target, send + + +@dataclass(slots=True) +class LogicalCPRingTransport: + """Implement the native KDA transport protocol with emulated ring cost.""" + + logical_cp: int + physical_rank: int + physical_world_size: int + ring_hops: int | None = field(default=None, kw_only=True) + events: list[tuple[str, torch.cuda.Event, torch.cuda.Event, torch.cuda.Event]] = field(default_factory=list) + local_payload: torch.Tensor | None = None + receive_payload: torch.Tensor | None = None + h_summaries: torch.Tensor | None = None + transition_summaries: torch.Tensor | None = None + + def clear_events(self) -> None: + self.events.clear() + + def _prepare(self, h: torch.Tensor, transition: torch.Tensor) -> None: + if h.dtype != torch.float32 or transition.dtype != torch.float32: + raise TypeError("logical KDA boundary summaries must remain FP32") + if h.device != transition.device: + raise ValueError("H and transition must be on the same device") + payload_shape = (h.shape[0], h.shape[-2] + transition.shape[-2], h.shape[-1]) + if self.local_payload is not None: + if self.local_payload.shape != payload_shape or self.local_payload.device != h.device: + raise ValueError("logical CP transport cannot be reused with a different H/M shape or device") + return + self.local_payload = torch.empty(payload_shape, dtype=torch.float32, device=h.device) + self.receive_payload = torch.empty_like(self.local_payload) + self.h_summaries = torch.empty((self.logical_cp - 1, *h.shape), dtype=torch.float32, device=h.device) + self.transition_summaries = torch.empty((self.logical_cp - 1, *transition.shape), dtype=torch.float32, device=h.device) + + @property + def workspace_bytes(self) -> int: + return sum( + tensor.numel() * tensor.element_size() + for tensor in ( + self.local_payload, + self.receive_payload, + self.h_summaries, + self.transition_summaries, + ) + if tensor is not None + ) + + def _ring(self) -> None: + assert self.local_payload is not None and self.receive_payload is not None + hops = self.logical_cp - 1 if self.ring_hops is None else self.ring_hops + if not 0 <= hops < self.logical_cp: + raise ValueError("ring_hops must be in [0, logical_cp)") + emulate_ring_hops( + self.local_payload, + self.receive_payload, + physical_rank=self.physical_rank, + physical_world_size=self.physical_world_size, + hops=hops, + ) + + def boundary_summaries( + self, + h: torch.Tensor, + transition: torch.Tensor, + *, + plan, + group, + reverse: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + del group + self._prepare(h, transition) + assert self.local_payload is not None + assert self.h_summaries is not None and self.transition_summaries is not None + self.local_payload[..., : h.shape[-2], :].copy_(h) + self.local_payload[..., h.shape[-2] :, :].copy_(transition) + start = torch.cuda.Event(enable_timing=True) + after_comm = torch.cuda.Event(enable_timing=True) + after_prepare = torch.cuda.Event(enable_timing=True) + start.record() + self._ring() + after_comm.record() + + pieces = plan.post_num_ranks if reverse else plan.pre_num_ranks + h_all = self.h_summaries[:pieces] + transition_all = self.transition_summaries[:pieces] + if pieces: + h_all.copy_(h.unsqueeze(0)) + transition_all.copy_(transition.unsqueeze(0)) + after_prepare.record() + self.events.append(("backward" if reverse else "forward", start, after_comm, after_prepare)) + return h_all, transition_all diff --git a/benchmarks/large_cp_emulation/workload.py b/benchmarks/large_cp_emulation/workload.py new file mode 100644 index 00000000..807009ad --- /dev/null +++ b/benchmarks/large_cp_emulation/workload.py @@ -0,0 +1,60 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Packed-THD workload descriptions shared by large-CP benchmarks.""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True, slots=True) +class PackedTHDWorkload: + global_cu_seqlens: tuple[int, ...] + + def __post_init__(self) -> None: + if len(self.global_cu_seqlens) < 2 or self.global_cu_seqlens[0] != 0: + raise ValueError("global_cu_seqlens must contain at least two entries and start at zero") + if any(left >= right for left, right in zip(self.global_cu_seqlens, self.global_cu_seqlens[1:])): + raise ValueError("global_cu_seqlens must be strictly increasing") + + @property + def global_tokens(self) -> int: + return self.global_cu_seqlens[-1] + + @property + def sequence_lengths(self) -> tuple[int, ...]: + return tuple(right - left for left, right in zip(self.global_cu_seqlens, self.global_cu_seqlens[1:])) + + @classmethod + def single_sequence(cls, global_tokens: int) -> PackedTHDWorkload: + if global_tokens < 1: + raise ValueError("global_tokens must be positive") + return cls((0, global_tokens)) + + @classmethod + def short_long_mix( + cls, + global_tokens: int, + *, + short_sequences: int = 20, + short_tokens: int = 1024, + long_sequences: int = 2, + ) -> PackedTHDWorkload: + for name, value in { + "global_tokens": global_tokens, + "short_sequences": short_sequences, + "short_tokens": short_tokens, + "long_sequences": long_sequences, + }.items(): + if value < 1: + raise ValueError(f"{name} must be positive") + short_total = short_sequences * short_tokens + remaining = global_tokens - short_total + if remaining <= 0 or remaining % long_sequences: + raise ValueError("global_tokens minus the short samples must split evenly across long samples") + lengths = [short_tokens] * short_sequences + [remaining // long_sequences] * long_sequences + boundaries = [0] + for length in lengths: + boundaries.append(boundaries[-1] + length) + return cls(tuple(boundaries)) diff --git a/cula/_metadata.py b/cula/_metadata.py new file mode 100644 index 00000000..850058fc --- /dev/null +++ b/cula/_metadata.py @@ -0,0 +1,59 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU-only packed-sequence and contiguous-shard metadata helpers.""" + +from __future__ import annotations + +from collections.abc import Sequence +from operator import index + +import torch + + +def as_boundaries(values: torch.Tensor | Sequence[int], *, name: str) -> tuple[int, ...]: + """Validate a strictly increasing, zero-based packed prefix on the host.""" + + if isinstance(values, torch.Tensor): + if values.device.type != "cpu": + raise ValueError(f"{name} must be a CPU tensor; copying CUDA metadata would synchronize the device") + if values.ndim != 1: + raise ValueError(f"{name} must be one-dimensional") + if values.dtype not in (torch.int32, torch.int64): + raise TypeError(f"{name} must have int32 or int64 dtype") + result = tuple(int(value) for value in values.tolist()) + else: + try: + result = tuple(index(value) for value in values) + except (TypeError, ValueError) as error: + raise TypeError(f"{name} must be an integer sequence") from error + if len(result) < 2: + raise ValueError(f"{name} must contain at least a start and an end") + if result[0] != 0: + raise ValueError(f"{name} must start at zero") + if any(left >= right for left, right in zip(result, result[1:])): + raise ValueError(f"{name} must be strictly increasing") + return result + + +def balanced_shards(total_tokens: int, world_size: int) -> tuple[int, ...]: + if total_tokens < world_size: + raise ValueError("packed context parallelism requires at least one token on every rank") + return tuple(shard * total_tokens // world_size for shard in range(world_size + 1)) + + +def validate_shards(shard_offsets: Sequence[int], *, total_tokens: int, world_size: int) -> tuple[int, ...]: + try: + offsets = tuple(index(value) for value in shard_offsets) + except (TypeError, ValueError) as error: + raise TypeError("shard_offsets must be an integer sequence") from error + if len(offsets) != world_size + 1: + raise ValueError(f"shard_offsets must contain world_size + 1 ({world_size + 1}) entries") + if offsets[0] != 0 or offsets[-1] != total_tokens: + raise ValueError("shard_offsets must cover the complete packed token range") + if any(left >= right for left, right in zip(offsets, offsets[1:])): + raise ValueError("shard_offsets must be strictly increasing; empty CP shards are unsupported") + return offsets + + +__all__ = ["as_boundaries", "balanced_shards", "validate_shards"] diff --git a/cula/distributed/__init__.py b/cula/distributed/__init__.py index 78aa7cb6..7930bbde 100644 --- a/cula/distributed/__init__.py +++ b/cula/distributed/__init__.py @@ -9,11 +9,25 @@ cudnn_kda_cp, ) from cula.distributed.plan import KimiCPPlan, build_kimi_cp_plan +from cula.distributed.thd_planner import ( + THDLayoutMetrics, + THDLayoutPlan, + THDPlannerConfig, + THDRankLayout, + THDSpan, + plan_thd_layout, +) __all__ = [ "AllGatherKimiCPTransport", "KimiCPBoundaryTransport", "KimiCPPlan", + "THDLayoutMetrics", + "THDLayoutPlan", + "THDPlannerConfig", + "THDRankLayout", + "THDSpan", "build_kimi_cp_plan", "cudnn_kda_cp", + "plan_thd_layout", ] diff --git a/cula/distributed/kda.py b/cula/distributed/kda.py index d9af7dc0..916583c5 100644 --- a/cula/distributed/kda.py +++ b/cula/distributed/kda.py @@ -19,20 +19,32 @@ @cache -def _load_cudnn_cp_ops(): +def _load_cudnn_kda_ops(): try: linear_attention = import_module("cudnn.linear_attention") ops = linear_attention - if not all( - hasattr(ops, name) - for name in ( - "kimi_delta_attention", - "kimi_delta_attention_summary", - "kimi_delta_attention_summary_bwd", - ) - ): + if not hasattr(ops, "kimi_delta_attention"): ops = import_module("cudnn.linear_attention.ops") kimi_delta_attention = ops.kimi_delta_attention + if "initial_state" not in signature(kimi_delta_attention).parameters: + raise AttributeError("KDA apply does not expose recurrent state inputs") + kimi_delta_attention_bwd = torch.ops.cudnn.kimi_delta_attention_bwd + except (ImportError, AttributeError) as error: + raise ImportError( + "cuDNN Kimi CP requires cuDNN Frontend KDA state APIs; " + "install cuda-linear-attention[cudnn-cp] and follow the compatible source-build instructions in USAGE.md" + ) from error + return kimi_delta_attention, kimi_delta_attention_bwd + + +@cache +def _load_cudnn_cp_ops(): + try: + kimi_delta_attention, kimi_delta_attention_bwd = _load_cudnn_kda_ops() + linear_attention = import_module("cudnn.linear_attention") + ops = linear_attention + if not all(hasattr(ops, name) for name in ("kimi_delta_attention_summary", "kimi_delta_attention_summary_bwd")): + ops = import_module("cudnn.linear_attention.ops") kimi_delta_attention_summary = ops.kimi_delta_attention_summary kimi_delta_attention_summary_bwd = ops.kimi_delta_attention_summary_bwd if "preceding_state_update" not in signature(kimi_delta_attention).parameters: @@ -40,11 +52,10 @@ def _load_cudnn_cp_ops(): summary_bwd_parameters = signature(kimi_delta_attention_summary_bwd).parameters if "output_context" not in summary_bwd_parameters or "checkpoint_every_n_tokens" not in summary_bwd_parameters: raise AttributeError("KDA backward summary does not expose the CP apply context") - kimi_delta_attention_bwd = torch.ops.cudnn.kimi_delta_attention_bwd except (ImportError, AttributeError) as error: raise ImportError( "cuDNN Kimi CP requires cuDNN Frontend with concatenated KDA summary apply APIs; " - "install cuda-linear-attention[cudnn-cp]" + "install cuda-linear-attention[cudnn-cp] and follow the compatible source-build instructions in USAGE.md" ) from error return ( kimi_delta_attention, @@ -368,6 +379,241 @@ def _validate_boundary_summaries( return summaries +def _compose_affine_summaries( + state_updates: torch.Tensor, + transitions: torch.Tensor, + *, + reverse: bool, +) -> torch.Tensor: + """Compose ordered FP32 state maps without changing their precision.""" + if state_updates.ndim != 4 or transitions.ndim != 4: + raise ValueError("ordered KDA summaries must have [ranks, heads, rows, cols] layout") + if state_updates.shape[0] != transitions.shape[0]: + raise ValueError("ordered KDA state updates and transitions must have the same rank count") + if state_updates.dtype != torch.float32 or transitions.dtype != torch.float32: + raise TypeError("ordered KDA summary composition requires FP32 inputs") + if state_updates.device != transitions.device: + raise ValueError("ordered KDA summaries must be on the same device") + state = state_updates.new_zeros(state_updates.shape[1:]) + indices = range(state_updates.shape[0] - 1, -1, -1) if reverse else range(state_updates.shape[0]) + for index in indices: + state = torch.matmul(state, transitions[index]) + state_updates[index] + return state.contiguous() + + +def _validate_hopper_state_probe_options( + options: _KdaOptions, + a_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, +) -> None: + """Reject transformations unsupported by the Hopper backward engine.""" + unsupported = [] + if options.use_qk_l2norm_in_kernel: + unsupported.append("use_qk_l2norm_in_kernel") + if options.use_beta_sigmoid_in_kernel: + unsupported.append("use_beta_sigmoid_in_kernel") + if options.safe_gate: + unsupported.append("safe_gate") + if a_log is not None: + unsupported.append("a_log") + if dt_bias is not None: + unsupported.append("dt_bias") + if unsupported: + names = ", ".join(unsupported) + raise ValueError( + "SM90 Kimi CP requires canonicalized q/k, beta, and log-space g because " + f"the cuDNN Hopper backward engine does not support in-kernel transforms: {names}" + ) + + +def _validate_ieee_fp32_composition(device: torch.device) -> None: + if device.type == "cuda" and torch.backends.cuda.matmul.allow_tf32: + raise RuntimeError( + "Kimi CP boundary composition requires IEEE FP32; set torch.backends.cuda.matmul.allow_tf32 = False" + ) + + +def _state_probe_summary( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + cu_seqlens: torch.Tensor, + *, + options: _KdaOptions, +) -> tuple[torch.Tensor, torch.Tensor]: + """Recover one local affine state map using the ordinary cuDNN KDA API. + + Hopper does not expose the FROST summary graph. KDA is affine in its + initial state, so evaluating a zero state gives H and evaluating an + embedded identity gives H + M. State tensors are FP32, hence the boundary + contract remains high precision. + """ + kimi_delta_attention, _ = _load_cudnn_kda_ops() + heads = max(q.shape[1], v.shape[1]) + key_dim = k.shape[-1] + value_dim = v.shape[-1] + if value_dim < key_dim: + raise ValueError( + f"SM90 Kimi CP state probing requires value_dim >= key_dim; got value_dim={value_dim}, key_dim={key_dim}" + ) + state_shape = (1, heads, value_dim, key_dim) + zero = torch.zeros(state_shape, dtype=torch.float32, device=q.device) + probe = zero.clone() + probe[:, :, :key_dim, :].copy_(torch.eye(key_dim, dtype=torch.float32, device=q.device)) + kwargs = options.kwargs(None, None) + + def final_state(initial_state: torch.Tensor) -> torch.Tensor: + outputs = kimi_delta_attention( + q, + k, + v, + g, + beta, + cu_seqlens, + scale=options.scale, + initial_state=initial_state, + output_final_state=True, + plan_name=options.plan_name, + **kwargs, + ) + return outputs[1].float() + + state_update = final_state(zero)[0].contiguous() + probed = final_state(probe)[0] + transition = (probed[:, :key_dim, :] - state_update[:, :key_dim, :]).contiguous() + _validate_local_boundary_summaries(state_update, transition) + return state_update, transition + + +class _CudnnKdaCPHopper(torch.autograd.Function): + """SM90 CP using ordinary Hopper state inputs instead of FROST summaries.""" + + @staticmethod + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + cu_seqlens: torch.Tensor, + plan: KimiCPPlan, + group: dist.ProcessGroup | None, + options: _KdaOptions, + boundary_transport: KimiCPBoundaryTransport, + ) -> torch.Tensor: + _validate_hopper_state_probe_options(options) + kimi_delta_attention, _ = _load_cudnn_kda_ops() + + last_start = plan.last_boundary_start + last_cu = cu_seqlens[-2:] - last_start + h_last, transition_last = _state_probe_summary( + q[last_start:], k[last_start:], v[last_start:], g[last_start:], beta[last_start:], last_cu, options=options + ) + preceding_h, preceding_transition = _validate_boundary_summaries( + boundary_transport.boundary_summaries(h_last, transition_last, plan=plan, group=group, reverse=False), + h_last, + transition_last, + rows=plan.pre_num_ranks, + ) + incoming_state = _compose_affine_summaries(preceding_h, preceding_transition, reverse=False) + initial_state = incoming_state.new_zeros((plan.num_local_sequences, *incoming_state.shape)) + initial_state[0].copy_(incoming_state) + + first_end = plan.first_boundary_end + if last_start == 0 and first_end == q.shape[0]: + transition_first = transition_last + else: + _, transition_first = _state_probe_summary( + q[:first_end], k[:first_end], v[:first_end], g[:first_end], beta[:first_end], cu_seqlens[:2], options=options + ) + + outputs = kimi_delta_attention( + q, + k, + v, + g, + beta, + cu_seqlens, + scale=options.scale, + initial_state=initial_state, + output_final_state=False, + plan_name=options.plan_name, + **options.attention_kwargs(None, None), + ) + output = outputs[0] + state_checkpoints = outputs[2] if len(outputs) > 2 else q.new_empty(0) + ctx.plan = plan + ctx.group = group + ctx.options = options + ctx.boundary_transport = boundary_transport + ctx.save_for_backward(q, k, v, g, beta, cu_seqlens, incoming_state, transition_first, state_checkpoints) + return output + + @staticmethod + @once_differentiable + def backward(ctx, grad_output: torch.Tensor): + _, kimi_delta_attention_bwd = _load_cudnn_kda_ops() + q, k, v, g, beta, cu_seqlens, incoming_state, transition_first, state_checkpoints = ctx.saved_tensors + plan = ctx.plan + options = ctx.options + _validate_ieee_fp32_composition(q.device) + initial_state = incoming_state.new_zeros((plan.num_local_sequences, *incoming_state.shape)) + initial_state[0].copy_(incoming_state) + grad_output = grad_output.contiguous() + scale = options.scale if options.scale is not None else q.shape[-1] ** -0.5 + + first_end = plan.first_boundary_end + zero_first = initial_state.new_zeros((1, *initial_state.shape[1:])) + first_grads = kimi_delta_attention_bwd( + grad_output[:first_end], + q[:first_end], + k[:first_end], + v[:first_end], + g[:first_end], + beta[:first_end], + cu_seqlens[:2], + float(scale), + initial_state=zero_first, + checkpoint_every_n_tokens=0, + batch_invariant=options.batch_invariant, + allow_neg_eigval=options.allow_neg_eigval, + plan_name=options.plan_name, + ) + g_local = first_grads[5][0].float().contiguous() + transition_backward = transition_first.transpose(-1, -2).contiguous() + following_h, following_transition = _validate_boundary_summaries( + ctx.boundary_transport.boundary_summaries(g_local, transition_backward, plan=plan, group=ctx.group, reverse=True), + g_local, + transition_backward, + rows=plan.post_num_ranks, + ) + outgoing_grad = _compose_affine_summaries(following_h, following_transition, reverse=True) + d_final_state = initial_state.new_zeros(initial_state.shape) + d_final_state[-1].copy_(outgoing_grad) + + dq, dk, dv, dg, dbeta, _, _, _ = kimi_delta_attention_bwd( + grad_output, + q, + k, + v, + g, + beta, + cu_seqlens, + float(scale), + initial_state=initial_state, + d_final_state=d_final_state, + state_checkpoints=state_checkpoints if state_checkpoints.numel() else None, + checkpoint_every_n_tokens=options.checkpoint_every_n_tokens, + batch_invariant=options.batch_invariant, + allow_neg_eigval=options.allow_neg_eigval, + plan_name=options.plan_name, + ) + return dq, dk, dv, dg, dbeta, None, None, None, None, None + + class _CudnnKdaCP(torch.autograd.Function): @staticmethod def forward( @@ -571,6 +817,19 @@ def _validate_inputs( raise ValueError("a_log and dt_bias are accepted only with safe_gate=True; canonicalize non-safe gates upstream") +def _resolve_kimi_cp_backend(backend: str, device: torch.device) -> str: + if backend not in ("auto", "frost-summary", "hopper-state-probe"): + raise ValueError("backend must be 'auto', 'frost-summary', or 'hopper-state-probe'") + if backend != "auto": + return backend + major, _ = torch.cuda.get_device_capability(device) + if major == 9: + return "hopper-state-probe" + if major >= 10: + return "frost-summary" + raise RuntimeError(f"cuDNN Kimi CP does not support CUDA compute capability {major}.x") + + def cudnn_kda_cp( q: torch.Tensor, k: torch.Tensor, @@ -592,9 +851,10 @@ def cudnn_kda_cp( dt_bias: torch.Tensor | None = None, batch_invariant: bool = True, checkpoint_every_n_tokens: int = 0, - plan_name: str | None = "kda_frost", - summary_plan_name: str | None = "kda_summary_frost", + plan_name: str | None = None, + summary_plan_name: str | None = None, boundary_transport: KimiCPBoundaryTransport | None = None, + backend: str = "auto", ) -> torch.Tensor: """Run differentiable cuDNN KDA over one contiguous packed-THD CP shard. @@ -624,6 +884,8 @@ def cudnn_kda_cp( if not isinstance(plan, KimiCPPlan): raise TypeError("plan must be a KimiCPPlan") + if backend not in ("auto", "frost-summary", "hopper-state-probe"): + raise ValueError("backend must be 'auto', 'frost-summary', or 'hopper-state-probe'") resolved_transport = _DEFAULT_BOUNDARY_TRANSPORT if boundary_transport is None else boundary_transport if not isinstance(resolved_transport, KimiCPBoundaryTransport): raise TypeError("boundary_transport must implement KimiCPBoundaryTransport") @@ -655,7 +917,7 @@ def cudnn_kda_cp( summary_plan_name=summary_plan_name, ) if plan.world_size == 1: - kimi_delta_attention, _, _, _ = _load_cudnn_cp_ops() + kimi_delta_attention, _ = _load_cudnn_kda_ops() attention_outputs = kimi_delta_attention( q, k, @@ -670,6 +932,22 @@ def cudnn_kda_cp( **options.attention_kwargs(a_log, dt_bias), ) return attention_outputs[0] + resolved_backend = _resolve_kimi_cp_backend(backend, q.device) + if resolved_backend == "hopper-state-probe": + _validate_hopper_state_probe_options(options, a_log, dt_bias) + _validate_ieee_fp32_composition(q.device) + return _CudnnKdaCPHopper.apply( + q.contiguous(), + k.contiguous(), + v.contiguous(), + g.contiguous(), + beta.contiguous(), + cu_seqlens.contiguous(), + plan, + group, + options, + resolved_transport, + ) return _CudnnKdaCP.apply( q.contiguous(), k.contiguous(), diff --git a/cula/distributed/plan.py b/cula/distributed/plan.py index b42b4275..3890e70d 100644 --- a/cula/distributed/plan.py +++ b/cula/distributed/plan.py @@ -14,53 +14,11 @@ from bisect import bisect_right from collections.abc import Sequence from dataclasses import dataclass -from operator import index import torch import torch.distributed as dist - -def _as_boundaries(values: torch.Tensor | Sequence[int], *, name: str) -> tuple[int, ...]: - if isinstance(values, torch.Tensor): - if values.device.type != "cpu": - raise ValueError(f"{name} must be a CPU tensor; copying CUDA metadata would synchronize the device") - if values.ndim != 1: - raise ValueError(f"{name} must be one-dimensional") - if values.dtype not in (torch.int32, torch.int64): - raise TypeError(f"{name} must have int32 or int64 dtype") - result = tuple(int(value) for value in values.tolist()) - else: - try: - result = tuple(index(value) for value in values) - except (TypeError, ValueError) as error: - raise TypeError(f"{name} must be an integer sequence") from error - if len(result) < 2: - raise ValueError(f"{name} must contain at least a start and an end") - if result[0] != 0: - raise ValueError(f"{name} must start at zero") - if any(left >= right for left, right in zip(result, result[1:])): - raise ValueError(f"{name} must be strictly increasing") - return result - - -def _balanced_shards(total_tokens: int, world_size: int) -> tuple[int, ...]: - if total_tokens < world_size: - raise ValueError("packed Kimi CP requires at least one token on every rank") - return tuple(index * total_tokens // world_size for index in range(world_size + 1)) - - -def _validate_shards(shard_offsets: Sequence[int], *, total_tokens: int, world_size: int) -> tuple[int, ...]: - try: - offsets = tuple(index(value) for value in shard_offsets) - except (TypeError, ValueError) as error: - raise TypeError("shard_offsets must be an integer sequence") from error - if len(offsets) != world_size + 1: - raise ValueError(f"shard_offsets must contain world_size + 1 ({world_size + 1}) entries") - if offsets[0] != 0 or offsets[-1] != total_tokens: - raise ValueError("shard_offsets must cover the complete packed token range") - if any(left >= right for left, right in zip(offsets, offsets[1:])): - raise ValueError("shard_offsets must be strictly increasing; empty CP shards are unsupported") - return offsets +from cula._metadata import as_boundaries, balanced_shards, validate_shards @dataclass(frozen=True, slots=True) @@ -152,12 +110,12 @@ def from_global_cu_seqlens( raise ValueError("world_size must be a positive integer") if not isinstance(rank, int) or not 0 <= rank < world_size: raise ValueError(f"rank must be in [0, {world_size})") - boundaries = _as_boundaries(global_cu_seqlens, name="global_cu_seqlens") + boundaries = as_boundaries(global_cu_seqlens, name="global_cu_seqlens") total_tokens = boundaries[-1] offsets = ( - _balanced_shards(total_tokens, world_size) + balanced_shards(total_tokens, world_size) if shard_offsets is None - else _validate_shards(shard_offsets, total_tokens=total_tokens, world_size=world_size) + else validate_shards(shard_offsets, total_tokens=total_tokens, world_size=world_size) ) start, end = offsets[rank], offsets[rank + 1] diff --git a/cula/distributed/thd_planner.py b/cula/distributed/thd_planner.py new file mode 100644 index 00000000..b7646fc0 --- /dev/null +++ b/cula/distributed/thd_planner.py @@ -0,0 +1,594 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU-only load planner for packed THD context parallelism. + +The planner balances two independent resources: token-linear work used by KDA +and position-dependent causal candidate work used by CSA. It deliberately +does not inspect CUDA tensors or kernel state, so one immutable plan can be +built at microbatch ingress and reused by the complete model stack. +""" + +from __future__ import annotations + +from collections import defaultdict +from collections.abc import Sequence +from dataclasses import dataclass +from math import ceil, isfinite +from operator import index + +import torch + +from cula._metadata import as_boundaries, balanced_shards + + +@dataclass(frozen=True, slots=True) +class THDSpan: + """One contiguous interval of a global packed sequence.""" + + sequence: int + token_start: int + token_end: int + sequence_start: int + sequence_end: int + owner: int + + @property + def tokens(self) -> int: + return self.token_end - self.token_start + + @property + def causal_offset(self) -> int: + return self.token_start - self.sequence_start + + +@dataclass(frozen=True, slots=True) +class THDRankLayout: + rank: int + spans: tuple[THDSpan, ...] + + @property + def local_tokens(self) -> int: + return sum(span.tokens for span in self.spans) + + @property + def local_cu_seqlens(self) -> tuple[int, ...]: + result = [0] + for span in self.spans: + result.append(result[-1] + span.tokens) + return tuple(result) + + @property + def causal_offsets(self) -> tuple[int, ...]: + return tuple(span.causal_offset for span in self.spans) + + +@dataclass(frozen=True, slots=True) +class THDLayoutMetrics: + token_loads: tuple[int, ...] + causal_loads: tuple[int, ...] + boundary_loads: tuple[int, ...] + token_imbalance: float + causal_imbalance: float + score: float + + +@dataclass(frozen=True, slots=True) +class THDLayoutPlan: + """Complete global span assignment selected by :func:`plan_thd_layout`.""" + + world_size: int + block_size: int + policy: str + ranks: tuple[THDRankLayout, ...] + metrics: THDLayoutMetrics + evaluated_scores: tuple[tuple[str, float], ...] + + def for_rank(self, rank: int) -> THDRankLayout: + if not isinstance(rank, int) or not 0 <= rank < self.world_size: + raise ValueError(f"rank must be in [0, {self.world_size})") + return self.ranks[rank] + + +@dataclass(frozen=True, slots=True) +class THDPlannerConfig: + """Shape-independent planner controls. + + ``max_chunks_per_rank`` controls the finest evaluated planning grid. The + planner compares every grid from one through this value, so a dominant + causal sequence can form an early/late zigzag without forcing the same + fragmentation on batches whose long sequences already balance each other. + """ + + block_size: int = 4 + min_split_tokens: int = 16 << 10 + max_chunks_per_rank: int = 6 + token_weight: float = 1.0 + causal_weight: float = 1.0 + boundary_weight: float = 0.02 + + +@dataclass(frozen=True, slots=True) +class _PlanningItem: + sequence: int + token_start: int + token_end: int + sequence_start: int + sequence_end: int + causal_work: int + + @property + def tokens(self) -> int: + return self.token_end - self.token_start + + +def _candidate_prefix(position: int, block_size: int) -> int: + quotient, remainder = divmod(position, block_size) + return block_size * quotient * (quotient - 1) // 2 + quotient * (remainder + 1) + + +def _candidate_work(start: int, tokens: int, block_size: int) -> int: + return _candidate_prefix(start + tokens, block_size) - _candidate_prefix(start, block_size) + + +def _validate_inputs( + global_cu_seqlens: torch.Tensor | Sequence[int], + *, + world_size: int, + config: THDPlannerConfig, +) -> tuple[int, ...]: + boundaries = as_boundaries(global_cu_seqlens, name="global_cu_seqlens") + if not isinstance(world_size, int) or isinstance(world_size, bool) or world_size < 1: + raise ValueError("world_size must be a positive integer") + for name in ("block_size", "min_split_tokens", "max_chunks_per_rank"): + raw_value = getattr(config, name) + if isinstance(raw_value, bool): + raise TypeError(f"{name} must be an integer") + value = index(raw_value) + if value < 1: + raise ValueError(f"{name} must be a positive integer") + if config.min_split_tokens % config.block_size: + raise ValueError("min_split_tokens must be block aligned") + if any((right - left) % config.block_size for left, right in zip(boundaries, boundaries[1:])): + raise ValueError("every sequence length must be block aligned") + if boundaries[-1] // config.block_size < world_size: + raise ValueError("packed context parallelism requires at least one block per rank") + for name in ("token_weight", "causal_weight", "boundary_weight"): + value = getattr(config, name) + if isinstance(value, bool) or not isinstance(value, int | float) or not isfinite(value) or value < 0: + raise ValueError(f"{name} must be finite and non-negative") + if config.token_weight == 0 and config.causal_weight == 0: + raise ValueError("at least one compute weight must be positive") + return boundaries + + +def _make_item( + sequence: int, + start: int, + end: int, + sequence_start: int, + sequence_end: int, + block_size: int, +) -> _PlanningItem: + return _PlanningItem( + sequence, + start, + end, + sequence_start, + sequence_end, + _candidate_work(start - sequence_start, end - start, block_size), + ) + + +def _contiguous_items(boundaries: tuple[int, ...], world_size: int, block_size: int) -> list[tuple[_PlanningItem, int]]: + total_blocks = boundaries[-1] // block_size + block_shards = balanced_shards(total_blocks, world_size) + shard_offsets = tuple(value * block_size for value in block_shards) + result = [] + for rank, (shard_start, shard_end) in enumerate(zip(shard_offsets, shard_offsets[1:])): + for sequence, (sequence_start, sequence_end) in enumerate(zip(boundaries, boundaries[1:])): + start = max(shard_start, sequence_start) + end = min(shard_end, sequence_end) + if start < end: + result.append((_make_item(sequence, start, end, sequence_start, sequence_end, block_size), rank)) + return result + + +def _adaptive_items( + boundaries: tuple[int, ...], + world_size: int, + config: THDPlannerConfig, + *, + chunks_per_rank: int, +) -> list[_PlanningItem]: + block_size = config.block_size + total_blocks = boundaries[-1] // block_size + requested_grid = world_size * chunks_per_rank + target_blocks = max(1, ceil(total_blocks / requested_grid)) + min_split_blocks = config.min_split_tokens // block_size + result = [] + for sequence, (sequence_start, sequence_end) in enumerate(zip(boundaries, boundaries[1:])): + length_blocks = (sequence_end - sequence_start) // block_size + pieces = min(requested_grid, max(1, ceil(length_blocks / target_blocks))) + # A sequence shorter than two useful spans stays whole. This is what + # prevents short THD samples from paying a max-CP fragmentation tax. + if length_blocks < 2 * min_split_blocks: + pieces = 1 + offsets = balanced_shards(length_blocks, pieces) + for left, right in zip(offsets, offsets[1:]): + start = sequence_start + left * block_size + end = sequence_start + right * block_size + result.append(_make_item(sequence, start, end, sequence_start, sequence_end, block_size)) + + # Keep all ranks non-empty even for small batches containing fewer than W + # sequences. The minimum span is a performance preference, not a semantic + # restriction, so only this feasibility repair may split below it. + while len(result) < world_size: + splittable = [item for item in result if item.tokens >= 2 * block_size] + if not splittable: + raise ValueError("block alignment leaves fewer non-empty spans than ranks") + item = max(splittable, key=lambda candidate: (candidate.tokens, candidate.causal_work)) + result.remove(item) + blocks = item.tokens // block_size + middle = item.token_start + (blocks // 2) * block_size + result.extend( + ( + _make_item( + item.sequence, + item.token_start, + middle, + item.sequence_start, + item.sequence_end, + block_size, + ), + _make_item( + item.sequence, + middle, + item.token_end, + item.sequence_start, + item.sequence_end, + block_size, + ), + ) + ) + return result + + +def _greedy_assign(items: list[_PlanningItem], world_size: int) -> list[tuple[_PlanningItem, int]]: + total_tokens = sum(item.tokens for item in items) + total_causal = sum(item.causal_work for item in items) + ideal_tokens = total_tokens / world_size + ideal_causal = total_causal / world_size + token_loads = [0] * world_size + causal_loads = [0] * world_size + item_counts = [0] * world_size + + ordered = sorted( + items, + key=lambda item: ( + max(item.tokens / ideal_tokens, item.causal_work / ideal_causal), + item.causal_work, + item.tokens, + -item.sequence, + -item.token_start, + ), + reverse=True, + ) + result = [] + for item in ordered: + owner = min( + range(world_size), + key=lambda rank: ( + max( + (token_loads[rank] + item.tokens) / ideal_tokens, + (causal_loads[rank] + item.causal_work) / ideal_causal, + ), + (token_loads[rank] + item.tokens) / ideal_tokens + (causal_loads[rank] + item.causal_work) / ideal_causal, + item_counts[rank], + rank, + ), + ) + result.append((item, owner)) + token_loads[owner] += item.tokens + causal_loads[owner] += item.causal_work + item_counts[owner] += 1 + return result + + +def _split_item( + item: _PlanningItem, + tokens: int, + *, + take_prefix: bool, + block_size: int, +) -> tuple[_PlanningItem, _PlanningItem]: + if tokens <= 0 or tokens >= item.tokens or tokens % block_size: + raise ValueError("split size must be a non-empty block-aligned item prefix or suffix") + middle = item.token_start + tokens if take_prefix else item.token_end - tokens + left = _make_item( + item.sequence, + item.token_start, + middle, + item.sequence_start, + item.sequence_end, + block_size, + ) + right = _make_item( + item.sequence, + middle, + item.token_end, + item.sequence_start, + item.sequence_end, + block_size, + ) + return (left, right) if take_prefix else (right, left) + + +def _rebalance_tokens( + assignments: list[tuple[_PlanningItem, int]], + *, + world_size: int, + block_size: int, +) -> list[tuple[_PlanningItem, int]]: + """Repair a compute-balanced assignment to exact balanced token bins.""" + + total_blocks = sum(item.tokens for item, _ in assignments) // block_size + capacity_offsets = balanced_shards(total_blocks, world_size) + capacities = [(right - left) * block_size for left, right in zip(capacity_offsets, capacity_offsets[1:])] + loads = [0] * world_size + causal_loads = [0] * world_size + for item, owner in assignments: + loads[owner] += item.tokens + causal_loads[owner] += item.causal_work + ideal_causal = sum(causal_loads) / world_size + repaired = list(assignments) + + while loads != capacities: + receivers = [rank for rank in range(world_size) if loads[rank] < capacities[rank]] + donors = [rank for rank in range(world_size) if loads[rank] > capacities[rank]] + if not receivers or not donors: + raise RuntimeError("internal error while repairing THD token capacities") + receiver = max(receivers, key=lambda rank: (capacities[rank] - loads[rank], -rank)) + donor = max(donors, key=lambda rank: (loads[rank] - capacities[rank], -rank)) + transfer_limit = min( + capacities[receiver] - loads[receiver], + loads[donor] - capacities[donor], + ) + + choices = [] + for assignment_index, (item, owner) in enumerate(repaired): + if owner != donor: + continue + transfer_tokens = min(item.tokens, transfer_limit) + if transfer_tokens == item.tokens: + variants = ((item, None),) + else: + variants = tuple( + _split_item( + item, + transfer_tokens, + take_prefix=take_prefix, + block_size=block_size, + ) + for take_prefix in (True, False) + ) + for moved, remainder in variants: + next_donor_causal = causal_loads[donor] - moved.causal_work + next_receiver_causal = causal_loads[receiver] + moved.causal_work + causal_pressure = max( + next_donor_causal / ideal_causal, + next_receiver_causal / ideal_causal, + *(causal_loads[rank] / ideal_causal for rank in range(world_size) if rank not in (donor, receiver)), + ) + choices.append( + ( + causal_pressure, + remainder is not None, + -moved.tokens, + moved.sequence, + moved.token_start, + assignment_index, + moved, + remainder, + ) + ) + if not choices: + raise RuntimeError("internal error: an overloaded rank has no movable THD span") + *_, assignment_index, moved, remainder = min(choices) + repaired.pop(assignment_index) + if remainder is not None: + repaired.append((remainder, donor)) + repaired.append((moved, receiver)) + loads[donor] -= moved.tokens + loads[receiver] += moved.tokens + causal_loads[donor] -= moved.causal_work + causal_loads[receiver] += moved.causal_work + return repaired + + +def _coalesce(assignments: list[tuple[_PlanningItem, int]]) -> list[THDSpan]: + ordered = sorted(assignments, key=lambda pair: (pair[0].sequence, pair[0].token_start)) + result: list[THDSpan] = [] + for item, owner in ordered: + if ( + result + and result[-1].sequence == item.sequence + and result[-1].owner == owner + and result[-1].token_end == item.token_start + ): + previous = result.pop() + result.append( + THDSpan( + previous.sequence, + previous.token_start, + item.token_end, + previous.sequence_start, + previous.sequence_end, + owner, + ) + ) + else: + result.append( + THDSpan( + item.sequence, + item.token_start, + item.token_end, + item.sequence_start, + item.sequence_end, + owner, + ) + ) + return result + + +def _metrics( + spans: list[THDSpan], + *, + world_size: int, + block_size: int, + config: THDPlannerConfig, +) -> THDLayoutMetrics: + token_loads = [0] * world_size + causal_loads = [0] * world_size + boundary_loads = [0] * world_size + by_sequence: dict[int, list[THDSpan]] = defaultdict(list) + for span in spans: + token_loads[span.owner] += span.tokens + causal_loads[span.owner] += _candidate_work(span.causal_offset, span.tokens, block_size) + by_sequence[span.sequence].append(span) + for sequence_spans in by_sequence.values(): + sequence_spans.sort(key=lambda span: span.token_start) + for left, right in zip(sequence_spans, sequence_spans[1:]): + if left.owner != right.owner: + boundary_loads[left.owner] += 1 + boundary_loads[right.owner] += 1 + + def imbalance(loads: list[int]) -> float: + average = sum(loads) / world_size + return max(loads) / average if average else 1.0 + + token_imbalance = imbalance(token_loads) + causal_imbalance = imbalance(causal_loads) + score = max( + config.token_weight * token_imbalance, + config.causal_weight * causal_imbalance, + ) + config.boundary_weight * max(boundary_loads, default=0) + return THDLayoutMetrics( + tuple(token_loads), + tuple(causal_loads), + tuple(boundary_loads), + token_imbalance, + causal_imbalance, + score, + ) + + +def _build_candidate( + assignments: list[tuple[_PlanningItem, int]], + *, + policy: str, + world_size: int, + config: THDPlannerConfig, +) -> tuple[str, list[THDSpan], THDLayoutMetrics] | None: + spans = _coalesce(assignments) + metrics = _metrics( + spans, + world_size=world_size, + block_size=config.block_size, + config=config, + ) + if any(load == 0 for load in metrics.token_loads): + return None + return policy, spans, metrics + + +def plan_thd_layout( + global_cu_seqlens: torch.Tensor | Sequence[int], + *, + world_size: int, + config: THDPlannerConfig | None = None, +) -> THDLayoutPlan: + """Plan a block-aligned THD layout and select the best modeled candidate. + + The result is guaranteed to have a score no worse than the evaluated + contiguous baseline under the declared cost model. This is a model-level + guarantee, not a claim that an analytic proxy exactly predicts every GPU. + """ + + config = THDPlannerConfig() if config is None else config + if not isinstance(config, THDPlannerConfig): + raise TypeError("config must be a THDPlannerConfig") + boundaries = _validate_inputs(global_cu_seqlens, world_size=world_size, config=config) + + contiguous = _build_candidate( + _contiguous_items(boundaries, world_size, config.block_size), + policy="contiguous", + world_size=world_size, + config=config, + ) + adaptive_candidates = [] + seen_assignments = set() + for chunks_per_rank in range(1, config.max_chunks_per_rank + 1): + adaptive_items = _adaptive_items( + boundaries, + world_size, + config, + chunks_per_rank=chunks_per_rank, + ) + adaptive = _build_candidate( + _rebalance_tokens( + _greedy_assign(adaptive_items, world_size), + world_size=world_size, + block_size=config.block_size, + ), + policy=f"cost-balanced-{chunks_per_rank}x", + world_size=world_size, + config=config, + ) + if adaptive is None: + continue + signature = tuple((span.sequence, span.token_start, span.token_end, span.owner) for span in adaptive[1]) + if signature not in seen_assignments: + adaptive_candidates.append(adaptive) + seen_assignments.add(signature) + candidates = [candidate for candidate in (contiguous, *adaptive_candidates) if candidate is not None] + if not candidates: + raise RuntimeError("no feasible THD layout candidate") + selected_policy, selected_spans, selected_metrics = min( + candidates, + key=lambda candidate: ( + candidate[2].score, + candidate[0] != "contiguous", + candidate[0], + ), + ) + ranks = tuple( + THDRankLayout( + rank, + tuple( + sorted( + (span for span in selected_spans if span.owner == rank), + key=lambda span: (span.sequence, span.token_start), + ) + ), + ) + for rank in range(world_size) + ) + return THDLayoutPlan( + world_size, + config.block_size, + selected_policy, + ranks, + selected_metrics, + tuple((policy, metrics.score) for policy, _, metrics in candidates), + ) + + +__all__ = [ + "THDLayoutMetrics", + "THDLayoutPlan", + "THDPlannerConfig", + "THDRankLayout", + "THDSpan", + "plan_thd_layout", +] diff --git a/cula/hybrid/__init__.py b/cula/hybrid/__init__.py new file mode 100644 index 00000000..708d6636 --- /dev/null +++ b/cula/hybrid/__init__.py @@ -0,0 +1,33 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""cuDNN-backed hybrid-attention integration primitives.""" + +from cula.hybrid.distributed_pool import cudnn_index_pool_cp +from cula.hybrid.dsa import ( + DEFAULT_CANDIDATE_BUDGET_BYTES, + MAX_BALANCED_CANDIDATE_BUDGET_BYTES, + MEMORY_EFFICIENT_CANDIDATE_BUDGET_BYTES, + DsaSelectionPlan, + DsaSelectionResult, + DsaSelectionWindow, + DsaSelectionWorkspace, + cudnn_dsa_select, +) +from cula.hybrid.index_pool import cudnn_index_pool +from cula.hybrid.plan import IndexPoolPlan, IndexPoolShardPlan + +__all__ = [ + "DEFAULT_CANDIDATE_BUDGET_BYTES", + "MAX_BALANCED_CANDIDATE_BUDGET_BYTES", + "MEMORY_EFFICIENT_CANDIDATE_BUDGET_BYTES", + "DsaSelectionPlan", + "DsaSelectionResult", + "DsaSelectionWindow", + "DsaSelectionWorkspace", + "IndexPoolPlan", + "IndexPoolShardPlan", + "cudnn_dsa_select", + "cudnn_index_pool", + "cudnn_index_pool_cp", +] diff --git a/cula/hybrid/distributed_pool.py b/cula/hybrid/distributed_pool.py new file mode 100644 index 00000000..fd6d91fe --- /dev/null +++ b/cula/hybrid/distributed_pool.py @@ -0,0 +1,275 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Pool context-parallel index keys before exchanging compressed blocks.""" + +from __future__ import annotations + +import torch +import torch.distributed as dist +from torch.autograd.function import once_differentiable + +from cula.hybrid.index_pool import _load_cudnn_index_pool_ops, _validate_inputs +from cula.hybrid.plan import IndexPoolShardPlan + + +def _reduce_scatter_tensor( + output: torch.Tensor, + tensor: torch.Tensor, + *, + group: dist.ProcessGroup | None, +) -> None: + """Use the public collective when available, retaining older torch support.""" + + if hasattr(dist, "reduce_scatter_single"): + dist.reduce_scatter_single(output, tensor, group=group) + elif hasattr(dist, "reduce_scatter_tensor"): + dist.reduce_scatter_tensor(output, tensor, group=group) + else: # pragma: no cover - compatibility with older PyTorch releases + dist._reduce_scatter_base(output, tensor, group=group) + + +def _all_gather_block_rows( + local: torch.Tensor, + *, + group: dist.ProcessGroup | None, + block_offsets: tuple[int, ...], + rank: int, +) -> torch.Tensor: + """Gather compressed rows while padding only an uneven compressed axis.""" + + counts = tuple(right - left for left, right in zip(block_offsets, block_offsets[1:])) + if all(count == counts[0] for count in counts[1:]): + output = local.new_empty((sum(counts), *local.shape[1:])) + if hasattr(dist, "all_gather_single"): + dist.all_gather_single(output, local.contiguous(), group=group) + else: # pragma: no cover - compatibility with older PyTorch releases + dist.all_gather_into_tensor(output, local.contiguous(), group=group) + return output + + max_count = max(counts) + padded = local.new_zeros((max_count, *local.shape[1:])) + padded[: counts[rank]].copy_(local) + gathered = local.new_empty((len(counts), max_count, *local.shape[1:])) + flat_gathered = gathered.flatten(0, 1) + if hasattr(dist, "all_gather_single"): + dist.all_gather_single(flat_gathered, padded, group=group) + else: # pragma: no cover - compatibility with older PyTorch releases + dist.all_gather_into_tensor(flat_gathered, padded, group=group) + + output = local.new_empty((sum(counts), *local.shape[1:])) + for owner, count in enumerate(counts): + output[block_offsets[owner] : block_offsets[owner + 1]].copy_(gathered[owner, :count]) + return output + + +class _CudnnIndexPoolCP(torch.autograd.Function): + """Distributed compressor with a precision-preserving reverse exchange.""" + + @staticmethod + def forward( + ctx, + kv: torch.Tensor, + score: torch.Tensor, + ape: torch.Tensor, + cu_seqlens: torch.Tensor, + cu_seqlens_comp: torch.Tensor, + plan: IndexPoolShardPlan, + group, + deterministic: bool, + ) -> torch.Tensor: + forward, _ = _load_cudnn_index_pool_ops() + local_plan = plan.local_plan + if local_plan is None: # Validated by the public entry point. + raise RuntimeError("the direct local pooling plan is unavailable") + result = forward( + kv, + score, + ape, + cu_seqlens, + cu_seqlens_comp, + ratio=local_plan.ratio, + head_dim=local_plan.head_dim, + coff=local_plan.coff, + total_comp=local_plan.total_comp, + ) + local_blocks = result["out"] + ctx.plan = plan + ctx.group = group + ctx.deterministic = deterministic + ctx.save_for_backward(kv, score, ape, cu_seqlens, cu_seqlens_comp) + if plan.world_size == 1: + return local_blocks + return _all_gather_block_rows( + local_blocks, + group=group, + block_offsets=plan.block_offsets, + rank=plan.rank, + ) + + @staticmethod + @once_differentiable + def backward(ctx, grad_output: torch.Tensor): + kv, score, ape, cu_seqlens, cu_seqlens_comp = ctx.saved_tensors + plan = ctx.plan + local_plan = plan.local_plan + if local_plan is None: # Validated by the public entry point. + raise RuntimeError("the direct local pooling plan is unavailable") + _, backward = _load_cudnn_index_pool_ops() + + if plan.world_size == 1: + result = backward( + kv, + score, + ape, + cu_seqlens, + cu_seqlens_comp, + grad_output.contiguous(), + ratio=local_plan.ratio, + head_dim=local_plan.head_dim, + coff=local_plan.coff, + deterministic=ctx.deterministic, + ) + return ( + result["grad_kv"], + result["grad_score"], + result["grad_ape"], + None, + None, + None, + None, + None, + ) + + # Each rank's loss supplies one gradient for every global compressed + # block. Route each owner its block slice without merging loss ranks. + # Keeping those BF16 slices separate is essential: merging them first + # would round the CP sum before the nonlinear compressor backward. + counts = plan.block_counts + local_count = counts[plan.rank] + received = grad_output.new_empty((plan.world_size * local_count, *grad_output.shape[1:])) + dist.all_to_all_single( + received, + grad_output.contiguous(), + output_split_sizes=[local_count] * plan.world_size, + input_split_sizes=list(counts), + group=ctx.group, + ) + received = received.view(plan.world_size, local_count, *grad_output.shape[1:]) + + grad_kv = torch.zeros_like(kv, dtype=torch.float32) + grad_score = torch.zeros_like(score, dtype=torch.float32) + grad_ape_by_loss = torch.empty((plan.world_size, *ape.shape), dtype=torch.float32, device=ape.device) + for loss_rank in range(plan.world_size): + result = backward( + kv, + score, + ape, + cu_seqlens, + cu_seqlens_comp, + received[loss_rank].contiguous(), + ratio=local_plan.ratio, + head_dim=local_plan.head_dim, + coff=local_plan.coff, + deterministic=ctx.deterministic, + ) + grad_kv.add_(result["grad_kv"]) + grad_score.add_(result["grad_score"]) + grad_ape_by_loss[loss_rank].copy_(result["grad_ape"]) + + # APE is replicated. Return this loss rank's gradient summed over all + # input owners, matching gather-then-pool autograd semantics before any + # framework-level data-parallel gradient synchronization. + grad_ape = torch.empty_like(ape, dtype=torch.float32) + _reduce_scatter_tensor( + grad_ape, + grad_ape_by_loss.flatten(0, 1).contiguous(), + group=ctx.group, + ) + return ( + grad_kv.to(kv.dtype), + grad_score.to(score.dtype), + grad_ape, + None, + None, + None, + None, + None, + ) + + +def cudnn_index_pool_cp( + kv: torch.Tensor, + score: torch.Tensor, + ape: torch.Tensor, + cu_seqlens: torch.Tensor, + cu_seqlens_comp: torch.Tensor, + *, + plan: IndexPoolShardPlan, + group: dist.ProcessGroup | None = None, + deterministic: bool = False, +) -> torch.Tensor: + """Pool local THD rows, then gather only compressed block rows. + + Complete ratio-sized blocks are owned by the rank containing their first + token. The current execution path requires every CP seam to already be a + block boundary; :class:`IndexPoolShardPlan` detects that condition without + reading device metadata. Misaligned shards are rejected until the explicit + neighbor-halo executor is selected, rather than silently pooling unrelated + tokens across a rank-local seam. + + The collective is differentiable. Backward routes compressed block-gradient + slices to their input owner without first merging BF16 loss-rank gradients, + runs the local cuDNN backward once per loss rank, accumulates token gradients + in FP32, and reduce-scatters FP32 APE gradients. This preserves the original + gather-then-pool numerical order instead of trading precision for one fewer + kernel launch. Rank block counts may differ; only the compressed axis uses + uneven collective splits. ``coff=2`` overlapping windows are intentionally + excluded because they need a left-window halo even when the owned block + boundary itself is aligned. + """ + + if not isinstance(plan, IndexPoolShardPlan): + raise TypeError("plan must be an IndexPoolShardPlan") + if not isinstance(deterministic, bool): + raise TypeError("deterministic must be a bool") + if plan.coff != 1 and plan.world_size > 1: + raise NotImplementedError( + "distributed pool-before-gather currently supports coff=1; coff=2 requires an overlapping-window halo" + ) + plan.validate_group(group) + if not plan.all_shards_direct or not plan.direct_local_pool: + raise NotImplementedError( + "this CP partition has a compressed block crossing a rank seam; " + "use the neighbor-halo executor or gather token rows before pooling" + ) + if any(count == 0 for count in plan.block_counts): + raise NotImplementedError( + "distributed pool-before-gather requires every rank to own at least one complete pooling block" + ) + if plan.local_plan is None: + raise RuntimeError("the direct local pooling plan is unavailable") + if kv.shape[0] != plan.local_tokens: + raise ValueError(f"kv must contain this rank's {plan.local_tokens} token rows, got {kv.shape[0]}") + + _validate_inputs( + kv, + score, + ape, + cu_seqlens, + cu_seqlens_comp, + plan=plan.local_plan, + ) + return _CudnnIndexPoolCP.apply( + kv, + score, + ape, + cu_seqlens, + cu_seqlens_comp, + plan, + group, + deterministic, + ) + + +__all__ = ["cudnn_index_pool_cp"] diff --git a/cula/hybrid/dsa.py b/cula/hybrid/dsa.py new file mode 100644 index 00000000..f3f0cf74 --- /dev/null +++ b/cula/hybrid/dsa.py @@ -0,0 +1,526 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Memory-bounded packed-THD block selection over cuDNN DSA kernels.""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from dataclasses import dataclass +from functools import cache +from importlib import import_module +from operator import index + +import torch + +from cula._metadata import as_boundaries + +MEMORY_EFFICIENT_CANDIDATE_BUDGET_BYTES = 512 << 20 +"""Lower-memory policy matching cuDNN QSA's serving-side selection budget.""" + +DEFAULT_CANDIDATE_BUDGET_BYTES = MEMORY_EFFICIENT_CANDIDATE_BUDGET_BYTES +"""Base FP32 compact-logit workspace used by the adaptive default policy.""" + +MAX_BALANCED_CANDIDATE_BUDGET_BYTES = 8 << 30 +"""Maximum FP32 compact-logit workspace selected by the adaptive policy.""" + +_CANDIDATE_BUDGET_POSITION_QUANTUM = 128 << 10 +_CANDIDATE_BUDGET_LOCAL_ROWS_QUANTUM = 64 << 10 +_MAX_LOCAL_ROWS_MULTIPLIER = 4 + + +def _as_nonnegative(values: torch.Tensor | Sequence[int], *, entries: int, name: str) -> tuple[int, ...]: + if isinstance(values, torch.Tensor): + if values.device.type != "cpu": + raise ValueError(f"{name} must be a CPU tensor; copying CUDA metadata would synchronize the device") + if values.ndim != 1: + raise ValueError(f"{name} must be one-dimensional") + if values.dtype not in (torch.int32, torch.int64): + raise TypeError(f"{name} must have int32 or int64 dtype") + result = tuple(int(value) for value in values.tolist()) + else: + try: + result = tuple(index(value) for value in values) + except (TypeError, ValueError) as error: + raise TypeError(f"{name} must be an integer sequence") from error + if len(result) != entries: + raise ValueError(f"{name} must contain {entries} entries") + if any(value < 0 for value in result): + raise ValueError(f"{name} entries must be non-negative") + if any(value >= 2**31 for value in result): + raise ValueError(f"{name} entries must fit int32") + return result + + +def _as_block_prefix(values: torch.Tensor | Sequence[int]) -> tuple[int, ...]: + """Read a host block prefix, permitting zero-block packed sequences.""" + + if isinstance(values, torch.Tensor): + if values.device.type != "cpu": + raise ValueError("cu_seqlens_k must be a CPU tensor; copying CUDA metadata would synchronize the device") + if values.ndim != 1: + raise ValueError("cu_seqlens_k must be one-dimensional") + if values.dtype not in (torch.int32, torch.int64): + raise TypeError("cu_seqlens_k must have int32 or int64 dtype") + result = tuple(int(value) for value in values.tolist()) + else: + try: + result = tuple(index(value) for value in values) + except (TypeError, ValueError) as error: + raise TypeError("cu_seqlens_k must be an integer sequence") from error + if len(result) < 2 or result[0] != 0: + raise ValueError("cu_seqlens_k must contain at least two entries and start at zero") + if any(left > right for left, right in zip(result, result[1:])): + raise ValueError("cu_seqlens_k must be non-decreasing") + if result[-1] == 0: + raise ValueError("the packed input contains no complete key block") + return result + + +def _candidate_prefix(position: int, block_size: int) -> int: + """Sum ``floor(i / block_size)`` for ``i`` in ``[0, position]``.""" + + quotient, remainder = divmod(position, block_size) + return block_size * quotient * (quotient - 1) // 2 + quotient * (remainder + 1) + + +def _candidate_floats(position: int, rows: int, block_size: int) -> int: + return _candidate_prefix(position + rows, block_size) - _candidate_prefix(position, block_size) + + +def _balanced_candidate_budget(cu_q: tuple[int, ...], offsets: tuple[int, ...]) -> int: + max_causal_end = max(offset + right - left for left, right, offset in zip(cu_q[:-1], cu_q[1:], offsets, strict=True)) + position_steps = max( + 1, + (max_causal_end + _CANDIDATE_BUDGET_POSITION_QUANTUM - 1) // _CANDIDATE_BUDGET_POSITION_QUANTUM, + ) + # Causal reach controls scratch per query row; total local rows control how + # many scorer/Top-K windows must launch. Preserve the 128K policy, then + # spend more workspace only where both dimensions make launch partitioning + # an observed throughput bottleneck. + total_local_rows = cu_q[-1] + local_rows_multiplier = 1 + if position_steps > 1: + local_rows_multiplier = min( + _MAX_LOCAL_ROWS_MULTIPLIER, + max( + 1, + (total_local_rows + _CANDIDATE_BUDGET_LOCAL_ROWS_QUANTUM - 1) // _CANDIDATE_BUDGET_LOCAL_ROWS_QUANTUM, + ), + ) + return min( + MAX_BALANCED_CANDIDATE_BUDGET_BYTES, + position_steps * DEFAULT_CANDIDATE_BUDGET_BYTES * local_rows_multiplier, + ) + + +@dataclass(frozen=True, slots=True) +class DsaSelectionWindow: + """One cuDNN THD launch over contiguous packed query and block rows.""" + + query_start: int + query_end: int + key_start: int + key_end: int + cu_seqlens_q: tuple[int, ...] + cu_seqlens_k: tuple[int, ...] + q_causal_offsets: tuple[int, ...] + candidate_offsets: tuple[int, ...] + + @property + def candidate_floats(self) -> int: + return self.candidate_offsets[-1] + + @property + def max_seqlen_q(self) -> int: + return max(right - left for left, right in zip(self.cu_seqlens_q, self.cu_seqlens_q[1:])) + + @property + def max_seqlen_k(self) -> int: + return max(right - left for left, right in zip(self.cu_seqlens_k, self.cu_seqlens_k[1:])) + + +@dataclass(frozen=True, slots=True) +class DsaSelectionPlan: + """Static, memory-bounded selection plan for packed indexer queries. + + ``cu_seqlens_k`` describes already-compressed block rows. The plan keeps + block ids local to each packed sequence, which is the native contract of + cuDNN QSA and avoids expanding 512 block ids into 2048 token ids. + + Long sequences are split into exact causal row windows whose compact-logit + scratch fits ``candidate_budget_bytes``. Windowing changes neither input + precision nor the selected set. Whole adjacent sequences are coalesced into + one launch when they fit the same budget. + """ + + block_size: int + top_k: int + cu_seqlens_q: tuple[int, ...] + cu_seqlens_k: tuple[int, ...] + q_causal_offsets: tuple[int, ...] + candidate_budget_bytes: int + windows: tuple[DsaSelectionWindow, ...] + + @property + def total_queries(self) -> int: + return self.cu_seqlens_q[-1] + + @property + def total_blocks(self) -> int: + return self.cu_seqlens_k[-1] + + @property + def num_sequences(self) -> int: + return len(self.cu_seqlens_q) - 1 + + @property + def candidate_workspace_floats(self) -> int: + return max(window.candidate_floats for window in self.windows) + + @classmethod + def from_cu_seqlens( + cls, + cu_seqlens_q: torch.Tensor | Sequence[int], + cu_seqlens_k: torch.Tensor | Sequence[int], + *, + top_k: int, + block_size: int = 4, + q_causal_offsets: torch.Tensor | Sequence[int] | None = None, + candidate_budget_bytes: int | None = None, + ) -> DsaSelectionPlan: + """Create a plan entirely from host-visible packed metadata. + + A ``None`` budget selects the balanced 512 MiB-per-128K causal-reach + policy. Beyond the first 128K of causal reach, local query rows add up + to a 4x throughput multiplier, capped at 8 GiB. Pass an explicit byte + count to enforce a fixed memory policy across ranks and microbatches. + """ + + cu_q = as_boundaries(cu_seqlens_q, name="cu_seqlens_q") + cu_k = _as_block_prefix(cu_seqlens_k) + if len(cu_q) != len(cu_k): + raise ValueError("cu_seqlens_q and cu_seqlens_k must describe the same number of sequences") + if not isinstance(block_size, int) or isinstance(block_size, bool) or block_size < 1: + raise ValueError("block_size must be a positive integer") + if not isinstance(top_k, int) or isinstance(top_k, bool) or not 1 <= top_k <= 2048: + raise ValueError("top_k must be an integer in [1, 2048]") + if candidate_budget_bytes is not None and ( + not isinstance(candidate_budget_bytes, int) + or isinstance(candidate_budget_bytes, bool) + or candidate_budget_bytes < torch.float32.itemsize + ): + raise ValueError("candidate_budget_bytes must hold at least one FP32 element") + + batch = len(cu_q) - 1 + if cu_q[-1] >= 2**31 or cu_k[-1] >= 2**31: + raise ValueError("packed query and block prefixes must fit int32") + offsets = ( + (0,) * batch + if q_causal_offsets is None + else _as_nonnegative(q_causal_offsets, entries=batch, name="q_causal_offsets") + ) + if candidate_budget_bytes is None: + candidate_budget_bytes = _balanced_candidate_budget(cu_q, offsets) + budget_floats = candidate_budget_bytes // torch.float32.itemsize + atoms: list[tuple[int, int, int, int, int, int, int]] = [] + for sequence in range(batch): + query_length = cu_q[sequence + 1] - cu_q[sequence] + key_length = cu_k[sequence + 1] - cu_k[sequence] + position = offsets[sequence] + if position + query_length >= 2**31: + raise ValueError(f"sequence {sequence} causal positions must fit int32") + required_blocks = (position + query_length) // block_size + if key_length < required_blocks: + raise ValueError( + f"sequence {sequence} has {key_length} compressed blocks but its query geometry requires " + f"at least {required_blocks}" + ) + max_row_candidates = (position + query_length) // block_size + if max_row_candidates > budget_floats: + raise ValueError( + f"candidate_budget_bytes is too small for one row of sequence {sequence}: " + f"need at least {max_row_candidates * torch.float32.itemsize} bytes" + ) + + row = 0 + while row < query_length: + low, high = row + 1, query_length + while low < high: + middle = (low + high + 1) // 2 + needed = _candidate_floats(position + row, middle - row, block_size) + if needed <= budget_floats: + low = middle + else: + high = middle - 1 + row_end = low + candidates = _candidate_floats(position + row, row_end - row, block_size) + atoms.append((sequence, row, row_end, key_length, position + row, candidates, query_length)) + row = row_end + + windows: list[DsaSelectionWindow] = [] + atom = 0 + while atom < len(atoms): + sequence, row_start, row_end, key_length, position, candidates, query_length = atoms[atom] + whole = row_start == 0 and row_end == query_length + group = [atoms[atom]] + atom += 1 + while whole and atom < len(atoms): + next_atom = atoms[atom] + next_sequence, next_start, next_end, _, _, next_candidates, next_length = next_atom + if ( + next_sequence != group[-1][0] + 1 + or next_start != 0 + or next_end != next_length + or sum(item[5] for item in group) + next_candidates > budget_floats + ): + break + group.append(next_atom) + atom += 1 + + first_sequence = group[0][0] + last_sequence = group[-1][0] + query_start = cu_q[first_sequence] + group[0][1] + query_end = cu_q[last_sequence] + group[-1][2] + key_start = cu_k[first_sequence] + key_end = cu_k[last_sequence + 1] + window_cu_q = [0] + window_cu_k = [0] + window_offsets = [] + candidate_offsets = [0] + for seq, start, end, seq_key_length, seq_position, seq_candidates, _ in group: + window_cu_q.append(window_cu_q[-1] + end - start) + window_cu_k.append(window_cu_k[-1] + seq_key_length) + window_offsets.append(seq_position) + candidate_offsets.append(candidate_offsets[-1] + seq_candidates) + windows.append( + DsaSelectionWindow( + query_start=query_start, + query_end=query_end, + key_start=key_start, + key_end=key_end, + cu_seqlens_q=tuple(window_cu_q), + cu_seqlens_k=tuple(window_cu_k), + q_causal_offsets=tuple(window_offsets), + candidate_offsets=tuple(candidate_offsets), + ) + ) + + return cls( + block_size=block_size, + top_k=top_k, + cu_seqlens_q=cu_q, + cu_seqlens_k=cu_k, + q_causal_offsets=offsets, + candidate_budget_bytes=candidate_budget_bytes, + windows=tuple(windows), + ) + + +@dataclass(frozen=True, slots=True) +class _DsaWindowMetadata: + cu_seqlens_q: torch.Tensor + cu_seqlens_k: torch.Tensor + q_causal_offsets: torch.Tensor | None + candidate_offsets: torch.Tensor + + +@dataclass(frozen=True, slots=True) +class DsaSelectionWorkspace: + """Caller-owned outputs, compact-logit scratch, and CUDA graph constants.""" + + plan: DsaSelectionPlan + candidates: torch.Tensor + block_ids: torch.Tensor + scores: torch.Tensor + windows: tuple[_DsaWindowMetadata, ...] + + @classmethod + def from_plan(cls, plan: DsaSelectionPlan, device: torch.device | str) -> DsaSelectionWorkspace: + if not isinstance(plan, DsaSelectionPlan): + raise TypeError("plan must be a DsaSelectionPlan") + resolved = torch.device(device) + if resolved.type != "cuda": + raise ValueError("cuDNN DSA selection workspace must be allocated on CUDA") + metadata = [] + for window in plan.windows: + q_offsets = None + if any(window.q_causal_offsets): + q_offsets = torch.tensor(window.q_causal_offsets, dtype=torch.int32, device=resolved) + metadata.append( + _DsaWindowMetadata( + cu_seqlens_q=torch.tensor(window.cu_seqlens_q, dtype=torch.int32, device=resolved), + cu_seqlens_k=torch.tensor(window.cu_seqlens_k, dtype=torch.int32, device=resolved), + q_causal_offsets=q_offsets, + candidate_offsets=torch.tensor(window.candidate_offsets, dtype=torch.int64, device=resolved), + ) + ) + return cls( + plan=plan, + candidates=torch.empty(plan.candidate_workspace_floats, dtype=torch.float32, device=resolved), + block_ids=torch.empty((plan.total_queries, plan.top_k), dtype=torch.int32, device=resolved), + scores=torch.empty((plan.total_queries, plan.top_k), dtype=torch.float32, device=resolved), + windows=tuple(metadata), + ) + + +@dataclass(frozen=True, slots=True) +class DsaSelectionResult: + """Workspace-owned local block ids and their FP32 selected logits. + + The `(block_id, score)` pairs form a deterministic set when selection uses + deterministic tie-breaking. Their radix output slot order is unspecified. + """ + + block_ids: torch.Tensor + scores: torch.Tensor + + +@cache +def _load_cudnn_dsa_select(): + try: + cudnn = import_module("cudnn") + namespace = getattr(cudnn, "DSA", None) + if namespace is not None and hasattr(namespace, "indexer_forward_top_k_wrapper"): + return namespace.indexer_forward_top_k_wrapper + module = import_module("cudnn.deepseek_sparse_attention") + return module.indexer_forward_top_k_wrapper + except (ImportError, AttributeError) as error: + raise ImportError( + "cuDNN DSA selection requires cuDNN Frontend with indexer_forward_top_k_wrapper; " + "install cuda-linear-attention[cudnn-hybrid]" + ) from error + + +def _validate_selection_inputs( + qi: torch.Tensor, + kbar: torch.Tensor, + weights: torch.Tensor, + *, + plan: DsaSelectionPlan, + workspace: DsaSelectionWorkspace, +) -> None: + if workspace.plan != plan: + raise ValueError("workspace was not created from the supplied DsaSelectionPlan") + if qi.device.type != "cuda" or kbar.device != qi.device or weights.device != qi.device: + raise ValueError("qi, kbar and weights must be CUDA tensors on one device") + if qi.ndim != 3 or qi.shape[0] != plan.total_queries: + raise ValueError(f"qi must have shape [{plan.total_queries}, H, D]") + if kbar.shape != (plan.total_blocks, 1, qi.shape[2]): + raise ValueError(f"kbar must have shape [{plan.total_blocks}, 1, {qi.shape[2]}], got {tuple(kbar.shape)}") + if weights.shape != qi.shape[:2]: + raise ValueError(f"weights must have shape {tuple(qi.shape[:2])}, got {tuple(weights.shape)}") + if qi.dtype != torch.bfloat16 or kbar.dtype != torch.bfloat16: + raise TypeError("qi and kbar must be bfloat16; cuLA does not silently quantize indexer inputs") + if weights.dtype not in (torch.bfloat16, torch.float32): + raise TypeError("weights must be bfloat16 or float32") + for name, tensor in (("qi", qi), ("kbar", kbar), ("weights", weights)): + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous; implicit hot-path layout copies are not allowed") + for name, tensor in ( + ("candidates", workspace.candidates), + ("block_ids", workspace.block_ids), + ("scores", workspace.scores), + ): + if tensor.device != qi.device: + raise ValueError(f"workspace {name} must be on {qi.device}") + if not tensor.is_contiguous(): + raise ValueError(f"workspace {name} must be contiguous") + if workspace.candidates.dtype != torch.float32 or workspace.candidates.numel() < plan.candidate_workspace_floats: + raise ValueError("workspace candidates must be a sufficiently large FP32 tensor") + if workspace.block_ids.dtype != torch.int32 or workspace.block_ids.shape != (plan.total_queries, plan.top_k): + raise ValueError("workspace block_ids has the wrong dtype or shape") + if workspace.scores.dtype != torch.float32 or workspace.scores.shape != (plan.total_queries, plan.top_k): + raise ValueError("workspace scores has the wrong dtype or shape") + if len(workspace.windows) != len(plan.windows): + raise ValueError("workspace window metadata does not match the plan") + for metadata, window in zip(workspace.windows, plan.windows): + entries = len(window.cu_seqlens_q) + for name, tensor, dtype in ( + ("cu_seqlens_q", metadata.cu_seqlens_q, torch.int32), + ("cu_seqlens_k", metadata.cu_seqlens_k, torch.int32), + ("candidate_offsets", metadata.candidate_offsets, torch.int64), + ): + if tensor.device != qi.device or tensor.dtype != dtype or tensor.shape != (entries,) or not tensor.is_contiguous(): + raise ValueError(f"workspace {name} metadata does not match its selection window") + if metadata.q_causal_offsets is not None and ( + metadata.q_causal_offsets.device != qi.device + or metadata.q_causal_offsets.dtype != torch.int32 + or metadata.q_causal_offsets.shape != (entries - 1,) + or not metadata.q_causal_offsets.is_contiguous() + ): + raise ValueError("workspace q_causal_offsets metadata does not match its selection window") + + +def cudnn_dsa_select( + qi: torch.Tensor, + kbar: torch.Tensor, + weights: torch.Tensor, + *, + plan: DsaSelectionPlan, + workspace: DsaSelectionWorkspace, + sm_scale: float = 1.0, + deterministic: bool = False, +) -> DsaSelectionResult: + """Select local compressed-block ids with bounded scratch and FP32 logits. + + The function consumes packed pooling output directly; it never pads blocks + to ``[B, max_blocks, D]``. Every long-sequence window reuses the same + caller-owned compact-logit buffer. Returned tensors alias ``workspace`` + and are overwritten by its next invocation. Non-deterministic selection is + the throughput-oriented default. Pass ``deterministic=True`` when exact-tie + reproducibility is required; cuDNN's radix output slot order remains + unspecified, so downstream code that requires bitwise accumulation order + must canonicalize the pairs explicitly. + """ + + if not isinstance(plan, DsaSelectionPlan): + raise TypeError("plan must be a DsaSelectionPlan") + if not isinstance(workspace, DsaSelectionWorkspace): + raise TypeError("workspace must be a DsaSelectionWorkspace") + if not isinstance(deterministic, bool): + raise TypeError("deterministic must be a bool") + if isinstance(sm_scale, bool) or not isinstance(sm_scale, int | float) or not math.isfinite(sm_scale): + raise ValueError("sm_scale must be a finite host scalar") + _validate_selection_inputs(qi, kbar, weights, plan=plan, workspace=workspace) + + select = _load_cudnn_dsa_select() + heads = int(qi.shape[1]) + for window, metadata in zip(plan.windows, workspace.windows): + query_slice = slice(window.query_start, window.query_end) + key_slice = slice(window.key_start, window.key_end) + select( + qi[query_slice], + kbar[key_slice], + weights[query_slice], + top_k=plan.top_k, + ratio=plan.block_size, + qhead_per_kv_head=heads, + sm_scale=float(sm_scale), + cu_seqlens_q=metadata.cu_seqlens_q, + cu_seqlens_k=metadata.cu_seqlens_k, + max_seqlen_q=window.max_seqlen_q, + max_seqlen_k=window.max_seqlen_k, + q_causal_offsets=metadata.q_causal_offsets, + topk_indices_global=False, + cand_buffer=workspace.candidates, + out_indices=workspace.block_ids[query_slice], + out_logits=workspace.scores[query_slice], + cand_batch_offsets=metadata.candidate_offsets, + return_softmax=False, + deterministic=deterministic, + ) + return DsaSelectionResult(block_ids=workspace.block_ids, scores=workspace.scores) + + +__all__ = [ + "DEFAULT_CANDIDATE_BUDGET_BYTES", + "MAX_BALANCED_CANDIDATE_BUDGET_BYTES", + "MEMORY_EFFICIENT_CANDIDATE_BUDGET_BYTES", + "DsaSelectionPlan", + "DsaSelectionResult", + "DsaSelectionWindow", + "DsaSelectionWorkspace", + "cudnn_dsa_select", +] diff --git a/cula/hybrid/index_pool.py b/cula/hybrid/index_pool.py new file mode 100644 index 00000000..d7ca3164 --- /dev/null +++ b/cula/hybrid/index_pool.py @@ -0,0 +1,162 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Differentiable cuDNN index pooling for packed THD inputs.""" + +from __future__ import annotations + +from functools import cache +from importlib import import_module + +import torch +from torch.autograd.function import once_differentiable + +from cula.hybrid.plan import IndexPoolPlan + + +@cache +def _load_cudnn_index_pool_ops(): + try: + cudnn = import_module("cudnn") + namespace = cudnn.CSA + forward = namespace.csa_compressor_forward_wrapper + backward = namespace.csa_compressor_backward_wrapper + except (ImportError, AttributeError) as error: + raise ImportError( + "cuDNN index pooling requires cuDNN Frontend with the CSA compressor APIs; install cuda-linear-attention[cudnn-cp]" + ) from error + return forward, backward + + +def _validate_metadata(tensor: torch.Tensor, *, name: str, entries: int, device: torch.device) -> None: + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if tensor.shape != (entries,): + raise ValueError(f"{name} must have shape [{entries}], got {tuple(tensor.shape)}") + if tensor.dtype != torch.int32: + raise TypeError(f"{name} must have int32 dtype") + if tensor.device != device: + raise ValueError(f"{name} must be on {device}") + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous") + + +def _validate_inputs( + kv: torch.Tensor, + score: torch.Tensor, + ape: torch.Tensor, + cu_seqlens: torch.Tensor, + cu_seqlens_comp: torch.Tensor, + *, + plan: IndexPoolPlan, +) -> None: + for name, tensor in (("kv", kv), ("score", score), ("ape", ape)): + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if kv.device.type != "cuda": + raise ValueError("cuDNN index pooling requires CUDA operands") + if score.device != kv.device or ape.device != kv.device: + raise ValueError("kv, score and ape must be on the same device") + if kv.shape != (plan.total_tokens, plan.input_width): + raise ValueError(f"kv must have shape [{plan.total_tokens}, {plan.input_width}], got {tuple(kv.shape)}") + if score.shape != kv.shape: + raise ValueError(f"score shape {tuple(score.shape)} must match kv shape {tuple(kv.shape)}") + if ape.shape != (plan.ratio, plan.input_width): + raise ValueError(f"ape must have shape [{plan.ratio}, {plan.input_width}], got {tuple(ape.shape)}") + if kv.dtype != torch.bfloat16 or score.dtype != torch.bfloat16: + raise TypeError("kv and score must have bfloat16 dtype") + if ape.dtype != torch.float32: + raise TypeError("ape must have float32 dtype") + for name, tensor in (("kv", kv), ("score", score), ("ape", ape)): + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous; implicit hot-path layout copies are not allowed") + entries = plan.num_sequences + 1 + _validate_metadata(cu_seqlens, name="cu_seqlens", entries=entries, device=kv.device) + _validate_metadata(cu_seqlens_comp, name="cu_seqlens_comp", entries=entries, device=kv.device) + if cu_seqlens.device.type == "cpu": + if tuple(int(value) for value in cu_seqlens.tolist()) != plan.cu_seqlens: + raise ValueError("cu_seqlens values do not match IndexPoolPlan") + if tuple(int(value) for value in cu_seqlens_comp.tolist()) != plan.cu_seqlens_comp: + raise ValueError("cu_seqlens_comp values do not match IndexPoolPlan") + + +class _CudnnIndexPool(torch.autograd.Function): + @staticmethod + def forward( + ctx, + kv: torch.Tensor, + score: torch.Tensor, + ape: torch.Tensor, + cu_seqlens: torch.Tensor, + cu_seqlens_comp: torch.Tensor, + plan: IndexPoolPlan, + deterministic: bool, + ) -> torch.Tensor: + forward, _ = _load_cudnn_index_pool_ops() + result = forward( + kv, + score, + ape, + cu_seqlens, + cu_seqlens_comp, + ratio=plan.ratio, + head_dim=plan.head_dim, + coff=plan.coff, + total_comp=plan.total_comp, + ) + ctx.plan = plan + ctx.deterministic = deterministic + ctx.save_for_backward(kv, score, ape, cu_seqlens, cu_seqlens_comp) + return result["out"] + + @staticmethod + @once_differentiable + def backward(ctx, grad_out: torch.Tensor): + _, backward = _load_cudnn_index_pool_ops() + kv, score, ape, cu_seqlens, cu_seqlens_comp = ctx.saved_tensors + plan = ctx.plan + result = backward( + kv, + score, + ape, + cu_seqlens, + cu_seqlens_comp, + grad_out.contiguous(), + ratio=plan.ratio, + head_dim=plan.head_dim, + coff=plan.coff, + deterministic=ctx.deterministic, + ) + return result["grad_kv"], result["grad_score"], result["grad_ape"], None, None, None, None + + +def cudnn_index_pool( + kv: torch.Tensor, + score: torch.Tensor, + ape: torch.Tensor, + cu_seqlens: torch.Tensor, + cu_seqlens_comp: torch.Tensor, + *, + plan: IndexPoolPlan, + deterministic: bool = False, +) -> torch.Tensor: + """Pool packed token rows with cuDNN's FP32-intermediate compressor. + + cuLA owns the static packed metadata and PyTorch autograd connection; + cuDNN Frontend owns the fused forward/backward kernels. Inputs must + already use the native contiguous THD rows. This function never performs + an implicit layout conversion or changes the kernel's numerical contract. + + Set ``deterministic=True`` to use cuDNN's fixed-order dAPE reduction. This + affects reproducibility and workspace scheduling, not arithmetic precision. + """ + + if not isinstance(plan, IndexPoolPlan): + raise TypeError("plan must be an IndexPoolPlan") + if not isinstance(deterministic, bool): + raise TypeError("deterministic must be a bool") + _validate_inputs(kv, score, ape, cu_seqlens, cu_seqlens_comp, plan=plan) + return _CudnnIndexPool.apply(kv, score, ape, cu_seqlens, cu_seqlens_comp, plan, deterministic) + + +__all__ = ["cudnn_index_pool"] diff --git a/cula/hybrid/plan.py b/cula/hybrid/plan.py new file mode 100644 index 00000000..4612092b --- /dev/null +++ b/cula/hybrid/plan.py @@ -0,0 +1,299 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Static packed-layout plans shared by cuDNN hybrid-attention operators.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from operator import index + +import torch +import torch.distributed as dist + +from cula._metadata import as_boundaries, balanced_shards, validate_shards + + +@dataclass(frozen=True, slots=True) +class IndexPoolPlan: + """Static THD metadata for one cuDNN index-pooling invocation. + + A plan is built from host metadata once per packed microbatch. It records + the compacted block axis without reading CUDA values back to the host. The + corresponding device tensors can then be materialized outside a captured + or allocation-free execution region and reused by every compatible layer. + + ``coff=1`` describes non-overlapping ``ratio``-token blocks. ``coff=2`` + selects cuDNN's overlapping compressor layout; it changes the projection + width, not the number of compacted rows. + """ + + ratio: int + coff: int + head_dim: int + cu_seqlens: tuple[int, ...] + cu_seqlens_comp: tuple[int, ...] + + @property + def total_tokens(self) -> int: + return self.cu_seqlens[-1] + + @property + def total_comp(self) -> int: + return self.cu_seqlens_comp[-1] + + @property + def num_sequences(self) -> int: + return len(self.cu_seqlens) - 1 + + @property + def input_width(self) -> int: + return self.coff * self.head_dim + + @classmethod + def from_cu_seqlens( + cls, + cu_seqlens: torch.Tensor | Sequence[int], + *, + head_dim: int, + ratio: int = 4, + coff: int = 1, + ) -> IndexPoolPlan: + """Build a plan without accessing CUDA-resident metadata.""" + + boundaries = as_boundaries(cu_seqlens, name="cu_seqlens") + if not isinstance(head_dim, int) or isinstance(head_dim, bool) or head_dim < 1: + raise ValueError("head_dim must be a positive integer") + if not isinstance(ratio, int) or isinstance(ratio, bool) or ratio < 1: + raise ValueError("ratio must be a positive integer") + if coff not in (1, 2): + raise ValueError("coff must be 1 (own-block) or 2 (overlapping)") + if boundaries[-1] >= 2**31: + raise ValueError("cu_seqlens values must fit int32") + + compact = [0] + for start, end in zip(boundaries, boundaries[1:]): + compact.append(compact[-1] + (end - start) // ratio) + if compact[-1] == 0: + raise ValueError("the packed input contains no complete pooling block") + return cls( + ratio=ratio, + coff=coff, + head_dim=head_dim, + cu_seqlens=boundaries, + cu_seqlens_comp=tuple(compact), + ) + + def materialize(self, device: torch.device | str) -> tuple[torch.Tensor, torch.Tensor]: + """Create reusable int32 metadata tensors on ``device``. + + This method allocates by design. Call it during plan construction or + warmup, not inside a CUDA Graph capture or the per-layer hot path. + """ + + resolved = torch.device(device) + cu_seqlens = torch.tensor(self.cu_seqlens, dtype=torch.int32, device=resolved) + cu_seqlens_comp = torch.tensor(self.cu_seqlens_comp, dtype=torch.int32, device=resolved) + return cu_seqlens, cu_seqlens_comp + + +def _ceil_div(value: int, divisor: int) -> int: + return (value + divisor - 1) // divisor + + +def _local_boundaries(global_cu_seqlens: tuple[int, ...], token_start: int, token_end: int) -> tuple[int, ...]: + return ( + (0,) + + tuple(value - token_start for value in global_cu_seqlens[1:-1] if token_start < value < token_end) + + (token_end - token_start,) + ) + + +def _interval_pool_geometry( + boundaries: tuple[int, ...], token_start: int, token_end: int, ratio: int +) -> tuple[int, int, int, bool]: + """Return owned blocks, left send, right receive, and direct-pool eligibility. + + A complete block is owned by the rank containing its first token. A rank + that starts in the middle of such a block contributes its prefix-adjacent + rows to the previous owner; a rank ending in the middle of an owned block + receives the remaining rows from the next rank. + """ + + owned_blocks = 0 + send_left_tokens = 0 + recv_right_tokens = 0 + direct = True + for sequence_start, sequence_end in zip(boundaries, boundaries[1:]): + local_start = max(token_start, sequence_start) + local_end = min(token_end, sequence_end) + if local_start >= local_end: + continue + sequence_length = sequence_end - sequence_start + complete_blocks = sequence_length // ratio + complete_end = complete_blocks * ratio + start_offset = local_start - sequence_start + end_offset = local_end - sequence_start + + first_owned = _ceil_div(start_offset, ratio) + last_owned = min(_ceil_div(end_offset, ratio), complete_blocks) + owned_blocks += max(last_owned - first_owned, 0) + + if start_offset % ratio and start_offset < complete_end: + direct = False + block_end = min(_ceil_div(start_offset, ratio) * ratio, complete_end) + send_left_tokens = max(send_left_tokens, max(block_end - start_offset, 0)) + # A sequence tail is intentionally dropped. It needs no halo and does + # not make direct local pooling invalid when the rank owns the tail. + if local_end < sequence_end and end_offset % ratio and end_offset < complete_end: + direct = False + owned_block_end = min(_ceil_div(end_offset, ratio) * ratio, complete_end) + recv_right_tokens = max(recv_right_tokens, max(owned_block_end - end_offset, 0)) + return owned_blocks, send_left_tokens, recv_right_tokens, direct + + +@dataclass(frozen=True, slots=True) +class IndexPoolShardPlan: + """Distributed ratio-block ownership and halo metadata for one CP rank. + + The plan assigns each complete compressed block to the rank that owns its + first token. Aligned shards can call cuDNN pooling directly and gather the + compact rows. Misaligned shards expose the bounded seam exchange needed + by a future/optional halo stitcher, avoiding a full raw-token all-gather. + """ + + rank: int + world_size: int + token_start: int + token_end: int + ratio: int + coff: int + head_dim: int + local_cu_seqlens: tuple[int, ...] + block_offsets: tuple[int, ...] + send_left_tokens: int + recv_right_tokens: int + direct_local_pool: bool + all_shards_direct: bool + halo_supported: bool + local_plan: IndexPoolPlan | None + + @property + def local_tokens(self) -> int: + return self.token_end - self.token_start + + @property + def local_blocks(self) -> int: + return self.block_offsets[self.rank + 1] - self.block_offsets[self.rank] + + @property + def global_blocks(self) -> int: + return self.block_offsets[-1] + + @property + def block_counts(self) -> tuple[int, ...]: + return tuple(right - left for left, right in zip(self.block_offsets, self.block_offsets[1:])) + + @property + def equal_block_counts(self) -> bool: + counts = self.block_counts + return all(count == counts[0] for count in counts[1:]) + + def validate_group(self, group: dist.ProcessGroup | None) -> None: + """Check that the runtime process group matches this rank-local plan.""" + + if not dist.is_available() or not dist.is_initialized(): + if group is not None or self.world_size != 1 or self.rank != 0: + raise RuntimeError("the index-pooling shard plan requires an initialized torch.distributed process group") + return + world_size = dist.get_world_size(group) + rank = dist.get_rank(group) + if (world_size, rank) != (self.world_size, self.rank): + raise ValueError( + "IndexPoolShardPlan process-group mismatch: " + f"plan has rank/world_size={self.rank}/{self.world_size}, " + f"runtime has {rank}/{world_size}" + ) + + @classmethod + def from_global_cu_seqlens( + cls, + global_cu_seqlens: torch.Tensor | Sequence[int], + *, + rank: int, + world_size: int, + head_dim: int, + ratio: int = 4, + coff: int = 1, + shard_offsets: Sequence[int] | None = None, + ) -> IndexPoolShardPlan: + """Build a static distributed plan from global host metadata.""" + + boundaries = as_boundaries(global_cu_seqlens, name="global_cu_seqlens") + # Reuse IndexPoolPlan's option validation and compacted global layout. + IndexPoolPlan.from_cu_seqlens(boundaries, head_dim=head_dim, ratio=ratio, coff=coff) + if not isinstance(world_size, int) or isinstance(world_size, bool) or world_size < 1: + raise ValueError("world_size must be a positive integer") + if not isinstance(rank, int) or isinstance(rank, bool) or not 0 <= rank < world_size: + raise ValueError(f"rank must be in [0, {world_size})") + offsets = ( + balanced_shards(boundaries[-1], world_size) + if shard_offsets is None + else validate_shards(shard_offsets, total_tokens=boundaries[-1], world_size=world_size) + ) + + counts = [] + direct_by_rank = [] + geometry = [] + for shard_rank in range(world_size): + values = _interval_pool_geometry( + boundaries, + offsets[shard_rank], + offsets[shard_rank + 1], + ratio, + ) + geometry.append(values) + counts.append(values[0]) + direct_by_rank.append(values[3]) + block_offsets = [0] + for count in counts: + block_offsets.append(block_offsets[-1] + count) + + start, end = offsets[rank], offsets[rank + 1] + local_cu = _local_boundaries(boundaries, start, end) + owned_blocks, send_left, recv_right, direct = geometry[rank] + local_plan = None + if direct and owned_blocks: + local_plan = IndexPoolPlan.from_cu_seqlens( + local_cu, + head_dim=head_dim, + ratio=ratio, + coff=coff, + ) + if local_plan.total_comp != owned_blocks: + raise RuntimeError("internal error while deriving the direct local index-pooling plan") + + # With at least ``ratio`` tokens per shard, one compressed block can + # cross at most one rank seam, so a single neighbor exchange suffices. + halo_supported = min(right - left for left, right in zip(offsets, offsets[1:])) >= ratio + return cls( + rank=index(rank), + world_size=index(world_size), + token_start=start, + token_end=end, + ratio=ratio, + coff=coff, + head_dim=head_dim, + local_cu_seqlens=local_cu, + block_offsets=tuple(block_offsets), + send_left_tokens=send_left, + recv_right_tokens=recv_right, + direct_local_pool=direct, + all_shards_direct=all(direct_by_rank), + halo_supported=halo_supported, + local_plan=local_plan, + ) + + +__all__ = ["IndexPoolPlan", "IndexPoolShardPlan"] diff --git a/docs/HYBRID_CUDNN_RUNTIME.md b/docs/HYBRID_CUDNN_RUNTIME.md new file mode 100644 index 00000000..1c479461 --- /dev/null +++ b/docs/HYBRID_CUDNN_RUNTIME.md @@ -0,0 +1,262 @@ +# cuDNN hybrid-attention runtime + +cuLA's hybrid layer is a layout-aware integration boundary for composing +cuDNN Frontend attention kernels. It is intentionally not another graph API: +cuDNN owns the kernels and their numerical contracts; cuLA owns static packed +metadata, bounded workspace, context-parallel layout, and framework-facing +PyTorch autograd. + +The first implemented path covers index pooling and DSA block selection: + +1. `cudnn_index_pool` consumes native contiguous THD rows and calls cuDNN CSA + compressor forward/backward. The FP32 intermediate and dAPE reduction + behavior are unchanged. +2. `DsaSelectionPlan` consumes host-visible packed query/block prefixes and + derives exact causal selection windows. +3. `cudnn_dsa_select` consumes the packed pooling output directly and calls + cuDNN's fused compact-logit + radix Top-K path. It does not materialize a + dense score tensor and does not pad blocks to `[B, max_blocks, D]`. + +The selected ids remain local compressed-block ids for each sequence. This is +the native QSA contract and avoids expanding `[T, 512]` block ids into a +`[T, 2048]` token-id tensor (128 MiB per rank at 16K local tokens). A future +QSA execution adapter should consume these block ids directly; cuLA should not +add an expansion solely to fit a token-index API. + +## Numerical contract + +The hybrid layer does not expose a precision downgrade knob: + +- index-pool KV/score inputs are BF16 and APE is FP32, matching cuDNN; +- DSA indexer Q/K are BF16, weights may be BF16 or FP32, and selected logits + are FP32; +- non-deterministic Top-K is the throughput-oriented default; callers can opt + into deterministic tie-breaking when exact ties at the K-th boundary must + keep the smaller block id; +- with deterministic selection, the `(block_id, logit)` set is reproducible, + not its radix slot order. CUDA Graph capture/replay preserved the sorted + pairs bitwise in validation, while the raw slot order changed; +- windowing replays the same causal rows and visible block range. It only + changes scratch lifetime and launch partitioning. + +On the remote SM100 validation environment, a 4096-token full launch and four +1024-row causal windows produced identical selected sets and bitwise-identical +selected logits (`max_abs = 0`). + +## Why the selection plan owns windows + +cuDNN's THD compact-logit path uses exact packed storage, but its scratch grows +with the number of visible `(query, block)` pairs. For a long causal sequence, +that is still quadratic scratch even though no dense score tensor is exposed. +`DsaSelectionPlan` splits rows at causal offsets and reuses one caller-owned +FP32 candidate buffer. It coalesces adjacent short packed sequences when they +fit the same budget. + +When `candidate_budget_bytes=None`, the balanced policy starts at 512 MiB and +adds 512 MiB for each 128K of absolute causal reach. Beyond the first 128K of +causal reach, each 64K of rank-local query rows adds up to a capped 4x +throughput multiplier; the global cap is 8 GiB. Thus the 128K packed workload +retains its 512 MiB workspace, the worst 1M/CP64 rank stays at its measured +4 GiB performance knee, and critical long-sequence ranks at 1M/CP4 select the +measured 8 GiB knee. The exported +`MEMORY_EFFICIENT_CANDIDATE_BUDGET_BYTES` forces a fixed 512 MiB policy, and +`MAX_BALANCED_CANDIDATE_BUDGET_BYTES` exposes the cap. This is a +memory/scheduling choice, not a numerical choice. Allocate one workspace per +concurrent execution stream and reuse it across compatible layers; do not +allocate one multi-GiB workspace per layer. + +### 128K heterogeneous THD result + +Input: 20 sequences of 1,024 tokens and two sequences of 55,296 tokens, +compression ratio 4, Top-K 512. + +| Selection form | Windows | Candidate scratch | Median time | +|---|---:|---:|---:| +| One packed launch | 1 | 2.8573 GiB | 5.3818 ms | +| 512 MiB memory policy | 7 | 512 MiB | 5.3246 ms | + +The memory policy reduced candidate scratch by 5.715x. The measured time was +effectively unchanged; the small apparent speedup is within benchmark noise. + +The same batch partitioned as four contiguous 32K-token CP shards shows the +rank imbalance that a global-only measurement hides: + +| CP rank | Full scratch | Bounded scratch | Windows | Full | Bounded | Delta | +|---:|---:|---:|---:|---:|---:|---:| +| 0 | 81.969 MiB | 81.969 MiB | 1 | 0.4274 ms | 0.4293 ms | +0.45% | +| 1 | 895.969 MiB | 511.970 MiB | 2 | 1.6170 ms | 1.6589 ms | +2.59% | +| 2 | 731.969 MiB | 489.990 MiB | 2 | 1.3802 ms | 1.4048 ms | +1.78% | +| 3 | 1,215.969 MiB | 511.985 MiB | 3 | 2.0281 ms | 2.1238 ms | +4.72% | + +The worst rank saves about 704 MiB of candidate scratch for a 4.72% local +selection-stage cost. This table excludes model tensors, outputs, and +collectives; it isolates the DSA compact-logit stage. + +### 1M context, simulated CP64 worst rank + +The simulation uses one CP64 rank's real shape: 16,384 query tokens at logical +rank 63, 262,144 globally visible compressed blocks, ratio 4 and Top-K 512. +Every entry executes the real cuDNN scorer/Top-K kernels; only the 64-rank +network is virtualized. + +| Candidate budget | Windows | Median selection time | Versus full | +|---:|---:|---:|---:| +| 256 MiB | 64 | 78.895 ms | 2.81x | +| 512 MiB | 32 | 72.247 ms | 2.58x | +| 1 GiB | 16 | 48.121 ms | 1.72x | +| 2 GiB | 8 | 37.752 ms | 1.35x | +| **4 GiB (adaptive CP64 choice)** | **4** | **33.519 ms** | **1.20x** | +| 8 GiB | 2 | 29.464 ms | 1.05x | +| 15.875 GiB full | 1 | 28.039 ms | 1.00x | + +Four GiB is the measured knee: it saves 11.875 GiB on the worst logical rank +while retaining about 84% of the full-scratch throughput. The 512 MiB policy +remains available when activation pressure matters more than selection latency. + +### 1M THD, real CP4 KDA + DSA indexer E2E + +Input: 20 sequences of 1,024 tokens plus two sequences of 514,048 tokens, +four physical CP ranks, KDA forward/backward, ratio-4 CSA pooling, compressed +block all-gather, and deterministic Top-K 512 selection. KDA and the independent +indexer branch run on separate CUDA streams. + +| Candidate policy | Critical GPU median | Versus old default | +|---|---:|---:| +| Old reach-only adaptive policy | 228.220 ms | 1.00x | +| Fixed 4 GiB | 218.031 ms | 1.05x | +| **Row-aware adaptive, up to 8 GiB** | **209.504 ms** | **1.09x** | +| Fixed 16 GiB | 210.363 ms | 1.08x | + +Eight GiB is the E2E knee. On the critical rank it uses 12 selection windows +and raises peak allocated memory from about 43.9 GiB to 49.7 GiB. Sixteen GiB +doubles candidate scratch without improving the critical path. The output and +all five KDA gradients remained finite; proxy loss, norms, selected-id sum, +and finite-logit sum were stable across seven repeats and matched the smaller +workspace policies. + +Nsight Systems measured only 1.318 ms of kernel-idle time in the 212.454 ms +critical-rank hot path (0.62%). CSA pooling itself took 0.046 ms and compressed +block communication about 0.2--0.3 ms. The remaining critical work is dominated +by 12 cuDNN score kernels (88.5 ms total) and 12 deterministic Top-K kernels +(90.1 ms total), overlapped with KDA and NCCL. High-priority streams and +20--100 ms staggered launches produced no stable gain, so cuLA does not retain +those scheduling knobs. + +An isolated NCU sample identifies the next optimization boundary inside the +cuDNN kernels. The score kernel used 112 registers per thread and 75.8 KiB of +dynamic shared memory per block; theoretical occupancy was 25% and achieved +occupancy 10.9%, with 45.2% compute and 35.3% maximum-memory throughput. The +Top-K kernel used 39 registers and 25.0 KiB of dynamic shared memory, reached +71.6% occupancy, but only 32.7% compute and 27.1% memory throughput. Further +material gains therefore require reducing score-kernel register/shared-memory +pressure and improving Top-K reduction/radix synchronization in cuDNN, rather +than adding more cuLA stream or caching policy. + +## Integration example + +```python +from cula.hybrid import ( + DsaSelectionPlan, + DsaSelectionWorkspace, + IndexPoolPlan, + cudnn_dsa_select, + cudnn_index_pool, +) + +pool_plan = IndexPoolPlan.from_cu_seqlens( + cu_seqlens_cpu, + head_dim=128, + ratio=4, +) +cu_tokens, cu_blocks = pool_plan.materialize(qi.device) +kbar = cudnn_index_pool( + raw_index_k, + pool_score, + ape, + cu_tokens, + cu_blocks, + plan=pool_plan, + deterministic=True, +) + +select_plan = DsaSelectionPlan.from_cu_seqlens( + cu_queries_cpu, + pool_plan.cu_seqlens_comp, + top_k=512, +) +select_workspace = DsaSelectionWorkspace.from_plan(select_plan, qi.device) +selected = cudnn_dsa_select( + qi, + kbar[:, None, :], + index_weights, + plan=select_plan, + workspace=select_workspace, +) +``` + +Construct plans and workspaces during setup/warmup, not inside CUDA Graph +capture. The execution call performs no implicit layout conversion. Returned +selection tensors alias the workspace and are overwritten on its next use. + +## Context-parallel ownership + +`IndexPoolShardPlan` assigns each complete ratio-4 block to the rank owning its +first token. `cudnn_index_pool_cp` executes the aligned fast path: each rank +pools its native THD rows first and all-gathers only compressed blocks. Unequal +block counts use uneven compressed-axis collective splits; raw token tensors +are never padded or gathered. + +Backward deliberately does not reduce BF16 block gradients before cuDNN. It +all-to-alls each loss rank's compressed slice to its input owner, invokes the +local cuDNN backward once per loss rank, accumulates KV/score gradients in +FP32, and reduce-scatters FP32 APE gradients. This has the same total token work +as the global reference and avoids an early BF16 CP merge. At 128K/CP4 the +forward, dKV, and dScore matched bitwise; dAPE differed only by FP32 reduction +order and was closer to the FP32 eager oracle than the reference. + +Measured on the four-GPU validation host with ratio 4, head dimension 128, and +the 20x1K + 2-long THD mix: + +| Global tokens | Raw / pooled forward payload | Forward | Forward + backward | +|---:|---:|---:|---:| +| 128K | 16 MiB / 2 MiB | 0.214 / 0.150 ms (1.43x) | 0.713 / 0.799 ms (0.89x) | +| 1M | 128 MiB / 16 MiB | 1.037 / 0.230 ms (4.51x) | 3.670 / 2.026 ms (1.81x) | + +The 1M forward+backward incremental peak allocation fell from 1,280 MiB to +640 MiB (2.0x). At 128K the precision-preserving four-launch backward costs +about 11% end-to-end for this operator; at 1M, reduced communication and +working set make the complete path 1.81x faster. Timings are medians of +interleaved, order-reversed pairs using each iteration's slowest rank. These +are isolated index-pooling measurements, not whole-model throughput claims. + +Misaligned shards still report a bounded neighbor halo (at most three raw +tokens for ratio 4), avoiding the need for a raw-token all-gather in the next +executor. The public function currently rejects those seams rather than +silently changing block membership. `coff=2` and ranks owning zero complete +blocks are also explicit unsupported boundaries. + +For large CP, the next execution layer should keep tokens stationary and move +only compressed blocks or QSA partial statistics. It must preserve FP32 merge +semantics; reducing collective payload by lowering merge precision is outside +this runtime's contract. + +## Current boundary + +Implemented and tested: + +- differentiable cuDNN index pooling; +- aligned CP pool-before-gather with precision-preserving backward; +- packed THD DSA selection with exact bounded workspace; +- deterministic local block ids and FP32 logits; +- CP block ownership and seam-halo planning. + +Not yet claimed as complete: + +- the distributed ratio-4 halo exchange executor; +- a public block-id-native QSA forward/backward composition; +- hybrid linear + MLA collective scheduling; +- end-to-end CUDA Graph capture tests of the complete hybrid block. + +These omissions are explicit because a reusable integration layer should not +hide layout copies or claim end-to-end correctness before those paths are +measured against the monolithic references. diff --git a/docs/KDA_CSA_1M_OPTIMIZATION.md b/docs/KDA_CSA_1M_OPTIMIZATION.md new file mode 100644 index 00000000..2fe14e7c --- /dev/null +++ b/docs/KDA_CSA_1M_OPTIMIZATION.md @@ -0,0 +1,81 @@ +# KDA + CSA 1M context optimization study + +This note records measured decisions for the packed-THD workload with 20 +sequences of 1,024 tokens and two long sequences. Unless stated otherwise, +the run uses four physical GPUs, 1,048,576 global tokens, BF16 operands, FP32 +Kimi CP merge, deterministic CSA selection, `top_k=512`, and block size 4. + +## Exact KDA checkpoint reuse + +cuDNN KDA can retain the state at every native 16-token chunk and consume it +directly in backward. cuLA exposes this as the explicit +`checkpoint_every_n_tokens` option. The default remains zero because saved +states live across the forward/backward interval for every layer, while the +recompute workspace can be reused between layers. + +| Workload | Checkpoint | Critical E2E | Critical KDA | Peak allocated | +|---|---:|---:|---:|---:| +| 128K / CP4 | 0 | 9.780 ms | 9.092 ms | 6.362 GiB | +| 128K / CP4 | 16 | 8.987 ms | 8.159 ms | 5.974 GiB | +| 1M / CP4 | 0 | 208.798 ms | 135.050 ms | 49.904 GiB | +| 1M / CP4 | 16 | 204.927 ms | 111.352 ms | 49.750 GiB | + +At 128K, checkpoint reuse improves the hybrid critical path by 8.1%. At 1M, +it improves the slowest KDA rank by 17.5%, but hybrid E2E improves only 1.9% +because deterministic CSA selection becomes the critical path. Output, proxy +loss, and all five input-gradient norms were exactly equal to the recompute +path on every rank in both measured workloads. + +## CSA scratch is not the remaining answer + +The default adaptive policy caps the compact-logit scratch at 8 GiB. Increasing +the cap reduces the number of exact query windows, but larger selector kernels +contend more strongly with KDA on the overlap stream. + +| Scratch cap | Max windows | Critical E2E | Critical select | Critical KDA | Peak allocated | +|---:|---:|---:|---:|---:|---:| +| 8 GiB | 12 | 204.927 ms | 204.524 ms | 111.352 ms | 49.750 GiB | +| 16 GiB | 6 | 205.090 ms | 204.742 ms | 194.178 ms | 59.701 GiB | +| 32 GiB | 3 | 204.129 ms | 180.216 ms | 203.184 ms | 73.649 GiB | +| 64 GiB | 2 | 204.970 ms | 172.687 ms | 204.300 ms | 105.649 GiB | + +The 32 GiB point buys less than 0.4% E2E for about 24 GiB more peak memory; +64 GiB regresses. The 8 GiB cap remains the better default Pareto point. + +## Why contiguous CP is imbalanced for CSA + +CSA selection work is proportional to the causal position of each query. A +contiguous token partition assigns later regions of each long sequence to a +small subset of ranks. For this workload the exact compact-candidate counts are: + +| Layout | Rank 0 | Rank 1 | Rank 2 | Rank 3 | Max / min | +|---|---:|---:|---:|---:|---:| +| Contiguous | 7.303B | 24.428B | 9.235B | 25.099B | 3.44x | +| Alternating by long sequence | 16.516B | 16.516B | 16.516B | 16.516B | 1.00x | + +The alternating plan keeps short samples whole. Each rank receives five of the +20 short samples, an ascending shard of the first long sample, and a descending +shard of the second long sample. Every rank still owns exactly 262,144 tokens; +the total arithmetic is unchanged, but early and late causal regions cancel. + +A layout-native DSA benchmark includes deterministic index pooling, all-gather, +global block reconstruction, discontiguous key remapping, and exact selection: +the published measurements below used four index heads. + +| Query layout | Critical DSA time | Peak allocated | +|---|---:|---:| +| Contiguous | 166.497 ms | 12.200 GiB | +| Alternating | 109.173 ms | 12.174 GiB | + +The measured critical path improves by 34.4%; remap plus communication stays +below 1.2 ms on every balanced rank. No arithmetic precision, Top-K semantics, +or deterministic mode changes. + +This is currently a layout study, not a production KDA E2E claim. A correct +integration needs a multi-span Kimi CP plan: cuDNN must receive one initial +state per local span, and forward/backward summaries must be routed by +`(sequence, segment-order)` rather than by physical rank order. The layout must +be produced before the model stack so layers do not pay token-layout conversion. +The CPU THD planner lives in `cula/distributed/thd_planner.py`; the layout +benchmark and its query-layout adapter live in `benchmarks/large_cp_emulation`, +separate from the contiguous KDA runtime. diff --git a/docs/KIMI_CP_RUNTIME.md b/docs/KIMI_CP_RUNTIME.md index 260ca71b..0f3f795c 100644 --- a/docs/KIMI_CP_RUNTIME.md +++ b/docs/KIMI_CP_RUNTIME.md @@ -3,10 +3,11 @@ ## Scope cuLA owns the distributed execution layer around cuDNN Frontend KDA. cuDNN -owns local KDA forward/backward, span summaries, and ordered FP32 affine-summary -composition inside the KDA apply operations. Frameworks own model-specific -tensor layouts, gate parameterization, and construction of the global packed -sequence metadata. +owns local KDA forward/backward and, on Blackwell, span summaries and ordered +FP32 affine-summary composition inside the KDA apply operations. Hopper uses +the same public layer through the ordinary cuDNN state ports described below. +Frameworks own model-specific tensor layouts, gate parameterization, and +construction of the global packed sequence metadata. The public surface intentionally has three KDA-specific layers: @@ -15,10 +16,12 @@ The public surface intentionally has three KDA-specific layers: with `AllGatherKimiCPTransport` as the exact prepared implementation. - `cudnn_kda_cp`: the differentiable distributed KDA call. -There is no generic backend hierarchy. The runtime is KDA-specific and uses -the cuDNN contracts directly. A future hierarchical or scan implementation can -replace only boundary-summary propagation without changing planning, local KDA, -autograd, or the framework adapter. +There is no generic backend hierarchy. The runtime has one narrow +architecture selector because cuDNN exposes different summary contracts on +Hopper and Blackwell, but both implementations remain KDA-specific and use +cuDNN state contracts directly. A future hierarchical or scan implementation +can replace only boundary-summary propagation without changing planning, +local KDA, autograd, or the framework adapter. ## Ownership boundary @@ -32,7 +35,7 @@ cuLA Kimi CP runtime | v cuDNN Frontend - KDA summaries, concatenated-summary composition, local KDA fwd/bwd + Blackwell summaries/composition or Hopper state ports, local KDA fwd/bwd ``` Framework-specific compatibility does not belong in the runtime. In @@ -83,7 +86,46 @@ release is published, use the tested source revision `9e056425989e75f406a5ce3ad58cb563e4dff8b5` as shown in `USAGE.md`; the runtime also checks the required call signatures and fails early on older builds. -## Current transport and intended evolution +### Hopper (SM90) compatibility backend + +cuDNN Hopper exposes ordinary FP32 `initial_state`, `final_state`, +`d_final_state`, and `d_initial_state` ports, but not the FROST summary graph. +`cudnn_kda_cp(backend="auto")` therefore selects a state-probe backend on +SM90 and the native summary backend on SM100 and newer. KDA is affine in its +initial state: a zero-state evaluation gives `H`, while an embedded-identity +evaluation gives `H + M`. cuLA subtracts the two FP32 final states, composes +the selected rank chain in FP32, and supplies the result through the ordinary +state ports. Backward obtains the local `G` from `d_initial_state`, uses +`M^T`, and composes following ranks in reverse causal order. +The external composition rejects an enabled PyTorch TF32 matmul policy; callers +must retain `torch.backends.cuda.matmul.allow_tf32 = False` so the advertised +FP32 boundary contract cannot silently change with a process-global setting. + +This is a correctness-first compatibility backend and performs two additional +boundary-state evaluations. It is not presented as the final SM90 performance +path. A dedicated pre-scan may replace only the summary provider after it has +equivalent compile-cache, CUDA Graph, accuracy, and fwd+bwd coverage. + +The Hopper backward engine requires canonical inputs: normalized q/k, +post-sigmoid beta, and log-space g. Framework adapters must apply q/k +normalization, beta sigmoid, and safe-gate conversion before calling cuLA so +PyTorch owns those outer chain-rule gradients. The SM90 backend rejects the +corresponding in-kernel flags instead of silently producing incomplete gate +gradients. + +The ordinary cuDNN Hopper backward state contract requires one FP32 state row +per local sequence; its `state_indices` indirection is forward-only. The SM90 +compatibility backend therefore materializes `initial_state`, `d_final_state`, +and the returned state gradient with shape `[N, H, V, K]` during backward. Each +live dense state tensor costs `4 * N * H * V * K` bytes (2 MiB per sequence for +`H=32`, `K=V=128`). Only the single-row incoming state is retained across the +forward/backward interval; the zero-padded initial state is rebuilt when +backward begins. Packed workloads with many local short samples must still +include the transient linear backward state storage in their memory plan. A +future compact backward state port in cuDNN could remove it without changing +cuLA's public distributed interface. + +## Current transport and rejected chain experiment The first implementation uses two typed `all_gather_into_tensor` calls per direction. Keeping H and M separate preserves the dense layout required by @@ -93,6 +135,28 @@ should decide whether a topology-aware exact-order transport is worth adding. A tree or scan must not silently reassociate FP32 products and change the numerical contract. +## 1M logical-CP64 KDA E2E emulation + +The four-GPU L20C environment cannot measure real 64-GPU scaling. The E2E +emulator therefore maps representative logical ranks `0/21/42/63` onto four +physical processes, uses 16,384 local tokens per logical rank, and executes 63 +real NCCL ring hops with the 4 MiB FP32 `(H,M)` payload in each direction. The +local summaries, cuDNN concatenated-summary KDA apply, backward summaries, and +local backward are real cuDNN operations. No DSA stage is included. + +The default shape is 1,048,576 global tokens, logical CP64, `H=32`, and +`K=V=128`. The harness reports total CUDA/wall time, transport and preparation +time, peak memory, proxy loss, output norm, gradient norms, and finiteness after +one warmup. These measurements model the serial CP64 dependency length on the +local four-GPU fabric; they do not model 64-GPU topology, cross-node contention, +or aggregate bisection bandwidth. The reproducible harness is +`benchmarks/large_cp_emulation/bench_kda.py`; its logical topology and +benchmark-only transport remain outside the production CP runtime. + +Results from the retired stand-alone merge API are intentionally not presented +as current numbers. Re-run the harness with the cuDNN Frontend version under +evaluation before using it for a performance decision. + ## Prepared large-CP transport The convenience path allocates its gathered H and M tensors on each forward and diff --git a/docs/THD_LAYOUT_PLANNER.md b/docs/THD_LAYOUT_PLANNER.md new file mode 100644 index 00000000..bb46035c --- /dev/null +++ b/docs/THD_LAYOUT_PLANNER.md @@ -0,0 +1,117 @@ +# Cost-driven THD context-parallel layout planner + +## Why the alternating prototype is not the production policy + +Alternating the shard direction of adjacent long samples exactly balances the +original `20 x 1K + 2 x 514048` experiment, but it encodes properties of that +batch into the policy. It is weak for a single dominant sequence, an odd +number of long sequences, and power-law mixtures. + +`plan_thd_layout` instead consumes arbitrary CPU `global_cu_seqlens` and +balances two independent resources: + +- token-linear KDA and index-pooling work; +- position-dependent causal candidate work in CSA/DSA. + +Short samples remain whole when exact token capacity permits it. Long samples +are considered on several planning grids from `1 x CP` through `6 x CP`. Equal-token pieces are assigned by a +deterministic multi-resource largest-processing-time scheduler, then adjacent +pieces with the same owner are coalesced. This naturally produces an +early/late zigzag for one long sequence and cross-sample compensation when +several long sequences are available. + +After compute scheduling, a capacity-repair pass moves whole spans where +possible and otherwise splits the cheapest span boundary. Every rank receives +its exact balanced block capacity (at most one block difference when the global +block count is not divisible by CP), so the plan remains compatible with +fixed-shape KDA and pooled all-gather collectives. + +## Objective and fallback guarantee + +For rank `r`, let `T_r` be its tokens and `C_r` its exact analytic causal +candidate count. The default robust objective is + +```text +max(max(T_r) / mean(T_r), max(C_r) / mean(C_r)) + + 0.02 * max(cross_rank_boundary_count_r) +``` + +The boundary term prevents fine planning grids from winning on an immaterial +compute difference while introducing extra KDA state chains. The weights and +maximum grid are explicit configuration rather than hidden thresholds. + +The planner always evaluates the ordinary contiguous CP partition with the +same objective and returns the lowest-scoring candidate. Therefore its result +is never worse than contiguous under the declared model. This is a model-level +guarantee; GPU profiling is still needed to calibrate how candidate work, +token work, and boundary communication overlap on a target system. + +## Representative CP4 distributions + +The CPU benchmark is `benchmarks/large_cp_emulation/bench_thd_planner.py`. +The following table uses 1M total tokens and the default cost model. + +| THD distribution | selected grid | token max/avg | causal max/avg | selected score | contiguous score | remote boundaries | +|---|---:|---:|---:|---:|---:|---:| +| 1024 equal short samples | contiguous | 1.000 | 1.000 | 1.000 | 1.000 | 0 | +| one 1M sequence | 2 x CP | 1.000 | 1.000 | 1.080 | 1.790 | 6 | +| 20 x 1K + 2 x 514048 | 2 x CP | 1.000 | 1.000 | 1.040 | 1.560 | 4 | +| 256 x 1K + 1 x 768K | 3 x CP | 1.000 | 1.037 | 1.157 | 2.261 | 10 | +| 32 x 4K + 128K + 256K + 512K | 4 x CP | 1.000 | 1.046 | 1.166 | 2.302 | 10 | +| power-law 1K through 511K | 4 x CP | 1.000 | 1.028 | 1.188 | 2.149 | 15 | + +The important result is not that one grid wins everywhere. Low-fragmentation +`2 x CP` is selected for the original and single-long cases, while irregular +mixtures spend additional boundaries only when they materially improve both +token and causal balance. + +At logical CP64, the planning grid scales below the fixed 16K short-sample +threshold for long sequences. A one-million-token sequence therefore still +selects `2 x CP`, with exact token and causal balance. The original two-long +mixture selects `5 x CP` and reduces modeled causal max/average from 2.007 to +1.072; the power-law mixture also selects `5 x CP` and reduces it from 2.770 +to 1.034. The threshold is used only to preserve genuinely short samples, not +to prevent a dominant long sequence from forming sub-rank zigzag pieces. + +## Four-GPU DSA validation + +The layout-native DSA benchmark executes real deterministic cuDNN index +pooling and DSA selection, including pooled-key all-gather and global block-ID +remap. Each number below is the median critical-path GPU time after two warmups +and three measured iterations. All cases contain 1,048,576 global tokens. +These published measurements used four index heads; current benchmark output +records `index_heads` explicitly because the default model shape is now 64. + +| THD distribution | contiguous | planned | speedup | planned causal max/avg | +|---|---:|---:|---:|---:| +| one 1M sequence | 418.51 ms | 239.16 ms | 42.9% | 1.000 | +| 20 x 1K + 2 x 514048 | 166.77 ms | 109.05 ms | 34.6% | 1.000 | +| 256 x 1K + 1 x 768K | 291.80 ms | 134.35 ms | 54.0% | 1.037 | + +For the original mixture, planned per-rank selection time is approximately +108 ms and remap plus collective time is 0.55--0.96 ms. Under contiguous CP, +the light ranks spend 107--120 ms waiting at the collective while the heavy +ranks finish selection. For the single-long sequence, contiguous candidate +work grows from 8.59 to 60.13 billion across ranks; the planned layout gives +every rank exactly 34.36 billion candidates. This confirms that the speedup is +removal of causal load bubbles rather than a precision or kernel change. + +## Plan contract and integration boundary + +`THDLayoutPlan` records the global spans, owner rank, local `cu_seqlens`, +causal offsets, per-resource loads, and all candidate scores. Planning is +CPU-only and deterministic. It must run once at microbatch ingress, before the +model layer stack, rather than permuting tokens inside every KDA or CSA layer. + +The current KDA runtime still assumes one contiguous global token shard. A +separate lowering step is required before `THDLayoutPlan` becomes its default: + +1. derive every sequence's ordered `(sequence, span ordinal)` chain; +2. merge forward and backward KDA state in sequence order, not physical-rank + order; +3. produce CSA global block IDs and key-remap metadata from the same spans; +4. apply the token permutation once and retain it across all model layers. + +The benchmark adapter already lets DSA consume the planned layout. Keeping the +planner independent from KDA transport is intentional: layout policy and the +collective implementation can be validated and evolved separately. diff --git a/pyproject.toml b/pyproject.toml index 1d289dae..97f2b161 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,6 +18,10 @@ dependencies = [ license = "Apache-2.0" [project.optional-dependencies] +cudnn-hybrid = [ + "nvidia-cutlass-dsl[cu13]>=4.8.0,<4.9", + "nvidia-cudnn-frontend>=1.33.0", +] cudnn-cp = [ "nvidia-cutlass-dsl[cu13]>=4.8.0,<4.9", ] diff --git a/tests/test_hybrid_distributed_pool.py b/tests/test_hybrid_distributed_pool.py new file mode 100644 index 00000000..289707da --- /dev/null +++ b/tests/test_hybrid_distributed_pool.py @@ -0,0 +1,177 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +from cula.hybrid import distributed_pool +from cula.hybrid.plan import IndexPoolShardPlan + + +def _fake_collectives(monkeypatch, remote_rows, *, rank=0): + calls = [] + + def all_gather(output, local, *, group): + calls.append(("gather", tuple(local.shape), group)) + gathered = output.view(2, *local.shape) + gathered[rank].copy_(local) + gathered[1 - rank].copy_(remote_rows) + + def reduce_scatter(output, tensor, *, group): + calls.append(("reduce_scatter", tuple(tensor.shape), group)) + chunks = tensor.view(2, *output.shape) + output.copy_(chunks[rank]) + + monkeypatch.setattr(distributed_pool.dist, "all_gather_into_tensor", all_gather) + monkeypatch.setattr(distributed_pool.dist, "reduce_scatter_tensor", reduce_scatter) + monkeypatch.delattr(distributed_pool.dist, "all_gather_single", raising=False) + monkeypatch.delattr(distributed_pool.dist, "reduce_scatter_single", raising=False) + return calls + + +def test_unequal_block_gather_pads_only_compressed_axis(monkeypatch): + local = torch.tensor([[1.0], [2.0], [3.0]]) + remote = torch.tensor([[10.0], [20.0], [0.0]]) + calls = _fake_collectives(monkeypatch, remote) + + gathered = distributed_pool._all_gather_block_rows( + local, + group="cp", + block_offsets=(0, 3, 5), + rank=0, + ) + + assert torch.equal(gathered, torch.tensor([[1.0], [2.0], [3.0], [10.0], [20.0]])) + assert calls == [("gather", (3, 1), "cp")] + + +def test_cp_pool_rejects_misaligned_seam_before_launch(monkeypatch): + plan = IndexPoolShardPlan.from_global_cu_seqlens( + [0, 16], + rank=0, + world_size=2, + head_dim=4, + shard_offsets=(0, 6, 16), + ) + monkeypatch.setattr(plan.__class__, "validate_group", lambda self, group: None) + monkeypatch.setattr( + distributed_pool, + "_load_cudnn_index_pool_ops", + lambda: pytest.fail("pooling must not launch"), + ) + + with pytest.raises(NotImplementedError, match="crossing a rank seam"): + distributed_pool.cudnn_index_pool_cp( + torch.empty(6, 4), + torch.empty(6, 4), + torch.empty(4, 4), + torch.tensor([0, 6], dtype=torch.int32), + torch.tensor([0, 1], dtype=torch.int32), + plan=plan, + group="cp", + ) + + +def test_cp_pool_runs_local_pool_before_block_gather(monkeypatch): + plan = IndexPoolShardPlan.from_global_cu_seqlens( + [0, 16], + rank=0, + world_size=2, + head_dim=4, + shard_offsets=(0, 8, 16), + ) + monkeypatch.setattr(plan.__class__, "validate_group", lambda self, group: None) + observed = [] + + def forward(kv, score, ape, cu, cb, **kwargs): + observed.append((kv.shape[0], kwargs["total_comp"], kwargs["coff"])) + return {"out": torch.arange(8, dtype=kv.dtype).reshape(2, 4)} + + monkeypatch.setattr( + distributed_pool, + "_load_cudnn_index_pool_ops", + lambda: (forward, pytest.fail), + ) + monkeypatch.setattr(distributed_pool, "_validate_inputs", lambda *args, **kwargs: None) + _fake_collectives( + monkeypatch, + torch.full((2, 4), 9.0), + ) + result = distributed_pool.cudnn_index_pool_cp( + torch.empty(8, 4), + torch.empty(8, 4), + torch.empty(4, 4), + torch.tensor([0, 8], dtype=torch.int32), + torch.tensor([0, 2], dtype=torch.int32), + plan=plan, + group="cp", + deterministic=True, + ) + + assert observed == [(8, 2, 1)] + assert torch.equal( + result, + torch.cat((torch.arange(8).reshape(2, 4), torch.full((2, 4), 9.0))), + ) + + +def test_cp_pool_backward_preserves_loss_rank_gradients(monkeypatch): + plan = IndexPoolShardPlan.from_global_cu_seqlens( + [0, 16], + rank=0, + world_size=2, + head_dim=4, + shard_offsets=(0, 8, 16), + ) + monkeypatch.setattr(plan.__class__, "validate_group", lambda self, group: None) + monkeypatch.setattr(distributed_pool, "_validate_inputs", lambda *args, **kwargs: None) + + backward_inputs = [] + + def forward(kv, score, ape, cu, cb, **kwargs): + return {"out": torch.zeros(2, 4, dtype=kv.dtype)} + + def backward(kv, score, ape, cu, cb, grad_out, **kwargs): + backward_inputs.append(grad_out.clone()) + value = grad_out.float().sum() + return { + "grad_kv": torch.full_like(kv, value), + "grad_score": torch.full_like(score, 2 * value), + "grad_ape": torch.full_like(ape, value), + } + + monkeypatch.setattr( + distributed_pool, + "_load_cudnn_index_pool_ops", + lambda: (forward, backward), + ) + _fake_collectives(monkeypatch, torch.full((2, 4), 9.0)) + + def all_to_all(output, tensor, **kwargs): + # Owner rank 0 receives its block gradients from loss ranks 0 and 1. + output.view(2, 2, 4)[0].copy_(tensor[:2]) + output.view(2, 2, 4)[1].fill_(3.0) + + monkeypatch.setattr(distributed_pool.dist, "all_to_all_single", all_to_all) + + kv = torch.ones(8, 4, requires_grad=True) + score = torch.ones(8, 4, requires_grad=True) + ape = torch.ones(4, 4, requires_grad=True) + output = distributed_pool.cudnn_index_pool_cp( + kv, + score, + ape, + torch.tensor([0, 8], dtype=torch.int32), + torch.tensor([0, 2], dtype=torch.int32), + plan=plan, + group="cp", + ) + output.backward(torch.cat((torch.ones(2, 4), torch.full((2, 4), 5.0)))) + + assert len(backward_inputs) == 2 + assert torch.equal(backward_inputs[0], torch.ones(2, 4)) + assert torch.equal(backward_inputs[1], torch.full((2, 4), 3.0)) + assert torch.equal(kv.grad, torch.full_like(kv, 32.0)) + assert torch.equal(score.grad, torch.full_like(score, 64.0)) + # The fake reduce-scatter selects this loss rank's owner-local APE chunk. + assert torch.equal(ape.grad, torch.full_like(ape, 8.0)) diff --git a/tests/test_hybrid_dsa_cudnn.py b/tests/test_hybrid_dsa_cudnn.py new file mode 100644 index 00000000..08b5f099 --- /dev/null +++ b/tests/test_hybrid_dsa_cudnn.py @@ -0,0 +1,54 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +from cula.hybrid import DsaSelectionPlan, DsaSelectionWorkspace, cudnn_dsa_select + + +@pytest.mark.sm100_only +@pytest.mark.cula_slow +def test_windowed_dsa_selection_is_cuda_graph_set_deterministic(): + pytest.importorskip("cudnn") + torch.manual_seed(43) + device = torch.device("cuda") + tokens = 4096 + blocks = tokens // 4 + top_k = 64 + plan = DsaSelectionPlan.from_cu_seqlens( + [0, tokens], + [0, blocks], + top_k=top_k, + candidate_budget_bytes=6 << 20, + ) + assert len(plan.windows) > 1 + workspace = DsaSelectionWorkspace.from_plan(plan, device) + qi = torch.randn(tokens, 4, 128, dtype=torch.bfloat16, device=device) + kbar = torch.randn(blocks, 1, 128, dtype=torch.bfloat16, device=device) + weights = torch.randn(tokens, 4, dtype=torch.float32, device=device) + + def run(): + return cudnn_dsa_select(qi, kbar, weights, plan=plan, workspace=workspace, deterministic=True) + + run() + torch.cuda.synchronize() + reference_ids = workspace.block_ids.clone() + reference_scores = workspace.scores.clone() + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + workspace.block_ids.fill_(-2) + workspace.scores.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + + reference_sorted, reference_order = torch.sort(reference_ids, dim=-1) + replay_sorted, replay_order = torch.sort(workspace.block_ids, dim=-1) + reference_logits = torch.gather(reference_scores, -1, reference_order) + replay_logits = torch.gather(workspace.scores, -1, replay_order) + valid = reference_sorted >= 0 + assert torch.equal(reference_sorted, replay_sorted) + assert torch.equal(reference_logits[valid], replay_logits[valid]) diff --git a/tests/test_hybrid_dsa_helpers.py b/tests/test_hybrid_dsa_helpers.py new file mode 100644 index 00000000..e3dce898 --- /dev/null +++ b/tests/test_hybrid_dsa_helpers.py @@ -0,0 +1,72 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace + +import torch + +from cula.hybrid import dsa + + +def test_cudnn_dsa_loader_prefers_public_namespace(monkeypatch): + select = object() + fake_cudnn = SimpleNamespace(DSA=SimpleNamespace(indexer_forward_top_k_wrapper=select)) + monkeypatch.setattr(dsa, "import_module", lambda name: fake_cudnn) + dsa._load_cudnn_dsa_select.cache_clear() + try: + assert dsa._load_cudnn_dsa_select() is select + finally: + dsa._load_cudnn_dsa_select.cache_clear() + + +def test_cudnn_dsa_selection_defaults_to_nondeterministic(monkeypatch): + observed = [] + + def select(*args, **kwargs): + del args + observed.append(kwargs["deterministic"]) + + window = dsa.DsaSelectionWindow( + query_start=0, + query_end=1, + key_start=0, + key_end=1, + cu_seqlens_q=(0, 1), + cu_seqlens_k=(0, 1), + q_causal_offsets=(0,), + candidate_offsets=(0, 1), + ) + plan = dsa.DsaSelectionPlan( + block_size=4, + top_k=1, + cu_seqlens_q=(0, 1), + cu_seqlens_k=(0, 1), + q_causal_offsets=(0,), + candidate_budget_bytes=4, + windows=(window,), + ) + metadata = SimpleNamespace( + cu_seqlens_q=torch.tensor([0, 1], dtype=torch.int32), + cu_seqlens_k=torch.tensor([0, 1], dtype=torch.int32), + q_causal_offsets=None, + candidate_offsets=torch.tensor([0, 1], dtype=torch.int64), + ) + workspace = dsa.DsaSelectionWorkspace( + plan=plan, + candidates=torch.empty(1), + block_ids=torch.empty((1, 1), dtype=torch.int32), + scores=torch.empty((1, 1)), + windows=(metadata,), + ) + monkeypatch.setattr(dsa, "_validate_selection_inputs", lambda *args, **kwargs: None) + monkeypatch.setattr(dsa, "_load_cudnn_dsa_select", lambda: select) + + dsa.cudnn_dsa_select( + torch.empty((1, 1, 1)), + torch.empty((1, 1, 1)), + torch.empty((1, 1)), + plan=plan, + workspace=workspace, + ) + + assert observed == [False] diff --git a/tests/test_hybrid_dsa_plan.py b/tests/test_hybrid_dsa_plan.py new file mode 100644 index 00000000..f4c1c1c6 --- /dev/null +++ b/tests/test_hybrid_dsa_plan.py @@ -0,0 +1,187 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +from cula.hybrid import ( + MAX_BALANCED_CANDIDATE_BUDGET_BYTES, + MEMORY_EFFICIENT_CANDIDATE_BUDGET_BYTES, + DsaSelectionPlan, +) + + +def test_selection_plan_coalesces_short_packed_sequences(): + plan = DsaSelectionPlan.from_cu_seqlens( + [0, 16, 36], + [0, 4, 9], + top_k=3, + candidate_budget_bytes=1 << 20, + ) + + assert plan.total_queries == 36 + assert plan.total_blocks == 9 + assert plan.num_sequences == 2 + assert len(plan.windows) == 1 + window = plan.windows[0] + assert window.cu_seqlens_q == (0, 16, 36) + assert window.cu_seqlens_k == (0, 4, 9) + assert window.q_causal_offsets == (0, 0) + assert window.candidate_offsets == (0, 28, 73) + + +def test_selection_plan_splits_one_sequence_to_bound_peak_scratch(): + plan = DsaSelectionPlan.from_cu_seqlens( + [0, 32], + [0, 8], + top_k=4, + candidate_budget_bytes=80 * 4, + ) + + assert len(plan.windows) == 2 + assert max(window.candidate_floats for window in plan.windows) <= 80 + assert [window.query_start for window in plan.windows] == [0, plan.windows[0].query_end] + assert plan.windows[-1].query_end == 32 + assert all(window.key_start == 0 and window.key_end == 8 for window in plan.windows) + assert plan.windows[1].q_causal_offsets == (plan.windows[1].query_start,) + + +def test_selection_plan_keeps_chunked_prefill_position_offset(): + plan = DsaSelectionPlan.from_cu_seqlens( + [0, 8], + [0, 6], + top_k=4, + q_causal_offsets=[16], + candidate_budget_bytes=1 << 20, + ) + + assert plan.windows[0].q_causal_offsets == (16,) + assert plan.windows[0].candidate_floats == 38 + + +def test_selection_plan_accepts_cpu_offset_tensor(): + plan = DsaSelectionPlan.from_cu_seqlens( + [0, 8], + [0, 6], + top_k=4, + q_causal_offsets=torch.tensor([16], dtype=torch.int32), + ) + + assert plan.q_causal_offsets == (16,) + + +def test_selection_plan_permits_a_sequence_without_a_complete_block(): + plan = DsaSelectionPlan.from_cu_seqlens( + [0, 3, 11], + [0, 0, 2], + top_k=2, + candidate_budget_bytes=1 << 20, + ) + + assert plan.windows[0].cu_seqlens_k == (0, 0, 2) + assert plan.windows[0].candidate_offsets == (0, 0, 6) + + +def test_selection_plan_covers_heterogeneous_128k_without_exceeding_budget(): + lengths = [1024] * 20 + [55296] * 2 + cu_q = [0] + cu_k = [0] + for length in lengths: + cu_q.append(cu_q[-1] + length) + cu_k.append(cu_k[-1] + length // 4) + + budget = 512 << 20 + plan = DsaSelectionPlan.from_cu_seqlens( + cu_q, + cu_k, + top_k=512, + candidate_budget_bytes=budget, + ) + + assert plan.total_queries == 131072 + assert plan.total_blocks == 32768 + assert len(plan.windows) == 7 + assert plan.candidate_workspace_floats * 4 <= budget + assert plan.windows[0].cu_seqlens_q == tuple(i * 1024 for i in range(21)) + assert plan.windows[0].candidate_floats * 4 < 10 << 20 + + +def test_adaptive_budget_keeps_128k_workload_memory_efficient(): + plan = DsaSelectionPlan.from_cu_seqlens( + [0, 1 << 17], + [0, 1 << 15], + top_k=512, + ) + + assert plan.candidate_budget_bytes == MEMORY_EFFICIENT_CANDIDATE_BUDGET_BYTES + assert plan.candidate_workspace_floats * 4 <= MEMORY_EFFICIENT_CANDIDATE_BUDGET_BYTES + + +def test_adaptive_budget_keeps_measured_4g_choice_at_1m_cp64_tail(): + local_tokens = (1 << 20) // 64 + query_offset = (1 << 20) - local_tokens + plan = DsaSelectionPlan.from_cu_seqlens( + [0, local_tokens], + [0, (1 << 20) // 4], + top_k=512, + q_causal_offsets=[query_offset], + ) + + assert plan.candidate_budget_bytes == 4 << 30 + assert len(plan.windows) == 4 + assert plan.candidate_workspace_floats * 4 <= 4 << 30 + + +def test_adaptive_budget_uses_throughput_multiplier_for_1m_cp4_tail(): + local_tokens = (1 << 20) // 4 + query_offset = (1 << 20) - local_tokens + plan = DsaSelectionPlan.from_cu_seqlens( + [0, local_tokens], + [0, (1 << 20) // 4], + top_k=512, + q_causal_offsets=[query_offset], + ) + + assert plan.candidate_budget_bytes == MAX_BALANCED_CANDIDATE_BUDGET_BYTES + assert plan.candidate_workspace_floats * 4 <= MAX_BALANCED_CANDIDATE_BUDGET_BYTES + + +def test_adaptive_budget_uses_8g_for_1m_thd_cp4_tail(): + plan = DsaSelectionPlan.from_cu_seqlens( + [0, 1 << 18], + [0, 514048 // 4], + top_k=512, + q_causal_offsets=[514048 - (1 << 18)], + ) + + assert plan.candidate_budget_bytes == MAX_BALANCED_CANDIDATE_BUDGET_BYTES + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"top_k": 0}, "top_k"), + ({"top_k": 3, "block_size": 0}, "block_size"), + ({"top_k": 3, "candidate_budget_bytes": 3}, "candidate_budget_bytes"), + ({"top_k": 3, "q_causal_offsets": [-1]}, "non-negative"), + ], +) +def test_invalid_selection_plans_are_rejected(kwargs, message): + with pytest.raises((TypeError, ValueError), match=message): + DsaSelectionPlan.from_cu_seqlens([0, 8], [0, 2], **kwargs) + + +def test_selection_plan_rejects_missing_visible_blocks(): + with pytest.raises(ValueError, match="requires at least 4"): + DsaSelectionPlan.from_cu_seqlens([0, 16], [0, 3], top_k=3) + + +def test_selection_plan_rejects_budget_smaller_than_one_causal_row(): + with pytest.raises(ValueError, match="too small for one row"): + DsaSelectionPlan.from_cu_seqlens( + [0, 16], + [0, 29], + top_k=3, + q_causal_offsets=[100], + candidate_budget_bytes=4, + ) diff --git a/tests/test_hybrid_index_pool_helpers.py b/tests/test_hybrid_index_pool_helpers.py new file mode 100644 index 00000000..f174bcc0 --- /dev/null +++ b/tests/test_hybrid_index_pool_helpers.py @@ -0,0 +1,23 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace + +from cula.hybrid import index_pool + + +def test_cudnn_index_pool_loader_uses_csa_namespace(monkeypatch): + forward = object() + backward = object() + fake_cudnn = SimpleNamespace( + CSA=SimpleNamespace( + csa_compressor_forward_wrapper=forward, + csa_compressor_backward_wrapper=backward, + ) + ) + monkeypatch.setattr(index_pool, "import_module", lambda name: fake_cudnn) + index_pool._load_cudnn_index_pool_ops.cache_clear() + try: + assert index_pool._load_cudnn_index_pool_ops() == (forward, backward) + finally: + index_pool._load_cudnn_index_pool_ops.cache_clear() diff --git a/tests/test_hybrid_index_pool_plan.py b/tests/test_hybrid_index_pool_plan.py new file mode 100644 index 00000000..67aa0ce5 --- /dev/null +++ b/tests/test_hybrid_index_pool_plan.py @@ -0,0 +1,163 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +import random + +import pytest +import torch + +from cula.hybrid import IndexPoolPlan, IndexPoolShardPlan + + +def test_index_pool_plan_compacts_each_packed_sequence_independently(): + plan = IndexPoolPlan.from_cu_seqlens( + torch.tensor([0, 5, 16, 24, 27], dtype=torch.int32), + head_dim=128, + ratio=4, + coff=1, + ) + + assert plan.cu_seqlens == (0, 5, 16, 24, 27) + assert plan.cu_seqlens_comp == (0, 1, 3, 5, 5) + assert plan.total_tokens == 27 + assert plan.total_comp == 5 + assert plan.num_sequences == 4 + assert plan.input_width == 128 + + +def test_overlapping_plan_changes_projection_width_not_block_count(): + plan = IndexPoolPlan.from_cu_seqlens([0, 8, 17], head_dim=512, ratio=4, coff=2) + + assert plan.input_width == 1024 + assert plan.cu_seqlens_comp == (0, 2, 4) + + +def test_plan_materializes_reusable_int32_metadata(): + plan = IndexPoolPlan.from_cu_seqlens([0, 5, 16], head_dim=128) + cu_seqlens, cu_seqlens_comp = plan.materialize("cpu") + + assert cu_seqlens.dtype == torch.int32 + assert cu_seqlens_comp.dtype == torch.int32 + assert torch.equal(cu_seqlens, torch.tensor([0, 5, 16], dtype=torch.int32)) + assert torch.equal(cu_seqlens_comp, torch.tensor([0, 1, 3], dtype=torch.int32)) + + +def test_aligned_cp_shards_pool_locally_and_gather_only_blocks(): + # Rank boundary 13 is sequence 1's block-2 boundary: sequence 1 starts at 5. + plans = [ + IndexPoolShardPlan.from_global_cu_seqlens( + [0, 5, 16, 24, 27], + rank=rank, + world_size=2, + head_dim=128, + shard_offsets=(0, 13, 27), + ) + for rank in range(2) + ] + + assert all(plan.all_shards_direct for plan in plans) + assert all(plan.direct_local_pool for plan in plans) + assert plans[0].block_offsets == (0, 3, 5) + assert plans[0].local_plan is not None + assert plans[0].local_plan.cu_seqlens == (0, 5, 13) + assert plans[0].local_plan.cu_seqlens_comp == (0, 1, 3) + assert plans[1].local_plan is not None + assert plans[1].local_plan.cu_seqlens == (0, 3, 11, 14) + assert plans[1].local_plan.cu_seqlens_comp == (0, 0, 2, 2) + assert plans[0].global_blocks == 5 + assert plans[0].block_counts == (3, 2) + assert not plans[0].equal_block_counts + + +def test_misaligned_cp_shard_derives_bounded_pooling_halo(): + plans = [ + IndexPoolShardPlan.from_global_cu_seqlens( + [0, 5, 16, 24, 27], + rank=rank, + world_size=2, + head_dim=128, + shard_offsets=(0, 12, 27), + ) + for rank in range(2) + ] + + assert not plans[0].all_shards_direct + assert plans[0].recv_right_tokens == 1 + assert plans[1].send_left_tokens == 1 + assert plans[0].block_offsets == (0, 3, 5) + assert plans[0].local_plan is None + assert plans[1].local_plan is None + assert all(plan.halo_supported for plan in plans) + + +def test_tiny_shards_declare_single_neighbor_halo_unsupported(): + plan = IndexPoolShardPlan.from_global_cu_seqlens( + [0, 8], + rank=0, + world_size=4, + head_dim=128, + shard_offsets=(0, 2, 4, 6, 8), + ) + + assert not plan.halo_supported + + +def test_cp_seam_inside_dropped_sequence_tail_needs_no_halo(): + plans = [ + IndexPoolShardPlan.from_global_cu_seqlens( + [0, 6, 14], + rank=rank, + world_size=2, + head_dim=128, + shard_offsets=(0, 5, 14), + ) + for rank in range(2) + ] + + assert plans[0].recv_right_tokens == 0 + assert plans[1].send_left_tokens == 0 + assert all(plan.all_shards_direct for plan in plans) + assert plans[0].local_plan is not None + assert plans[1].local_plan is not None + + +def test_random_shards_assign_every_complete_block_exactly_once(): + generator = random.Random(47) + for _ in range(200): + lengths = [generator.randint(1, 64) for _ in range(generator.randint(1, 12))] + cu_seqlens = [0] + for length in lengths: + cu_seqlens.append(cu_seqlens[-1] + length) + world_size = generator.randint(1, min(8, cu_seqlens[-1])) + cuts = sorted(generator.sample(range(1, cu_seqlens[-1]), world_size - 1)) + shard_offsets = (0, *cuts, cu_seqlens[-1]) + plans = [ + IndexPoolShardPlan.from_global_cu_seqlens( + cu_seqlens, + rank=rank, + world_size=world_size, + head_dim=128, + shard_offsets=shard_offsets, + ) + for rank in range(world_size) + ] + + expected_blocks = sum(length // 4 for length in lengths) + assert plans[0].block_offsets[-1] == expected_blocks + assert sum(plan.local_blocks for plan in plans) == expected_blocks + for rank in range(world_size - 1): + assert plans[rank].recv_right_tokens == plans[rank + 1].send_left_tokens + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"head_dim": 0}, "head_dim"), + ({"head_dim": 128, "ratio": 0}, "ratio"), + ({"head_dim": 128, "coff": 3}, "coff"), + ({"head_dim": 128, "ratio": 32}, "no complete"), + ], +) +def test_invalid_index_pool_plans_are_rejected(kwargs, message): + with pytest.raises(ValueError, match=message): + IndexPoolPlan.from_cu_seqlens([0, 8], **kwargs) diff --git a/tests/test_kimi_cp_distributed.py b/tests/test_kimi_cp_distributed.py index 670e003f..72d8102e 100644 --- a/tests/test_kimi_cp_distributed.py +++ b/tests/test_kimi_cp_distributed.py @@ -17,7 +17,6 @@ def _relative_rms(actual, expected): @pytest.mark.distributed -@pytest.mark.sm100_only def test_cudnn_kda_cp_matches_monolithic_forward_backward(): if int(os.environ.get("WORLD_SIZE", "1")) < 2: pytest.skip("launch this deterministic fixture with at least two processes") @@ -27,13 +26,19 @@ def test_cudnn_kda_cp_matches_monolithic_forward_backward(): try: cudnn_la = pytest.importorskip("cudnn.linear_attention") try: - supports_concatenated_summaries = ( - "preceding_state_update" in inspect.signature(cudnn_la.kimi_delta_attention).parameters - ) + supports_initial_state = "initial_state" in inspect.signature(cudnn_la.kimi_delta_attention).parameters except (TypeError, ValueError): - supports_concatenated_summaries = False - if not supports_concatenated_summaries: - pytest.skip("cuDNN Frontend build does not include concatenated KDA summary apply") + supports_initial_state = False + if not supports_initial_state: + pytest.skip("cuDNN Frontend build does not include KDA state inputs") + capability_major = torch.cuda.get_device_capability(local_rank)[0] + if capability_major < 9: + pytest.skip("cuDNN Kimi CP requires SM90 or newer") + if ( + capability_major >= 10 + and "preceding_state_update" not in inspect.signature(cudnn_la.kimi_delta_attention).parameters + ): + pytest.skip("the Blackwell build does not include concatenated KDA summary apply") def stage(name): print(f"rank={local_rank} stage={name}", flush=True) @@ -77,7 +82,7 @@ def randn(*shape, dtype=torch.bfloat16): global_cu, output_final_state=True, batch_invariant=True, - plan_name="kda_frost", + plan_name=None, ) stage("monolithic-forward-done") full_output.backward(grad_output) diff --git a/tests/test_kimi_cp_helpers.py b/tests/test_kimi_cp_helpers.py index d68735a5..932f48e4 100644 --- a/tests/test_kimi_cp_helpers.py +++ b/tests/test_kimi_cp_helpers.py @@ -50,7 +50,7 @@ def test_attention_kwargs_keep_checkpoints_out_of_summary_calls(): def test_cudnn_loader_accepts_current_linear_attention_exports(monkeypatch): - def attention(*args, preceding_state_update=None, preceding_transition=None, **kwargs): + def attention(*args, initial_state=None, preceding_state_update=None, preceding_transition=None, **kwargs): return args, preceding_state_update, preceding_transition, kwargs def summary(*args, **kwargs): @@ -76,16 +76,32 @@ def import_fake(name): monkeypatch.setattr(cp, "import_module", import_fake) monkeypatch.setattr(cp.torch.ops.cudnn, "kimi_delta_attention_bwd", backward, raising=False) + cp._load_cudnn_kda_ops.cache_clear() cp._load_cudnn_cp_ops.cache_clear() try: loaded = cp._load_cudnn_cp_ops() assert loaded == (attention, summary, summary_bwd, backward) finally: cp._load_cudnn_cp_ops.cache_clear() + cp._load_cudnn_kda_ops.cache_clear() + + +def test_cudnn_kda_loader_does_not_require_frost_summary_exports(monkeypatch): + def attention(*args, initial_state=None, **kwargs): + return args, initial_state, kwargs + + backward = object() + monkeypatch.setattr(cp, "import_module", lambda name: SimpleNamespace(kimi_delta_attention=attention)) + monkeypatch.setattr(cp.torch.ops.cudnn, "kimi_delta_attention_bwd", backward, raising=False) + cp._load_cudnn_kda_ops.cache_clear() + try: + assert cp._load_cudnn_kda_ops() == (attention, backward) + finally: + cp._load_cudnn_kda_ops.cache_clear() def test_cudnn_loader_rejects_pre_concatenated_summary_api(monkeypatch): - def old_attention(*args, **kwargs): + def old_attention(*args, initial_state=None, **kwargs): return args, kwargs modules = { @@ -103,6 +119,7 @@ def import_fake(name): raise ModuleNotFoundError(name) from error monkeypatch.setattr(cp, "import_module", import_fake) + cp._load_cudnn_kda_ops.cache_clear() cp._load_cudnn_cp_ops.cache_clear() try: try: @@ -113,6 +130,7 @@ def import_fake(name): raise AssertionError("the retired stand-alone merge API must not be accepted") finally: cp._load_cudnn_cp_ops.cache_clear() + cp._load_cudnn_kda_ops.cache_clear() def test_local_boundary_summary_contract_preserves_fp32(): @@ -166,6 +184,133 @@ def test_local_boundary_summary_contract_rejects_precision_downgrade(): raise AssertionError("BF16 boundary state must be rejected") +def test_affine_summary_composition_uses_causal_order_in_both_directions(): + updates = torch.tensor([2.0, 4.0], dtype=torch.float32).reshape(2, 1, 1, 1) + transitions = torch.tensor([5.0, 7.0], dtype=torch.float32).reshape(2, 1, 1, 1) + + forward = cp._compose_affine_summaries(updates, transitions, reverse=False) + reverse = cp._compose_affine_summaries(updates, transitions, reverse=True) + + assert forward.item() == 18.0 # (0 * 5 + 2) * 7 + 4 + assert reverse.item() == 22.0 # (0 * 7 + 4) * 5 + 2 + + +def test_ieee_composition_validation_precedes_cuda_collectives(): + original = cp.torch.backends.cuda.matmul.allow_tf32 + try: + cp.torch.backends.cuda.matmul.allow_tf32 = True + try: + cp._validate_ieee_fp32_composition(torch.device("cuda", 0)) + except RuntimeError as error: + assert "IEEE FP32" in str(error) + else: + raise AssertionError("TF32 boundary composition was accepted") + finally: + cp.torch.backends.cuda.matmul.allow_tf32 = original + + cp._validate_ieee_fp32_composition(torch.device("cpu")) + + +def test_hopper_state_probe_recovers_fp32_affine_map(monkeypatch): + heads, key_dim, value_dim = 2, 4, 6 + state_update = torch.arange(heads * value_dim * key_dim, dtype=torch.float32).reshape(1, heads, value_dim, key_dim) / 32 + transition = torch.eye(key_dim, dtype=torch.float32).expand(heads, key_dim, key_dim).clone() + transition[:, 0, 1] = 0.25 + observed_kwargs = {} + + def attention(q, k, v, g, beta, cu, *, initial_state, **kwargs): + observed_kwargs.update(kwargs) + final = torch.matmul(initial_state, transition) + state_update + return torch.zeros_like(v), final, torch.empty(0), torch.empty(0) + + monkeypatch.setattr(cp, "_load_cudnn_kda_ops", lambda: (attention, object())) + options = cp._KdaOptions( + scale=None, + use_qk_l2norm_in_kernel=False, + use_beta_sigmoid_in_kernel=False, + allow_neg_eigval=False, + safe_gate=False, + lower_bound=None, + gate_domain="log", + batch_invariant=True, + checkpoint_every_n_tokens=0, + plan_name=None, + summary_plan_name=None, + ) + q = torch.zeros(8, heads, key_dim) + v = torch.zeros(8, heads, value_dim) + beta = torch.zeros(8, heads) + cu = torch.tensor([0, 8], dtype=torch.int32) + + actual_h, actual_transition = cp._state_probe_summary(q, q, v, q, beta, cu, options=options) + + assert actual_h.dtype == actual_transition.dtype == torch.float32 + torch.testing.assert_close(actual_h, state_update[0]) + torch.testing.assert_close(actual_transition, transition) + assert "checkpoint_every_n_tokens" not in observed_kwargs + + +def test_hopper_state_probe_rejects_in_kernel_backward_transforms(): + base = dict( + scale=None, + use_qk_l2norm_in_kernel=False, + use_beta_sigmoid_in_kernel=False, + allow_neg_eigval=False, + safe_gate=False, + lower_bound=None, + gate_domain="log", + batch_invariant=True, + checkpoint_every_n_tokens=0, + plan_name=None, + summary_plan_name=None, + ) + for option in ("use_qk_l2norm_in_kernel", "use_beta_sigmoid_in_kernel", "safe_gate"): + values = dict(base) + values[option] = True + try: + cp._validate_hopper_state_probe_options(cp._KdaOptions(**values)) + except ValueError as error: + assert option in str(error) + else: + raise AssertionError(f"SM90 accepted unsupported Hopper option {option}") + + options = cp._KdaOptions(**base) + for name, kwargs in ( + ("a_log", {"a_log": torch.zeros(1)}), + ("dt_bias", {"dt_bias": torch.zeros(1, 1)}), + ): + try: + cp._validate_hopper_state_probe_options(options, **kwargs) + except ValueError as error: + assert name in str(error) + else: + raise AssertionError(f"SM90 accepted unsupported Hopper input {name}") + + +def test_kimi_cp_backend_resolution(monkeypatch): + device = torch.device("cuda", 0) + monkeypatch.setattr(cp.torch.cuda, "get_device_capability", lambda unused: (9, 0)) + assert cp._resolve_kimi_cp_backend("auto", device) == "hopper-state-probe" + monkeypatch.setattr(cp.torch.cuda, "get_device_capability", lambda unused: (10, 0)) + assert cp._resolve_kimi_cp_backend("auto", device) == "frost-summary" + monkeypatch.setattr(cp.torch.cuda, "get_device_capability", lambda unused: (8, 0)) + try: + cp._resolve_kimi_cp_backend("auto", device) + except RuntimeError as error: + assert "8.x" in str(error) + else: + raise AssertionError("unsupported architectures must not select a Kimi CP backend") + + for backend in ("frost-summary", "hopper-state-probe"): + assert cp._resolve_kimi_cp_backend(backend, device) == backend + try: + cp._resolve_kimi_cp_backend("unknown", device) + except ValueError as error: + assert "backend" in str(error) + else: + raise AssertionError("an unknown Kimi CP backend was accepted") + + def _transport_plan(*, rank=2, pre=2, post=1): return KimiCPPlan( rank=rank, @@ -462,7 +607,7 @@ def attention(*args, **kwargs): return torch.ones_like(v), torch.empty(0) monkeypatch.setattr(cp, "_validate_inputs", lambda *args, **kwargs: None) - monkeypatch.setattr(cp, "_load_cudnn_cp_ops", lambda: (attention, object(), object(), object())) + monkeypatch.setattr(cp, "_load_cudnn_kda_ops", lambda: (attention, object())) result = cp.cudnn_kda_cp(q, q, v, q, beta, cu, plan=plan) diff --git a/tests/test_large_cp_emulation.py b/tests/test_large_cp_emulation.py new file mode 100644 index 00000000..656d9b3c --- /dev/null +++ b/tests/test_large_cp_emulation.py @@ -0,0 +1,73 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +from benchmarks.large_cp_emulation import ( + LogicalCPEmulationPlan, + PackedTHDWorkload, + bench_hybrid, + sample_chain_hops, +) +from benchmarks.large_cp_emulation.bench_dsa_layout import _validate_block_aligned_workload + + +def test_sample_chain_hops_for_one_megatoken_thd_mix(): + workload = PackedTHDWorkload.short_long_mix(1 << 20) + + assert workload.sequence_lengths == (1024,) * 20 + (514048, 514048) + assert sample_chain_hops(workload.global_cu_seqlens, local_tokens=16384) == 31 + + +def test_sample_chain_hops_for_rank_local_samples(): + assert sample_chain_hops((0, 1024, 2048, 4096), local_tokens=4096) == 0 + + +def test_emulation_plan_separates_logical_and_physical_identity(): + workload = PackedTHDWorkload.short_long_mix(1 << 20) + + plan = LogicalCPEmulationPlan.build( + workload.global_cu_seqlens, + logical_cp=64, + physical_rank=2, + physical_world_size=4, + transport_mode="sample-chain", + ) + + assert plan.logical_rank == 42 + assert (plan.logical_plan.rank, plan.logical_plan.world_size) == (42, 64) + assert (plan.runtime_plan.rank, plan.runtime_plan.world_size) == (2, 4) + assert plan.runtime_plan.local_cu_seqlens == plan.logical_plan.local_cu_seqlens + assert plan.transport_hops == 31 + + +def test_dsa_layout_rejects_non_block_aligned_sequences(): + workload = PackedTHDWorkload((0, 1023, 2048)) + + with pytest.raises(ValueError, match="every packed sequence length"): + _validate_block_aligned_workload(workload, 4) + + +def test_logical_dsa_transport_preserves_local_blocks_before_ring_reuses_send(monkeypatch): + branch = object.__new__(bench_hybrid._DsaIndexerBranch) + branch.mode = "logical" + branch.physical_rank = 0 + branch.physical_world_size = 2 + branch.transport_hops = 2 + branch.global_block_start = 2 + branch.local_blocks = 2 + branch.global_kbar = torch.full((6, 3), -1.0) + branch.ring_receive = torch.empty((2, 3)) + local_kbar = torch.arange(6, dtype=torch.float32).reshape(2, 3) + expected = local_kbar.clone() + + def overwrite_send(source, receive, **kwargs): + del receive, kwargs + source.fill_(99) + + monkeypatch.setattr(bench_hybrid, "emulate_ring_hops", overwrite_send) + branch._communicate(local_kbar) + + assert torch.equal(branch.global_kbar[2:4], expected) + assert torch.equal(local_kbar, torch.full_like(local_kbar, 99)) diff --git a/tests/test_large_cp_query_layout.py b/tests/test_large_cp_query_layout.py new file mode 100644 index 00000000..622dca55 --- /dev/null +++ b/tests/test_large_cp_query_layout.py @@ -0,0 +1,66 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +from benchmarks.large_cp_emulation.query_layout import ( + build_alternating_sequence_query_layout, + build_contiguous_query_layout, + build_planned_query_layout, +) + + +def _workload_boundaries(): + lengths = [1024] * 20 + [514048] * 2 + result = [0] + for length in lengths: + result.append(result[-1] + length) + return result + + +def test_alternating_sequence_layout_keeps_short_samples_whole_and_tokens_balanced(): + boundaries = _workload_boundaries() + layouts = [build_alternating_sequence_query_layout(boundaries, rank=rank, world_size=4) for rank in range(4)] + + assert {layout.local_tokens for layout in layouts} == {1 << 18} + assert all(len(layout.spans) == 7 for layout in layouts) + assert all(sum(span.sequence < 20 for span in layout.spans) == 5 for layout in layouts) + assert all(sum(span.sequence >= 20 for span in layout.spans) == 2 for layout in layouts) + + covered = sorted((span.sequence, span.token_start, span.token_end) for layout in layouts for span in layout.spans) + for sequence, (start, end) in enumerate(zip(boundaries, boundaries[1:])): + sequence_spans = [span for span in covered if span[0] == sequence] + assert sequence_spans[0][1] == start + assert sequence_spans[-1][2] == end + assert all(left[2] == right[1] for left, right in zip(sequence_spans, sequence_spans[1:])) + + +def test_alternating_sequence_layout_balances_causal_candidates_better_than_contiguous(): + boundaries = _workload_boundaries() + contiguous = [build_contiguous_query_layout(boundaries, rank=rank, world_size=4) for rank in range(4)] + alternating = [build_alternating_sequence_query_layout(boundaries, rank=rank, world_size=4) for rank in range(4)] + + contiguous_costs = [layout.candidate_floats for layout in contiguous] + alternating_costs = [layout.candidate_floats for layout in alternating] + assert max(contiguous_costs) / min(contiguous_costs) > 2.5 + assert max(alternating_costs) / min(alternating_costs) < 1.01 + + +def test_planned_layout_balances_reference_workload_without_special_case(): + boundaries = _workload_boundaries() + layouts = [build_planned_query_layout(boundaries, rank=rank, world_size=4) for rank in range(4)] + + assert {layout.local_tokens for layout in layouts} == {1 << 18} + assert max(layout.candidate_floats for layout in layouts) == min(layout.candidate_floats for layout in layouts) + + +def test_contiguous_layout_uses_block_aligned_shards_for_non_power_of_two_cp(): + boundaries = [0, 1024, 4096] + layouts = [build_contiguous_query_layout(boundaries, rank=rank, world_size=3) for rank in range(3)] + + assert sum(layout.local_tokens for layout in layouts) == boundaries[-1] + assert all(span.token_start % 4 == 0 and span.token_end % 4 == 0 for layout in layouts for span in layout.spans) + covered = sorted((span.sequence, span.token_start, span.token_end) for layout in layouts for span in layout.spans) + for sequence, (start, end) in enumerate(zip(boundaries, boundaries[1:])): + sequence_spans = [span for span in covered if span[0] == sequence] + assert sequence_spans[0][1] == start + assert sequence_spans[-1][2] == end + assert all(left[2] == right[1] for left, right in zip(sequence_spans, sequence_spans[1:])) diff --git a/tests/test_thd_planner.py b/tests/test_thd_planner.py new file mode 100644 index 00000000..963cd543 --- /dev/null +++ b/tests/test_thd_planner.py @@ -0,0 +1,116 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +import random +from collections import defaultdict + +import pytest + +from cula.distributed.thd_planner import THDPlannerConfig, plan_thd_layout + + +def _boundaries(lengths): + result = [0] + for length in lengths: + result.append(result[-1] + length) + return result + + +def _assert_exact_cover(plan, boundaries): + by_sequence = defaultdict(list) + for rank_layout in plan.ranks: + assert rank_layout.local_tokens > 0 + for span in rank_layout.spans: + assert span.owner == rank_layout.rank + by_sequence[span.sequence].append(span) + for sequence, (start, end) in enumerate(zip(boundaries, boundaries[1:])): + spans = sorted(by_sequence[sequence], key=lambda span: span.token_start) + assert spans[0].token_start == start + assert spans[-1].token_end == end + assert all(left.token_end == right.token_start for left, right in zip(spans, spans[1:])) + + +@pytest.mark.parametrize( + "lengths", + [ + [1024] * 1024, + [1024] * 20 + [514048] * 2, + [1 << 20], + [1024] * 256 + [786432], + [4096] * 32 + [131072, 262144, 524288], + [1024, 4096, 16384, 65536, 131072, 307200, 523264], + ], + ids=("all-short", "two-long", "single-long", "long-tail", "multi-long", "power-law"), +) +def test_planner_handles_representative_thd_distributions(lengths): + boundaries = _boundaries(lengths) + plan = plan_thd_layout(boundaries, world_size=4) + + _assert_exact_cover(plan, boundaries) + scores = dict(plan.evaluated_scores) + assert plan.metrics.score <= scores["contiguous"] + assert max(plan.metrics.token_loads) - min(plan.metrics.token_loads) <= plan.block_size + + +def test_single_long_sequence_forms_balanced_multi_span_layout(): + plan = plan_thd_layout([0, 1 << 20], world_size=4) + + assert plan.policy == "cost-balanced-2x" + assert plan.metrics.token_imbalance == pytest.approx(1.0) + assert plan.metrics.causal_imbalance < 1.02 + assert all(1 <= len(rank.spans) <= 2 for rank in plan.ranks) + + +def test_short_sequences_are_not_fragmented_by_cost_balanced_candidate(): + lengths = [1024] * 20 + [514048] * 2 + boundaries = _boundaries(lengths) + plan = plan_thd_layout(boundaries, world_size=4) + + assert plan.policy.startswith("cost-balanced-") + short_counts = defaultdict(int) + for rank in plan.ranks: + for span in rank.spans: + if span.sequence < 20: + short_counts[span.sequence] += 1 + assert set(short_counts.values()) == {1} + + +def test_boundary_penalty_can_select_contiguous_fallback(): + config = THDPlannerConfig(boundary_weight=100.0) + plan = plan_thd_layout([0, 1 << 20], world_size=4, config=config) + + assert plan.policy == "contiguous" + + +def test_cp64_long_sequence_uses_scaled_zigzag_chunks(): + plan = plan_thd_layout([0, 1 << 20], world_size=64) + + assert plan.policy == "cost-balanced-2x" + assert plan.metrics.token_imbalance == pytest.approx(1.0) + assert plan.metrics.causal_imbalance == pytest.approx(1.0) + assert all(1 <= len(rank.spans) <= 2 for rank in plan.ranks) + + +@pytest.mark.parametrize( + ("boundaries", "message"), + [ + ([0, 3, 8], "block aligned"), + ([0, 8], "at least one block per rank"), + ], +) +def test_planner_rejects_invalid_partitions(boundaries, message): + with pytest.raises(ValueError, match=message): + plan_thd_layout(boundaries, world_size=8) + + +def test_random_thd_plans_preserve_coverage_and_never_lose_to_contiguous(): + generator = random.Random(20261104) + for _ in range(100): + sequence_count = generator.randint(1, 128) + lengths = [4 * generator.randint(1, 1 << 15) for _ in range(sequence_count)] + boundaries = _boundaries(lengths) + plan = plan_thd_layout(boundaries, world_size=4) + + _assert_exact_cover(plan, boundaries) + assert plan.metrics.score <= dict(plan.evaluated_scores)["contiguous"] + assert max(plan.metrics.token_loads) - min(plan.metrics.token_loads) <= plan.block_size