Core API

def set_context_parallel_group(self, cp_group: ProcessGroup | None) -> None:
    """Enable or disable context parallelism for ring attention.

    Args:
        cp_group: Process group for context parallel; use None to disable.
    """
    if cp_group is None:
        self.device_mesh = None
    else:
        self.device_mesh = DeviceMesh.from_group(cp_group, device_type="cuda")

        # Need to disable load balance for torch context parallel to work.
        from torch.distributed.tensor.experimental._attention import (
            _cp_options,
            set_rotate_method,
        )

        _cp_options.enable_load_balance = False
        set_rotate_method("allgather")

def is_context_parallel_enabled(self) -> bool:
    """Return True if context parallelism is active."""
    return self.device_mesh is not None

def context_parallel_size(self) -> int:
    """Return the context parallel world size, or 1 if disabled."""
    return self.device_mesh.size() if self.device_mesh is not None else 1

def forward(self, query: Tensor, key: Tensor, value: Tensor) -> Tensor:
    """Run context-parallel SDPA (or single-rank SDPA when CP is disabled).

    Args:
        query: Query tensor in configured ``qkv_format``.
        key: Key tensor in configured ``qkv_format``.
        value: Value tensor in configured ``qkv_format``.

    Returns:
        Attention output in the same format as inputs.
    """
    # SDPA / low-level ops expect (B, H, S, D). "bshd" is (B, S, H, D) → transpose once.
    if self.qkv_format == "bshd":
        query = query.transpose(1, 2)
        key = key.transpose(1, 2)
        value = value.transpose(1, 2)
    out = self._impl(query=query, key=key, value=value)
    if self.qkv_format == "bshd":
        out = out.transpose(1, 2)
    return out

@classmethod
def from_tensor(cls, k: Tensor, v: Tensor, seq_dim: int) -> Self:
    """Build a single-chunk cache pre-filled with the given key and value tensors."""
    cache = cls(
        k_shape=k.shape,
        v_shape=v.shape,
        seq_dim=seq_dim,
        chunk_size=k.shape[seq_dim],
        window_size=k.shape[seq_dim],
        device=k.device,
        dtype=k.dtype,
    )
    cache.before_update(0)
    cache.update(k, v)
    cache.after_update(0)
    cache._curr_chunk_idx = 0
    return cache

def is_steady_state(self) -> bool:
    """Return True if the cache is full (steady-state phase)."""
    assert self._curr_chunk_idx is not None, (
        "Must call before_update() before is_steady_state()"
    )
    total_size = self._k.shape[self.seq_dim]
    is_full = total_size == self._n_cached
    is_overlapping_with_sink = (
        self.sink_size > 0
        and self._curr_chunk_idx * self.chunk_size
        < self.sink_size  # start < sink_size
    )
    return is_full and not is_overlapping_with_sink

def before_update(self, chunk_idx: int) -> None:
    """
    Prepare the cache before writing new tokens.

    If ``chunk_idx`` equals the previous chunk index, this is a no-op. Otherwise,
    we expect the ``chunk_idx`` to be +1 from the previous chunk index. In this case,
    we will roll the local window left if the cache is in steady-state, or no op
    if the cache is in filling phase.

    Args:
        chunk_idx: Chunk index of the new chunk in the full sequence.
    """
    assert self._curr_chunk_idx is None, (
        "Must call after_update() before before_update()"
    )
    self._curr_chunk_idx = chunk_idx

    if chunk_idx == self._prev_chunk_idx:
        return

    assert chunk_idx == self._prev_chunk_idx + 1, (
        "Expected the new chunk_idx to be +1 from the previous chunk_idx, "
        f"got {chunk_idx} != {self._prev_chunk_idx} + 1"
    )
    total_size = self._k.shape[self.seq_dim]
    if not self._current_chunk_overlaps_sink():
        overflow = self._n_cached + self.chunk_size - total_size
        if overflow > 0:
            self._roll_local_window_left(overflow)

