From d7f30d71248042bb23a204044319ffbbac5c9085 Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Wed, 11 Nov 2026 01:01:10 +0800 Subject: [PATCH 1/2] feat(kda): add cuDNN concatenated-summary Kimi CP --- REPO_LAYOUT.md | 3 + USAGE.md | 44 ++ cula/distributed/__init__.py | 19 + cula/distributed/kda.py | 693 ++++++++++++++++++++++++++++++ cula/distributed/plan.py | 218 ++++++++++ cula/kda/__init__.py | 10 + docs/KIMI_CP_RUNTIME.md | 154 +++++++ pyproject.toml | 4 + tests/conftest.py | 1 + tests/test_kimi_cp_distributed.py | 109 +++++ tests/test_kimi_cp_helpers.py | 470 ++++++++++++++++++++ tests/test_kimi_cp_plan.py | 105 +++++ 12 files changed, 1830 insertions(+) create mode 100644 cula/distributed/__init__.py create mode 100644 cula/distributed/kda.py create mode 100644 cula/distributed/plan.py create mode 100644 docs/KIMI_CP_RUNTIME.md create mode 100644 tests/test_kimi_cp_distributed.py create mode 100644 tests/test_kimi_cp_helpers.py create mode 100644 tests/test_kimi_cp_plan.py diff --git a/REPO_LAYOUT.md b/REPO_LAYOUT.md index e45517a3..d9df540e 100644 --- a/REPO_LAYOUT.md +++ b/REPO_LAYOUT.md @@ -18,6 +18,9 @@ cuLA/ │ │ ├── prefill.py # Public FlashKDA prefill validation and routing │ │ ├── flashkda.py # SM90 CuTeDSL K1+K2 wrapper, autograd shell, and CP selection │ │ └── chunk_intra.py # Pure-CuTeDSL SM100 intra/WU composition +│ ├── distributed/ # Optional cuDNN KDA Kimi CP runtime +│ │ ├── plan.py # CPU-only packed-THD boundary planning +│ │ └── kda.py # Collectives and distributed autograd │ │ │ ├── lightning/ # [non-KDA] Lightning Attention public wrappers │ │ diff --git a/USAGE.md b/USAGE.md index e3d40dbd..abf03ea6 100644 --- a/USAGE.md +++ b/USAGE.md @@ -4,6 +4,50 @@ This document provides usage notes and examples for cuLA kernels. --- +## cuDNN KDA Kimi context parallelism + +The optional distributed runtime executes cuDNN KDA over contiguous packed-THD +shards. Build one CPU plan per packed microbatch and reuse it across layers: + +```python +from cula.distributed import build_kimi_cp_plan, cudnn_kda_cp + +cp_plan = build_kimi_cp_plan( + global_cu_seqlens_cpu, + group=cp_group, +) +local_cu_seqlens = torch.tensor( + cp_plan.local_cu_seqlens, + dtype=torch.int32, + device=q.device, +) + +output = cudnn_kda_cp( + q, k, v, g, beta, + local_cu_seqlens, + plan=cp_plan, + group=cp_group, + safe_gate=True, + a_log=a_log, # [heads] + dt_bias=dt_bias, # canonical [heads, key_dim] +) +``` + +Inputs use cuDNN's canonical THD layout. Framework adapters must reshape +model-specific parameter layouts before the call. The runtime summarizes only +the two local boundary spans and passes their ordered FP32 summaries to cuDNN's +context-parallel KDA apply. +See [`docs/KIMI_CP_RUNTIME.md`](docs/KIMI_CP_RUNTIME.md) for the ownership and +validation contracts. + +Install the optional CUDA 13 dependencies with: + +```bash +pip install 'cuda-linear-attention[cudnn-cp]' +``` + +--- + ## KDA cuLA exposes several KDA entry points targeting different GPU architectures: diff --git a/cula/distributed/__init__.py b/cula/distributed/__init__.py new file mode 100644 index 00000000..78aa7cb6 --- /dev/null +++ b/cula/distributed/__init__.py @@ -0,0 +1,19 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Distributed linear-attention execution layers.""" + +from cula.distributed.kda import ( + AllGatherKimiCPTransport, + KimiCPBoundaryTransport, + cudnn_kda_cp, +) +from cula.distributed.plan import KimiCPPlan, build_kimi_cp_plan + +__all__ = [ + "AllGatherKimiCPTransport", + "KimiCPBoundaryTransport", + "KimiCPPlan", + "build_kimi_cp_plan", + "cudnn_kda_cp", +] diff --git a/cula/distributed/kda.py b/cula/distributed/kda.py new file mode 100644 index 00000000..d9af7dc0 --- /dev/null +++ b/cula/distributed/kda.py @@ -0,0 +1,693 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Autograd-aware cuDNN KDA context parallelism for packed THD shards.""" + +from __future__ import annotations + +from dataclasses import dataclass +from functools import cache +from importlib import import_module +from inspect import signature +from typing import Any, Protocol, runtime_checkable + +import torch +import torch.distributed as dist +from torch.autograd.function import once_differentiable + +from cula.distributed.plan import KimiCPPlan + + +@cache +def _load_cudnn_cp_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", + ) + ): + ops = import_module("cudnn.linear_attention.ops") + kimi_delta_attention = ops.kimi_delta_attention + 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: + raise AttributeError("KDA apply does not accept concatenated CP summaries") + summary_bwd_parameters = signature(kimi_delta_attention_summary_bwd).parameters + if "output_context" not in summary_bwd_parameters or "checkpoint_every_n_tokens" not in summary_bwd_parameters: + raise AttributeError("KDA backward summary does not expose the CP apply context") + kimi_delta_attention_bwd = torch.ops.cudnn.kimi_delta_attention_bwd + except (ImportError, AttributeError) as error: + raise ImportError( + "cuDNN Kimi CP requires cuDNN Frontend with concatenated KDA summary apply APIs; " + "install cuda-linear-attention[cudnn-cp]" + ) from error + return ( + kimi_delta_attention, + kimi_delta_attention_summary, + kimi_delta_attention_summary_bwd, + kimi_delta_attention_bwd, + ) + + +@dataclass(frozen=True, slots=True) +class _KdaOptions: + scale: float | None + use_qk_l2norm_in_kernel: bool + use_beta_sigmoid_in_kernel: bool + allow_neg_eigval: bool + safe_gate: bool + lower_bound: float | None + gate_domain: str + batch_invariant: bool + checkpoint_every_n_tokens: int + plan_name: str | None + summary_plan_name: str | None + + def kwargs(self, a_log: torch.Tensor | None, dt_bias: torch.Tensor | None) -> dict[str, Any]: + return { + "use_qk_l2norm_in_kernel": self.use_qk_l2norm_in_kernel, + "use_beta_sigmoid_in_kernel": self.use_beta_sigmoid_in_kernel, + "allow_neg_eigval": self.allow_neg_eigval, + "safe_gate": self.safe_gate, + "gate_lower_bound": self.lower_bound, + "gate_domain": self.gate_domain, + "a_log": a_log, + "dt_bias": dt_bias, + "batch_invariant": self.batch_invariant, + } + + def attention_kwargs(self, a_log: torch.Tensor | None, dt_bias: torch.Tensor | None) -> dict[str, Any]: + kwargs = self.kwargs(a_log, dt_bias) + kwargs["checkpoint_every_n_tokens"] = self.checkpoint_every_n_tokens + return kwargs + + +def _validate_checkpoint_every_n_tokens(value: int) -> int: + if not isinstance(value, int) or isinstance(value, bool): + raise TypeError("checkpoint_every_n_tokens must be an integer") + if value not in (0, 16): + raise ValueError("checkpoint_every_n_tokens must be 0 or the exact native chunk cadence 16") + return value + + +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: + return tensor.unsqueeze(0) + gathered = torch.empty( + (world_size * tensor.shape[0], *tensor.shape[1:]), + dtype=tensor.dtype, + device=tensor.device, + ) + _all_gather_into(gathered, tensor.contiguous(), group=group) + return gathered.view(world_size, *tensor.shape) + + +def _validate_local_boundary_summaries(h: torch.Tensor, transition: torch.Tensor) -> None: + """Validate the two dense FP32 tensors emitted by a local summary op.""" + if h.ndim != 3 or transition.ndim != 3: + raise ValueError("KDA boundary summaries must have [heads, dim, key_dim] layout") + if h.dtype != torch.float32 or transition.dtype != torch.float32: + raise TypeError("KDA boundary summaries must remain FP32") + if h.shape[0] != transition.shape[0] or h.shape[-1] != transition.shape[-1]: + raise ValueError("KDA H and transition boundary shapes are incompatible") + if not h.is_contiguous() or not transition.is_contiguous(): + raise ValueError("local KDA boundary summaries must be contiguous") + + +def _select_boundary_summaries( + h_all: torch.Tensor, + transition_all: torch.Tensor, + *, + rank: int, + num_remote_ranks: int, + reverse: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + """Select the ordered summaries consumed by cuDNN's CP-aware KDA apply.""" + if reverse: + start, stop = rank + 1, rank + num_remote_ranks + 1 + else: + start, stop = rank - num_remote_ranks, rank + return h_all[start:stop], transition_all[start:stop] + + +@dataclass(slots=True) +class _KimiCPCollectiveWorkspace: + """Reusable FP32 payload storage for one Kimi CP collective stream. + + The workspace is temporary transport storage and is not retained by the + local KDA autograd graph. One instance can therefore be reused across KDA + layers executed sequentially on the same CUDA stream. Concurrent streams + require separate instances. + """ + + world_size: int + heads: int + key_dim: int + value_dim: int + gathered_h: torch.Tensor + gathered_transition: torch.Tensor + + @classmethod + def allocate( + cls, + *, + world_size: int, + heads: int, + key_dim: int, + value_dim: int, + device: torch.device | str | int, + ) -> _KimiCPCollectiveWorkspace: + dimensions = { + "world_size": world_size, + "heads": heads, + "key_dim": key_dim, + "value_dim": value_dim, + } + for name, value in dimensions.items(): + if not isinstance(value, int) or isinstance(value, bool) or value < 1: + raise ValueError(f"{name} must be a positive integer") + resolved = torch.device("cuda", device) if isinstance(device, int) else torch.device(device) + gathered_h = torch.empty((world_size * heads, value_dim, key_dim), dtype=torch.float32, device=resolved) + gathered_transition = torch.empty((world_size * heads, key_dim, key_dim), dtype=torch.float32, device=resolved) + return cls(world_size, heads, key_dim, value_dim, gathered_h, gathered_transition) + + @property + def nbytes(self) -> int: + return sum(tensor.numel() * tensor.element_size() for tensor in (self.gathered_h, self.gathered_transition)) + + def validate(self, h: torch.Tensor, transition: torch.Tensor, *, world_size: int) -> None: + expected_h = (self.heads, self.value_dim, self.key_dim) + expected_transition = (self.heads, self.key_dim, self.key_dim) + expected_gathered_h = (self.world_size * self.heads, self.value_dim, self.key_dim) + expected_gathered_transition = (self.world_size * self.heads, self.key_dim, self.key_dim) + if world_size != self.world_size: + raise ValueError(f"workspace world_size={self.world_size} does not match KimiCPPlan world_size={world_size}") + if h.dtype != torch.float32 or transition.dtype != torch.float32: + raise TypeError("Kimi CP boundary summaries must remain FP32") + if transition.device != h.device: + raise ValueError("Kimi CP H and transition summaries must be on the same device") + if tuple(h.shape) != expected_h or tuple(transition.shape) != expected_transition: + raise ValueError( + "workspace H/M shape mismatch: " + f"expected {expected_h} and {expected_transition}, got {tuple(h.shape)} and {tuple(transition.shape)}" + ) + for name, tensor, shape in ( + ("gathered_h", self.gathered_h, expected_gathered_h), + ("gathered_transition", self.gathered_transition, expected_gathered_transition), + ): + if tensor.dtype != torch.float32 or tensor.device != h.device or tuple(tensor.shape) != shape: + raise ValueError(f"workspace {name} must be FP32 on {h.device} with shape {shape}") + if not tensor.is_contiguous(): + raise ValueError(f"workspace {name} must be contiguous") + + +@runtime_checkable +class KimiCPBoundaryTransport(Protocol): + """KDA-specific contract for ordered cross-rank summary propagation.""" + + def boundary_summaries( + self, + h: torch.Tensor, + transition: torch.Tensor, + *, + plan: KimiCPPlan, + group: dist.ProcessGroup | None, + reverse: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: ... + + +@dataclass(slots=True) +class AllGatherKimiCPTransport: + """Exact ordered Kimi CP transport backed by typed H/M all-gathers. + + With no workspace, storage is allocated lazily on each call. Use + :meth:`prepare` for stable large-CP memory and CUDA Graph integration. + """ + + workspace: _KimiCPCollectiveWorkspace | None = None + + def __post_init__(self) -> None: + if self.workspace is not None and not isinstance(self.workspace, _KimiCPCollectiveWorkspace): + raise TypeError("workspace must be created by AllGatherKimiCPTransport.prepare") + + @staticmethod + def required_workspace_bytes( + plan: KimiCPPlan, + *, + heads: int, + key_dim: int, + value_dim: int, + ) -> int: + """Return local plus gathered FP32 payload bytes without allocating.""" + if not isinstance(plan, KimiCPPlan): + raise TypeError("plan must be a KimiCPPlan") + for name, value in {"heads": heads, "key_dim": key_dim, "value_dim": value_dim}.items(): + if not isinstance(value, int) or isinstance(value, bool) or value < 1: + raise ValueError(f"{name} must be a positive integer") + return plan.world_size * heads * (value_dim + key_dim) * key_dim * torch.float32.itemsize + + @classmethod + def prepare( + cls, + plan: KimiCPPlan, + *, + heads: int, + key_dim: int, + value_dim: int, + device: torch.device | str | int, + ) -> AllGatherKimiCPTransport: + if not isinstance(plan, KimiCPPlan): + raise TypeError("plan must be a KimiCPPlan") + return cls( + _KimiCPCollectiveWorkspace.allocate( + world_size=plan.world_size, + heads=heads, + key_dim=key_dim, + value_dim=value_dim, + device=device, + ) + ) + + @property + def workspace_bytes(self) -> int: + return 0 if self.workspace is None else self.workspace.nbytes + + def boundary_summaries( + self, + h: torch.Tensor, + transition: torch.Tensor, + *, + plan: KimiCPPlan, + group: dist.ProcessGroup | None, + reverse: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + _validate_local_boundary_summaries(h, transition) + if self.workspace is None: + h_all = _all_gather(h, group, plan.world_size) + transition_all = _all_gather(transition, group, plan.world_size) + else: + workspace = self.workspace + workspace.validate(h, transition, world_size=plan.world_size) + if plan.world_size == 1: + h_all = h.unsqueeze(0) + transition_all = transition.unsqueeze(0) + else: + _all_gather_into( + workspace.gathered_h, + h, + group=group, + ) + _all_gather_into( + workspace.gathered_transition, + transition, + group=group, + ) + h_all = workspace.gathered_h.view(plan.world_size, *h.shape) + transition_all = workspace.gathered_transition.view(plan.world_size, *transition.shape) + return _select_boundary_summaries( + h_all, + transition_all, + rank=plan.rank, + num_remote_ranks=plan.post_num_ranks if reverse else plan.pre_num_ranks, + reverse=reverse, + ) + + +_DEFAULT_BOUNDARY_TRANSPORT = AllGatherKimiCPTransport() + + +def _validate_boundary_summaries( + summaries: tuple[torch.Tensor, torch.Tensor], + h_reference: torch.Tensor, + transition_reference: torch.Tensor, + *, + rows: int, +) -> tuple[torch.Tensor, torch.Tensor]: + if not isinstance(summaries, tuple) or len(summaries) != 2: + raise TypeError("KimiCPBoundaryTransport must return an (H, transition) tuple") + expected = ( + (rows, *h_reference.shape), + (rows, *transition_reference.shape), + ) + for name, tensor, reference, shape in zip( + ("H", "transition"), + summaries, + (h_reference, transition_reference), + expected, + strict=True, + ): + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"KimiCPBoundaryTransport {name} must be a torch.Tensor") + if tensor.dtype != torch.float32: + raise TypeError(f"KimiCPBoundaryTransport {name} must remain FP32") + if tensor.device != reference.device or tuple(tensor.shape) != shape: + raise ValueError( + f"KimiCPBoundaryTransport {name} mismatch: expected shape/device " + f"{shape} on {reference.device}, got {tuple(tensor.shape)} on {tensor.device}" + ) + if not tensor.is_contiguous(): + raise ValueError(f"KimiCPBoundaryTransport {name} must be contiguous") + return summaries + + +class _CudnnKdaCP(torch.autograd.Function): + @staticmethod + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + cu_seqlens: torch.Tensor, + a_log: torch.Tensor | None, + dt_bias: torch.Tensor | None, + plan: KimiCPPlan, + group: dist.ProcessGroup | None, + options: _KdaOptions, + boundary_transport: KimiCPBoundaryTransport, + ) -> torch.Tensor: + kimi_delta_attention, kimi_delta_attention_summary, _, _ = _load_cudnn_cp_ops() + kwargs = options.kwargs(a_log, dt_bias) + + # A contiguous shard can carry state only through its last local span. + # Restricting the summary to that span avoids rescanning short samples. + last_start = plan.last_boundary_start + last_cu = cu_seqlens[-2:] - last_start + h_local, transition_local = kimi_delta_attention_summary( + k[last_start:], + v[last_start:], + g[last_start:], + beta[last_start:], + last_cu, + plan_name=options.summary_plan_name, + **kwargs, + ) + preceding_h, preceding_transition = _validate_boundary_summaries( + boundary_transport.boundary_summaries( + h_local[0], + transition_local[0], + plan=plan, + group=group, + reverse=False, + ), + h_local[0], + transition_local[0], + rows=plan.pre_num_ranks, + ) + + # cuDNN owns the ordered FP32 summary composition. Passing an empty + # leading dimension on the first rank keeps one stable return contract. + attention_outputs = kimi_delta_attention( + q, + k, + v, + g, + beta, + cu_seqlens, + scale=options.scale, + output_final_state=False, + plan_name=options.plan_name, + preceding_state_update=preceding_h, + preceding_transition=preceding_transition, + **options.attention_kwargs(a_log, dt_bias), + ) + output, _, state_checkpoints, composed_initial_state = attention_outputs + + ctx.plan = plan + ctx.group = group + ctx.options = options + ctx.boundary_transport = boundary_transport + ctx.a_log_present = a_log is not None + ctx.dt_bias_present = dt_bias is not None + ctx.save_for_backward( + q, + k, + v, + g, + beta, + cu_seqlens, + state_checkpoints, + composed_initial_state, + *(tensor for tensor in (a_log, dt_bias) if tensor is not None), + ) + return output + + @staticmethod + @once_differentiable + def backward(ctx, grad_output: torch.Tensor): + _, _, kimi_delta_attention_summary_bwd, kimi_delta_attention_bwd = _load_cudnn_cp_ops() + saved = iter(ctx.saved_tensors) + q = next(saved) + k = next(saved) + v = next(saved) + g = next(saved) + beta = next(saved) + cu_seqlens = next(saved) + state_checkpoints = next(saved) + composed_initial_state = next(saved) + a_log = next(saved) if ctx.a_log_present else None + dt_bias = next(saved) if ctx.dt_bias_present else None + plan = ctx.plan + options = ctx.options + grad_output = grad_output.contiguous() + + # The reverse carry is determined only by the first local span. cuDNN + # backward apply consumes the summary transition in its stored M^T domain. + first_end = plan.first_boundary_end + g_local, transition_backward, backward_context = kimi_delta_attention_summary_bwd( + q[:first_end], + k[:first_end], + g[:first_end], + beta[:first_end], + grad_output[:first_end], + cu_seqlens[:2], + scale=options.scale, + output_transition=True, + output_context=True, + checkpoint_every_n_tokens=options.checkpoint_every_n_tokens, + plan_name=options.summary_plan_name, + **options.kwargs(a_log, dt_bias), + ) + following_h, following_transition = _validate_boundary_summaries( + ctx.boundary_transport.boundary_summaries( + g_local[0], + transition_backward[0], + plan=plan, + group=ctx.group, + reverse=True, + ), + g_local[0], + transition_backward[0], + rows=plan.post_num_ranks, + ) + + scale = options.scale if options.scale is not None else q.shape[-1] ** -0.5 + dq, dk, dv, dg, dbeta, _, da_log, ddt_bias = kimi_delta_attention_bwd( + grad_output, + q, + k, + v, + g, + beta, + cu_seqlens, + float(scale), + state_checkpoints=state_checkpoints if state_checkpoints.numel() else None, + checkpoint_every_n_tokens=options.checkpoint_every_n_tokens, + use_qk_l2norm_in_kernel=options.use_qk_l2norm_in_kernel, + batch_invariant=options.batch_invariant, + use_beta_sigmoid_in_kernel=options.use_beta_sigmoid_in_kernel, + allow_neg_eigval=options.allow_neg_eigval, + safe_gate=options.safe_gate, + gate_lower_bound=options.lower_bound, + gate_domain=options.gate_domain, + a_log=a_log, + dt_bias=dt_bias, + plan_name=options.plan_name, + following_state_grad=following_h, + following_transition=following_transition, + composed_initial_state=(composed_initial_state if composed_initial_state.numel() else None), + context=backward_context, + ) + da_log = da_log if ctx.a_log_present else None + ddt_bias = ddt_bias if ctx.dt_bias_present else None + return dq, dk, dv, dg, dbeta, None, da_log, ddt_bias, None, None, None, None + + +def _validate_inputs( + 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, + safe_gate: bool, + a_log: torch.Tensor | None, + dt_bias: torch.Tensor | None, +) -> None: + operands = {"q": q, "k": k, "v": v, "g": g, "beta": beta} + if any(not isinstance(tensor, torch.Tensor) for tensor in operands.values()): + raise TypeError("all KDA operands must be torch.Tensor instances") + if any(tensor.ndim != 3 for tensor in (q, k, v, g)) or beta.ndim != 2: + raise ValueError("cuDNN Kimi CP expects q/k/v/g in THD and beta in TH layout") + if any(tensor.shape[0] != q.shape[0] for tensor in (k, v, g, beta)): + raise ValueError("all KDA operands must have the same local token count") + if q.shape[0] != plan.local_tokens: + raise ValueError(f"KDA operands contain {q.shape[0]} local tokens but KimiCPPlan expects {plan.local_tokens}") + if any(tensor.device != q.device for tensor in (k, v, g, beta)): + raise ValueError("all KDA operands must be on the same device") + if q.device.type != "cuda": + raise ValueError("cuDNN Kimi CP requires CUDA operands") + plan.validate_group(group) + plan.validate_local_cu_seqlens(cu_seqlens, device=q.device) + if safe_gate: + heads, key_dim = g.shape[-2:] + if a_log is not None and tuple(a_log.shape) != (heads,): + raise ValueError(f"safe-gate a_log must have shape [{heads}]") + if dt_bias is not None and tuple(dt_bias.shape) != (heads, key_dim): + raise ValueError(f"safe-gate dt_bias must have shape [{heads}, {key_dim}]") + elif a_log is not None or dt_bias is not None: + raise ValueError("a_log and dt_bias are accepted only with safe_gate=True; canonicalize non-safe gates upstream") + + +def cudnn_kda_cp( + 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 = None, + scale: float | None = None, + use_qk_l2norm_in_kernel: bool = False, + use_beta_sigmoid_in_kernel: bool = False, + allow_neg_eigval: bool = False, + safe_gate: bool = False, + lower_bound: float | None = None, + gate_domain: str = "log", + a_log: torch.Tensor | None = None, + 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", + boundary_transport: KimiCPBoundaryTransport | None = None, +) -> torch.Tensor: + """Run differentiable cuDNN KDA over one contiguous packed-THD CP shard. + + cuLA owns boundary planning, collectives, ordered summary propagation and + distributed autograd. cuDNN Frontend owns local KDA summaries, ordered + FP32 summary composition, and local forward/backward. The summary dtype is + an invariant rather than a user option so large-CP behavior cannot silently + degrade. + + ``plan`` should be built once per packed microbatch and reused by all KDA + layers. Inputs use cuDNN's canonical THD contract; framework-specific + layout and gate conversions belong in the framework adapter. + + ``boundary_transport`` defaults to exact ordered summary all-gather with lazy + storage. Large-CP callers should prepare one + :class:`AllGatherKimiCPTransport` per concurrent CUDA stream and reuse it + across layers to avoid repeated collective-buffer allocations. + + ``checkpoint_every_n_tokens=0`` lets cuDNN regenerate the chunk-state + series during backward. Passing 16 retains every native chunk state for + exact backward reuse. Coarser upstream cadences still recompute seeded + spans, were slower in the 1M study, and are intentionally not exposed. + Because exact checkpoints are large and live across the forward/backward + interval, cuLA never enables them implicitly; callers should select 16 + only from measured whole-model memory/latency tradeoffs. + """ + + if not isinstance(plan, KimiCPPlan): + raise TypeError("plan must be a KimiCPPlan") + 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") + checkpoint_every_n_tokens = _validate_checkpoint_every_n_tokens(checkpoint_every_n_tokens) + _validate_inputs( + q, + k, + v, + g, + beta, + cu_seqlens, + plan=plan, + group=group, + safe_gate=safe_gate, + a_log=a_log, + dt_bias=dt_bias, + ) + options = _KdaOptions( + scale=scale, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel, + allow_neg_eigval=allow_neg_eigval, + safe_gate=safe_gate, + lower_bound=lower_bound, + gate_domain=gate_domain, + batch_invariant=batch_invariant, + checkpoint_every_n_tokens=checkpoint_every_n_tokens, + plan_name=plan_name, + summary_plan_name=summary_plan_name, + ) + if plan.world_size == 1: + kimi_delta_attention, _, _, _ = _load_cudnn_cp_ops() + attention_outputs = kimi_delta_attention( + q, + k, + v, + g, + beta, + cu_seqlens, + scale=scale, + initial_state=None, + output_final_state=False, + plan_name=plan_name, + **options.attention_kwargs(a_log, dt_bias), + ) + return attention_outputs[0] + return _CudnnKdaCP.apply( + q.contiguous(), + k.contiguous(), + v.contiguous(), + g.contiguous(), + beta.contiguous(), + cu_seqlens.contiguous(), + a_log, + dt_bias, + plan, + group, + options, + resolved_transport, + ) + + +__all__ = [ + "AllGatherKimiCPTransport", + "KimiCPBoundaryTransport", + "cudnn_kda_cp", +] diff --git a/cula/distributed/plan.py b/cula/distributed/plan.py new file mode 100644 index 00000000..b42b4275 --- /dev/null +++ b/cula/distributed/plan.py @@ -0,0 +1,218 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Packed-THD shard planning for Kimi context parallelism. + +The planner is deliberately CPU-only. A framework supplies global sequence +boundaries and the contiguous token partition; the resulting immutable plan is +reusable by every KDA layer in the microbatch. No CUDA tensor is inspected, so +building a plan never introduces a device-to-host synchronization. +""" + +from __future__ import annotations + +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 + + +@dataclass(frozen=True, slots=True) +class KimiCPPlan: + """Immutable metadata for one contiguous packed-THD CP shard. + + Only the first and last local sequences can cross rank boundaries. The + plan records their token spans and the ordered rank chains that contribute + incoming forward state or outgoing backward state. Tensor shapes and + process-group identity are checked again at execution time. + """ + + rank: int + world_size: int + token_start: int + token_end: int + first_sequence: int + last_sequence: int + local_cu_seqlens: tuple[int, ...] + first_boundary_end: int + last_boundary_start: int + pre_num_ranks: int + post_num_ranks: int + + @property + def local_tokens(self) -> int: + return self.token_end - self.token_start + + @property + def num_local_sequences(self) -> int: + return len(self.local_cu_seqlens) - 1 + + @property + def first_boundary_tokens(self) -> int: + return self.first_boundary_end + + @property + def last_boundary_tokens(self) -> int: + return self.local_tokens - self.last_boundary_start + + 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 Kimi CP 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( + "KimiCPPlan process-group mismatch: " + f"plan has rank/world_size={self.rank}/{self.world_size}, " + f"runtime has {rank}/{world_size}" + ) + + def validate_local_cu_seqlens(self, cu_seqlens: torch.Tensor, *, device: torch.device) -> None: + if not isinstance(cu_seqlens, torch.Tensor): + raise TypeError("cu_seqlens must be a torch.Tensor") + if cu_seqlens.ndim != 1: + raise ValueError("cu_seqlens must be one-dimensional") + if cu_seqlens.dtype not in (torch.int32, torch.int64): + raise TypeError("cu_seqlens must have int32 or int64 dtype") + if cu_seqlens.device != device: + raise ValueError("cu_seqlens and KDA operands must be on the same device") + if cu_seqlens.numel() != len(self.local_cu_seqlens): + raise ValueError( + "cu_seqlens does not match KimiCPPlan: " + f"expected {len(self.local_cu_seqlens)} entries, got {cu_seqlens.numel()}" + ) + if cu_seqlens.device.type == "cpu" and tuple(int(value) for value in cu_seqlens.tolist()) != self.local_cu_seqlens: + raise ValueError("CPU cu_seqlens values do not match KimiCPPlan") + + @classmethod + def from_global_cu_seqlens( + cls, + global_cu_seqlens: torch.Tensor | Sequence[int], + *, + rank: int, + world_size: int, + shard_offsets: Sequence[int] | None = None, + ) -> KimiCPPlan: + """Build a plan for a contiguous rank-ordered token partition. + + ``shard_offsets`` defaults to a balanced partition of the total packed + tokens. Frameworks with a different contiguous partition can pass the + exact global token offsets for all ranks. + """ + + 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(f"rank must be in [0, {world_size})") + boundaries = _as_boundaries(global_cu_seqlens, name="global_cu_seqlens") + total_tokens = boundaries[-1] + offsets = ( + _balanced_shards(total_tokens, world_size) + if shard_offsets is None + else _validate_shards(shard_offsets, total_tokens=total_tokens, world_size=world_size) + ) + start, end = offsets[rank], offsets[rank + 1] + + first_sequence = bisect_right(boundaries, start) - 1 + last_sequence = bisect_right(boundaries, end - 1) - 1 + local_boundaries = (0,) + tuple(value - start for value in boundaries[1:-1] if start < value < end) + (end - start,) + + first_sequence_rank = bisect_right(offsets, boundaries[first_sequence]) - 1 + last_sequence_rank = bisect_right(offsets, boundaries[last_sequence + 1] - 1) - 1 + pre_num_ranks = rank - first_sequence_rank + post_num_ranks = last_sequence_rank - rank + if pre_num_ranks < 0 or post_num_ranks < 0: + raise RuntimeError("internal error while deriving Kimi CP boundary rank ranges") + + return cls( + rank=rank, + world_size=world_size, + token_start=start, + token_end=end, + first_sequence=first_sequence, + last_sequence=last_sequence, + local_cu_seqlens=local_boundaries, + first_boundary_end=local_boundaries[1], + last_boundary_start=local_boundaries[-2], + pre_num_ranks=pre_num_ranks, + post_num_ranks=post_num_ranks, + ) + + +def build_kimi_cp_plan( + global_cu_seqlens: torch.Tensor | Sequence[int], + *, + group: dist.ProcessGroup | None = None, + rank: int | None = None, + world_size: int | None = None, + shard_offsets: Sequence[int] | None = None, +) -> KimiCPPlan: + """Build :class:`KimiCPPlan`, deriving rank identity from ``group`` by default.""" + + if (rank is None) != (world_size is None): + raise ValueError("rank and world_size must be supplied together") + if rank is None: + if dist.is_available() and dist.is_initialized(): + rank = dist.get_rank(group) + world_size = dist.get_world_size(group) + elif group is None: + rank, world_size = 0, 1 + else: + raise RuntimeError("a process group was supplied but torch.distributed is not initialized") + return KimiCPPlan.from_global_cu_seqlens( + global_cu_seqlens, + rank=rank, + world_size=world_size, + shard_offsets=shard_offsets, + ) + + +__all__ = ["KimiCPPlan", "build_kimi_cp_plan"] diff --git a/cula/kda/__init__.py b/cula/kda/__init__.py index 91cc4417..cda26803 100644 --- a/cula/kda/__init__.py +++ b/cula/kda/__init__.py @@ -28,12 +28,22 @@ "fused_sigmoid_gating_delta_rule_update", "flashkda_prefill", "kda_prefill", + "AllGatherKimiCPTransport", + "KimiCPBoundaryTransport", + "KimiCPPlan", + "build_kimi_cp_plan", + "cudnn_kda_cp", ] _LAZY = { "chunk_kda_fwd_intra": ("cula.kda.chunk_intra", "chunk_kda_fwd_intra"), "flashkda_prefill": ("cula.kda.flashkda", "cula_kda_prefill"), "kda_prefill": ("cula.kda.prefill", "kda_prefill"), + "AllGatherKimiCPTransport": ("cula.distributed", "AllGatherKimiCPTransport"), + "KimiCPBoundaryTransport": ("cula.distributed", "KimiCPBoundaryTransport"), + "KimiCPPlan": ("cula.distributed", "KimiCPPlan"), + "build_kimi_cp_plan": ("cula.distributed", "build_kimi_cp_plan"), + "cudnn_kda_cp": ("cula.distributed", "cudnn_kda_cp"), "kda_decode": ("cula.ops.kda.decode.cute", "kda_decode"), "kda_decode_mtp": ("cula.ops.kda.decode.mtp", "kda_decode_mtp"), "kda_decode_mtp_recurrent": ("cula.ops.kda.decode.mtp", "kda_decode_mtp_recurrent"), diff --git a/docs/KIMI_CP_RUNTIME.md b/docs/KIMI_CP_RUNTIME.md new file mode 100644 index 00000000..0a349891 --- /dev/null +++ b/docs/KIMI_CP_RUNTIME.md @@ -0,0 +1,154 @@ +# cuDNN KDA Kimi CP runtime + +## 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. + +The public surface intentionally has three KDA-specific layers: + +- `KimiCPPlan`: immutable CPU metadata for one contiguous packed-THD shard. +- `KimiCPBoundaryTransport`: the narrow ordered-summary propagation contract, + 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. + +## Ownership boundary + +```text +framework adapter + layout/gate canonicalization, global cu_seqlens, process group + | + v +cuLA Kimi CP runtime + shard plan, boundary-only summaries, boundary transport, distributed autograd + | + v +cuDNN Frontend + KDA summaries, concatenated-summary composition, local KDA fwd/bwd +``` + +Framework-specific compatibility does not belong in the runtime. In +particular, flattened Megatron `dt_bias`, non-safe gate preprocessing, and +`[1,T,H,D]` wrappers must be canonicalized before calling cuLA. + +## Packed shard plan + +`KimiCPPlan.from_global_cu_seqlens` accepts CPU global boundaries and either a +balanced partition or explicit global shard offsets. It derives: + +- clipped local `cu_seqlens`; +- first and last local sequence IDs; +- the first and last boundary token spans; +- the number of preceding and following ranks in those sequence chains. + +Plan construction is CPU-only and must happen once per packed microbatch, not +once per layer. The same plan is reusable by every KDA layer. + +Only the first and last local spans can cross rank boundaries under the +contiguous-shard contract. Forward therefore summarizes only the last local +span; backward summarizes only the first. This is the boundary-only strategy +validated by the Ant Megatron 256K workload (20 short samples plus two long +samples), where it removed the full-local-summary overhead. + +## Forward and backward + +Forward: + +1. Produce dense FP32 `(H, M)` for the last local span. +2. All-gather H and M into separate dense typed buffers. +3. Select only preceding ranks belonging to the first local sequence. +4. Pass them as `preceding_state_update` and `preceding_transition`. +5. Let cuDNN compose the rows and seed the first local sequence. + +Backward performs the symmetric reverse chain: + +1. Produce `(G, M^T)` and reusable backward context for the first local span. +2. All-gather the transition in its stored-domain orientation. +3. Select following ranks belonging to the last local sequence. +4. Pass them as `following_state_grad` and `following_transition`. +5. Let cuDNN compose the reverse chain and compute local tensor and gate gradients. + +The state payload and transition remain FP32. cuLA exposes no merge-precision +switch; the cuDNN apply API owns the fixed high-precision state chain. This +runtime requires cuDNN Frontend 1.33 or a newer develop build containing the +concatenated-summary KDA API. + +## Current transport and intended evolution + +The first implementation uses two typed `all_gather_into_tensor` calls per +direction. Keeping H and M separate preserves the dense layout required by +cuDNN without a CP-sized repack buffer; total communication bytes are unchanged. +It still materializes every rank's summary on every rank, so CP32/CP64 profiling +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. + +## Prepared large-CP transport + +The convenience path allocates its gathered H and M tensors on each forward and +backward call. At large CP this creates avoidable allocator traffic: for +`H=32`, `K=V=128`, and CP64, the two gathered tensors total 256 MiB. + +`AllGatherKimiCPTransport.prepare` allocates those FP32 buffers once. The same +transport can be reused by every KDA layer that executes sequentially on one +CUDA stream because the buffers are not retained by the local autograd graph. +Use one prepared transport per concurrently executing stream. Shape, device, +world-size, dtype, and contiguity are validated before every collective. +`required_workspace_bytes` reports the exact allocation before materializing +it; the standard CP64 shape above requires 256 MiB. + +```python +from cula.distributed import ( + AllGatherKimiCPTransport, + build_kimi_cp_plan, + cudnn_kda_cp, +) + +plan = build_kimi_cp_plan(global_cu_seqlens_cpu, group=cp_group) +transport = AllGatherKimiCPTransport.prepare( + plan, + heads=32, + key_dim=128, + value_dim=128, + device=q.device, +) + +# Reuse `plan` and `transport` across compatible KDA layers. +output = cudnn_kda_cp( + q, + k, + v, + g, + beta, + local_cu_seqlens, + plan=plan, + group=cp_group, + boundary_transport=transport, +) +``` + +The transport contract returns only the ordered FP32 `(state_update, +transition)` rows belonging to the boundary sequence. It does not compose the +state or expose framework layouts and local KDA details. This is the intended +extension point for future topology-aware CP64 implementations. + +CUDA Graph is an execution mode, not a CP algorithm. The prepared transport +provides stable collective addresses; full graph capture also depends on the +process group and cuDNN summary-apply implementation supporting capture. + +## Validation ladder + +1. CPU plan tests, including the 256K 20-short/2-long partition. +2. Helper tests for boundary selection and mandatory FP32 payloads. +3. Two-rank distributed comparison with monolithic cuDNN forward and gradients. +4. CP4/CP8 128K and 256K packed-THD accuracy/performance regression. +5. CP16/CP32/CP64 accuracy, peak memory, collective scaling, and loss checks. +6. Warm CUDA Graph capture/replay after reusable workspace is introduced. diff --git a/pyproject.toml b/pyproject.toml index e2fb2e9c..4da6b6c6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,6 +18,10 @@ dependencies = [ license = "Apache-2.0" [project.optional-dependencies] +cudnn-cp = [ + "nvidia-cutlass-dsl[cu13]>=4.8.0,<4.9", + "nvidia-cudnn-frontend>=1.33.0", +] dev = [ "pytest", "matplotlib", diff --git a/tests/conftest.py b/tests/conftest.py index 541deb55..f0a9df4d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -41,6 +41,7 @@ def _prepare_markexpr(markexpr: str) -> tuple[str, bool, bool]: def pytest_configure(config): + config.addinivalue_line("markers", "distributed: requires launch through torchrun") config.addinivalue_line("markers", "sm100_only: only run on SM100 devices") config.addinivalue_line("markers", "sm90_only: only run on SM90 devices") config.addinivalue_line("markers", "benchmark: long-running benchmark-shaped coverage") diff --git a/tests/test_kimi_cp_distributed.py b/tests/test_kimi_cp_distributed.py new file mode 100644 index 00000000..670e003f --- /dev/null +++ b/tests/test_kimi_cp_distributed.py @@ -0,0 +1,109 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Run with two or more ranks under ``torchrun``.""" + +import inspect +import os + +import pytest +import torch +import torch.distributed as dist + + +def _relative_rms(actual, expected): + delta = actual.float() - expected.float() + return delta.norm() / expected.float().norm().clamp_min(1e-8) + + +@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") + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group("nccl") + try: + cudnn_la = pytest.importorskip("cudnn.linear_attention") + try: + supports_concatenated_summaries = ( + "preceding_state_update" 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") + + def stage(name): + print(f"rank={local_rank} stage={name}", flush=True) + + from cula.distributed import AllGatherKimiCPTransport, build_kimi_cp_plan, cudnn_kda_cp + + device = torch.device("cuda", local_rank) + generator = torch.Generator(device=device).manual_seed(20261102) + heads, key_dim, value_dim = 8, 128, 128 + lengths = (48, 112) + total = sum(lengths) + global_cu_cpu = torch.tensor([0, lengths[0], total], dtype=torch.int32) + global_cu = global_cu_cpu.to(device) + plan = build_kimi_cp_plan(global_cu_cpu, group=dist.group.WORLD) + boundary_transport = AllGatherKimiCPTransport.prepare( + plan, + heads=heads, + key_dim=key_dim, + value_dim=value_dim, + device=device, + ) + local_cu = torch.tensor(plan.local_cu_seqlens, dtype=torch.int32, device=device) + start, end = plan.token_start, plan.token_end + + def randn(*shape, dtype=torch.bfloat16): + return torch.randn(*shape, dtype=dtype, device=device, generator=generator) + + q = torch.nn.functional.normalize(randn(total, heads, key_dim).float(), dim=-1).to(torch.bfloat16) + k = torch.nn.functional.normalize(randn(total, heads, key_dim).float(), dim=-1).to(torch.bfloat16) + v = randn(total, heads, value_dim) + g = -(0.01 + 0.04 * torch.rand(total, heads, key_dim, dtype=torch.float32, device=device, generator=generator)) + beta = torch.sigmoid(randn(total, heads, dtype=torch.float32)) + grad_output = randn(total, heads, value_dim) + + full_inputs = [tensor.detach().clone().requires_grad_(True) for tensor in (q, k, v, g, beta)] + local_inputs = [tensor[start:end].detach().clone().requires_grad_(True) for tensor in (q, k, v, g, beta)] + + stage("monolithic-forward-start") + full_output, _ = cudnn_la.kimi_delta_attention( + *full_inputs, + global_cu, + output_final_state=True, + batch_invariant=True, + plan_name="kda_frost", + ) + stage("monolithic-forward-done") + full_output.backward(grad_output) + stage("monolithic-backward-done") + stage("cp-forward-start") + local_output = cudnn_kda_cp( + *local_inputs, + local_cu, + plan=plan, + group=dist.group.WORLD, + boundary_transport=boundary_transport, + ) + stage("cp-forward-done") + local_output.backward(grad_output[start:end]) + stage("cp-backward-done") + + torch.testing.assert_close(local_output, full_output[start:end], rtol=2e-2, atol=1e-2) + for actual, expected in zip(local_inputs, full_inputs): + assert torch.isfinite(actual.grad).all() + delta = (actual.grad.float() - expected.grad[start:end].float()).abs() + assert delta.max().item() <= 1.25e-1 + assert _relative_rms(actual.grad, expected.grad[start:end]).item() < 5e-2 + dist.barrier() + finally: + dist.destroy_process_group() + + +if __name__ == "__main__": + test_cudnn_kda_cp_matches_monolithic_forward_backward() diff --git a/tests/test_kimi_cp_helpers.py b/tests/test_kimi_cp_helpers.py new file mode 100644 index 00000000..d68735a5 --- /dev/null +++ b/tests/test_kimi_cp_helpers.py @@ -0,0 +1,470 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace + +import torch + +from cula.distributed import kda as cp +from cula.distributed.plan import KimiCPPlan + + +def test_checkpoint_cadence_validation(): + assert cp._validate_checkpoint_every_n_tokens(0) == 0 + assert cp._validate_checkpoint_every_n_tokens(16) == 16 + + for invalid in (-16, 1, 24, 32, 64, 128, 256): + try: + cp._validate_checkpoint_every_n_tokens(invalid) + except ValueError as error: + assert "native chunk cadence 16" in str(error) + else: + raise AssertionError(f"invalid checkpoint cadence {invalid} was accepted") + + for invalid in (True, 16.0, "16"): + try: + cp._validate_checkpoint_every_n_tokens(invalid) + except TypeError as error: + assert "integer" in str(error) + else: + raise AssertionError(f"non-integer checkpoint cadence {invalid!r} was accepted") + + +def test_attention_kwargs_keep_checkpoints_out_of_summary_calls(): + 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=16, + plan_name="kda_frost", + summary_plan_name="kda_summary_frost", + ) + + assert "checkpoint_every_n_tokens" not in options.kwargs(None, None) + assert options.attention_kwargs(None, None)["checkpoint_every_n_tokens"] == 16 + + +def test_cudnn_loader_accepts_current_linear_attention_exports(monkeypatch): + def attention(*args, preceding_state_update=None, preceding_transition=None, **kwargs): + return args, preceding_state_update, preceding_transition, kwargs + + def summary(*args, **kwargs): + return args, kwargs + + def summary_bwd(*args, output_context=False, checkpoint_every_n_tokens=0, **kwargs): + return args, output_context, checkpoint_every_n_tokens, kwargs + + backward = object() + modules = { + "cudnn.linear_attention": SimpleNamespace( + kimi_delta_attention=attention, + kimi_delta_attention_summary=summary, + kimi_delta_attention_summary_bwd=summary_bwd, + ), + } + + def import_fake(name): + try: + return modules[name] + except KeyError as error: + raise ModuleNotFoundError(name) from error + + monkeypatch.setattr(cp, "import_module", import_fake) + monkeypatch.setattr(cp.torch.ops.cudnn, "kimi_delta_attention_bwd", backward, raising=False) + 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() + + +def test_cudnn_loader_rejects_pre_concatenated_summary_api(monkeypatch): + def old_attention(*args, **kwargs): + return args, kwargs + + modules = { + "cudnn.linear_attention": SimpleNamespace( + kimi_delta_attention=old_attention, + kimi_delta_attention_summary=object(), + kimi_delta_attention_summary_bwd=object(), + ), + } + + def import_fake(name): + try: + return modules[name] + except KeyError as error: + raise ModuleNotFoundError(name) from error + + monkeypatch.setattr(cp, "import_module", import_fake) + cp._load_cudnn_cp_ops.cache_clear() + try: + try: + cp._load_cudnn_cp_ops() + except ImportError as error: + assert "concatenated KDA summary" in str(error) + else: + raise AssertionError("the retired stand-alone merge API must not be accepted") + finally: + cp._load_cudnn_cp_ops.cache_clear() + + +def test_local_boundary_summary_contract_preserves_fp32(): + h = torch.arange(24, dtype=torch.float32).reshape(2, 3, 4) + transition = torch.arange(32, dtype=torch.float32).reshape(2, 4, 4) + + cp._validate_local_boundary_summaries(h, transition) + + +def test_boundary_chain_selects_only_matching_remote_ranks(): + h_all = torch.arange(1, 5, dtype=torch.float32).reshape(4, 1, 1, 1) + transition_all = torch.zeros(4, 1, 1, 1) + + incoming = cp._select_boundary_summaries( + h_all, + transition_all, + rank=2, + num_remote_ranks=2, + reverse=False, + ) + outgoing = cp._select_boundary_summaries( + h_all, + transition_all, + rank=1, + num_remote_ranks=2, + reverse=True, + ) + no_remote = cp._select_boundary_summaries( + h_all, + transition_all, + rank=0, + num_remote_ranks=0, + reverse=False, + ) + + assert incoming[0].flatten().tolist() == [1.0, 2.0] + assert outgoing[0].flatten().tolist() == [3.0, 4.0] + assert incoming[1].shape == outgoing[1].shape == (2, 1, 1, 1) + assert no_remote[0].shape == no_remote[1].shape == (0, 1, 1, 1) + + +def test_local_boundary_summary_contract_rejects_precision_downgrade(): + h = torch.zeros(2, 3, 4, dtype=torch.bfloat16) + transition = torch.zeros(2, 4, 4, dtype=torch.float32) + + try: + cp._validate_local_boundary_summaries(h, transition) + except TypeError as error: + assert "FP32" in str(error) + else: + raise AssertionError("BF16 boundary state must be rejected") + + +def _transport_plan(*, rank=2, pre=2, post=1): + return KimiCPPlan( + rank=rank, + world_size=4, + token_start=8, + token_end=12, + first_sequence=0, + last_sequence=0, + local_cu_seqlens=(0, 4), + first_boundary_end=4, + last_boundary_start=0, + pre_num_ranks=pre, + post_num_ranks=post, + ) + + +def test_collective_workspace_has_stable_fp32_storage(): + workspace = cp._KimiCPCollectiveWorkspace.allocate( + world_size=4, + heads=2, + key_dim=4, + value_dim=3, + device="cpu", + ) + + assert workspace.gathered_h.shape == (8, 3, 4) + assert workspace.gathered_transition.shape == (8, 4, 4) + assert workspace.gathered_h.dtype == workspace.gathered_transition.dtype == torch.float32 + assert workspace.nbytes == (8 * 3 * 4 + 8 * 4 * 4) * 4 + cp64_plan = KimiCPPlan( + rank=0, + world_size=64, + token_start=0, + token_end=1, + first_sequence=0, + last_sequence=0, + local_cu_seqlens=(0, 1), + first_boundary_end=1, + last_boundary_start=0, + pre_num_ranks=0, + post_num_ranks=63, + ) + required = cp.AllGatherKimiCPTransport.required_workspace_bytes( + cp64_plan, + heads=32, + key_dim=128, + value_dim=128, + ) + assert required == 256 << 20 + + try: + cp.AllGatherKimiCPTransport(workspace=object()) + except TypeError as error: + assert "prepare" in str(error) + else: + raise AssertionError("transport must reject an unrelated workspace object") + + +def test_prepared_all_gather_transport_reuses_workspace(monkeypatch): + plan = _transport_plan() + transport = cp.AllGatherKimiCPTransport.prepare( + plan, + heads=1, + key_dim=1, + value_dim=1, + device="cpu", + ) + workspace = transport.workspace + assert workspace is not None + h_pointer = workspace.gathered_h.data_ptr() + transition_pointer = workspace.gathered_transition.data_ptr() + calls = 0 + + def fake_all_gather(output, tensor, *, group): + nonlocal calls + assert group == "cp" + gathered = output.view(4, *tensor.shape) + for rank in range(4): + gathered[rank].copy_(tensor) + if calls % 2 == 0: + gathered[rank, :, 0, :].fill_(rank + 1) + calls += 1 + + monkeypatch.setattr(cp, "_all_gather_into", fake_all_gather) + h = torch.ones(1, 1, 1, dtype=torch.float32) + transition = torch.zeros(1, 1, 1, dtype=torch.float32) + + incoming = transport.boundary_summaries(h, transition, plan=plan, group="cp", reverse=False) + outgoing = transport.boundary_summaries( + h, + transition, + plan=_transport_plan(rank=1, pre=1, post=2), + group="cp", + reverse=True, + ) + + assert incoming[0].flatten().tolist() == [1.0, 2.0] + assert outgoing[0].flatten().tolist() == [3.0, 4.0] + assert incoming[1].flatten().tolist() == outgoing[1].flatten().tolist() == [0.0, 0.0] + assert workspace.gathered_h.data_ptr() == h_pointer + assert workspace.gathered_transition.data_ptr() == transition_pointer + + +def test_collective_workspace_rejects_shape_or_world_mismatch(): + workspace = cp._KimiCPCollectiveWorkspace.allocate( + world_size=4, + heads=2, + key_dim=4, + value_dim=3, + device="cpu", + ) + h = torch.zeros(2, 3, 4, dtype=torch.float32) + transition = torch.zeros(2, 4, 4, dtype=torch.float32) + + workspace.validate(h, transition, world_size=4) + for bad_h, world_size, message in ( + (h[:1], 4, "shape mismatch"), + (h, 8, "world_size"), + ): + try: + workspace.validate(bad_h, transition, world_size=world_size) + except ValueError as error: + assert message in str(error) + else: + raise AssertionError("workspace mismatch must be rejected") + + try: + workspace.validate(h.to(torch.bfloat16), transition, world_size=4) + except TypeError as error: + assert "FP32" in str(error) + else: + raise AssertionError("workspace must not silently cast BF16 summaries") + + +def test_boundary_transport_protocol_is_structural(): + class CustomTransport: + def boundary_summaries(self, h, transition, *, plan, group, reverse): + return h.unsqueeze(0), transition.unsqueeze(0) + + assert isinstance(CustomTransport(), cp.KimiCPBoundaryTransport) + assert not isinstance(object(), cp.KimiCPBoundaryTransport) + + +def test_boundary_transport_summary_contract_rejects_precision_or_shape_changes(): + h_reference = torch.zeros(2, 3, 4, dtype=torch.float32) + transition_reference = torch.zeros(2, 4, 4, dtype=torch.float32) + valid = (h_reference.unsqueeze(0), transition_reference.unsqueeze(0)) + assert ( + cp._validate_boundary_summaries( + valid, + h_reference, + transition_reference, + rows=1, + ) + is valid + ) + + noncontiguous = torch.zeros(1, 3, 2, 4, dtype=torch.float32).transpose(1, 2) + for summaries, error_type, message in ( + (None, TypeError, "tuple"), + ((valid[0].to(torch.bfloat16), valid[1]), TypeError, "FP32"), + ((valid[0][:0], valid[1]), ValueError, "mismatch"), + ((noncontiguous, valid[1]), ValueError, "contiguous"), + ): + try: + cp._validate_boundary_summaries( + summaries, + h_reference, + transition_reference, + rows=1, + ) + except error_type as error: + assert message in str(error) + else: + raise AssertionError("invalid transport summaries must be rejected") + + +def test_distributed_autograd_delegates_summary_composition_to_cudnn(monkeypatch): + plan = _transport_plan(rank=2, pre=2, post=1) + heads, key_dim, value_dim = 2, 4, 3 + q = torch.zeros(4, heads, key_dim) + k = torch.zeros_like(q) + v = torch.zeros(4, heads, value_dim) + g = torch.zeros_like(q) + beta = torch.zeros(4, heads) + cu = torch.tensor([0, 4], dtype=torch.int32) + 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=16, + plan_name="kda_frost", + summary_plan_name="kda_summary_frost", + ) + observed = {} + + class Transport: + def boundary_summaries(self, h, transition, *, plan, group, reverse): + rows = plan.post_num_ranks if reverse else plan.pre_num_ranks + observed["transport_reverse"] = reverse + return h.unsqueeze(0).expand(rows, *h.shape).contiguous(), transition.unsqueeze(0).expand( + rows, *transition.shape + ).contiguous() + + class Context: + def save_for_backward(self, *tensors): + self.saved_tensors = tensors + + def summary(k, v, g, beta, cu, **kwargs): + return torch.ones(1, heads, value_dim, key_dim), torch.full((1, heads, key_dim, key_dim), 2.0) + + def attention(q, k, v, g, beta, cu, **kwargs): + observed["preceding_h"] = kwargs["preceding_state_update"].clone() + observed["preceding_transition"] = kwargs["preceding_transition"].clone() + return ( + torch.ones_like(v), + torch.empty(0), + torch.full((1, heads, value_dim, key_dim), 3.0), + torch.full((1, heads, value_dim, key_dim), 4.0), + ) + + def summary_bwd(q, k, g, beta, d_o, cu, **kwargs): + assert kwargs["output_context"] is True + assert kwargs["checkpoint_every_n_tokens"] == 16 + return ( + torch.full((1, heads, value_dim, key_dim), 5.0), + torch.full((1, heads, key_dim, key_dim), 6.0), + torch.tensor([7], dtype=torch.uint8), + ) + + def backward(d_o, q, k, v, g, beta, cu, scale, **kwargs): + observed["following_h"] = kwargs["following_state_grad"].clone() + observed["following_transition"] = kwargs["following_transition"].clone() + observed["composed"] = kwargs["composed_initial_state"].clone() + observed["context"] = kwargs["context"].clone() + observed["checkpoints"] = kwargs["state_checkpoints"].clone() + return ( + torch.ones_like(q), + torch.ones_like(k), + torch.ones_like(v), + torch.ones_like(g), + torch.ones_like(beta), + torch.empty(0), + torch.empty(0), + torch.empty(0), + ) + + monkeypatch.setattr(cp, "_load_cudnn_cp_ops", lambda: (attention, summary, summary_bwd, backward)) + ctx = Context() + output = cp._CudnnKdaCP.forward( + ctx, + q, + k, + v, + g, + beta, + cu, + None, + None, + plan, + "cp", + options, + Transport(), + ) + gradients = cp._CudnnKdaCP.backward(ctx, torch.ones_like(output)) + + assert observed["preceding_h"].shape[0] == 2 + assert observed["preceding_transition"].shape[0] == 2 + assert observed["following_h"].shape[0] == 1 + assert torch.equal(observed["following_transition"], torch.full((1, heads, key_dim, key_dim), 6.0)) + assert torch.equal(observed["composed"], torch.full((1, heads, value_dim, key_dim), 4.0)) + assert torch.equal(observed["context"], torch.tensor([7], dtype=torch.uint8)) + assert torch.equal(observed["checkpoints"], torch.full((1, heads, value_dim, key_dim), 3.0)) + assert len(gradients) == 12 + assert gradients[6:8] == (None, None) + + +def test_single_rank_fast_path_calls_attention_export(monkeypatch): + plan = KimiCPPlan.from_global_cu_seqlens([0, 2], rank=0, world_size=1) + q = torch.zeros(2, 1, 4) + v = torch.zeros(2, 1, 3) + beta = torch.zeros(2, 1) + cu = torch.tensor([0, 2], dtype=torch.int32) + observed = {} + + def attention(*args, **kwargs): + observed["args"] = args + observed["kwargs"] = 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())) + + result = cp.cudnn_kda_cp(q, q, v, q, beta, cu, plan=plan) + + assert torch.equal(result, torch.ones_like(v)) + assert all(actual is expected for actual, expected in zip(observed["args"][:5], (q, q, v, q, beta), strict=True)) diff --git a/tests/test_kimi_cp_plan.py b/tests/test_kimi_cp_plan.py new file mode 100644 index 00000000..f6c6c791 --- /dev/null +++ b/tests/test_kimi_cp_plan.py @@ -0,0 +1,105 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +from cula.distributed.plan import KimiCPPlan, build_kimi_cp_plan + + +def _global_cu(lengths): + values = [0] + for length in lengths: + values.append(values[-1] + length) + return torch.tensor(values, dtype=torch.int32) + + +def test_256k_twenty_short_two_long_boundary_plans(): + global_cu = _global_cu([1024] * 20 + [120832, 120832]) + plans = [KimiCPPlan.from_global_cu_seqlens(global_cu, rank=rank, world_size=8) for rank in range(8)] + + assert all(plan.local_tokens == 32768 for plan in plans) + assert plans[0].local_cu_seqlens == tuple(range(0, 20 * 1024 + 1, 1024)) + (32768,) + assert (plans[0].first_sequence, plans[0].last_sequence) == (0, 20) + assert (plans[0].pre_num_ranks, plans[0].post_num_ranks) == (0, 4) + + # Rank 4 ends the first long sample and starts the second one. + assert (plans[4].first_sequence, plans[4].last_sequence) == (20, 21) + assert plans[4].local_cu_seqlens == (0, 10240, 32768) + assert (plans[4].first_boundary_end, plans[4].last_boundary_start) == (10240, 10240) + assert (plans[4].pre_num_ranks, plans[4].post_num_ranks) == (4, 3) + + assert (plans[7].first_sequence, plans[7].last_sequence) == (21, 21) + assert (plans[7].pre_num_ranks, plans[7].post_num_ranks) == (3, 0) + + +def test_custom_uneven_shards_derive_ordered_boundary_chains(): + global_cu = _global_cu([5, 11, 7]) + offsets = (0, 4, 10, 17, 23) + plans = [ + KimiCPPlan.from_global_cu_seqlens( + global_cu, + rank=rank, + world_size=4, + shard_offsets=offsets, + ) + for rank in range(4) + ] + + assert plans[0].local_cu_seqlens == (0, 4) + assert (plans[0].pre_num_ranks, plans[0].post_num_ranks) == (0, 1) + assert plans[1].local_cu_seqlens == (0, 1, 6) + assert (plans[1].first_sequence, plans[1].last_sequence) == (0, 1) + assert (plans[1].pre_num_ranks, plans[1].post_num_ranks) == (1, 1) + assert plans[2].local_cu_seqlens == (0, 6, 7) + assert (plans[2].pre_num_ranks, plans[2].post_num_ranks) == (1, 1) + assert plans[3].local_cu_seqlens == (0, 6) + assert (plans[3].pre_num_ranks, plans[3].post_num_ranks) == (1, 0) + + +def test_build_single_rank_plan_without_distributed_runtime(): + plan = build_kimi_cp_plan([0, 3, 8]) + + assert plan.rank == 0 + assert plan.world_size == 1 + assert plan.local_cu_seqlens == (0, 3, 8) + assert plan.pre_num_ranks == plan.post_num_ranks == 0 + plan.validate_group(None) + plan.validate_local_cu_seqlens(torch.tensor([0, 3, 8], dtype=torch.int32), device=torch.device("cpu")) + + +@pytest.mark.parametrize( + ("global_cu", "message"), + [ + (torch.tensor([1, 4], dtype=torch.int32), "start at zero"), + (torch.tensor([0, 4, 4], dtype=torch.int32), "strictly increasing"), + (torch.tensor([[0, 4]], dtype=torch.int32), "one-dimensional"), + (torch.tensor([0.0, 4.0]), "int32 or int64"), + ], +) +def test_invalid_global_boundaries_fail_without_cuda_sync(global_cu, message): + with pytest.raises((TypeError, ValueError), match=message): + KimiCPPlan.from_global_cu_seqlens(global_cu, rank=0, world_size=1) + + +def test_invalid_or_empty_shards_are_rejected(): + with pytest.raises(TypeError, match="integer sequence"): + KimiCPPlan.from_global_cu_seqlens([0.0, 8.0], rank=0, world_size=1) + with pytest.raises(ValueError, match="at least one token"): + KimiCPPlan.from_global_cu_seqlens([0, 2], rank=0, world_size=4) + with pytest.raises(ValueError, match="strictly increasing"): + KimiCPPlan.from_global_cu_seqlens( + [0, 8], + rank=0, + world_size=2, + shard_offsets=(0, 0, 8), + ) + + +def test_cpu_local_boundaries_must_match_plan(): + plan = KimiCPPlan.from_global_cu_seqlens([0, 4, 8], rank=0, world_size=2) + with pytest.raises(ValueError, match="do not match"): + plan.validate_local_cu_seqlens( + torch.tensor([0, 3], dtype=torch.int32), + device=torch.device("cpu"), + ) From 42592faa3ac9e2de8cc076403cf22f4e0de606ef Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Wed, 11 Nov 2026 01:24:55 +0800 Subject: [PATCH 2/2] fix(packaging): document cuDNN source dependency --- USAGE.md | 8 ++++++-- docs/KIMI_CP_RUNTIME.md | 8 +++++--- pyproject.toml | 1 - 3 files changed, 11 insertions(+), 6 deletions(-) diff --git a/USAGE.md b/USAGE.md index abf03ea6..e269c64d 100644 --- a/USAGE.md +++ b/USAGE.md @@ -40,10 +40,14 @@ context-parallel KDA apply. See [`docs/KIMI_CP_RUNTIME.md`](docs/KIMI_CP_RUNTIME.md) for the ownership and validation contracts. -Install the optional CUDA 13 dependencies with: +The public cuDNN Frontend 1.32.2 wheel predates the concatenated-summary API. +Until a compatible release is published, install the CUDA 13 dependencies and +the tested cuDNN Frontend source revision explicitly: ```bash -pip install 'cuda-linear-attention[cudnn-cp]' +python -m pip install '.[cudnn-cp]' +python -m pip install -v \ + 'git+https://github.com/NVIDIA/cudnn-frontend.git@9e056425989e75f406a5ce3ad58cb563e4dff8b5' ``` --- diff --git a/docs/KIMI_CP_RUNTIME.md b/docs/KIMI_CP_RUNTIME.md index 0a349891..260ca71b 100644 --- a/docs/KIMI_CP_RUNTIME.md +++ b/docs/KIMI_CP_RUNTIME.md @@ -77,9 +77,11 @@ Backward performs the symmetric reverse chain: 5. Let cuDNN compose the reverse chain and compute local tensor and gate gradients. The state payload and transition remain FP32. cuLA exposes no merge-precision -switch; the cuDNN apply API owns the fixed high-precision state chain. This -runtime requires cuDNN Frontend 1.33 or a newer develop build containing the -concatenated-summary KDA API. +switch; the cuDNN apply API owns the fixed high-precision state chain. The +public cuDNN Frontend 1.32.2 wheel does not contain this API. Until a compatible +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 diff --git a/pyproject.toml b/pyproject.toml index 4da6b6c6..1d289dae 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,7 +20,6 @@ license = "Apache-2.0" [project.optional-dependencies] cudnn-cp = [ "nvidia-cutlass-dsl[cu13]>=4.8.0,<4.9", - "nvidia-cudnn-frontend>=1.33.0", ] dev = [ "pytest",