Diff Coverage

# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""RD pointwise kernels; matrix multiplication stays on the existing path."""
import triton
import triton.language as tl


@triton.jit
def kda_rd_finish_round_kernel(
    product: tl.tensor, previous: tl.tensor, matrix: tl.tensor,
    row_count: tl.constexpr, state_width: tl.constexpr, packed_width: tl.constexpr,
    block_rows: tl.constexpr, save_matrix: tl.constexpr,
) -> None:

        packed_width: Columns in the concatenated S/M row.
        block_rows: Flattened matrix rows per program, unrelated to token chunk size.
        save_matrix: Whether a later round consumes compact M.
    """
    rows = tl.program_id(0) * block_rows + tl.arange(0, block_rows)
    cols = tl.arange(0, state_width)
    offsets = rows[:, None] * packed_width + cols[None, :]
    left = tl.load(product + offsets, rows[:, None] < row_count, 0)
    old = tl.load(previous + offsets, rows[:, None] < row_count, 0)
    tl.store(product + offsets, left + old, rows[:, None] < row_count)
    if save_matrix:
        right = tl.load(product + offsets + state_width, rows[:, None] < row_count, 0)
        tl.store(matrix + rows[:, None] * state_width + cols[None, :], right, rows[:, None] < row_count)


@triton.jit
def kda_rd_finish_to_state_kernel(
    product: tl.tensor, previous: tl.tensor, matrix: tl.tensor, state: tl.tensor,
    row_count: tl.constexpr, state_width: tl.constexpr, packed_width: tl.constexpr,
    block_rows: tl.constexpr, save_matrix: tl.constexpr,
) -> None:

        packed_width: Columns in the concatenated S/M row.
        block_rows: Flattened matrix rows per program, unrelated to token chunk size.
        save_matrix: Whether a later round consumes compact M.
    """
    rows = tl.program_id(0) * block_rows + tl.arange(0, block_rows)
    cols = tl.arange(0, state_width)
    packed = rows[:, None] * packed_width + cols[None, :]
    compact = rows[:, None] * state_width + cols[None, :]
    left = tl.load(product + packed, rows[:, None] < row_count, 0)
    old = tl.load(previous + packed, rows[:, None] < row_count, 0)
    result = left + old
    tl.store(product + packed, result, rows[:, None] < row_count)
    tl.store(state + compact, result, rows[:, None] < row_count)
    if save_matrix:
        right = tl.load(product + packed + state_width, rows[:, None] < row_count, 0)
        tl.store(matrix + compact, right, rows[:, None] < row_count)