def update(self, k: Tensor, v: Tensor) -> None:
    """
    Write the new chunk's keys and values into the cache.

    Must be called after ``before_update()`` and before ``after_update()``.

    Args:
        k: Keys; shape must match cached keys except at seq_dim, where length must be chunk_size.
        v: Values; shape must match cached values except at seq_dim, where length must be chunk_size.
    """
    assert self._curr_chunk_idx is not None, (
        "Must call before_update() before update()"
    )

    chunk_size_k = k.shape[self.seq_dim]
    chunk_size_v = v.shape[self.seq_dim]
    assert chunk_size_k == self.chunk_size, (
        f"Expected input k to have chunk_size ({chunk_size_k}) at seq_dim ({self.seq_dim}), "
        f"got {chunk_size_k} != {self.chunk_size}"
    )
    assert chunk_size_v == self.chunk_size, (
        f"Expected input v to have chunk_size ({chunk_size_v}) at seq_dim ({self.seq_dim}), "
        f"got {chunk_size_v} != {self.chunk_size}"
    )
    self._write_current_chunk(k, v)

def after_update(self, chunk_idx: int) -> None:
    """
    Finalize bookkeeping after writing new tokens.

    Updates ``_prev_chunk_idx`` and, in filling phase, ``_n_cached``.

    Args:
        chunk_idx: The index of the new chunk in the full sequence.
    """
    assert chunk_idx == self._curr_chunk_idx, (
        f"Expected chunk_idx to be {self._curr_chunk_idx}, got {chunk_idx}"
    )

    if self._curr_chunk_idx == self._prev_chunk_idx + 1:
        if self.is_steady_state():
            pass
        else:
            total_size = self._k.shape[self.seq_dim]
            self._n_cached = min(self._n_cached + self.chunk_size, total_size)
        self._prev_chunk_idx += 1
    elif self._curr_chunk_idx == self._prev_chunk_idx:
        pass
    else:
        raise ValueError(
            f"{self._curr_chunk_idx=} should be either {self._prev_chunk_idx + 1} or {self._prev_chunk_idx}."
        )

    self._curr_chunk_idx = None

def cached_k(self) -> Tensor:
    """
    Return cached keys for attention (valid prefix in filling phase, full buffer in steady-state).
    """
    return self._k[self._seq_slice(0, self._visible_end())]

def cached_v(self) -> Tensor:
    """
    Return cached values for attention (valid prefix in filling phase, full buffer in steady-state).
    """
    return self._v[self._seq_slice(0, self._visible_end())]

def reset(self) -> None:
    """Reset bookkeeping while preserving the allocated tensor storage."""
    self._prev_chunk_idx = -1
    self._curr_chunk_idx = None
    self._n_cached = 0

def clone_kv(self) -> tuple[Tensor, Tensor]:
    """Return clones of the full physical K/V buffers."""
    return self._k.clone(), self._v.clone()

def overwrite_kv_(self, k: Tensor, v: Tensor) -> None:
    """Overwrite the full K/V buffers without changing their addresses.

    Args:
        k: Replacement keys with the exact cache shape.
        v: Replacement values with the exact cache shape.
    """
    if k.shape != self._k.shape or v.shape != self._v.shape:
        raise ValueError(
            "overwrite_kv_ shape mismatch: "
            f"got k {tuple(k.shape)} / v {tuple(v.shape)}, "
            f"cache holds k {tuple(self._k.shape)} / v {tuple(self._v.shape)}"
        )
    self._k.copy_(k)
    self._v.copy_(v)

def get_name(self) -> str:
    """Return the marketing name reported by NVML for this device."""
    return pynvml.nvmlDeviceGetName(self.handle)

def get_cpu_affinity(self) -> list[int]:
    """Return the indices of CPUs ideally affined to this GPU per NVML."""
    affinity_string = ""
    for j in pynvml.nvmlDeviceGetCpuAffinity(
        self.handle, Device._nvml_affinity_elements
    ):
        # NVML returns a sequence of 64-bit affinity bitmasks, low word first.
        affinity_string = "{:064b}".format(j) + affinity_string
    affinity_list = [int(x) for x in affinity_string]
    affinity_list.reverse()  # so core 0 is in the 0th element of the list
    return [i for i, e in enumerate(affinity_list) if e != 0]