From 0a6066971b05949f48bff3d6f1b4b20ea3940666 Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Mon, 2 Nov 2026 17:13:58 +0800 Subject: [PATCH 01/18] feat(hybrid): add packed cuDNN attention runtime Keep index-pooling output in packed THD form, add a differentiable CSA wrapper, and bound DSA compact-logit scratch with exact causal windows. Include CP block ownership metadata, deterministic-set CUDA coverage, and 128K heterogeneous/CP4 measurements. --- README.md | 3 + REPO_LAYOUT.md | 5 + USAGE.md | 15 + cula/_metadata.py | 59 +++ cula/distributed/plan.py | 50 +-- cula/hybrid/__init__.py | 27 ++ cula/hybrid/dsa.py | 484 ++++++++++++++++++++++++ cula/hybrid/index_pool.py | 167 ++++++++ cula/hybrid/plan.py | 271 +++++++++++++ docs/HYBRID_CUDNN_RUNTIME.md | 163 ++++++++ pyproject.toml | 4 + tests/test_hybrid_dsa_cudnn.py | 52 +++ tests/test_hybrid_dsa_helpers.py | 17 + tests/test_hybrid_dsa_plan.py | 132 +++++++ tests/test_hybrid_index_pool_helpers.py | 23 ++ tests/test_hybrid_index_pool_plan.py | 160 ++++++++ 16 files changed, 1586 insertions(+), 46 deletions(-) create mode 100644 cula/_metadata.py create mode 100644 cula/hybrid/__init__.py create mode 100644 cula/hybrid/dsa.py create mode 100644 cula/hybrid/index_pool.py create mode 100644 cula/hybrid/plan.py create mode 100644 docs/HYBRID_CUDNN_RUNTIME.md create mode 100644 tests/test_hybrid_dsa_cudnn.py create mode 100644 tests/test_hybrid_dsa_helpers.py create mode 100644 tests/test_hybrid_dsa_plan.py create mode 100644 tests/test_hybrid_index_pool_helpers.py create mode 100644 tests/test_hybrid_index_pool_plan.py 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..135d0ca2 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 │ │ @@ -61,6 +65,7 @@ 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). | | `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/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/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/hybrid/__init__.py b/cula/hybrid/__init__.py new file mode 100644 index 00000000..c955561f --- /dev/null +++ b/cula/hybrid/__init__.py @@ -0,0 +1,27 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""cuDNN-backed hybrid-attention integration primitives.""" + +from cula.hybrid.dsa import ( + DEFAULT_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", + "DsaSelectionPlan", + "DsaSelectionResult", + "DsaSelectionWindow", + "DsaSelectionWorkspace", + "IndexPoolPlan", + "IndexPoolShardPlan", + "cudnn_dsa_select", + "cudnn_index_pool", +] diff --git a/cula/hybrid/dsa.py b/cula/hybrid/dsa.py new file mode 100644 index 00000000..82a15e54 --- /dev/null +++ b/cula/hybrid/dsa.py @@ -0,0 +1,484 @@ +# 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 + +DEFAULT_CANDIDATE_BUDGET_BYTES = 512 << 20 +"""Default peak FP32 compact-logit workspace, matching cuDNN's QSA budget.""" + + +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) + + +@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; deterministic tie-breaking remains enabled + by default at execution. 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 = DEFAULT_CANDIDATE_BUDGET_BYTES, + ) -> DsaSelectionPlan: + """Create a plan entirely from host-visible packed metadata.""" + + 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 ( + 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") + ) + 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 = True, +) -> 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. ``deterministic=True`` makes + the selected set reproducible; 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", + "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..310c39fb --- /dev/null +++ b/cula/hybrid/index_pool.py @@ -0,0 +1,167 @@ +# 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..dae45820 --- /dev/null +++ b/cula/hybrid/plan.py @@ -0,0 +1,271 @@ +# 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 + +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] + + @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..5ae65751 --- /dev/null +++ b/docs/HYBRID_CUDNN_RUNTIME.md @@ -0,0 +1,163 @@ +# 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; +- deterministic Top-K is the default, so exact ties at the K-th boundary keep + the smaller block id; +- deterministic selection guarantees the `(block_id, logit)` set, not 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. + +The default budget is 512 MiB, matching cuDNN QSA's serving-side selection +budget. It is a memory/scheduling choice, not a numerical choice. Callers can +set another explicit budget when constructing the plan. + +### 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 | +| cuLA bounded plan | 7 | 512 MiB | 5.3246 ms | + +The bounded plan 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. + +## 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. Aligned shards pool locally and exchange only compressed rows. +Misaligned shards report a bounded neighbor halo (at most three raw tokens for +ratio 4), avoiding a raw-token all-gather. The current implementation exposes +and tests this ownership/halo plan; the distributed halo executor is not yet +part of the public path. + +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; +- 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/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_dsa_cudnn.py b/tests/test_hybrid_dsa_cudnn.py new file mode 100644 index 00000000..a3bd80c4 --- /dev/null +++ b/tests/test_hybrid_dsa_cudnn.py @@ -0,0 +1,52 @@ +# 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) + + 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() + 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..bd348269 --- /dev/null +++ b/tests/test_hybrid_dsa_helpers.py @@ -0,0 +1,17 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace + +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() diff --git a/tests/test_hybrid_dsa_plan.py b/tests/test_hybrid_dsa_plan.py new file mode 100644 index 00000000..42536abc --- /dev/null +++ b/tests/test_hybrid_dsa_plan.py @@ -0,0 +1,132 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +from cula.hybrid import 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 + + +@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..d87bab35 --- /dev/null +++ b/tests/test_hybrid_index_pool_plan.py @@ -0,0 +1,160 @@ +# 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) + + +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) From e6ae6bc054ba371ecbe181bec9621a0483a25e1a Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Mon, 2 Nov 2026 22:06:25 +0800 Subject: [PATCH 02/18] perf(hybrid): adapt DSA workspace for million-token CP Scale the default candidate budget by absolute causal reach up to the measured 4 GiB CP64 knee while preserving the 512 MiB 128K policy. Add a reproducible DSA/ring emulator and document why the exact ordered P2P KDA chain was rejected. --- benchmarks/bench_hybrid_cp_emulation.py | 225 ++++++++++++++++++++++++ cula/hybrid/__init__.py | 4 + cula/hybrid/dsa.py | 35 +++- docs/HYBRID_CUDNN_RUNTIME.md | 37 +++- docs/KIMI_CP_RUNTIME.md | 2 +- tests/test_hybrid_dsa_plan.py | 32 +++- 6 files changed, 323 insertions(+), 12 deletions(-) create mode 100644 benchmarks/bench_hybrid_cp_emulation.py diff --git a/benchmarks/bench_hybrid_cp_emulation.py b/benchmarks/bench_hybrid_cp_emulation.py new file mode 100644 index 00000000..933cb759 --- /dev/null +++ b/benchmarks/bench_hybrid_cp_emulation.py @@ -0,0 +1,225 @@ +# 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().parent.parent)) + +from cula.hybrid import DsaSelectionPlan, DsaSelectionWorkspace, cudnn_dsa_select + + +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, + "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( + { + "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) + next_rank = (rank + 1) % world_size + previous_rank = (rank - 1 + world_size) % world_size + + def execute(): + source, target = send, receive + for _ in range(steps): + requests = dist.batch_isend_irecv( + [ + dist.P2POp(dist.isend, source, next_rank), + dist.P2POp(dist.irecv, target, previous_rank), + ] + ) + for request in requests: + request.wait() + source, target = target, source + + 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=4) + 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/cula/hybrid/__init__.py b/cula/hybrid/__init__.py index c955561f..fca53c71 100644 --- a/cula/hybrid/__init__.py +++ b/cula/hybrid/__init__.py @@ -5,6 +5,8 @@ from cula.hybrid.dsa import ( DEFAULT_CANDIDATE_BUDGET_BYTES, + MAX_BALANCED_CANDIDATE_BUDGET_BYTES, + MEMORY_EFFICIENT_CANDIDATE_BUDGET_BYTES, DsaSelectionPlan, DsaSelectionResult, DsaSelectionWindow, @@ -16,6 +18,8 @@ __all__ = [ "DEFAULT_CANDIDATE_BUDGET_BYTES", + "MAX_BALANCED_CANDIDATE_BUDGET_BYTES", + "MEMORY_EFFICIENT_CANDIDATE_BUDGET_BYTES", "DsaSelectionPlan", "DsaSelectionResult", "DsaSelectionWindow", diff --git a/cula/hybrid/dsa.py b/cula/hybrid/dsa.py index 82a15e54..2e07e4c1 100644 --- a/cula/hybrid/dsa.py +++ b/cula/hybrid/dsa.py @@ -16,8 +16,16 @@ from cula._metadata import as_boundaries -DEFAULT_CANDIDATE_BUDGET_BYTES = 512 << 20 -"""Default peak FP32 compact-logit workspace, matching cuDNN's QSA budget.""" +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 = 4 << 30 +"""Maximum FP32 compact-logit workspace selected by the adaptive policy.""" + +_CANDIDATE_BUDGET_POSITION_QUANTUM = 128 << 10 def _as_nonnegative( @@ -81,6 +89,14 @@ 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) + ) + steps = max(1, (max_causal_end + _CANDIDATE_BUDGET_POSITION_QUANTUM - 1) // _CANDIDATE_BUDGET_POSITION_QUANTUM) + return min(MAX_BALANCED_CANDIDATE_BUDGET_BYTES, steps * DEFAULT_CANDIDATE_BUDGET_BYTES) + + @dataclass(frozen=True, slots=True) class DsaSelectionWindow: """One cuDNN THD launch over contiguous packed query and block rows.""" @@ -155,9 +171,14 @@ def from_cu_seqlens( top_k: int, block_size: int = 4, q_causal_offsets: torch.Tensor | Sequence[int] | None = None, - candidate_budget_bytes: int = DEFAULT_CANDIDATE_BUDGET_BYTES, + candidate_budget_bytes: int | None = None, ) -> DsaSelectionPlan: - """Create a plan entirely from host-visible packed metadata.""" + """Create a plan entirely from host-visible packed metadata. + + A ``None`` budget selects the balanced 512 MiB-per-128K causal-reach + policy, capped at 4 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) @@ -167,7 +188,7 @@ def from_cu_seqlens( 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 ( + 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 @@ -182,6 +203,8 @@ def from_cu_seqlens( 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): @@ -476,6 +499,8 @@ def cudnn_dsa_select( __all__ = [ "DEFAULT_CANDIDATE_BUDGET_BYTES", + "MAX_BALANCED_CANDIDATE_BUDGET_BYTES", + "MEMORY_EFFICIENT_CANDIDATE_BUDGET_BYTES", "DsaSelectionPlan", "DsaSelectionResult", "DsaSelectionWindow", diff --git a/docs/HYBRID_CUDNN_RUNTIME.md b/docs/HYBRID_CUDNN_RUNTIME.md index 5ae65751..adc8d6cc 100644 --- a/docs/HYBRID_CUDNN_RUNTIME.md +++ b/docs/HYBRID_CUDNN_RUNTIME.md @@ -51,9 +51,15 @@ that is still quadratic scratch even though no dense score tensor is exposed. FP32 candidate buffer. It coalesces adjacent short packed sequences when they fit the same budget. -The default budget is 512 MiB, matching cuDNN QSA's serving-side selection -budget. It is a memory/scheduling choice, not a numerical choice. Callers can -set another explicit budget when constructing the plan. +When `candidate_budget_bytes=None`, the balanced policy starts at 512 MiB and +adds 512 MiB for each 128K of absolute causal reach, capped at 4 GiB. Thus the +128K packed workload retains its 512 MiB workspace while the worst 1M/CP64 +rank reaches the measured 4 GiB performance 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 @@ -63,9 +69,9 @@ compression ratio 4, Top-K 512. | Selection form | Windows | Candidate scratch | Median time | |---|---:|---:|---:| | One packed launch | 1 | 2.8573 GiB | 5.3818 ms | -| cuLA bounded plan | 7 | 512 MiB | 5.3246 ms | +| 512 MiB memory policy | 7 | 512 MiB | 5.3246 ms | -The bounded plan reduced candidate scratch by 5.715x. The measured time was +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 @@ -82,6 +88,27 @@ 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 cap)** | **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. + ## Integration example ```python diff --git a/docs/KIMI_CP_RUNTIME.md b/docs/KIMI_CP_RUNTIME.md index 260ca71b..a533fd56 100644 --- a/docs/KIMI_CP_RUNTIME.md +++ b/docs/KIMI_CP_RUNTIME.md @@ -83,7 +83,7 @@ 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 +## 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 diff --git a/tests/test_hybrid_dsa_plan.py b/tests/test_hybrid_dsa_plan.py index 42536abc..d3407b48 100644 --- a/tests/test_hybrid_dsa_plan.py +++ b/tests/test_hybrid_dsa_plan.py @@ -4,7 +4,11 @@ import pytest import torch -from cula.hybrid import DsaSelectionPlan +from cula.hybrid import ( + MAX_BALANCED_CANDIDATE_BUDGET_BYTES, + MEMORY_EFFICIENT_CANDIDATE_BUDGET_BYTES, + DsaSelectionPlan, +) def test_selection_plan_coalesces_short_packed_sequences(): @@ -102,6 +106,32 @@ def test_selection_plan_covers_heterogeneous_128k_without_exceeding_budget(): 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_reaches_measured_cap_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 == MAX_BALANCED_CANDIDATE_BUDGET_BYTES + assert len(plan.windows) == 4 + assert plan.candidate_workspace_floats * 4 <= MAX_BALANCED_CANDIDATE_BUDGET_BYTES + + @pytest.mark.parametrize( ("kwargs", "message"), [ From f17a7f41cc31771b98736c9f21ccf65208f5a23f Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Tue, 3 Nov 2026 12:20:32 +0800 Subject: [PATCH 03/18] feat(benchmarks): add hybrid KDA DSA CP64 emulator Benchmark serial and overlapped KDA plus deterministic DSA index selection on real or logical CP groups. Preserve a conservative full-ring baseline and add packed-sample-aware transport spans for THD workloads, with exact profiler markers and focused coverage. --- benchmarks/bench_hybrid_kda_dsa_e2e.py | 585 +++++++++++++++++++++++++ benchmarks/bench_kimi_cp_e2e.py | 311 +++++++++++++ tests/test_hybrid_cp_benchmark.py | 14 + 3 files changed, 910 insertions(+) create mode 100644 benchmarks/bench_hybrid_kda_dsa_e2e.py create mode 100644 benchmarks/bench_kimi_cp_e2e.py create mode 100644 tests/test_hybrid_cp_benchmark.py diff --git a/benchmarks/bench_hybrid_kda_dsa_e2e.py b/benchmarks/bench_hybrid_kda_dsa_e2e.py new file mode 100644 index 00000000..034689af --- /dev/null +++ b/benchmarks/bench_hybrid_kda_dsa_e2e.py @@ -0,0 +1,585 @@ +# 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, replace + +import torch +import torch.distributed as dist + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) + +from benchmarks.bench_kimi_cp_e2e import ( # noqa: E402 + _LogicalCPRingTransport, + _representative_logical_rank, +) +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 + + +def _packed_boundaries(args: argparse.Namespace) -> tuple[int, ...]: + if not args.thd_20s2l: + return (0, args.global_tokens) + short_total = args.short_sequences * args.short_tokens + remaining = args.global_tokens - short_total + if remaining <= 0 or remaining % args.long_sequences: + raise ValueError("global_tokens minus the short samples must split evenly across the long samples") + lengths = [args.short_tokens] * args.short_sequences + [remaining // args.long_sequences] * args.long_sequences + boundaries = [0] + for length in lengths: + boundaries.append(boundaries[-1] + length) + return tuple(boundaries) + + +def _sample_chain_hops(global_cu_seqlens: tuple[int, ...], local_tokens: int) -> int: + """Return the largest number of rank-to-rank edges crossed by one sample.""" + max_ranks = 1 + for left, right in zip(global_cu_seqlens, global_cu_seqlens[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(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 _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": + dist.all_gather_into_tensor(self.global_kbar, local_kbar.contiguous()) + return + + source, target = local_kbar, self.ring_receive + next_rank = (self.physical_rank + 1) % self.physical_world_size + previous_rank = (self.physical_rank - 1 + self.physical_world_size) % self.physical_world_size + for _ in range(self.transport_hops): + requests = dist.batch_isend_irecv( + [ + dist.P2POp(dist.isend, source, next_rank), + dist.P2POp(dist.irecv, target, previous_rank), + ] + ) + for request in requests: + request.wait() + source, target = target, source + self.global_kbar[self.global_block_start : self.global_block_start + self.local_blocks].copy_(local_kbar) + + 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") + + logical_cp = world_size if args.mode == "real" else args.logical_cp + if logical_cp < world_size: + raise ValueError("logical_cp must be at least the physical 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 + if args.mode == "real" + else _representative_logical_rank(rank, world_size, logical_cp) + ) + global_cu_seqlens = _packed_boundaries(args) + transport_hops = logical_cp - 1 + if args.mode == "logical" and args.logical_transport == "sample-chain": + transport_hops = _sample_chain_hops(global_cu_seqlens, local_tokens) + 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: + logical_shard_plan = KimiCPPlan.from_global_cu_seqlens( + global_cu_seqlens, + rank=logical_rank, + world_size=logical_cp, + ) + # cuDNN needs the logical shard's packed boundaries, while runtime + # group validation must still describe the four physical processes. + kda_plan = replace( + logical_shard_plan, + rank=rank, + world_size=world_size, + token_start=rank * local_tokens, + token_end=(rank + 1) * local_tokens, + ) + 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, + ) + else: + kda_transport = _LogicalCPRingTransport( + logical_cp, + logical_rank, + 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 the remote logical block ranges once. Timed iterations + # update the local logical slice after emulating its 63-hop transport. + 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): + kda_transport.clear_events() + start = torch.cuda.Event(enable_timing=True) + kda_start = 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) + 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, + ) + 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() + + result = dsa.last_result + assert result is not None + sample: dict[str, object] = { + "cuda_total_ms": start.elapsed_time(end) if measure else None, + "kda_ms": kda_start.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_merge) + ) + for direction, begin, after_comm, after_merge in kda_transport.events + for stage in ("comm", "merge") + } + 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]) + 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]) + wall_start = time.perf_counter() + with torch.cuda.nvtx.range(f"hybrid.{schedule}"): + sample = execute(schedule, True) + sample["wall_ms"] = (time.perf_counter() - wall_start) * 1000 + 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), + } + 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", + "transport_hops": transport_hops, + "physical_gpus": world_size, + "local_tokens": local_tokens, + "sequence_lengths": [ + right - left for left, right in zip(global_cu_seqlens, global_cu_seqlens[1:]) + ], + "schedules": summary, + "records": 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) + 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") + 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=4) + 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( + "--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("--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") + return parser.parse_args() + + +if __name__ == "__main__": + _run(_parse_args()) diff --git a/benchmarks/bench_kimi_cp_e2e.py b/benchmarks/bench_kimi_cp_e2e.py new file mode 100644 index 00000000..beb977c9 --- /dev/null +++ b/benchmarks/bench_kimi_cp_e2e.py @@ -0,0 +1,311 @@ +# 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 IEEE merge, 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/bench_kimi_cp_e2e.py +""" + +from __future__ import annotations + +import argparse +import inspect +import json +import os +import pathlib +import statistics +import sys +import time +from dataclasses import dataclass, field, replace + +import torch +import torch.distributed as dist + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) + +from cula.distributed import KimiCPBoundaryTransport, KimiCPPlan, cudnn_kda_cp + + +def _positive_int(value: str) -> int: + result = int(value) + if result < 1: + raise argparse.ArgumentTypeError("value must be positive") + return result + + +def _load_merge(): + from cudnn.linear_attention import kimi_cp_merge + + if "precision" not in inspect.signature(kimi_cp_merge).parameters: + return kimi_cp_merge + + def ieee_merge(*args, **kwargs): + return kimi_cp_merge(*args, precision="ieee", **kwargs) + + return ieee_merge + + +@dataclass(slots=True) +class _LogicalCPRingTransport: + """Benchmark-only logical-CP transport with real NCCL hop count.""" + + logical_cp: int + logical_rank: int + physical_rank: int + physical_world_size: int + ring_hops: int | None = field(default=None, kw_only=True) + merge: object = field(default_factory=_load_merge) + 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: + 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 + source, target = self.local_payload, self.receive_payload + next_rank = (self.physical_rank + 1) % self.physical_world_size + previous_rank = (self.physical_rank - 1 + self.physical_world_size) % self.physical_world_size + 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)") + for _ in range(hops): + requests = dist.batch_isend_irecv( + [ + dist.P2POp(dist.isend, source, next_rank), + dist.P2POp(dist.irecv, target, previous_rank), + ] + ) + for request in requests: + request.wait() + source, target = target, source + + def boundary_state( + self, + h: torch.Tensor, + transition: torch.Tensor, + *, + plan, + group, + reverse: bool, + ) -> 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_merge = 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 + if pieces == 0: + state = torch.zeros_like(h) + else: + h_all = self.h_summaries[:pieces] + transition_all = self.transition_summaries[:pieces] + h_all.copy_(h.unsqueeze(0)) + transition_all.copy_(transition.unsqueeze(0)) + state = self.merge(h_all, transition_all, reverse=reverse)["final_state"][0] + after_merge.record() + self.events.append(("backward" if reverse else "forward", start, after_comm, after_merge)) + return state + + +def _representative_logical_rank(physical_rank: int, physical_world_size: int, logical_cp: int) -> int: + if physical_world_size == 1: + return logical_cp - 1 + return physical_rank * (logical_cp - 1) // (physical_world_size - 1) + + +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") + if args.logical_cp < world_size: + raise ValueError("logical_cp must be at least the physical world size") + if args.global_tokens % args.logical_cp: + raise ValueError("global_tokens must be divisible by logical_cp") + + local_tokens = args.global_tokens // args.logical_cp + logical_rank = _representative_logical_rank(rank, world_size, args.logical_cp) + 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() + + # The physical plan supplies the local packed-THD shape. The benchmark + # transport replaces only boundary propagation with logical-CP semantics. + logical_plan = KimiCPPlan.from_global_cu_seqlens( + [0, args.global_tokens], + rank=logical_rank, + world_size=args.logical_cp, + ) + plan = replace( + logical_plan, + rank=rank, + world_size=world_size, + token_start=rank * local_tokens, + token_end=(rank + 1) * local_tokens, + ) + cu_seqlens = torch.tensor(plan.local_cu_seqlens, dtype=torch.int32, device=device) + transport = _LogicalCPRingTransport(args.logical_cp, logical_rank, rank, world_size) + 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_merge in transport.events: + transport_times[f"comm_{direction}_ms"] = begin.elapsed_time(after_comm) + transport_times[f"merge_{direction}_ms"] = after_comm.elapsed_time(after_merge) + 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() + + +if __name__ == "__main__": + _run(_parse_args()) diff --git a/tests/test_hybrid_cp_benchmark.py b/tests/test_hybrid_cp_benchmark.py new file mode 100644 index 00000000..d511024b --- /dev/null +++ b/tests/test_hybrid_cp_benchmark.py @@ -0,0 +1,14 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +from benchmarks.bench_hybrid_kda_dsa_e2e import _sample_chain_hops + + +def test_sample_chain_hops_for_one_megatoken_thd_mix(): + boundaries = tuple([index * 1024 for index in range(21)] + [534528, 1048576]) + + assert _sample_chain_hops(boundaries, 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 From 668b54cff24304320ed50410239b6fea1e9df991 Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Tue, 3 Nov 2026 15:28:24 +0800 Subject: [PATCH 04/18] refactor(benchmarks): isolate large CP emulation layer Separate logical topology, benchmark-only transports, and packed THD workloads from native CP runtime code. Keep the previous benchmark paths as thin compatibility entry points and document the one-way dependency boundary. --- REPO_LAYOUT.md | 8 + benchmarks/bench_hybrid_cp_emulation.py | 213 +------ benchmarks/bench_hybrid_kda_dsa_e2e.py | 575 +----------------- benchmarks/bench_kimi_cp_e2e.py | 301 +-------- benchmarks/large_cp_emulation/README.md | 33 + benchmarks/large_cp_emulation/__init__.py | 17 + .../large_cp_emulation/bench_components.py | 220 +++++++ benchmarks/large_cp_emulation/bench_hybrid.py | 563 +++++++++++++++++ benchmarks/large_cp_emulation/bench_kda.py | 188 ++++++ benchmarks/large_cp_emulation/topology.py | 109 ++++ benchmarks/large_cp_emulation/transport.py | 151 +++++ benchmarks/large_cp_emulation/workload.py | 62 ++ docs/KIMI_CP_RUNTIME.md | 36 ++ tests/test_hybrid_cp_benchmark.py | 14 - tests/test_large_cp_emulation.py | 37 ++ 15 files changed, 1435 insertions(+), 1092 deletions(-) create mode 100644 benchmarks/large_cp_emulation/README.md create mode 100644 benchmarks/large_cp_emulation/__init__.py create mode 100644 benchmarks/large_cp_emulation/bench_components.py create mode 100644 benchmarks/large_cp_emulation/bench_hybrid.py create mode 100644 benchmarks/large_cp_emulation/bench_kda.py create mode 100644 benchmarks/large_cp_emulation/topology.py create mode 100644 benchmarks/large_cp_emulation/transport.py create mode 100644 benchmarks/large_cp_emulation/workload.py delete mode 100644 tests/test_hybrid_cp_benchmark.py create mode 100644 tests/test_large_cp_emulation.py diff --git a/REPO_LAYOUT.md b/REPO_LAYOUT.md index 135d0ca2..d3c53d74 100644 --- a/REPO_LAYOUT.md +++ b/REPO_LAYOUT.md @@ -52,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 @@ -66,6 +73,7 @@ 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/benchmarks/bench_hybrid_cp_emulation.py b/benchmarks/bench_hybrid_cp_emulation.py index 933cb759..078b009d 100644 --- a/benchmarks/bench_hybrid_cp_emulation.py +++ b/benchmarks/bench_hybrid_cp_emulation.py @@ -1,224 +1,17 @@ # Copyright 2025-2026 Ant Group Co., Ltd. # SPDX-License-Identifier: Apache-2.0 +# ruff: noqa: E402, I001 -"""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. -""" +"""Compatibility entry point for large-CP component emulation.""" 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().parent.parent)) -from cula.hybrid import DsaSelectionPlan, DsaSelectionWorkspace, cudnn_dsa_select - - -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, - "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( - { - "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) - next_rank = (rank + 1) % world_size - previous_rank = (rank - 1 + world_size) % world_size - - def execute(): - source, target = send, receive - for _ in range(steps): - requests = dist.batch_isend_irecv( - [ - dist.P2POp(dist.isend, source, next_rank), - dist.P2POp(dist.irecv, target, previous_rank), - ] - ) - for request in requests: - request.wait() - source, target = target, source - - 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=4) - 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) +from benchmarks.large_cp_emulation.bench_components import main if __name__ == "__main__": diff --git a/benchmarks/bench_hybrid_kda_dsa_e2e.py b/benchmarks/bench_hybrid_kda_dsa_e2e.py index 034689af..63d26bbb 100644 --- a/benchmarks/bench_hybrid_kda_dsa_e2e.py +++ b/benchmarks/bench_hybrid_kda_dsa_e2e.py @@ -1,585 +1,18 @@ # Copyright 2025-2026 Ant Group Co., Ltd. # SPDX-License-Identifier: Apache-2.0 +# ruff: noqa: E402, I001 -"""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. -""" +"""Compatibility entry point for the large-CP hybrid emulator.""" from __future__ import annotations -import argparse -import json -import os import pathlib -import statistics import sys -import time -from dataclasses import dataclass, replace - -import torch -import torch.distributed as dist sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) -from benchmarks.bench_kimi_cp_e2e import ( # noqa: E402 - _LogicalCPRingTransport, - _representative_logical_rank, -) -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 - - -def _packed_boundaries(args: argparse.Namespace) -> tuple[int, ...]: - if not args.thd_20s2l: - return (0, args.global_tokens) - short_total = args.short_sequences * args.short_tokens - remaining = args.global_tokens - short_total - if remaining <= 0 or remaining % args.long_sequences: - raise ValueError("global_tokens minus the short samples must split evenly across the long samples") - lengths = [args.short_tokens] * args.short_sequences + [remaining // args.long_sequences] * args.long_sequences - boundaries = [0] - for length in lengths: - boundaries.append(boundaries[-1] + length) - return tuple(boundaries) - - -def _sample_chain_hops(global_cu_seqlens: tuple[int, ...], local_tokens: int) -> int: - """Return the largest number of rank-to-rank edges crossed by one sample.""" - max_ranks = 1 - for left, right in zip(global_cu_seqlens, global_cu_seqlens[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(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 _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": - dist.all_gather_into_tensor(self.global_kbar, local_kbar.contiguous()) - return - - source, target = local_kbar, self.ring_receive - next_rank = (self.physical_rank + 1) % self.physical_world_size - previous_rank = (self.physical_rank - 1 + self.physical_world_size) % self.physical_world_size - for _ in range(self.transport_hops): - requests = dist.batch_isend_irecv( - [ - dist.P2POp(dist.isend, source, next_rank), - dist.P2POp(dist.irecv, target, previous_rank), - ] - ) - for request in requests: - request.wait() - source, target = target, source - self.global_kbar[self.global_block_start : self.global_block_start + self.local_blocks].copy_(local_kbar) - - 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") - - logical_cp = world_size if args.mode == "real" else args.logical_cp - if logical_cp < world_size: - raise ValueError("logical_cp must be at least the physical 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 - if args.mode == "real" - else _representative_logical_rank(rank, world_size, logical_cp) - ) - global_cu_seqlens = _packed_boundaries(args) - transport_hops = logical_cp - 1 - if args.mode == "logical" and args.logical_transport == "sample-chain": - transport_hops = _sample_chain_hops(global_cu_seqlens, local_tokens) - 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: - logical_shard_plan = KimiCPPlan.from_global_cu_seqlens( - global_cu_seqlens, - rank=logical_rank, - world_size=logical_cp, - ) - # cuDNN needs the logical shard's packed boundaries, while runtime - # group validation must still describe the four physical processes. - kda_plan = replace( - logical_shard_plan, - rank=rank, - world_size=world_size, - token_start=rank * local_tokens, - token_end=(rank + 1) * local_tokens, - ) - 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, - ) - else: - kda_transport = _LogicalCPRingTransport( - logical_cp, - logical_rank, - 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 the remote logical block ranges once. Timed iterations - # update the local logical slice after emulating its 63-hop transport. - 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): - kda_transport.clear_events() - start = torch.cuda.Event(enable_timing=True) - kda_start = 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) - 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, - ) - 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() - - result = dsa.last_result - assert result is not None - sample: dict[str, object] = { - "cuda_total_ms": start.elapsed_time(end) if measure else None, - "kda_ms": kda_start.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_merge) - ) - for direction, begin, after_comm, after_merge in kda_transport.events - for stage in ("comm", "merge") - } - 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]) - 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]) - wall_start = time.perf_counter() - with torch.cuda.nvtx.range(f"hybrid.{schedule}"): - sample = execute(schedule, True) - sample["wall_ms"] = (time.perf_counter() - wall_start) * 1000 - 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), - } - 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", - "transport_hops": transport_hops, - "physical_gpus": world_size, - "local_tokens": local_tokens, - "sequence_lengths": [ - right - left for left, right in zip(global_cu_seqlens, global_cu_seqlens[1:]) - ], - "schedules": summary, - "records": 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) - 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") - 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=4) - 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( - "--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("--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") - return parser.parse_args() +from benchmarks.large_cp_emulation.bench_hybrid import main if __name__ == "__main__": - _run(_parse_args()) + main() diff --git a/benchmarks/bench_kimi_cp_e2e.py b/benchmarks/bench_kimi_cp_e2e.py index beb977c9..47962e8c 100644 --- a/benchmarks/bench_kimi_cp_e2e.py +++ b/benchmarks/bench_kimi_cp_e2e.py @@ -1,311 +1,18 @@ # Copyright 2025-2026 Ant Group Co., Ltd. # SPDX-License-Identifier: Apache-2.0 +# ruff: noqa: E402, I001 -"""Emulate a large logical Kimi CP group on fewer physical GPUs. - -The benchmark executes the complete cuLA KDA path: local summary, boundary -transport and IEEE merge, 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/bench_kimi_cp_e2e.py -""" +"""Compatibility entry point for the large-CP KDA emulator.""" from __future__ import annotations -import argparse -import inspect -import json -import os import pathlib -import statistics import sys -import time -from dataclasses import dataclass, field, replace - -import torch -import torch.distributed as dist sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) -from cula.distributed import KimiCPBoundaryTransport, KimiCPPlan, cudnn_kda_cp - - -def _positive_int(value: str) -> int: - result = int(value) - if result < 1: - raise argparse.ArgumentTypeError("value must be positive") - return result - - -def _load_merge(): - from cudnn.linear_attention import kimi_cp_merge - - if "precision" not in inspect.signature(kimi_cp_merge).parameters: - return kimi_cp_merge - - def ieee_merge(*args, **kwargs): - return kimi_cp_merge(*args, precision="ieee", **kwargs) - - return ieee_merge - - -@dataclass(slots=True) -class _LogicalCPRingTransport: - """Benchmark-only logical-CP transport with real NCCL hop count.""" - - logical_cp: int - logical_rank: int - physical_rank: int - physical_world_size: int - ring_hops: int | None = field(default=None, kw_only=True) - merge: object = field(default_factory=_load_merge) - 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: - 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 - source, target = self.local_payload, self.receive_payload - next_rank = (self.physical_rank + 1) % self.physical_world_size - previous_rank = (self.physical_rank - 1 + self.physical_world_size) % self.physical_world_size - 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)") - for _ in range(hops): - requests = dist.batch_isend_irecv( - [ - dist.P2POp(dist.isend, source, next_rank), - dist.P2POp(dist.irecv, target, previous_rank), - ] - ) - for request in requests: - request.wait() - source, target = target, source - - def boundary_state( - self, - h: torch.Tensor, - transition: torch.Tensor, - *, - plan, - group, - reverse: bool, - ) -> 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_merge = 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 - if pieces == 0: - state = torch.zeros_like(h) - else: - h_all = self.h_summaries[:pieces] - transition_all = self.transition_summaries[:pieces] - h_all.copy_(h.unsqueeze(0)) - transition_all.copy_(transition.unsqueeze(0)) - state = self.merge(h_all, transition_all, reverse=reverse)["final_state"][0] - after_merge.record() - self.events.append(("backward" if reverse else "forward", start, after_comm, after_merge)) - return state - - -def _representative_logical_rank(physical_rank: int, physical_world_size: int, logical_cp: int) -> int: - if physical_world_size == 1: - return logical_cp - 1 - return physical_rank * (logical_cp - 1) // (physical_world_size - 1) - - -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") - if args.logical_cp < world_size: - raise ValueError("logical_cp must be at least the physical world size") - if args.global_tokens % args.logical_cp: - raise ValueError("global_tokens must be divisible by logical_cp") - - local_tokens = args.global_tokens // args.logical_cp - logical_rank = _representative_logical_rank(rank, world_size, args.logical_cp) - 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() - - # The physical plan supplies the local packed-THD shape. The benchmark - # transport replaces only boundary propagation with logical-CP semantics. - logical_plan = KimiCPPlan.from_global_cu_seqlens( - [0, args.global_tokens], - rank=logical_rank, - world_size=args.logical_cp, - ) - plan = replace( - logical_plan, - rank=rank, - world_size=world_size, - token_start=rank * local_tokens, - token_end=(rank + 1) * local_tokens, - ) - cu_seqlens = torch.tensor(plan.local_cu_seqlens, dtype=torch.int32, device=device) - transport = _LogicalCPRingTransport(args.logical_cp, logical_rank, rank, world_size) - 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_merge in transport.events: - transport_times[f"comm_{direction}_ms"] = begin.elapsed_time(after_comm) - transport_times[f"merge_{direction}_ms"] = after_comm.elapsed_time(after_merge) - 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() +from benchmarks.large_cp_emulation.bench_kda import main if __name__ == "__main__": - _run(_parse_args()) + main() diff --git a/benchmarks/large_cp_emulation/README.md b/benchmarks/large_cp_emulation/README.md new file mode 100644 index 00000000..d02184f5 --- /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 KDA merge precision 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..e85d1083 --- /dev/null +++ b/benchmarks/large_cp_emulation/bench_components.py @@ -0,0 +1,220 @@ +# 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, + "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( + { + "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=4) + 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_hybrid.py b/benchmarks/large_cp_emulation/bench_hybrid.py new file mode 100644 index 00000000..e87050e0 --- /dev/null +++ b/benchmarks/large_cp_emulation/bench_hybrid.py @@ -0,0 +1,563 @@ +# 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 _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": + dist.all_gather_into_tensor(self.global_kbar, local_kbar.contiguous()) + return + + emulate_ring_hops( + local_kbar, + self.ring_receive, + physical_rank=self.physical_rank, + physical_world_size=self.physical_world_size, + hops=self.transport_hops, + ) + self.global_kbar[self.global_block_start : self.global_block_start + self.local_blocks].copy_(local_kbar) + + 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, + ) + 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): + kda_transport.clear_events() + start = torch.cuda.Event(enable_timing=True) + kda_start = 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) + 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, + ) + 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() + + result = dsa.last_result + assert result is not None + sample: dict[str, object] = { + "cuda_total_ms": start.elapsed_time(end) if measure else None, + "kda_ms": kda_start.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_merge) + ) + for direction, begin, after_comm, after_merge in kda_transport.events + for stage in ("comm", "merge") + } + 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]) + 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]) + wall_start = time.perf_counter() + with torch.cuda.nvtx.range(f"hybrid.{schedule}"): + sample = execute(schedule, True) + sample["wall_ms"] = (time.perf_counter() - wall_start) * 1000 + 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), + } + 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", + "transport_hops": transport_hops, + "physical_gpus": world_size, + "local_tokens": local_tokens, + "sequence_lengths": workload.sequence_lengths, + "schedules": summary, + "records": 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) + 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") + 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=4) + 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( + "--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("--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") + 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..2dcd262d --- /dev/null +++ b/benchmarks/large_cp_emulation/bench_kda.py @@ -0,0 +1,188 @@ +# 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 IEEE merge, 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_merge in transport.events: + transport_times[f"comm_{direction}_ms"] = begin.elapsed_time(after_comm) + transport_times[f"merge_{direction}_ms"] = after_comm.elapsed_time(after_merge) + 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/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..73fc6f23 --- /dev/null +++ b/benchmarks/large_cp_emulation/transport.py @@ -0,0 +1,151 @@ +# 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 + +import inspect +from dataclasses import dataclass, field + +import torch +import torch.distributed as dist + + +def _load_merge(): + from cudnn.linear_attention import kimi_cp_merge + + if "precision" not in inspect.signature(kimi_cp_merge).parameters: + return kimi_cp_merge + + def ieee_merge(*args, **kwargs): + return kimi_cp_merge(*args, precision="ieee", **kwargs) + + return ieee_merge + + +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) + merge: object = field(default_factory=_load_merge) + 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_state( + self, + h: torch.Tensor, + transition: torch.Tensor, + *, + plan, + group, + reverse: bool, + ) -> 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_merge = 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 + if pieces == 0: + state = torch.zeros_like(h) + else: + h_all = self.h_summaries[:pieces] + transition_all = self.transition_summaries[:pieces] + h_all.copy_(h.unsqueeze(0)) + transition_all.copy_(transition.unsqueeze(0)) + state = self.merge(h_all, transition_all, reverse=reverse)["final_state"][0] + after_merge.record() + self.events.append(("backward" if reverse else "forward", start, after_comm, after_merge)) + return state diff --git a/benchmarks/large_cp_emulation/workload.py b/benchmarks/large_cp_emulation/workload.py new file mode 100644 index 00000000..c2776969 --- /dev/null +++ b/benchmarks/large_cp_emulation/workload.py @@ -0,0 +1,62 @@ +# 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/docs/KIMI_CP_RUNTIME.md b/docs/KIMI_CP_RUNTIME.md index a533fd56..4f42f4bf 100644 --- a/docs/KIMI_CP_RUNTIME.md +++ b/docs/KIMI_CP_RUNTIME.md @@ -93,6 +93,42 @@ 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 summary, IEEE merge, KDA forward, backward summary, reverse merge, and +local autograd are real cuDNN operations. No DSA stage is included. + +Shape: 1,048,576 global tokens, logical CP64, `H=32`, `K=V=128`. One complete +warmup was excluded, followed by three steady-state samples. + +| E2E stage | Critical-stage median | +|---|---:| +| Forward summary | 0.903 ms | +| Forward 63-hop communication | 3.008 ms | +| Forward FP32 merge | 0.293 ms | +| Local KDA forward | 1.054 ms | +| Backward summary | 1.083 ms | +| Backward 63-hop communication | 3.216 ms | +| Backward FP32 merge | 0.403 ms | +| Local KDA backward | 3.261 ms | + +The measured critical-path median was **12.953 ms GPU time** and **16.102 ms +wall time**. Peak allocated memory was 3.049 GiB per physical process; peak +reserved memory was 3.15–3.27 GiB. Outputs and all five local input gradients +were finite on every representative rank. Proxy loss, output norm, and gradient +norms were identical across all three repeats on each rank. + +Communication accounts for roughly 6.22 ms of the stage envelope and is the +largest optimization target. These numbers 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`. The logical topology +and benchmark-only transport live beside it rather than in the production CP +runtime. ## Prepared large-CP transport The convenience path allocates its gathered H and M tensors on each forward and diff --git a/tests/test_hybrid_cp_benchmark.py b/tests/test_hybrid_cp_benchmark.py deleted file mode 100644 index d511024b..00000000 --- a/tests/test_hybrid_cp_benchmark.py +++ /dev/null @@ -1,14 +0,0 @@ -# Copyright 2025-2026 Ant Group Co., Ltd. -# SPDX-License-Identifier: Apache-2.0 - -from benchmarks.bench_hybrid_kda_dsa_e2e import _sample_chain_hops - - -def test_sample_chain_hops_for_one_megatoken_thd_mix(): - boundaries = tuple([index * 1024 for index in range(21)] + [534528, 1048576]) - - assert _sample_chain_hops(boundaries, 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 diff --git a/tests/test_large_cp_emulation.py b/tests/test_large_cp_emulation.py new file mode 100644 index 00000000..ca762592 --- /dev/null +++ b/tests/test_large_cp_emulation.py @@ -0,0 +1,37 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +from benchmarks.large_cp_emulation import ( + LogicalCPEmulationPlan, + PackedTHDWorkload, + sample_chain_hops, +) + + +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 From e705c84b47aff6b92eac8af6eb0604c2bb37ebcb Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Wed, 4 Nov 2026 00:44:12 +0800 Subject: [PATCH 05/18] perf(hybrid): pool index blocks before CP gather Keep raw THD rows rank-local and exchange only compressed blocks. Preserve the gather-then-pool backward contract by routing loss-rank gradients separately and accumulating token and APE gradients in FP32. --- benchmarks/bench_index_pool_cp.py | 371 ++++++++++++++++++++++++++ cula/hybrid/__init__.py | 2 + cula/hybrid/distributed_pool.py | 285 ++++++++++++++++++++ cula/hybrid/plan.py | 35 +++ docs/HYBRID_CUDNN_RUNTIME.md | 39 ++- tests/test_hybrid_distributed_pool.py | 177 ++++++++++++ tests/test_hybrid_index_pool_plan.py | 3 + 7 files changed, 907 insertions(+), 5 deletions(-) create mode 100644 benchmarks/bench_index_pool_cp.py create mode 100644 cula/hybrid/distributed_pool.py create mode 100644 tests/test_hybrid_distributed_pool.py diff --git a/benchmarks/bench_index_pool_cp.py b/benchmarks/bench_index_pool_cp.py new file mode 100644 index 00000000..b8e83346 --- /dev/null +++ b/benchmarks/bench_index_pool_cp.py @@ -0,0 +1,371 @@ +# 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/cula/hybrid/__init__.py b/cula/hybrid/__init__.py index fca53c71..708d6636 100644 --- a/cula/hybrid/__init__.py +++ b/cula/hybrid/__init__.py @@ -3,6 +3,7 @@ """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, @@ -28,4 +29,5 @@ "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..ea84920c --- /dev/null +++ b/cula/hybrid/distributed_pool.py @@ -0,0 +1,285 @@ +# 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/plan.py b/cula/hybrid/plan.py index dae45820..d333583c 100644 --- a/cula/hybrid/plan.py +++ b/cula/hybrid/plan.py @@ -10,6 +10,7 @@ from operator import index import torch +import torch.distributed as dist from cula._metadata import as_boundaries, balanced_shards, validate_shards @@ -188,6 +189,40 @@ def local_tokens(self) -> int: 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, diff --git a/docs/HYBRID_CUDNN_RUNTIME.md b/docs/HYBRID_CUDNN_RUNTIME.md index adc8d6cc..2766e12a 100644 --- a/docs/HYBRID_CUDNN_RUNTIME.md +++ b/docs/HYBRID_CUDNN_RUNTIME.md @@ -158,11 +158,39 @@ 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. Aligned shards pool locally and exchange only compressed rows. -Misaligned shards report a bounded neighbor halo (at most three raw tokens for -ratio 4), avoiding a raw-token all-gather. The current implementation exposes -and tests this ownership/halo plan; the distributed halo executor is not yet -part of the public path. +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 @@ -174,6 +202,7 @@ this runtime's contract. 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. 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_index_pool_plan.py b/tests/test_hybrid_index_pool_plan.py index d87bab35..67aa0ce5 100644 --- a/tests/test_hybrid_index_pool_plan.py +++ b/tests/test_hybrid_index_pool_plan.py @@ -64,6 +64,9 @@ def test_aligned_cp_shards_pool_locally_and_gather_only_blocks(): 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(): From 9565612c6b32be5e83869bb0ad93af9b9d1bc0ac Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Wed, 4 Nov 2026 01:18:42 +0800 Subject: [PATCH 06/18] perf(hybrid): tune 1M CP workspace policy Scale compact-logit scratch by causal reach and rank-local query rows so real CP4 uses the measured 8 GiB knee while CP64 retains its 4 GiB choice. Add stage/resource/correctness summaries for reproducible E2E profiling and migrate KDA gathers to the current PyTorch collective API. --- benchmarks/large_cp_emulation/bench_hybrid.py | 92 ++++++++++++++++++- cula/distributed/kda.py | 12 +++ cula/hybrid/dsa.py | 35 ++++++- docs/HYBRID_CUDNN_RUNTIME.md | 50 +++++++++- tests/test_hybrid_dsa_plan.py | 29 +++++- 5 files changed, 203 insertions(+), 15 deletions(-) diff --git a/benchmarks/large_cp_emulation/bench_hybrid.py b/benchmarks/large_cp_emulation/bench_hybrid.py index e87050e0..b9e84f57 100644 --- a/benchmarks/large_cp_emulation/bench_hybrid.py +++ b/benchmarks/large_cp_emulation/bench_hybrid.py @@ -209,7 +209,10 @@ def workspace_bytes(self) -> int: def _communicate(self, local_kbar: torch.Tensor) -> None: if self.mode == "real": - dist.all_gather_into_tensor(self.global_kbar, local_kbar.contiguous()) + 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 emulate_ring_hops( @@ -372,6 +375,7 @@ def execute(schedule: str, measure: bool) -> dict[str, object]: 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() @@ -402,11 +406,13 @@ def execute(schedule: str, measure: bool) -> dict[str, object]: 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, "dsa_ms": dsa_events.timings() if measure else None, "proxy_loss": (output.float() * grad_output.float()).sum().item(), @@ -445,10 +451,8 @@ def execute(schedule: str, measure: bool) -> dict[str, object]: samples = [] for _ in range(args.repeats): dist.barrier(device_ids=[local_rank]) - wall_start = time.perf_counter() with torch.cuda.nvtx.range(f"hybrid.{schedule}"): sample = execute(schedule, True) - sample["wall_ms"] = (time.perf_counter() - wall_start) * 1000 samples.append(sample) if args.cuda_profiler_capture: torch.cuda.cudart().cudaProfilerStop() @@ -486,6 +490,23 @@ def execute(schedule: str, measure: bool) -> dict[str, object]: "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] + ), + **{ + 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"] @@ -503,7 +524,65 @@ def execute(schedule: str, measure: bool) -> dict[str, object]: "local_tokens": local_tokens, "sequence_lengths": workload.sequence_lengths, "schedules": summary, - "records": records, + "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, @@ -552,6 +631,11 @@ def _parse_args() -> argparse.Namespace: 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() diff --git a/cula/distributed/kda.py b/cula/distributed/kda.py index d9af7dc0..8efbf07d 100644 --- a/cula/distributed/kda.py +++ b/cula/distributed/kda.py @@ -107,6 +107,18 @@ def _all_gather_into( dist.all_gather_into_tensor(output, tensor, group=group) +def _all_gather_into( + output: torch.Tensor, + tensor: torch.Tensor, + *, + group: dist.ProcessGroup | None, +) -> None: + if hasattr(dist, "all_gather_single"): + dist.all_gather_single(output, tensor, group=group) + else: # pragma: no cover - compatibility with older PyTorch releases + dist.all_gather_into_tensor(output, tensor, group=group) + + def _all_gather(tensor: torch.Tensor, group: dist.ProcessGroup | None, world_size: int) -> torch.Tensor: """Gather one equal-shaped boundary payload per rank.""" if world_size == 1: diff --git a/cula/hybrid/dsa.py b/cula/hybrid/dsa.py index 2e07e4c1..d9e2ad28 100644 --- a/cula/hybrid/dsa.py +++ b/cula/hybrid/dsa.py @@ -22,10 +22,12 @@ 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 = 4 << 30 +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( @@ -93,8 +95,30 @@ def _balanced_candidate_budget(cu_q: tuple[int, ...], offsets: tuple[int, ...]) max_causal_end = max( offset + right - left for left, right, offset in zip(cu_q[:-1], cu_q[1:], offsets, strict=True) ) - steps = max(1, (max_causal_end + _CANDIDATE_BUDGET_POSITION_QUANTUM - 1) // _CANDIDATE_BUDGET_POSITION_QUANTUM) - return min(MAX_BALANCED_CANDIDATE_BUDGET_BYTES, steps * DEFAULT_CANDIDATE_BUDGET_BYTES) + 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) @@ -176,8 +200,9 @@ def from_cu_seqlens( """Create a plan entirely from host-visible packed metadata. A ``None`` budget selects the balanced 512 MiB-per-128K causal-reach - policy, capped at 4 GiB. Pass an explicit byte count to enforce a fixed - memory policy across ranks and microbatches. + 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") diff --git a/docs/HYBRID_CUDNN_RUNTIME.md b/docs/HYBRID_CUDNN_RUNTIME.md index 2766e12a..0baf2542 100644 --- a/docs/HYBRID_CUDNN_RUNTIME.md +++ b/docs/HYBRID_CUDNN_RUNTIME.md @@ -52,9 +52,12 @@ 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, capped at 4 GiB. Thus the -128K packed workload retains its 512 MiB workspace while the worst 1M/CP64 -rank reaches the measured 4 GiB performance knee. The exported +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 @@ -101,7 +104,7 @@ network is virtualized. | 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 cap)** | **4** | **33.519 ms** | **1.20x** | +| **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 | @@ -109,6 +112,45 @@ 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 diff --git a/tests/test_hybrid_dsa_plan.py b/tests/test_hybrid_dsa_plan.py index d3407b48..f4c1c1c6 100644 --- a/tests/test_hybrid_dsa_plan.py +++ b/tests/test_hybrid_dsa_plan.py @@ -117,7 +117,7 @@ def test_adaptive_budget_keeps_128k_workload_memory_efficient(): assert plan.candidate_workspace_floats * 4 <= MEMORY_EFFICIENT_CANDIDATE_BUDGET_BYTES -def test_adaptive_budget_reaches_measured_cap_at_1m_cp64_tail(): +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( @@ -127,11 +127,36 @@ def test_adaptive_budget_reaches_measured_cap_at_1m_cp64_tail(): q_causal_offsets=[query_offset], ) - assert plan.candidate_budget_bytes == MAX_BALANCED_CANDIDATE_BUDGET_BYTES + 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"), [ From 6e5503388c96250433f4639acfed25b25affc93c Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Wed, 11 Nov 2026 01:03:09 +0800 Subject: [PATCH 07/18] perf(benchmarks): model exact KDA checkpoint reuse --- benchmarks/large_cp_emulation/bench_hybrid.py | 106 +++++++++++++++++- benchmarks/large_cp_emulation/transport.py | 15 +-- 2 files changed, 108 insertions(+), 13 deletions(-) diff --git a/benchmarks/large_cp_emulation/bench_hybrid.py b/benchmarks/large_cp_emulation/bench_hybrid.py index b9e84f57..867152e0 100644 --- a/benchmarks/large_cp_emulation/bench_hybrid.py +++ b/benchmarks/large_cp_emulation/bench_hybrid.py @@ -5,7 +5,7 @@ Two modes deliberately answer different questions: -* ``real`` runs a real CP group. KDA uses cuLA's prepared all-gather +* ``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 @@ -86,6 +86,58 @@ def timings(self) -> dict[str, float]: } +@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_state( + self, + h: torch.Tensor, + transition: torch.Tensor, + *, + plan: KimiCPPlan, + group: dist.ProcessGroup | None, + reverse: bool, + ) -> torch.Tensor: + begin = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + begin.record() + result = self.transport.boundary_state( + 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 @@ -327,6 +379,8 @@ def randn(*shape, dtype=torch.bfloat16): value_dim=args.value_dim, device=device, ) + if args.profile_kda_transport: + kda_transport = _TimedKimiCPTransport.wrap(kda_transport) else: kda_transport = LogicalCPRingTransport( logical_cp, @@ -362,10 +416,11 @@ def randn(*shape, dtype=torch.bfloat16): 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): + 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: @@ -395,7 +450,9 @@ def execute(schedule: str, measure: bool) -> dict[str, object]: 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": @@ -414,6 +471,8 @@ def execute(schedule: str, measure: bool) -> 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(), @@ -434,6 +493,8 @@ def execute(schedule: str, measure: bool) -> dict[str, object]: for direction, begin, after_comm, after_merge in kda_transport.events for stage in ("comm", "merge") } + 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,) @@ -443,7 +504,8 @@ def execute(schedule: str, measure: bool) -> dict[str, object]: memory_by_schedule = {} for schedule in schedules: dist.barrier(device_ids=[local_rank]) - execute(schedule, False) + 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: @@ -497,6 +559,23 @@ def execute(schedule: str, measure: bool) -> dict[str, object]: "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] @@ -519,9 +598,11 @@ def execute(schedule: str, measure: bool) -> dict[str, object]: "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, + "kda_checkpoint_tokens": args.kda_checkpoint_tokens, "sequence_lengths": workload.sequence_lengths, "schedules": summary, "rank_resources": [ @@ -595,6 +676,11 @@ def _parse_args() -> argparse.Namespace: 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) @@ -605,6 +691,7 @@ def _parse_args() -> argparse.Namespace: 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) @@ -613,6 +700,13 @@ def _parse_args() -> argparse.Namespace: 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, @@ -620,6 +714,12 @@ def _parse_args() -> argparse.Namespace: 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( diff --git a/benchmarks/large_cp_emulation/transport.py b/benchmarks/large_cp_emulation/transport.py index 73fc6f23..5b2f235c 100644 --- a/benchmarks/large_cp_emulation/transport.py +++ b/benchmarks/large_cp_emulation/transport.py @@ -5,23 +5,18 @@ from __future__ import annotations -import inspect from dataclasses import dataclass, field import torch import torch.distributed as dist +from cula.distributed.kda import _load_cudnn_cp_ops -def _load_merge(): - from cudnn.linear_attention import kimi_cp_merge - - if "precision" not in inspect.signature(kimi_cp_merge).parameters: - return kimi_cp_merge - def ieee_merge(*args, **kwargs): - return kimi_cp_merge(*args, precision="ieee", **kwargs) - - return ieee_merge +def _load_merge(): + # Keep the benchmark emulator on exactly the same export compatibility + # and fixed-IEEE path as the production CP runtime. + return _load_cudnn_cp_ops()[0] def emulate_ring_hops( From a8d76054319618273ce949a5a8d7308161ef9eb7 Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Wed, 4 Nov 2026 12:58:39 +0800 Subject: [PATCH 08/18] perf(hybrid): prototype balanced causal query layout Keep short packed samples whole while alternating long-sequence shard direction across ranks. Add a layout-native deterministic DSA benchmark and document the measured 1M CP4 compute, memory, and scratch tradeoffs without changing the contiguous production runtime. --- .../large_cp_emulation/bench_dsa_layout.py | 308 ++++++++++++++++++ benchmarks/large_cp_emulation/query_layout.py | 186 +++++++++++ docs/KDA_CSA_1M_OPTIMIZATION.md | 79 +++++ tests/test_large_cp_query_layout.py | 56 ++++ 4 files changed, 629 insertions(+) create mode 100644 benchmarks/large_cp_emulation/bench_dsa_layout.py create mode 100644 benchmarks/large_cp_emulation/query_layout.py create mode 100644 docs/KDA_CSA_1M_OPTIMIZATION.md create mode 100644 tests/test_large_cp_query_layout.py 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..6b8b1e07 --- /dev/null +++ b/benchmarks/large_cp_emulation/bench_dsa_layout.py @@ -0,0 +1,308 @@ +# 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, +) +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 _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 = ( + build_contiguous_query_layout + if name == "contiguous" + else build_alternating_sequence_query_layout + ) + 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 _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() + + workload = PackedTHDWorkload.short_long_mix( + args.global_tokens, + short_sequences=args.short_sequences, + short_tokens=args.short_tokens, + long_sequences=args.long_sequences, + ) + boundaries = workload.global_cu_seqlens + 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": args.global_tokens, + "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"), required=True) + parser.add_argument("--global-tokens", type=_positive_int, default=1 << 20) + 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=4) + 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/query_layout.py b/benchmarks/large_cp_emulation/query_layout.py new file mode 100644 index 00000000..c2c145e0 --- /dev/null +++ b/benchmarks/large_cp_emulation/query_layout.py @@ -0,0 +1,186 @@ +# 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 + + +@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, + ) + shard_offsets = balanced_shards(boundaries[-1], world_size) + shard_start, shard_end = shard_offsets[rank : rank + 2] + if shard_start % block_size or shard_end % block_size: + raise ValueError("contiguous CP shard boundaries must be block aligned") + 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) + + +__all__ = [ + "QueryLayout", + "QuerySpan", + "build_alternating_sequence_query_layout", + "build_contiguous_query_layout", +] diff --git a/docs/KDA_CSA_1M_OPTIMIZATION.md b/docs/KDA_CSA_1M_OPTIMIZATION.md new file mode 100644 index 00000000..27773451 --- /dev/null +++ b/docs/KDA_CSA_1M_OPTIMIZATION.md @@ -0,0 +1,79 @@ +# 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: + +| 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 benchmark and CPU planner live in `benchmarks/large_cp_emulation` to keep +this not-yet-production policy separate from the contiguous runtime. diff --git a/tests/test_large_cp_query_layout.py b/tests/test_large_cp_query_layout.py new file mode 100644 index 00000000..d402ec2e --- /dev/null +++ b/tests/test_large_cp_query_layout.py @@ -0,0 +1,56 @@ +# 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, +) + + +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 From b2235c41b35a443580919d2b18bbdb4de04dde03 Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Wed, 4 Nov 2026 13:51:39 +0800 Subject: [PATCH 09/18] feat(distributed): add cost-driven THD layout planner Balance token-linear KDA work and position-dependent CSA candidate work across arbitrary packed THD batches. Evaluate multiple planning grids with a contiguous fallback, enforce exact balanced token capacity, and expose deterministic per-rank span metadata for future KDA/CSA lowering. --- .../large_cp_emulation/bench_dsa_layout.py | 48 +- .../large_cp_emulation/bench_thd_planner.py | 41 ++ benchmarks/large_cp_emulation/query_layout.py | 29 + cula/distributed/__init__.py | 14 + cula/distributed/thd_planner.py | 605 ++++++++++++++++++ docs/THD_LAYOUT_PLANNER.md | 115 ++++ tests/test_large_cp_query_layout.py | 14 + tests/test_thd_planner.py | 116 ++++ 8 files changed, 969 insertions(+), 13 deletions(-) create mode 100644 benchmarks/large_cp_emulation/bench_thd_planner.py create mode 100644 cula/distributed/thd_planner.py create mode 100644 docs/THD_LAYOUT_PLANNER.md create mode 100644 tests/test_thd_planner.py diff --git a/benchmarks/large_cp_emulation/bench_dsa_layout.py b/benchmarks/large_cp_emulation/bench_dsa_layout.py index 6b8b1e07..d5c7ea89 100644 --- a/benchmarks/large_cp_emulation/bench_dsa_layout.py +++ b/benchmarks/large_cp_emulation/bench_dsa_layout.py @@ -21,6 +21,7 @@ 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 @@ -39,6 +40,16 @@ def _positive_int(value: str) -> int: 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:]): @@ -54,11 +65,11 @@ def _layout_for_rank( world_size: int, block_size: int, ) -> QueryLayout: - builder = ( - build_contiguous_query_layout - if name == "contiguous" - else build_alternating_sequence_query_layout - ) + builder = { + "contiguous": build_contiguous_query_layout, + "alternating": build_alternating_sequence_query_layout, + "planned": build_planned_query_layout, + }[name] return builder( boundaries, rank=rank, @@ -94,12 +105,18 @@ def _run(args: argparse.Namespace) -> None: rank = dist.get_rank() world_size = dist.get_world_size() - 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.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 layouts = tuple( _layout_for_rank( @@ -243,7 +260,7 @@ def execute(measure: bool) -> dict[str, object]: json.dumps( { "layout": args.layout, - "global_tokens": args.global_tokens, + "global_tokens": workload.global_tokens, "physical_gpus": world_size, "sequence_lengths": workload.sequence_lengths, "critical_samples_ms": critical, @@ -288,8 +305,13 @@ def execute(measure: bool) -> dict[str, object]: def _parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--layout", choices=("contiguous", "alternating"), required=True) + 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) 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 index c2c145e0..476aabc7 100644 --- a/benchmarks/large_cp_emulation/query_layout.py +++ b/benchmarks/large_cp_emulation/query_layout.py @@ -8,6 +8,7 @@ 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) @@ -178,9 +179,37 @@ def build_alternating_sequence_query_layout( 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/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/thd_planner.py b/cula/distributed/thd_planner.py new file mode 100644 index 00000000..4ea842ec --- /dev/null +++ b/cula/distributed/thd_planner.py @@ -0,0 +1,605 @@ +# 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 +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") + if boundaries[-1] < world_size: + raise ValueError("packed context parallelism requires at least one token per rank") + 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") + for name in ("token_weight", "causal_weight", "boundary_weight"): + value = getattr(config, name) + if not isinstance(value, int | float) or value < 0: + raise ValueError(f"{name} must be 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/docs/THD_LAYOUT_PLANNER.md b/docs/THD_LAYOUT_PLANNER.md new file mode 100644 index 00000000..ca6a4ef7 --- /dev/null +++ b/docs/THD_LAYOUT_PLANNER.md @@ -0,0 +1,115 @@ +# 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. + +| 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/tests/test_large_cp_query_layout.py b/tests/test_large_cp_query_layout.py index d402ec2e..2407a1c5 100644 --- a/tests/test_large_cp_query_layout.py +++ b/tests/test_large_cp_query_layout.py @@ -4,6 +4,7 @@ from benchmarks.large_cp_emulation.query_layout import ( build_alternating_sequence_query_layout, build_contiguous_query_layout, + build_planned_query_layout, ) @@ -54,3 +55,16 @@ def test_alternating_sequence_layout_balances_causal_candidates_better_than_cont 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 + ) diff --git a/tests/test_thd_planner.py b/tests/test_thd_planner.py new file mode 100644 index 00000000..e222ef06 --- /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, 4], "at least one token 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 From bfe73dcc952b0f148cf207cabedffb0dabd46c1b Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Wed, 4 Nov 2026 14:05:17 +0800 Subject: [PATCH 10/18] fix(distributed): validate THD planner block capacity Reject layouts with fewer aligned blocks than CP ranks and require finite numeric cost weights before candidate construction. --- cula/distributed/thd_planner.py | 10 +++++----- tests/test_thd_planner.py | 2 +- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/cula/distributed/thd_planner.py b/cula/distributed/thd_planner.py index 4ea842ec..07877ac5 100644 --- a/cula/distributed/thd_planner.py +++ b/cula/distributed/thd_planner.py @@ -14,7 +14,7 @@ from collections import defaultdict from collections.abc import Sequence from dataclasses import dataclass -from math import ceil +from math import ceil, isfinite from operator import index import torch @@ -140,8 +140,6 @@ def _validate_inputs( 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") - if boundaries[-1] < world_size: - raise ValueError("packed context parallelism requires at least one token per rank") for name in ("block_size", "min_split_tokens", "max_chunks_per_rank"): raw_value = getattr(config, name) if isinstance(raw_value, bool): @@ -153,10 +151,12 @@ def _validate_inputs( 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 not isinstance(value, int | float) or value < 0: - raise ValueError(f"{name} must be non-negative") + 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 diff --git a/tests/test_thd_planner.py b/tests/test_thd_planner.py index e222ef06..963cd543 100644 --- a/tests/test_thd_planner.py +++ b/tests/test_thd_planner.py @@ -95,7 +95,7 @@ def test_cp64_long_sequence_uses_scaled_zigzag_chunks(): ("boundaries", "message"), [ ([0, 3, 8], "block aligned"), - ([0, 4], "at least one token per rank"), + ([0, 8], "at least one block per rank"), ], ) def test_planner_rejects_invalid_partitions(boundaries, message): From 36b5027ee54960cd121ddba1bd622f1550ad74c9 Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Wed, 11 Nov 2026 01:03:32 +0800 Subject: [PATCH 11/18] fix(benchmarks): use cuDNN concatenated KDA summaries --- benchmarks/bench_index_pool_cp.py | 39 +++---------- benchmarks/large_cp_emulation/README.md | 4 +- .../large_cp_emulation/bench_dsa_layout.py | 26 +++------ benchmarks/large_cp_emulation/bench_hybrid.py | 58 ++++++------------- benchmarks/large_cp_emulation/bench_kda.py | 21 +++---- benchmarks/large_cp_emulation/query_layout.py | 20 ++----- benchmarks/large_cp_emulation/transport.py | 34 ++++------- benchmarks/large_cp_emulation/workload.py | 4 +- cula/distributed/thd_planner.py | 21 ++----- cula/hybrid/distributed_pool.py | 22 ++----- cula/hybrid/dsa.py | 18 ++---- cula/hybrid/index_pool.py | 11 +--- cula/hybrid/plan.py | 13 +---- tests/test_large_cp_query_layout.py | 30 ++-------- 14 files changed, 89 insertions(+), 232 deletions(-) diff --git a/benchmarks/bench_index_pool_cp.py b/benchmarks/bench_index_pool_cp.py index b8e83346..fb240cd4 100644 --- a/benchmarks/bench_index_pool_cp.py +++ b/benchmarks/bench_index_pool_cp.py @@ -118,20 +118,14 @@ def _local_backward_oracle(kv, score, ape, grad_blocks, plan): 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.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_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) - ) + 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)) @@ -190,15 +184,11 @@ def main() -> None: 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) - ) + 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 - ) + global_score = _GatherReferenceRows.apply(score_, dist.group.WORLD, token_offsets, rank) return cudnn_index_pool( global_kv, global_score, @@ -251,16 +241,9 @@ def candidate(kv_, score_, ape_): "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() - ), + "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() - ), + "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]), @@ -343,13 +326,9 @@ def candidate_forward_backward(): "global_tokens": args.global_tokens, "world_size": world_size, "ratio": args.ratio, - "sequence_lengths": [ - right - left for left, right in zip(boundaries, boundaries[1:]) - ], + "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() - ), + "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, diff --git a/benchmarks/large_cp_emulation/README.md b/benchmarks/large_cp_emulation/README.md index d02184f5..6f4d4c8d 100644 --- a/benchmarks/large_cp_emulation/README.md +++ b/benchmarks/large_cp_emulation/README.md @@ -20,8 +20,8 @@ cross-node congestion, or aggregate bandwidth. - `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 KDA merge precision or -DSA arithmetic. +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: diff --git a/benchmarks/large_cp_emulation/bench_dsa_layout.py b/benchmarks/large_cp_emulation/bench_dsa_layout.py index d5c7ea89..598edadb 100644 --- a/benchmarks/large_cp_emulation/bench_dsa_layout.py +++ b/benchmarks/large_cp_emulation/bench_dsa_layout.py @@ -80,9 +80,7 @@ def _layout_for_rank( 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) + block for span in layout.spans for block in range(span.token_start // block_size, span.token_end // block_size) ) @@ -91,9 +89,7 @@ def _selection_key_ids( 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]) + block for span in layout.spans for block in range(block_prefix[span.sequence], block_prefix[span.sequence + 1]) ) @@ -149,9 +145,7 @@ def _run(args: argparse.Namespace) -> None: ) 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] - ) + 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, @@ -253,8 +247,7 @@ def execute(measure: bool) -> dict[str, object]: 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) + max(records[owner]["samples"][sample]["total_ms"] for owner in range(world_size)) for sample in range(args.repeats) ] print( json.dumps( @@ -279,19 +272,14 @@ def execute(measure: bool) -> dict[str, object]: ) } | { - f"{stage}_median_ms": statistics.median( - sample[stage] for sample in record["samples"] - ) + 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"] + 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"] - ) + "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 diff --git a/benchmarks/large_cp_emulation/bench_hybrid.py b/benchmarks/large_cp_emulation/bench_hybrid.py index 867152e0..a9d4823b 100644 --- a/benchmarks/large_cp_emulation/bench_hybrid.py +++ b/benchmarks/large_cp_emulation/bench_hybrid.py @@ -104,7 +104,7 @@ def workspace_bytes(self) -> int: def clear_events(self) -> None: self.events.clear() - def boundary_state( + def boundary_summaries( self, h: torch.Tensor, transition: torch.Tensor, @@ -112,11 +112,11 @@ def boundary_state( plan: KimiCPPlan, group: dist.ProcessGroup | None, reverse: bool, - ) -> torch.Tensor: + ) -> 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_state( + result = self.transport.boundary_summaries( h, transition, plan=plan, @@ -209,9 +209,7 @@ def create( 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] - ) + 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( @@ -352,12 +350,8 @@ def _run(args: argparse.Namespace) -> None: 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 - ) + 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)) @@ -486,12 +480,10 @@ def execute(schedule: str, measure: bool) -> dict[str, object]: 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_merge) + begin.elapsed_time(after_comm) if stage == "comm" else after_comm.elapsed_time(after_prepare) ) - for direction, begin, after_comm, after_merge in kda_transport.events - for stage in ("comm", "merge") + 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() @@ -544,8 +536,7 @@ def execute(schedule: str, measure: bool) -> dict[str, object]: 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) + 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, @@ -556,9 +547,7 @@ def execute(schedule: str, measure: bool) -> dict[str, object]: { "physical_rank": r, "logical_rank": records[r]["logical_rank"], - "kda": statistics.median( - sample["kda_ms"] for sample in records[r]["samples"][schedule] - ), + "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] ), @@ -568,8 +557,7 @@ def execute(schedule: str, measure: bool) -> dict[str, object]: **( { f"kda_boundary_{stage}": statistics.median( - sample["kda_boundary_ms"][stage] - for sample in records[r]["samples"][schedule] + sample["kda_boundary_ms"][stage] for sample in records[r]["samples"][schedule] ) for stage in ("forward_ms", "reverse_ms", "total_ms") } @@ -578,8 +566,7 @@ def execute(schedule: str, measure: bool) -> dict[str, object]: ), **{ f"dsa_{stage}": statistics.median( - sample["dsa_ms"][stage] - for sample in records[r]["samples"][schedule] + sample["dsa_ms"][stage] for sample in records[r]["samples"][schedule] ) for stage in ("pool_ms", "comm_ms", "select_ms", "total_ms") }, @@ -610,13 +597,9 @@ def execute(schedule: str, measure: bool) -> dict[str, object]: "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_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" - ], + "kda_transport_workspace_gib": records[r]["kda_transport_workspace_gib"], "memory_by_schedule": records[r]["memory_by_schedule"], } for r in range(world_size) @@ -626,19 +609,14 @@ def execute(schedule: str, measure: bool) -> dict[str, object]: "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] + 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] + 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] + sample[key] == records[r]["samples"][schedules[0]][0][key] for key in ( "proxy_loss", "output_norm", diff --git a/benchmarks/large_cp_emulation/bench_kda.py b/benchmarks/large_cp_emulation/bench_kda.py index 2dcd262d..35bfe7c3 100644 --- a/benchmarks/large_cp_emulation/bench_kda.py +++ b/benchmarks/large_cp_emulation/bench_kda.py @@ -4,7 +4,7 @@ """Emulate a large logical Kimi CP group on fewer physical GPUs. The benchmark executes the complete cuLA KDA path: local summary, boundary -transport and IEEE merge, local forward, and distributed backward. It maps one +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. @@ -62,12 +62,8 @@ def _run(args) -> None: 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 - ) + 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)) @@ -105,17 +101,16 @@ def execute(measure: bool): torch.cuda.synchronize(device) transport_times = {} if measure: - for direction, begin, after_comm, after_merge in transport.events: + for direction, begin, after_comm, after_prepare in transport.events: transport_times[f"comm_{direction}_ms"] = begin.elapsed_time(after_comm) - transport_times[f"merge_{direction}_ms"] = after_comm.elapsed_time(after_merge) + 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), + "finite": bool(torch.isfinite(output).all()) and all(bool(torch.isfinite(tensor.grad).all()) for tensor in inputs), } dist.barrier(device_ids=[local_rank]) @@ -141,7 +136,9 @@ def execute(measure: bool): 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_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( diff --git a/benchmarks/large_cp_emulation/query_layout.py b/benchmarks/large_cp_emulation/query_layout.py index 476aabc7..6f0165da 100644 --- a/benchmarks/large_cp_emulation/query_layout.py +++ b/benchmarks/large_cp_emulation/query_layout.py @@ -62,9 +62,7 @@ def _candidate_prefix(position: int, block_size: int) -> int: 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 - ) + return _candidate_prefix(start + span.tokens, block_size) - _candidate_prefix(start, block_size) def _validate_layout_inputs( @@ -104,9 +102,7 @@ def build_contiguous_query_layout( if shard_start % block_size or shard_end % block_size: raise ValueError("contiguous CP shard boundaries must be block aligned") spans = [] - for sequence, (sequence_start, sequence_end) in enumerate( - zip(boundaries, boundaries[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: @@ -149,9 +145,7 @@ def build_alternating_sequence_query_layout( assigned_tokens = [0] * world_size long_index = 0 short_sequences = [] - for sequence, (sequence_start, sequence_end) in enumerate( - zip(boundaries, boundaries[1:]) - ): + 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)) @@ -162,17 +156,13 @@ def build_alternating_sequence_query_layout( 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) - ) + 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) - ) + 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))) diff --git a/benchmarks/large_cp_emulation/transport.py b/benchmarks/large_cp_emulation/transport.py index 5b2f235c..225bcb7d 100644 --- a/benchmarks/large_cp_emulation/transport.py +++ b/benchmarks/large_cp_emulation/transport.py @@ -10,14 +10,6 @@ import torch import torch.distributed as dist -from cula.distributed.kda import _load_cudnn_cp_ops - - -def _load_merge(): - # Keep the benchmark emulator on exactly the same export compatibility - # and fixed-IEEE path as the production CP runtime. - return _load_cudnn_cp_ops()[0] - def emulate_ring_hops( source: torch.Tensor, @@ -57,7 +49,6 @@ class LogicalCPRingTransport: physical_rank: int physical_world_size: int ring_hops: int | None = field(default=None, kw_only=True) - merge: object = field(default_factory=_load_merge) 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 @@ -80,9 +71,7 @@ def _prepare(self, h: torch.Tensor, transition: torch.Tensor) -> None: 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 - ) + self.transition_summaries = torch.empty((self.logical_cp - 1, *transition.shape), dtype=torch.float32, device=h.device) @property def workspace_bytes(self) -> int: @@ -110,7 +99,7 @@ def _ring(self) -> None: hops=hops, ) - def boundary_state( + def boundary_summaries( self, h: torch.Tensor, transition: torch.Tensor, @@ -118,7 +107,7 @@ def boundary_state( plan, group, reverse: bool, - ) -> torch.Tensor: + ) -> tuple[torch.Tensor, torch.Tensor]: del group self._prepare(h, transition) assert self.local_payload is not None @@ -127,20 +116,17 @@ def boundary_state( self.local_payload[..., h.shape[-2] :, :].copy_(transition) start = torch.cuda.Event(enable_timing=True) after_comm = torch.cuda.Event(enable_timing=True) - after_merge = 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 - if pieces == 0: - state = torch.zeros_like(h) - else: - h_all = self.h_summaries[:pieces] - transition_all = self.transition_summaries[:pieces] + 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)) - state = self.merge(h_all, transition_all, reverse=reverse)["final_state"][0] - after_merge.record() - self.events.append(("backward" if reverse else "forward", start, after_comm, after_merge)) - return state + 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 index c2776969..807009ad 100644 --- a/benchmarks/large_cp_emulation/workload.py +++ b/benchmarks/large_cp_emulation/workload.py @@ -24,9 +24,7 @@ def global_tokens(self) -> int: @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:]) - ) + 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: diff --git a/cula/distributed/thd_planner.py b/cula/distributed/thd_planner.py index 07877ac5..b7646fc0 100644 --- a/cula/distributed/thd_planner.py +++ b/cula/distributed/thd_planner.py @@ -190,9 +190,7 @@ def _contiguous_items(boundaries: tuple[int, ...], world_size: int, block_size: 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) - ) + result.append((_make_item(sequence, start, end, sequence_start, sequence_end, block_size), rank)) return result @@ -220,9 +218,7 @@ def _adaptive_items( 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) - ) + 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 @@ -287,8 +283,7 @@ def _greedy_assign(items: list[_PlanningItem], world_size: int) -> list[tuple[_P (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, + (token_loads[rank] + item.tokens) / ideal_tokens + (causal_loads[rank] + item.causal_work) / ideal_causal, item_counts[rank], rank, ), @@ -339,10 +334,7 @@ def _rebalance_tokens( 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:]) - ] + 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: @@ -555,10 +547,7 @@ def plan_thd_layout( ) if adaptive is None: continue - signature = tuple( - (span.sequence, span.token_start, span.token_end, span.owner) - for span in adaptive[1] - ) + 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) diff --git a/cula/hybrid/distributed_pool.py b/cula/hybrid/distributed_pool.py index ea84920c..fd6d91fe 100644 --- a/cula/hybrid/distributed_pool.py +++ b/cula/hybrid/distributed_pool.py @@ -59,9 +59,7 @@ def _all_gather_block_rows( 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] - ) + output[block_offsets[owner] : block_offsets[owner + 1]].copy_(gathered[owner, :count]) return output @@ -149,9 +147,7 @@ def backward(ctx, grad_output: torch.Tensor): # 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:]) - ) + received = grad_output.new_empty((plan.world_size * local_count, *grad_output.shape[1:])) dist.all_to_all_single( received, grad_output.contiguous(), @@ -163,9 +159,7 @@ def backward(ctx, grad_output: torch.Tensor): 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 - ) + 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, @@ -241,8 +235,7 @@ def cudnn_index_pool_cp( 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" + "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: @@ -252,15 +245,12 @@ def cudnn_index_pool_cp( ) 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" + "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]}" - ) + raise ValueError(f"kv must contain this rank's {plan.local_tokens} token rows, got {kv.shape[0]}") _validate_inputs( kv, diff --git a/cula/hybrid/dsa.py b/cula/hybrid/dsa.py index d9e2ad28..f27f5a8c 100644 --- a/cula/hybrid/dsa.py +++ b/cula/hybrid/dsa.py @@ -30,9 +30,7 @@ _MAX_LOCAL_ROWS_MULTIPLIER = 4 -def _as_nonnegative( - values: torch.Tensor | Sequence[int], *, entries: int, name: str -) -> tuple[int, ...]: +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") @@ -92,13 +90,10 @@ def _candidate_floats(position: int, rows: int, block_size: int) -> int: 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) - ) + 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, + (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 @@ -111,8 +106,7 @@ def _balanced_candidate_budget(cu_q: tuple[int, ...], offsets: tuple[int, ...]) _MAX_LOCAL_ROWS_MULTIPLIER, max( 1, - (total_local_rows + _CANDIDATE_BUDGET_LOCAL_ROWS_QUANTUM - 1) - // _CANDIDATE_BUDGET_LOCAL_ROWS_QUANTUM, + (total_local_rows + _CANDIDATE_BUDGET_LOCAL_ROWS_QUANTUM - 1) // _CANDIDATE_BUDGET_LOCAL_ROWS_QUANTUM, ), ) return min( @@ -415,9 +409,7 @@ def _validate_selection_inputs( 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)}" - ) + 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: diff --git a/cula/hybrid/index_pool.py b/cula/hybrid/index_pool.py index 310c39fb..d7ca3164 100644 --- a/cula/hybrid/index_pool.py +++ b/cula/hybrid/index_pool.py @@ -23,8 +23,7 @@ def _load_cudnn_index_pool_ops(): 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]" + "cuDNN index pooling requires cuDNN Frontend with the CSA compressor APIs; install cuda-linear-attention[cudnn-cp]" ) from error return forward, backward @@ -59,15 +58,11 @@ def _validate_inputs( 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)}" - ) + 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)}" - ) + 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: diff --git a/cula/hybrid/plan.py b/cula/hybrid/plan.py index d333583c..4612092b 100644 --- a/cula/hybrid/plan.py +++ b/cula/hybrid/plan.py @@ -102,9 +102,7 @@ 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, ...]: +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) @@ -195,9 +193,7 @@ def global_blocks(self) -> int: @property def block_counts(self) -> tuple[int, ...]: - return tuple( - right - left for left, right in zip(self.block_offsets, self.block_offsets[1:]) - ) + return tuple(right - left for left, right in zip(self.block_offsets, self.block_offsets[1:])) @property def equal_block_counts(self) -> bool: @@ -209,10 +205,7 @@ def validate_group(self, group: dist.ProcessGroup | None) -> None: 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" - ) + 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) diff --git a/tests/test_large_cp_query_layout.py b/tests/test_large_cp_query_layout.py index 2407a1c5..a7e637ce 100644 --- a/tests/test_large_cp_query_layout.py +++ b/tests/test_large_cp_query_layout.py @@ -18,21 +18,14 @@ def _workload_boundaries(): 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) - ] + 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 - ) + 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 @@ -42,14 +35,8 @@ def test_alternating_sequence_layout_keeps_short_samples_whole_and_tokens_balanc 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 = [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] @@ -59,12 +46,7 @@ def test_alternating_sequence_layout_balances_causal_candidates_better_than_cont 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) - ] + 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 - ) + assert max(layout.candidate_floats for layout in layouts) == min(layout.candidate_floats for layout in layouts) From bccce11fd64d8bdc4b02177b07e48543ca7de238 Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Wed, 11 Nov 2026 01:04:32 +0800 Subject: [PATCH 12/18] docs(kda): retire standalone merge benchmark results --- docs/KIMI_CP_RUNTIME.md | 46 ++++++++++++++--------------------------- 1 file changed, 16 insertions(+), 30 deletions(-) diff --git a/docs/KIMI_CP_RUNTIME.md b/docs/KIMI_CP_RUNTIME.md index 4f42f4bf..c5284b92 100644 --- a/docs/KIMI_CP_RUNTIME.md +++ b/docs/KIMI_CP_RUNTIME.md @@ -99,36 +99,22 @@ 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 summary, IEEE merge, KDA forward, backward summary, reverse merge, and -local autograd are real cuDNN operations. No DSA stage is included. - -Shape: 1,048,576 global tokens, logical CP64, `H=32`, `K=V=128`. One complete -warmup was excluded, followed by three steady-state samples. - -| E2E stage | Critical-stage median | -|---|---:| -| Forward summary | 0.903 ms | -| Forward 63-hop communication | 3.008 ms | -| Forward FP32 merge | 0.293 ms | -| Local KDA forward | 1.054 ms | -| Backward summary | 1.083 ms | -| Backward 63-hop communication | 3.216 ms | -| Backward FP32 merge | 0.403 ms | -| Local KDA backward | 3.261 ms | - -The measured critical-path median was **12.953 ms GPU time** and **16.102 ms -wall time**. Peak allocated memory was 3.049 GiB per physical process; peak -reserved memory was 3.15–3.27 GiB. Outputs and all five local input gradients -were finite on every representative rank. Proxy loss, output norm, and gradient -norms were identical across all three repeats on each rank. - -Communication accounts for roughly 6.22 ms of the stage envelope and is the -largest optimization target. These numbers 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`. The logical topology -and benchmark-only transport live beside it rather than in the production CP -runtime. +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 From 181fc8566a93d6bd07ec692352539772d7b8b106 Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Wed, 11 Nov 2026 01:28:27 +0800 Subject: [PATCH 13/18] fix(benchmarks): preserve block-aligned CP emulation --- .../large_cp_emulation/bench_dsa_layout.py | 6 ++++ benchmarks/large_cp_emulation/bench_hybrid.py | 4 ++- benchmarks/large_cp_emulation/query_layout.py | 6 ++-- docs/KDA_CSA_1M_OPTIMIZATION.md | 5 +-- tests/test_hybrid_dsa_cudnn.py | 2 ++ tests/test_large_cp_emulation.py | 36 +++++++++++++++++++ tests/test_large_cp_query_layout.py | 14 ++++++++ 7 files changed, 66 insertions(+), 7 deletions(-) diff --git a/benchmarks/large_cp_emulation/bench_dsa_layout.py b/benchmarks/large_cp_emulation/bench_dsa_layout.py index 598edadb..afc4fdff 100644 --- a/benchmarks/large_cp_emulation/bench_dsa_layout.py +++ b/benchmarks/large_cp_emulation/bench_dsa_layout.py @@ -93,6 +93,11 @@ def _selection_key_ids( ) +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) @@ -114,6 +119,7 @@ def _run(args: argparse.Namespace) -> None: 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, diff --git a/benchmarks/large_cp_emulation/bench_hybrid.py b/benchmarks/large_cp_emulation/bench_hybrid.py index a9d4823b..9b689c27 100644 --- a/benchmarks/large_cp_emulation/bench_hybrid.py +++ b/benchmarks/large_cp_emulation/bench_hybrid.py @@ -265,6 +265,9 @@ def _communicate(self, local_kbar: torch.Tensor) -> None: 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, @@ -272,7 +275,6 @@ def _communicate(self, local_kbar: torch.Tensor) -> None: physical_world_size=self.physical_world_size, hops=self.transport_hops, ) - self.global_kbar[self.global_block_start : self.global_block_start + self.local_blocks].copy_(local_kbar) def execute(self) -> _DsaEvents: events = _DsaEvents(*(torch.cuda.Event(enable_timing=True) for _ in range(4))) diff --git a/benchmarks/large_cp_emulation/query_layout.py b/benchmarks/large_cp_emulation/query_layout.py index 6f0165da..92292d45 100644 --- a/benchmarks/large_cp_emulation/query_layout.py +++ b/benchmarks/large_cp_emulation/query_layout.py @@ -97,10 +97,8 @@ def build_contiguous_query_layout( world_size=world_size, block_size=block_size, ) - shard_offsets = balanced_shards(boundaries[-1], world_size) - shard_start, shard_end = shard_offsets[rank : rank + 2] - if shard_start % block_size or shard_end % block_size: - raise ValueError("contiguous CP shard boundaries must be block aligned") + 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) diff --git a/docs/KDA_CSA_1M_OPTIMIZATION.md b/docs/KDA_CSA_1M_OPTIMIZATION.md index 27773451..b113f4c0 100644 --- a/docs/KDA_CSA_1M_OPTIMIZATION.md +++ b/docs/KDA_CSA_1M_OPTIMIZATION.md @@ -75,5 +75,6 @@ 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 benchmark and CPU planner live in `benchmarks/large_cp_emulation` to keep -this not-yet-production policy separate from the contiguous runtime. +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/tests/test_hybrid_dsa_cudnn.py b/tests/test_hybrid_dsa_cudnn.py index a3bd80c4..abe7376b 100644 --- a/tests/test_hybrid_dsa_cudnn.py +++ b/tests/test_hybrid_dsa_cudnn.py @@ -40,6 +40,8 @@ def run(): 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() diff --git a/tests/test_large_cp_emulation.py b/tests/test_large_cp_emulation.py index ca762592..656d9b3c 100644 --- a/tests/test_large_cp_emulation.py +++ b/tests/test_large_cp_emulation.py @@ -1,11 +1,16 @@ # 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(): @@ -35,3 +40,34 @@ def test_emulation_plan_separates_logical_and_physical_identity(): 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 index a7e637ce..622dca55 100644 --- a/tests/test_large_cp_query_layout.py +++ b/tests/test_large_cp_query_layout.py @@ -50,3 +50,17 @@ def test_planned_layout_balances_reference_workload_without_special_case(): 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:])) From 9283528936875529ad3405c743cb2b5202697ea9 Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Wed, 11 Nov 2026 01:32:49 +0800 Subject: [PATCH 14/18] fix(hybrid): default DSA selection to nondeterministic --- cula/hybrid/dsa.py | 12 +++---- docs/HYBRID_CUDNN_RUNTIME.md | 11 ++++--- tests/test_hybrid_dsa_cudnn.py | 2 +- tests/test_hybrid_dsa_helpers.py | 55 ++++++++++++++++++++++++++++++++ 4 files changed, 68 insertions(+), 12 deletions(-) diff --git a/cula/hybrid/dsa.py b/cula/hybrid/dsa.py index f27f5a8c..f3f0cf74 100644 --- a/cula/hybrid/dsa.py +++ b/cula/hybrid/dsa.py @@ -151,9 +151,8 @@ class DsaSelectionPlan: 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; deterministic tie-breaking remains enabled - by default at execution. Whole adjacent sequences are coalesced into one - launch when they fit the same budget. + precision nor the selected set. Whole adjacent sequences are coalesced into + one launch when they fit the same budget. """ block_size: int @@ -462,15 +461,16 @@ def cudnn_dsa_select( plan: DsaSelectionPlan, workspace: DsaSelectionWorkspace, sm_scale: float = 1.0, - deterministic: bool = True, + 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. ``deterministic=True`` makes - the selected set reproducible; cuDNN's radix output slot order remains + 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. """ diff --git a/docs/HYBRID_CUDNN_RUNTIME.md b/docs/HYBRID_CUDNN_RUNTIME.md index 0baf2542..1c479461 100644 --- a/docs/HYBRID_CUDNN_RUNTIME.md +++ b/docs/HYBRID_CUDNN_RUNTIME.md @@ -30,11 +30,12 @@ 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; -- deterministic Top-K is the default, so exact ties at the K-th boundary keep - the smaller block id; -- deterministic selection guarantees the `(block_id, logit)` set, not radix - slot order. CUDA Graph capture/replay preserved the sorted pairs bitwise in - validation, while the raw slot order changed; +- 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. diff --git a/tests/test_hybrid_dsa_cudnn.py b/tests/test_hybrid_dsa_cudnn.py index abe7376b..08b5f099 100644 --- a/tests/test_hybrid_dsa_cudnn.py +++ b/tests/test_hybrid_dsa_cudnn.py @@ -29,7 +29,7 @@ def test_windowed_dsa_selection_is_cuda_graph_set_deterministic(): weights = torch.randn(tokens, 4, dtype=torch.float32, device=device) def run(): - return cudnn_dsa_select(qi, kbar, weights, plan=plan, workspace=workspace) + return cudnn_dsa_select(qi, kbar, weights, plan=plan, workspace=workspace, deterministic=True) run() torch.cuda.synchronize() diff --git a/tests/test_hybrid_dsa_helpers.py b/tests/test_hybrid_dsa_helpers.py index bd348269..e3dce898 100644 --- a/tests/test_hybrid_dsa_helpers.py +++ b/tests/test_hybrid_dsa_helpers.py @@ -3,6 +3,8 @@ from types import SimpleNamespace +import torch + from cula.hybrid import dsa @@ -15,3 +17,56 @@ def test_cudnn_dsa_loader_prefers_public_namespace(monkeypatch): 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] From 4a9c1a1dbc84b77c026e0d795d45b14b1b543386 Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Wed, 11 Nov 2026 13:23:05 +0800 Subject: [PATCH 15/18] feat(distributed): support Kimi CP on SM90 Select a Hopper state-probe backend that derives FP32 affine boundary summaries from cuDNN recurrent-state ports, propagates reverse state gradients, and preserves the Blackwell native-summary path. Validate canonical Hopper inputs and IEEE FP32 composition, cover the architecture-neutral distributed path, and align SM90 indexer benchmark defaults. --- .../large_cp_emulation/bench_components.py | 2 +- .../large_cp_emulation/bench_dsa_layout.py | 2 +- benchmarks/large_cp_emulation/bench_hybrid.py | 2 +- cula/distributed/kda.py | 305 ++++++++++++++++-- docs/KIMI_CP_RUNTIME.md | 48 ++- tests/test_kimi_cp_distributed.py | 13 +- tests/test_kimi_cp_helpers.py | 95 +++++- 7 files changed, 419 insertions(+), 48 deletions(-) diff --git a/benchmarks/large_cp_emulation/bench_components.py b/benchmarks/large_cp_emulation/bench_components.py index e85d1083..a4b47af3 100644 --- a/benchmarks/large_cp_emulation/bench_components.py +++ b/benchmarks/large_cp_emulation/bench_components.py @@ -196,7 +196,7 @@ def parse_args(): 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=4) + 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) diff --git a/benchmarks/large_cp_emulation/bench_dsa_layout.py b/benchmarks/large_cp_emulation/bench_dsa_layout.py index afc4fdff..5b5e09dc 100644 --- a/benchmarks/large_cp_emulation/bench_dsa_layout.py +++ b/benchmarks/large_cp_emulation/bench_dsa_layout.py @@ -311,7 +311,7 @@ def _parse_args() -> argparse.Namespace: 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=4) + 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) diff --git a/benchmarks/large_cp_emulation/bench_hybrid.py b/benchmarks/large_cp_emulation/bench_hybrid.py index 9b689c27..3f8baf1d 100644 --- a/benchmarks/large_cp_emulation/bench_hybrid.py +++ b/benchmarks/large_cp_emulation/bench_hybrid.py @@ -676,7 +676,7 @@ def _parse_args() -> argparse.Namespace: 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=4) + 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) diff --git a/cula/distributed/kda.py b/cula/distributed/kda.py index 8efbf07d..a0688295 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]" + ) 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,7 +52,6 @@ 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; " @@ -107,18 +118,6 @@ def _all_gather_into( dist.all_gather_into_tensor(output, tensor, group=group) -def _all_gather_into( - output: torch.Tensor, - tensor: torch.Tensor, - *, - group: dist.ProcessGroup | None, -) -> None: - if hasattr(dist, "all_gather_single"): - dist.all_gather_single(output, tensor, group=group) - else: # pragma: no cover - compatibility with older PyTorch releases - dist.all_gather_into_tensor(output, tensor, group=group) - - def _all_gather(tensor: torch.Tensor, group: dist.ProcessGroup | None, world_size: int) -> torch.Tensor: """Gather one equal-shaped boundary payload per rank.""" if world_size == 1: @@ -380,6 +379,235 @@ 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") + if state_updates.is_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" + ) + 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) -> 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 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 _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( + "SM90 Kimi CP state probing requires value_dim >= key_dim; " + f"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.attention_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, initial_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, initial_state, transition_first, state_checkpoints = ctx.saved_tensors + plan = ctx.plan + options = ctx.options + 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( @@ -604,9 +832,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. @@ -636,6 +865,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") @@ -667,7 +898,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, @@ -682,6 +913,30 @@ def cudnn_kda_cp( **options.attention_kwargs(a_log, dt_bias), ) return attention_outputs[0] + resolved_backend = backend + if resolved_backend == "auto": + major, _ = torch.cuda.get_device_capability(q.device) + if major == 9: + resolved_backend = "hopper-state-probe" + elif major >= 10: + resolved_backend = "frost-summary" + else: + raise RuntimeError(f"cuDNN Kimi CP does not support CUDA compute capability {major}.x") + if resolved_backend == "hopper-state-probe": + if a_log is not None or dt_bias is not None: + raise ValueError("SM90 Kimi CP requires safe-gate preprocessing in the framework adapter") + 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/docs/KIMI_CP_RUNTIME.md b/docs/KIMI_CP_RUNTIME.md index c5284b92..72a34cb1 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,6 +86,33 @@ 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. +### 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. + ## Current transport and rejected chain experiment The first implementation uses two typed `all_gather_into_tensor` calls per diff --git a/tests/test_kimi_cp_distributed.py b/tests/test_kimi_cp_distributed.py index 670e003f..56b74dd0 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,11 @@ 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") def stage(name): print(f"rank={local_rank} stage={name}", flush=True) @@ -77,7 +74,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..40d093e6 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,77 @@ 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_hopper_state_probe_recovers_fp32_affine_map(monkeypatch): + heads, dim = 2, 4 + state_update = torch.arange(heads * dim * dim, dtype=torch.float32).reshape(1, heads, dim, dim) / 32 + transition = torch.eye(dim, dtype=torch.float32).expand(heads, dim, dim).clone() + transition[:, 0, 1] = 0.25 + + def attention(q, k, v, g, beta, cu, *, initial_state, **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, dim) + beta = torch.zeros(8, heads) + cu = torch.tensor([0, 8], dtype=torch.int32) + + actual_h, actual_transition = cp._state_probe_summary(q, q, q, 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) + + +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}") + + def _transport_plan(*, rank=2, pre=2, post=1): return KimiCPPlan( rank=rank, @@ -462,7 +551,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) From e09dfaa297ce33e9ba401f8bad0fea6c2805f2f1 Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Wed, 11 Nov 2026 13:55:47 +0800 Subject: [PATCH 16/18] fix(distributed): harden SM90 Kimi CP integration Avoid checkpoint retention in boundary state probes, record index-head shapes in benchmark output, document Hopper state-memory scaling, and cover backend dispatch and wider value states. Apply repository formatting so the stacked PR passes Ruff CI. --- .../large_cp_emulation/bench_components.py | 2 + .../large_cp_emulation/bench_dsa_layout.py | 1 + benchmarks/large_cp_emulation/bench_hybrid.py | 1 + cula/distributed/kda.py | 59 ++++++++++--------- docs/KDA_CSA_1M_OPTIMIZATION.md | 1 + docs/KIMI_CP_RUNTIME.md | 9 +++ docs/THD_LAYOUT_PLANNER.md | 2 + tests/test_kimi_cp_distributed.py | 8 +++ tests/test_kimi_cp_helpers.py | 50 ++++++++++++++-- 9 files changed, 100 insertions(+), 33 deletions(-) diff --git a/benchmarks/large_cp_emulation/bench_components.py b/benchmarks/large_cp_emulation/bench_components.py index a4b47af3..9671b2a2 100644 --- a/benchmarks/large_cp_emulation/bench_components.py +++ b/benchmarks/large_cp_emulation/bench_components.py @@ -91,6 +91,7 @@ def run_dsa(args) -> None: "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, } @@ -117,6 +118,7 @@ def run_dsa(args) -> None: 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, diff --git a/benchmarks/large_cp_emulation/bench_dsa_layout.py b/benchmarks/large_cp_emulation/bench_dsa_layout.py index 5b5e09dc..d9f8b72a 100644 --- a/benchmarks/large_cp_emulation/bench_dsa_layout.py +++ b/benchmarks/large_cp_emulation/bench_dsa_layout.py @@ -260,6 +260,7 @@ def execute(measure: bool) -> dict[str, object]: { "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, diff --git a/benchmarks/large_cp_emulation/bench_hybrid.py b/benchmarks/large_cp_emulation/bench_hybrid.py index 3f8baf1d..a9fa45dc 100644 --- a/benchmarks/large_cp_emulation/bench_hybrid.py +++ b/benchmarks/large_cp_emulation/bench_hybrid.py @@ -591,6 +591,7 @@ def execute(schedule: str, measure: bool) -> dict[str, object]: "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, diff --git a/cula/distributed/kda.py b/cula/distributed/kda.py index a0688295..4633a6f1 100644 --- a/cula/distributed/kda.py +++ b/cula/distributed/kda.py @@ -31,8 +31,7 @@ def _load_cudnn_kda_ops(): 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]" + "cuDNN Kimi CP requires cuDNN Frontend KDA state APIs; install cuda-linear-attention[cudnn-cp]" ) from error return kimi_delta_attention, kimi_delta_attention_bwd @@ -396,8 +395,7 @@ def _compose_affine_summaries( raise ValueError("ordered KDA summaries must be on the same device") if state_updates.is_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" + "Kimi CP boundary composition requires IEEE FP32; set torch.backends.cuda.matmul.allow_tf32 = False" ) 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]) @@ -406,7 +404,11 @@ def _compose_affine_summaries( return state.contiguous() -def _validate_hopper_state_probe_options(options: _KdaOptions) -> None: +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: @@ -415,6 +417,10 @@ def _validate_hopper_state_probe_options(options: _KdaOptions) -> None: 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( @@ -446,14 +452,13 @@ def _state_probe_summary( value_dim = v.shape[-1] if value_dim < key_dim: raise ValueError( - "SM90 Kimi CP state probing requires value_dim >= key_dim; " - f"got value_dim={value_dim}, key_dim={key_dim}" + 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.attention_kwargs(None, None) + kwargs = options.kwargs(None, None) def final_state(initial_state: torch.Tensor) -> torch.Tensor: outputs = kimi_delta_attention( @@ -504,9 +509,7 @@ def forward( 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 - ), + boundary_transport.boundary_summaries(h_last, transition_last, plan=plan, group=group, reverse=False), h_last, transition_last, rows=plan.pre_num_ranks, @@ -542,9 +545,7 @@ def forward( ctx.group = group ctx.options = options ctx.boundary_transport = boundary_transport - ctx.save_for_backward( - q, k, v, g, beta, cu_seqlens, initial_state, transition_first, state_checkpoints - ) + ctx.save_for_backward(q, k, v, g, beta, cu_seqlens, initial_state, transition_first, state_checkpoints) return output @staticmethod @@ -577,9 +578,7 @@ def backward(ctx, grad_output: torch.Tensor): 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 - ), + 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, @@ -811,6 +810,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, @@ -913,18 +925,9 @@ def cudnn_kda_cp( **options.attention_kwargs(a_log, dt_bias), ) return attention_outputs[0] - resolved_backend = backend - if resolved_backend == "auto": - major, _ = torch.cuda.get_device_capability(q.device) - if major == 9: - resolved_backend = "hopper-state-probe" - elif major >= 10: - resolved_backend = "frost-summary" - else: - raise RuntimeError(f"cuDNN Kimi CP does not support CUDA compute capability {major}.x") + resolved_backend = _resolve_kimi_cp_backend(backend, q.device) if resolved_backend == "hopper-state-probe": - if a_log is not None or dt_bias is not None: - raise ValueError("SM90 Kimi CP requires safe-gate preprocessing in the framework adapter") + _validate_hopper_state_probe_options(options, a_log, dt_bias) return _CudnnKdaCPHopper.apply( q.contiguous(), k.contiguous(), diff --git a/docs/KDA_CSA_1M_OPTIMIZATION.md b/docs/KDA_CSA_1M_OPTIMIZATION.md index b113f4c0..2fe14e7c 100644 --- a/docs/KDA_CSA_1M_OPTIMIZATION.md +++ b/docs/KDA_CSA_1M_OPTIMIZATION.md @@ -60,6 +60,7 @@ 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 | |---|---:|---:| diff --git a/docs/KIMI_CP_RUNTIME.md b/docs/KIMI_CP_RUNTIME.md index 72a34cb1..87b6d161 100644 --- a/docs/KIMI_CP_RUNTIME.md +++ b/docs/KIMI_CP_RUNTIME.md @@ -113,6 +113,15 @@ 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 stores `initial_state`, `d_final_state`, and +the returned state gradient with shape `[N, H, V, K]`. Each live state tensor +costs `4 * N * H * V * K` bytes (2 MiB per sequence for `H=32`, `K=V=128`). +Packed workloads with many local short samples must include this linear 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 diff --git a/docs/THD_LAYOUT_PLANNER.md b/docs/THD_LAYOUT_PLANNER.md index ca6a4ef7..bb46035c 100644 --- a/docs/THD_LAYOUT_PLANNER.md +++ b/docs/THD_LAYOUT_PLANNER.md @@ -79,6 +79,8 @@ 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 | |---|---:|---:|---:|---:| diff --git a/tests/test_kimi_cp_distributed.py b/tests/test_kimi_cp_distributed.py index 56b74dd0..72d8102e 100644 --- a/tests/test_kimi_cp_distributed.py +++ b/tests/test_kimi_cp_distributed.py @@ -31,6 +31,14 @@ def test_cudnn_kda_cp_matches_monolithic_forward_backward(): 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) diff --git a/tests/test_kimi_cp_helpers.py b/tests/test_kimi_cp_helpers.py index 40d093e6..f4054b82 100644 --- a/tests/test_kimi_cp_helpers.py +++ b/tests/test_kimi_cp_helpers.py @@ -196,12 +196,14 @@ def test_affine_summary_composition_uses_causal_order_in_both_directions(): def test_hopper_state_probe_recovers_fp32_affine_map(monkeypatch): - heads, dim = 2, 4 - state_update = torch.arange(heads * dim * dim, dtype=torch.float32).reshape(1, heads, dim, dim) / 32 - transition = torch.eye(dim, dtype=torch.float32).expand(heads, dim, dim).clone() + 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) @@ -219,15 +221,17 @@ def attention(q, k, v, g, beta, cu, *, initial_state, **kwargs): plan_name=None, summary_plan_name=None, ) - q = torch.zeros(8, heads, dim) + 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, q, q, beta, cu, options=options) + 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(): @@ -254,6 +258,42 @@ def test_hopper_state_probe_rejects_in_kernel_backward_transforms(): 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( From cfd9051a43b7c54e65b0b299d2fbdb805fd45a25 Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Wed, 11 Nov 2026 14:08:28 +0800 Subject: [PATCH 17/18] fix(distributed): reduce retained Hopper state Retain only the incoming boundary row across forward and backward, rebuild the cuDNN-required dense state at backward entry, and validate IEEE composition before collectives. Improve compatible cuDNN build guidance while keeping global TF32 settings untouched. --- cula/distributed/kda.py | 24 ++++++++++++++++-------- docs/KIMI_CP_RUNTIME.md | 15 +++++++++------ tests/test_kimi_cp_helpers.py | 12 ++++++++++++ 3 files changed, 37 insertions(+), 14 deletions(-) diff --git a/cula/distributed/kda.py b/cula/distributed/kda.py index 4633a6f1..916583c5 100644 --- a/cula/distributed/kda.py +++ b/cula/distributed/kda.py @@ -31,7 +31,8 @@ def _load_cudnn_kda_ops(): 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]" + "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 @@ -54,7 +55,7 @@ def _load_cudnn_cp_ops(): 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, @@ -393,10 +394,6 @@ def _compose_affine_summaries( 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") - if state_updates.is_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" - ) 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: @@ -429,6 +426,13 @@ def _validate_hopper_state_probe_options( ) +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, @@ -545,16 +549,19 @@ def forward( ctx.group = group ctx.options = options ctx.boundary_transport = boundary_transport - ctx.save_for_backward(q, k, v, g, beta, cu_seqlens, initial_state, transition_first, state_checkpoints) + 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, initial_state, transition_first, state_checkpoints = ctx.saved_tensors + 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 @@ -928,6 +935,7 @@ def cudnn_kda_cp( 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(), diff --git a/docs/KIMI_CP_RUNTIME.md b/docs/KIMI_CP_RUNTIME.md index 87b6d161..0f3f795c 100644 --- a/docs/KIMI_CP_RUNTIME.md +++ b/docs/KIMI_CP_RUNTIME.md @@ -115,12 +115,15 @@ 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 stores `initial_state`, `d_final_state`, and -the returned state gradient with shape `[N, H, V, K]`. Each live state tensor -costs `4 * N * H * V * K` bytes (2 MiB per sequence for `H=32`, `K=V=128`). -Packed workloads with many local short samples must include this linear state -storage in their memory plan. A future compact backward state port in cuDNN -could remove it without changing cuLA's public distributed interface. +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 diff --git a/tests/test_kimi_cp_helpers.py b/tests/test_kimi_cp_helpers.py index f4054b82..142986fd 100644 --- a/tests/test_kimi_cp_helpers.py +++ b/tests/test_kimi_cp_helpers.py @@ -195,6 +195,18 @@ def test_affine_summary_composition_uses_causal_order_in_both_directions(): assert reverse.item() == 22.0 # (0 * 7 + 4) * 5 + 2 +def test_ieee_composition_validation_precedes_cuda_collectives(monkeypatch): + monkeypatch.setattr(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") + + 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 From d6e2f8c5ee31631c1ba85524cf052b5dfe70dad5 Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Wed, 11 Nov 2026 14:09:43 +0800 Subject: [PATCH 18/18] test(distributed): restore dynamic TF32 backend state Save and restore the PyTorch cuBLAS TF32 property explicitly so the review regression works with dynamic backend attributes on the SM90 validation environment. --- tests/test_kimi_cp_helpers.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/tests/test_kimi_cp_helpers.py b/tests/test_kimi_cp_helpers.py index 142986fd..932f48e4 100644 --- a/tests/test_kimi_cp_helpers.py +++ b/tests/test_kimi_cp_helpers.py @@ -195,14 +195,18 @@ def test_affine_summary_composition_uses_causal_order_in_both_directions(): assert reverse.item() == 22.0 # (0 * 7 + 4) * 5 + 2 -def test_ieee_composition_validation_precedes_cuda_collectives(monkeypatch): - monkeypatch.setattr(cp.torch.backends.cuda.matmul, "allow_tf32", True) +def test_ieee_composition_validation_precedes_cuda_collectives(): + original = cp.torch.backends.cuda.matmul.allow_tf32 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") + 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"))