def _prepare_packed_kda_summary(kg, w, u, gate, sequence_info, state_cp, send_next, chunk_size):
    """Skip summaries with no consumer when the invocation supports absent tensors."""
    if state_cp is None and not send_next:
        return None, None
    state_shape = (kg.shape[0], u.shape[2], kg.shape[-1], u.shape[-1])
    transition = None
    if getattr(state_cp, "needs_state_summary", True):
        tail = sequence_info.tail

            cp_size: Number of state partitions.
            state_cp: Executor returning an invocation cache; None selects P2P.
        """
        ops = get_fla_kda_staged_ops()
        bind_sequence = getattr(state_cp, "bind_sequence", None)
        if sequence_info is not None and bind_sequence is not None:
            state_cp = bind_sequence(
                sequence_info, state_shape=(query.shape[0], value.shape[2], query.shape[-1], value.shape[-1]))
        plan = None if sequence_info is None else sequence_info.materialize(query.device, chunk_size)
        varlen = {} if plan is None else plan.kwargs()
        recurrent_varlen = dict(varlen)

            disable_recompute=True,
            **varlen,
        )
        # The gated query is unused here and is recomputed in backward.
        del _
        state_ext = None
        transition = None
        if sequence_info is not None:
            state_ext, transition = _prepare_packed_kda_summary(
                kg, w, u, gate, sequence_info, state_cp, send_next, chunk_size,
            )
        else:
            if state_cp is None and cp_rank == 0 and cp_size > 1:

            recv_handle = None

        send_handle = None
        send_state = None
        state_cache = (transition,) if transition is not None else ()
        if state_cp is not None:
            initial_state, state_cache = state_cp.forward(state_ext, transition)
            transition = None
        elif send_next:
            summary_initial = initial_state if transition is not None else None
            final_state = apply_kda_state_summary(state_ext, transition, summary_initial)
            send_state = final_state.contiguous()

            output_final_state=False, chunk_size=chunk_size,
            **recurrent_varlen,
        )
        # Local h was their last forward consumer on this stream.
        del w, u, kg, _
        output = ops.chunk_gla_fwd_o_gk(
            q=query, v=value_new, g=gate, A=attention_qk, h=states, scale=scale, chunk_size=chunk_size,
            **varlen,
        )

        ctx.next_rank = next_rank
        ctx.cp_rank = cp_rank
        ctx.cp_size = cp_size
        ctx.state_cp = state_cp
        ctx.state_cache_count = len(state_cache)
        ctx.receive_previous = receive_previous
        ctx.send_next = send_next
        ctx.sequence_info = sequence_info
        return output.to(value.dtype)

        (
            query, query_rstd, key, key_rstd, value, gate_raw, beta_raw, a_log, dt_bias,
            attention_qk, attention_kk, saved_initial_state, *saved_tail,
        ) = ctx.saved_tensors
        state_cache = tuple(saved_tail[:ctx.state_cache_count])
        packed_tensors = saved_tail[ctx.state_cache_count:]
        sequence_info = ctx.sequence_info
        cu_seqlens = chunk_indices = chunk_offsets = None
        if sequence_info is not None:
            cu_seqlens, chunk_indices, chunk_offsets = packed_tensors

        )
        grad_state_ext = None
        if ctx.state_cp is not None or ctx.receive_previous:
            if sequence_info is not None and not sequence_info.continues_previous:
                if not getattr(ctx.state_cp, "allows_absent_summaries", False):
                    grad_state_ext = query.new_zeros(
                        (query.shape[0], value.shape[2], query.shape[-1], value.shape[-1]), dtype=torch.float32)
            else:
                head = slice(None) if sequence_info is None else sequence_info.head
                grad_state_ext = kda_state_gradient_summary_from_prepared(

            recv_handle = None
        send_handle = None
        send_state_gradient = None
        if ctx.state_cp is not None:
            grad_final_state = ctx.state_cp.backward(grad_state_ext, state_cache)
        elif ctx.receive_previous:
            transition = state_cache[0] if state_cache else None
            summary_final_gradient = grad_final_state if transition is not None else None
            grad_initial_state = apply_kda_state_gradient_summary(grad_state_ext, transition, summary_final_gradient)
            send_state_gradient = grad_initial_state.contiguous()
            send_handle = dist.isend(send_state_gradient, dst=ctx.prev_rank, group=ctx.cp_group)

        self.cp_size = self.cp_mesh.size()
        self.cp_rank = self.cp_mesh.get_local_rank()
        self.cp_group = self.cp_mesh.get_group()
        if backend != "triton" and state_cp_method != "p2p":
            raise ValueError("KDA non-P2P state exchange requires backend='triton'.")
        self.state_cp = build_kda_state_cp(self.cp_mesh, state_cp_method, allgather_group_size)
        self.prev_rank = _global_peer_rank(
            self.cp_mesh,
            max(self.cp_rank - 1, 0),

        device_mesh = _ensure_1d(device_mesh)
        if state_cp_method not in ("p2p", "allgather", "grouped_allgather_p2p", "recursive_doubling"):
            raise ValueError("KDA requires a supported state_cp_method.")
        if backend.lower() != "triton" and state_cp_method != "p2p":
            raise ValueError("KDA non-P2P state exchange requires backend='triton'.")
        _validate_cp_degrees(device_mesh.size(), ulysses_degree, allgather_group_size)
        if state_cp_method != "grouped_allgather_p2p" and allgather_group_size != 1:
            raise ValueError("allgather_group_size is only valid for grouped_allgather_p2p.")
        if lower_bound >= 0:

@lru_cache(maxsize=1)
def _rd_kernels() -> Any:
    """Load optional Ascend kernels only when an NPU epilogue is executed."""
    return importlib.import_module(
        "hyper_parallel.components.functional._triton.kimi_delta_attention.rd_pointwise"
    )

        product: Fresh contiguous FP32 packed matrix product [B, H, 128, 256].
        previous: Previous immutable packed S/M tensor with the same layout.
        save_m: Whether a future round consumes the new contiguous M.
    """
    state_width, packed_width = product.shape[-2:]
    matrix = (torch.empty((*product.shape[:-1], state_width), dtype=product.dtype, device=product.device)
              if save_m else None)
    kernels = _rd_kernels()

    rows = product.numel() // packed_width
    block_rows = _RD_EPILOGUE_BLOCK_ROWS
    kernels.kda_rd_finish_round_kernel[((rows + block_rows - 1) // block_rows,)](
        product, previous, matrix if save_m else product,
        row_count=rows, state_width=state_width, packed_width=packed_width,
        block_rows=block_rows, save_matrix=save_m, num_warps=4)
    return matrix


def _finish_to_state(
    product: torch.Tensor, previous: torch.Tensor, save_m: bool,

        product: Fresh contiguous FP32 packed matrix product [B, H, 128, 256].
        previous: Previous immutable packed S/M tensor with the same layout.
        save_m: Whether a future round consumes the new contiguous M.
    """
    state_width, packed_width = product.shape[-2:]
    state = torch.empty((*product.shape[:-1], state_width), dtype=product.dtype, device=product.device)
    matrix = torch.empty_like(state) if save_m else None
    kernels = _rd_kernels()

    rows = product.numel() // packed_width
    block_rows = _RD_EPILOGUE_BLOCK_ROWS
    kernels.kda_rd_finish_to_state_kernel[((rows + block_rows - 1) // block_rows,)](
        product, previous, matrix if save_m else product, state,
        row_count=rows, state_width=state_width, packed_width=packed_width,
        block_rows=block_rows, save_matrix=save_m, num_warps=4)
    return matrix, state


class RoundExchange:
    """Keep immutable payload alive until send completion; outputs belong to the call."""

    def __init__(self, group: Any, rank: int, ranks: tuple[int, ...]) -> None:
        """Resolve ordered peers and the process group's protocol device."""
        if not ranks or len(set(ranks)) != len(ranks) or not 0 <= rank < len(ranks):
            raise ValueError("RD requires unique global peers and a valid local rank")
        if tuple(dist.get_process_group_ranks(group)) != ranks:
            raise ValueError("RD mesh sequence order must match its communication group order")
        self.group, self.rank, self.size, self.peers = group, rank, len(ranks), ranks
        if dist.get_backend(group) != "hccl":

        if dist.get_backend(group) != "hccl":
            raise ValueError("KDA recursive doubling requires an HCCL process group")
        # Receive-only TND ranks have no local summary to infer a device from.
        # Load the optional extension here and bind buffers to the caller's NPU.
        torch_npu = importlib.import_module("torch_npu")
        self.device = torch.device("npu", torch_npu.npu.current_device())

    def __call__(
        self, payload: torch.Tensor | None, destination: int | None, source: int | None,
        receive_shape: tuple[int, ...] | None = None,

                raise ValueError('RD send requires a payload')
            if source is not None and receive_shape is None:
                raise ValueError('RD receive without a payload requires a shape')
        elif payload.device != self.device or not payload.is_contiguous():
            raise ValueError('RD payload must be contiguous on the protocol device')
        if any(peer is not None and not 0 <= peer < self.size for peer in (source, destination)):
            raise ValueError('RD peer is outside the process group')
        output = None
        if source is not None:

) -> None:
    """Reject incompatible invocation cache layouts before any peer exchange."""
    if gradient is None:
        if state_shape is None:
            raise ValueError("RD gradient may be absent only with invocation shape metadata")
    else:
        _validate_summary(gradient, device, state_shape)
        state_shape = gradient.shape
    if not isinstance(cache, tuple) or len(cache) != expected:

def _compact_transfer(matrix: torch.Tensor) -> torch.Tensor:
    """Keep a contiguous M view from retaining a larger producer allocation."""
    if matrix.untyped_storage().nbytes() != matrix.numel() * matrix.element_size():
        return matrix.clone()
    return matrix


class KDAStateRecursiveDoubling:

            sequence_info: This call's packed document offsets and local interval.
            state_shape: B/H/K/V dimensions, without retaining a tensor allocation.
        """
        if sequence_info.start != self.rank * sequence_info.length:
            raise ValueError("RD packed interval must follow the state process-group order")
        if state_shape is not None:
            state_shape = tuple(state_shape)
            if (len(state_shape) != 4 or state_shape[-2:] != (128, 128)
                    or any(not isinstance(dim, int) or dim < 1 for dim in state_shape)):

        """
        rounds = self.rounds if plan is None else plan.execution_rounds(self.rank)
        destination, source = self._shift_peers(plan)
        if state_shape is not None and plan is None:
            raise ValueError("RD optional summaries require a packed invocation plan")
        if state is None:
            if state_shape is None or matrix is not None or rounds or destination is not None:
                raise ValueError("RD state may be absent only when no exit-state consumer exists")
            return self.exchange(None, None, source, receive_shape=state_shape), ()

            received = self.exchange(narrow_state if final else packed, destination, source)
            if source is not None:
                cache.append(before_m)
                if final:
                    narrow_state = before_m @ received + narrow_state
                else:
                    packed, before_m, narrow_state = self._compose_packed_round(
                        step, before_m, received, packed)
        if narrow_state is None:

        expected = (sum(self.rank >= distance for distance in self.distances)
                    if plan is None else plan.cache_count(self.rank))
        source, destination = self._shift_peers(plan)
        if state_shape is not None and plan is None:
            raise ValueError("RD optional summaries require a packed invocation plan")
        if gradient is None and destination is not None:
            raise ValueError("RD gradient may be absent only when no head-gradient consumer exists")
        _validate_backward_summary(gradient, cache, expected, self.device, state_shape)
        received = (self.exchange(gradient, destination, source) if gradient is not None

                payload = cache[index].transpose(-2, -1) @ adjoint
            received = self.exchange(payload, destination, source)
            if received is not None:
                # Adjoint is invocation-owned and every send has completed.
                adjoint.add_(received)
        return adjoint


@dataclass(frozen=True)

        """Reuse local bounds or locate another rank for an audit query."""
        if rank == self.local_rank and self._local_interval is not None:
            return self._local_interval
        if not 0 <= rank < self.size:
            raise ValueError("RD execution-plan rank is outside the state process group")
        return _rank_interval(self.offsets, self.local_length, rank)

    def peers(self, stage: int, rank: int) -> tuple[int | None, int | None]:
        """Return forward destination/source, both in process-group order."""
        return self._interval(rank).peers(rank, self.distances[stage])

    def shift_peers(self, rank: int) -> tuple[int | None, int | None]:
        """Return the tail's consumer and the head's producer independently."""
        interval = self._interval(rank)

    """Share document validation across rank-specific plans with bounded storage."""
    if (size < 1 or local_length < 1 or len(offsets) < 2 or offsets[0] != 0
            or offsets[-1] != size * local_length
            or any(left >= right for left, right in zip(offsets, offsets[1:]))):
        raise ValueError("RD reset planning requires ordered documents covering equal state-rank intervals")


@lru_cache(maxsize=32)
def build_reset_plan(

        An immutable CPU plan reusable across layers with identical metadata.
        Omitting rank retains only metadata for subsequent diagnostic queries.
    """
    if rank is not None and not 0 <= rank < size:
        raise ValueError("RD execution-plan rank is outside the state process group")
    if isinstance(offsets, PackedDocumentInfo):
        offsets = offsets.offsets
        if size < 1 or local_length < 1 or offsets[-1] != size * local_length:
            raise ValueError("RD reset planning requires ordered documents covering equal state-rank intervals")

    Returns:
        State executor with an invocation-owned cache, or None for the P2P/local path.
    """
    if protocol not in ("p2p", "allgather", "grouped_allgather_p2p", "recursive_doubling"):
        raise ValueError(f"Unknown KDA state_cp_method: {protocol!r}.")
    if not isinstance(allgather_group_size, int) or isinstance(allgather_group_size, bool) or allgather_group_size < 1:
        raise ValueError("KDA allgather_group_size must be a positive integer.")
    if protocol != "grouped_allgather_p2p" and allgather_group_size != 1:

    if protocol == "p2p":
        return None
    ranks = tuple(mesh.rank_list)
    rank, size = mesh.get_local_rank(), mesh.size()
    if protocol == "recursive_doubling":
        return KDAStateRecursiveDoubling(mesh.get_group(), rank, ranks) if size > 1 else None
    if protocol == "allgather":
        _validate_mesh_order(mesh)
        return KDAStateAllGather(mesh.get_group(), rank, size)
    if size % allgather_group_size:

            This rank's incoming recurrent state and a one-tensor local M cache.
        """
        _validate_summary(state, matrix)
        maps = _gather_maps(state, matrix, self.group, self.size)
        return merge_affine_summaries(maps[..., :128], maps[..., 128:], self.rank, True), (matrix,)

    def backward(self, gradient: torch.Tensor, cache: tuple[torch.Tensor, ...]) -> torch.Tensor:
        """Collect new G with the saved local M and recover the outgoing gradient.

        if result is None:
            result, _ = _scan_maps(maps, incoming, range(self.local), True)
        if send_handle is not None:
            send_handle.wait()
        return result, (matrix,)

    def backward(self, gradient: torch.Tensor, cache: tuple[torch.Tensor, ...]) -> torch.Tensor:
        """Regather G/M and reverse the owner chain with ordered state feedback.