llm.core — Building Blocks

Reusable, framework-agnostic building blocks shared by training and serving. Everything here is pure PyTorch (no FastAPI, no trainer plumbing) so it can be reused in notebooks and other runners.

KV Cache

kv_cache

KV Cache manager for efficient autoregressive generation.

This module provides a pre-allocated cache that avoids repeated memory allocations during token generation, significantly improving inference performance.

KVCache

Pre-allocated Key-Value cache for efficient autoregressive generation.

Instead of using torch.cat to concatenate new K/V with past K/V (which allocates new memory each step), this class pre-allocates buffers and updates them in-place.

参数：

名称	类型	描述	默认
max_batch_size	int	Maximum batch size to support.	必需
max_seq_len	int	Maximum sequence length to cache.	必需
num_kv_heads	int	Number of key-value heads (for GQA, this may differ from num_heads).	必需
head_dim	int	Dimension of each attention head.	必需
device	device | str | None	Device to allocate buffers on.	None
dtype	dtype | None	Data type for cache buffers.	None

Example

cache = KVCache(max_batch_size=2, max_seq_len=16, num_kv_heads=4, head_dim=64, ... device="cpu", dtype=torch.float32) k_new = torch.randn(2, 4, 8, 64) # (batch, kv_heads, seq, head_dim) v_new = torch.randn(2, 4, 8, 64) _ = cache.update(k_new, v_new) # returns (k_cache, v_cache) views

源代码位于： src/llm/core/kv_cache.py

class KVCache:
    """Pre-allocated Key-Value cache for efficient autoregressive generation.

    Instead of using torch.cat to concatenate new K/V with past K/V (which allocates
    new memory each step), this class pre-allocates buffers and updates them in-place.

    Args:
        max_batch_size: Maximum batch size to support.
        max_seq_len: Maximum sequence length to cache.
        num_kv_heads: Number of key-value heads (for GQA, this may differ from num_heads).
        head_dim: Dimension of each attention head.
        device: Device to allocate buffers on.
        dtype: Data type for cache buffers.

    Example:
        >>> cache = KVCache(max_batch_size=2, max_seq_len=16, num_kv_heads=4, head_dim=64,
        ...                 device="cpu", dtype=torch.float32)
        >>> k_new = torch.randn(2, 4, 8, 64)   # (batch, kv_heads, seq, head_dim)
        >>> v_new = torch.randn(2, 4, 8, 64)
        >>> _ = cache.update(k_new, v_new)  # returns (k_cache, v_cache) views
    """

    def __init__(
        self,
        max_batch_size: int,
        max_seq_len: int,
        num_kv_heads: int,
        head_dim: int,
        device: torch.device | str | None = None,
        dtype: torch.dtype | None = None,
    ) -> None:
        self.max_batch_size = max_batch_size
        self.max_seq_len = max_seq_len
        self.num_kv_heads = num_kv_heads
        self.head_dim = head_dim

        # Pre-allocate buffers: [B, N_kv, S_max, D]
        self.k_cache = torch.zeros(max_batch_size, num_kv_heads, max_seq_len, head_dim, device=device, dtype=dtype)
        self.v_cache = torch.zeros(max_batch_size, num_kv_heads, max_seq_len, head_dim, device=device, dtype=dtype)
        self._seq_len = 0

    @property
    def seq_len(self) -> int:
        """Current cached sequence length."""
        return self._seq_len

    @property
    def device(self) -> torch.device:
        """Device of the cache buffers."""
        return self.k_cache.device

    @property
    def dtype(self) -> torch.dtype:
        """Data type of the cache buffers."""
        return self.k_cache.dtype

    def update(self, k_new: Tensor, v_new: Tensor) -> tuple[Tensor, Tensor]:
        """Update cache with new key-value tensors and return full cache view.

        Args:
            k_new: New key tensor of shape [B, N_kv, S_new, D].
            v_new: New value tensor of shape [B, N_kv, S_new, D].

        Returns:
            Tuple of (k_cached, v_cached) containing all cached keys and values
            up to and including the new tokens. Shape: [B, N_kv, S_total, D].

        Raises:
            ValueError: If update would exceed max_seq_len.
        """
        batch_size = k_new.size(0)
        new_tokens = k_new.size(2)
        new_seq_len = self._seq_len + new_tokens

        if new_seq_len > self.max_seq_len:
            raise ValueError(
                f"Cache overflow: trying to cache {new_seq_len} tokens, but max_seq_len is {self.max_seq_len}"
            )

        # In-place update (no memory allocation)
        self.k_cache[:batch_size, :, self._seq_len : new_seq_len] = k_new
        self.v_cache[:batch_size, :, self._seq_len : new_seq_len] = v_new
        self._seq_len = new_seq_len

        # Return view of valid cache region
        return (
            self.k_cache[:batch_size, :, :new_seq_len],
            self.v_cache[:batch_size, :, :new_seq_len],
        )

    def update_at_indices(
        self,
        batch_indices: Tensor,
        k_new: Tensor,
        v_new: Tensor,
        start_pos: Tensor | int,
    ) -> tuple[Tensor, Tensor]:
        """Update cache at specific batch indices and positions.

        This is used for continuous batching or when batch slots are managed explicitly.

        Args:
            batch_indices: Tensor of shape [B_curr] containing the cache slot indices.
            k_new: New key tensor of shape [B_curr, N_kv, S_new, D].
            v_new: New value tensor of shape [B_curr, N_kv, S_new, D].
            start_pos: Starting position to write to. Can be an int (broadcast) or
                Tensor of shape ``[B_curr, S_new]`` (one position per batch slot
                and token). When a per-batch tensor is provided, the same value
                is written for every S_new position; we rely on the caller to
                pass position_ids so we never need a host-device ``.item()``
                sync on the happy path.

        Returns:
             Tuple of (k_out, v_out) for the current batch.
             Note: This returns the *full* valid context for the *current batch indices*.
             Shape: [B_curr, N_kv, Max_Context_Len, D].
             Since different sequences have different lengths, we return up to max(start_pos + S_new).
             Correct handling usually implies the model knows how to mask using attention mask.
        """
        seq_len_new = k_new.size(2)

        if isinstance(start_pos, int):
            # Scalar start_pos: every batch slot writes the same contiguous range.
            # This is the typical prefill path (all slots start at 0).
            pos_end = start_pos + seq_len_new
            if pos_end > self.max_seq_len:
                raise ValueError(
                    f"Cache overflow: trying to cache {pos_end} tokens, but max_seq_len is {self.max_seq_len}"
                )
            self.k_cache[batch_indices, :, start_pos:pos_end] = k_new
            self.v_cache[batch_indices, :, start_pos:pos_end] = v_new
        elif seq_len_new == 1:
            # Decode path: one position per batch slot, advanced indexing is
            # already a single fused op with no host sync. ``start_pos`` is
            # ``[B, 1]`` (position_ids); flatten to ``[B]`` so each slot is
            # written at its own position. The unflattened shape would
            # broadcast with ``batch_indices`` into a [B, B] index grid,
            # writing every slot's K/V at every batch position and silently
            # corrupting unrelated cache entries.
            start_pos = start_pos.reshape(-1)
            self.k_cache[batch_indices, :, start_pos] = k_new.squeeze(2)
            self.v_cache[batch_indices, :, start_pos] = v_new.squeeze(2)
        else:
            # Mixed batch prefill path: each slot writes its own contiguous
            # range starting at start_pos[b, 0]. The previous implementation
            # used a Python-level ``for`` loop with ``.item()`` to materialize
            # one scalar start position per batch — that stalled the pipeline
            # on every step. Here we keep the whole thing on-device with one
            # advanced-indexed assignment per cache (k, v).
            #
            # ``start_pos`` is ``position_ids`` of shape ``[B, S_new]``.
            # Continuous batching left-pads each row's real positions with 0
            # (decode rows carry only their single real position), so the
            # *real* positions of a row form its leading run of strictly
            # increasing-by-one values starting at ``start_pos[b, 0]`` (e.g.
            # a prefill row is ``[0, 1, 2, 0, 0]`` -> run length 3; a decode
            # row is ``[p, 0, 0, 0]`` -> run length 1).  Everything after that
            # run is a pad and must never be written to the cache.
            #
            # Two consequences we must handle that the previous
            # deduplication-only approach got wrong:
            #  * a decode row has NO real write at position 0 this step (its
            #    real write is at ``start_pos[b,0] > 0``), so its padded
            #    position-0 columns must not clobber the genuine position-0
            #    K/V cached during the row's original prefill — filtering to
            #    the real run removes them entirely;
            #  * the overflow check must use each row's *own* real write end
            #    (``start_pos[b,0] + real_len``), not the batch-max
            #    ``seq_len_new``, otherwise a decode row near ``max_seq_len``
            #    batched with a longer prefill row raises spuriously.
            batch_starts = start_pos[:, 0]  # [B]
            col_positions = torch.arange(seq_len_new, device=start_pos.device)
            # Real iff ``start_pos[b, j] == start_pos[b, 0] + j`` — exactly the
            # leading contiguous run (pads are 0 and stop the progression).
            real_mask = start_pos == (batch_starts[:, None] + col_positions[None, :])
            real_lengths = real_mask.sum(dim=1)  # [B] real write count per row

            overflow_mask = (batch_starts + real_lengths) > self.max_seq_len
            if overflow_mask.any():
                overflow_slots = batch_indices[overflow_mask].tolist()
                raise ValueError(
                    f"Cache overflow for slots {overflow_slots} (start_pos + real_len > max_seq_len={self.max_seq_len})"
                )

            b_curr = batch_indices.size(0)
            n_kv = k_new.size(1)
            d_dim = k_new.size(3)

            # Flatten [B, S] -> [B*S], then drop the padded (non-real) entries.
            # After filtering, every slot's real position set is strictly
            # increasing (positions are equal to ``start_pos[b,0] + column``),
            # so keys are unique and no dedup is needed.
            b_idx = batch_indices.view(b_curr, 1).expand(b_curr, seq_len_new).reshape(-1)
            s_idx = start_pos.reshape(-1)
            real_flat = real_mask.reshape(-1)
            b_idx = b_idx[real_flat]
            s_idx = s_idx[real_flat]

            # Permute [B, N_kv, S, D] -> [B, S, N_kv, D] then flatten -> [B*S, N_kv, D]
            # The reshape is a view when memory is contiguous (it is here because
            # permute + reshape is followed by assignment, not by a graph op).
            k_flat = k_new.permute(0, 2, 1, 3).reshape(b_curr * seq_len_new, n_kv, d_dim)[real_flat]
            v_flat = v_new.permute(0, 2, 1, 3).reshape(b_curr * seq_len_new, n_kv, d_dim)[real_flat]

            # Advanced indexing with broadcasting on the head dim. The whole
            # write happens in one scatter-style kernel per cache; no host sync.
            n_kv_idx = torch.arange(n_kv, device=k_new.device)
            self.k_cache[b_idx[:, None], n_kv_idx[None, :], s_idx[:, None]] = k_flat
            self.v_cache[b_idx[:, None], n_kv_idx[None, :], s_idx[:, None]] = v_flat

        # Return the updated cache for these indices.
        # We return the full pre-allocated buffer [B_curr, N_kv, max_seq_len, D]
        # to ensure compatibility with the global attention masks used in continuous batching.
        return self.k_cache[batch_indices], self.v_cache[batch_indices]

    def reset(self) -> None:
        """Reset cache to empty state (does not deallocate memory)."""
        self._seq_len = 0

    def get_usable_length(self, new_tokens: int) -> int:
        """Get the usable cache length after adding new tokens."""
        return self._seq_len + new_tokens

    @classmethod
    def from_model_config(
        cls,
        max_batch_size: int,
        max_seq_len: int,
        num_layers: int,
        num_kv_heads: int,
        head_dim: int,
        device: torch.device | str | None = None,
        dtype: torch.dtype | None = None,
    ) -> list[KVCache]:
        """Create a list of KVCache objects, one per transformer layer.

        Args:
            max_batch_size: Maximum batch size.
            max_seq_len: Maximum sequence length.
            num_layers: Number of transformer layers.
            num_kv_heads: Number of KV heads per layer.
            head_dim: Dimension per head.
            device: Target device.
            dtype: Data type.

        Returns:
            List of KVCache objects, one for each layer.
        """
        return [cls(max_batch_size, max_seq_len, num_kv_heads, head_dim, device, dtype) for _ in range(num_layers)]

seq_len property

seq_len

Current cached sequence length.

device property

device

Device of the cache buffers.

dtype property

dtype

Data type of the cache buffers.

update

update(k_new, v_new)

Update cache with new key-value tensors and return full cache view.

参数：

名称	类型	描述	默认
k_new	Tensor	New key tensor of shape [B, N_kv, S_new, D].	必需
v_new	Tensor	New value tensor of shape [B, N_kv, S_new, D].	必需

返回：

类型	描述
Tensor	Tuple of (k_cached, v_cached) containing all cached keys and values
Tensor	up to and including the new tokens. Shape: [B, N_kv, S_total, D].

引发：

类型	描述
ValueError	If update would exceed max_seq_len.

源代码位于： src/llm/core/kv_cache.py

def update(self, k_new: Tensor, v_new: Tensor) -> tuple[Tensor, Tensor]:
    """Update cache with new key-value tensors and return full cache view.

    Args:
        k_new: New key tensor of shape [B, N_kv, S_new, D].
        v_new: New value tensor of shape [B, N_kv, S_new, D].

    Returns:
        Tuple of (k_cached, v_cached) containing all cached keys and values
        up to and including the new tokens. Shape: [B, N_kv, S_total, D].

    Raises:
        ValueError: If update would exceed max_seq_len.
    """
    batch_size = k_new.size(0)
    new_tokens = k_new.size(2)
    new_seq_len = self._seq_len + new_tokens

    if new_seq_len > self.max_seq_len:
        raise ValueError(
            f"Cache overflow: trying to cache {new_seq_len} tokens, but max_seq_len is {self.max_seq_len}"
        )

    # In-place update (no memory allocation)
    self.k_cache[:batch_size, :, self._seq_len : new_seq_len] = k_new
    self.v_cache[:batch_size, :, self._seq_len : new_seq_len] = v_new
    self._seq_len = new_seq_len

    # Return view of valid cache region
    return (
        self.k_cache[:batch_size, :, :new_seq_len],
        self.v_cache[:batch_size, :, :new_seq_len],
    )

update_at_indices

update_at_indices(batch_indices, k_new, v_new, start_pos)

Update cache at specific batch indices and positions.

This is used for continuous batching or when batch slots are managed explicitly.

参数：

名称	类型	描述	默认
batch_indices	Tensor	Tensor of shape [B_curr] containing the cache slot indices.	必需
k_new	Tensor	New key tensor of shape [B_curr, N_kv, S_new, D].	必需
v_new	Tensor	New value tensor of shape [B_curr, N_kv, S_new, D].	必需
start_pos	Tensor | int	Starting position to write to. Can be an int (broadcast) or Tensor of shape [B_curr, S_new] (one position per batch slot and token). When a per-batch tensor is provided, the same value is written for every S_new position; we rely on the caller to pass position_ids so we never need a host-device .item() sync on the happy path.	必需

返回：

名称	类型	描述
	Tensor	Tuple of (k_out, v_out) for the current batch.
Note	Tensor	This returns the full valid context for the current batch indices.
Shape	tuple[Tensor, Tensor]	[B_curr, N_kv, Max_Context_Len, D].
	tuple[Tensor, Tensor]	Since different sequences have different lengths, we return up to max(start_pos + S_new).
	tuple[Tensor, Tensor]	Correct handling usually implies the model knows how to mask using attention mask.

源代码位于： src/llm/core/kv_cache.py

def update_at_indices(
    self,
    batch_indices: Tensor,
    k_new: Tensor,
    v_new: Tensor,
    start_pos: Tensor | int,
) -> tuple[Tensor, Tensor]:
    """Update cache at specific batch indices and positions.

    This is used for continuous batching or when batch slots are managed explicitly.

    Args:
        batch_indices: Tensor of shape [B_curr] containing the cache slot indices.
        k_new: New key tensor of shape [B_curr, N_kv, S_new, D].
        v_new: New value tensor of shape [B_curr, N_kv, S_new, D].
        start_pos: Starting position to write to. Can be an int (broadcast) or
            Tensor of shape ``[B_curr, S_new]`` (one position per batch slot
            and token). When a per-batch tensor is provided, the same value
            is written for every S_new position; we rely on the caller to
            pass position_ids so we never need a host-device ``.item()``
            sync on the happy path.

    Returns:
         Tuple of (k_out, v_out) for the current batch.
         Note: This returns the *full* valid context for the *current batch indices*.
         Shape: [B_curr, N_kv, Max_Context_Len, D].
         Since different sequences have different lengths, we return up to max(start_pos + S_new).
         Correct handling usually implies the model knows how to mask using attention mask.
    """
    seq_len_new = k_new.size(2)

    if isinstance(start_pos, int):
        # Scalar start_pos: every batch slot writes the same contiguous range.
        # This is the typical prefill path (all slots start at 0).
        pos_end = start_pos + seq_len_new
        if pos_end > self.max_seq_len:
            raise ValueError(
                f"Cache overflow: trying to cache {pos_end} tokens, but max_seq_len is {self.max_seq_len}"
            )
        self.k_cache[batch_indices, :, start_pos:pos_end] = k_new
        self.v_cache[batch_indices, :, start_pos:pos_end] = v_new
    elif seq_len_new == 1:
        # Decode path: one position per batch slot, advanced indexing is
        # already a single fused op with no host sync. ``start_pos`` is
        # ``[B, 1]`` (position_ids); flatten to ``[B]`` so each slot is
        # written at its own position. The unflattened shape would
        # broadcast with ``batch_indices`` into a [B, B] index grid,
        # writing every slot's K/V at every batch position and silently
        # corrupting unrelated cache entries.
        start_pos = start_pos.reshape(-1)
        self.k_cache[batch_indices, :, start_pos] = k_new.squeeze(2)
        self.v_cache[batch_indices, :, start_pos] = v_new.squeeze(2)
    else:
        # Mixed batch prefill path: each slot writes its own contiguous
        # range starting at start_pos[b, 0]. The previous implementation
        # used a Python-level ``for`` loop with ``.item()`` to materialize
        # one scalar start position per batch — that stalled the pipeline
        # on every step. Here we keep the whole thing on-device with one
        # advanced-indexed assignment per cache (k, v).
        #
        # ``start_pos`` is ``position_ids`` of shape ``[B, S_new]``.
        # Continuous batching left-pads each row's real positions with 0
        # (decode rows carry only their single real position), so the
        # *real* positions of a row form its leading run of strictly
        # increasing-by-one values starting at ``start_pos[b, 0]`` (e.g.
        # a prefill row is ``[0, 1, 2, 0, 0]`` -> run length 3; a decode
        # row is ``[p, 0, 0, 0]`` -> run length 1).  Everything after that
        # run is a pad and must never be written to the cache.
        #
        # Two consequences we must handle that the previous
        # deduplication-only approach got wrong:
        #  * a decode row has NO real write at position 0 this step (its
        #    real write is at ``start_pos[b,0] > 0``), so its padded
        #    position-0 columns must not clobber the genuine position-0
        #    K/V cached during the row's original prefill — filtering to
        #    the real run removes them entirely;
        #  * the overflow check must use each row's *own* real write end
        #    (``start_pos[b,0] + real_len``), not the batch-max
        #    ``seq_len_new``, otherwise a decode row near ``max_seq_len``
        #    batched with a longer prefill row raises spuriously.
        batch_starts = start_pos[:, 0]  # [B]
        col_positions = torch.arange(seq_len_new, device=start_pos.device)
        # Real iff ``start_pos[b, j] == start_pos[b, 0] + j`` — exactly the
        # leading contiguous run (pads are 0 and stop the progression).
        real_mask = start_pos == (batch_starts[:, None] + col_positions[None, :])
        real_lengths = real_mask.sum(dim=1)  # [B] real write count per row

        overflow_mask = (batch_starts + real_lengths) > self.max_seq_len
        if overflow_mask.any():
            overflow_slots = batch_indices[overflow_mask].tolist()
            raise ValueError(
                f"Cache overflow for slots {overflow_slots} (start_pos + real_len > max_seq_len={self.max_seq_len})"
            )

        b_curr = batch_indices.size(0)
        n_kv = k_new.size(1)
        d_dim = k_new.size(3)

        # Flatten [B, S] -> [B*S], then drop the padded (non-real) entries.
        # After filtering, every slot's real position set is strictly
        # increasing (positions are equal to ``start_pos[b,0] + column``),
        # so keys are unique and no dedup is needed.
        b_idx = batch_indices.view(b_curr, 1).expand(b_curr, seq_len_new).reshape(-1)
        s_idx = start_pos.reshape(-1)
        real_flat = real_mask.reshape(-1)
        b_idx = b_idx[real_flat]
        s_idx = s_idx[real_flat]

        # Permute [B, N_kv, S, D] -> [B, S, N_kv, D] then flatten -> [B*S, N_kv, D]
        # The reshape is a view when memory is contiguous (it is here because
        # permute + reshape is followed by assignment, not by a graph op).
        k_flat = k_new.permute(0, 2, 1, 3).reshape(b_curr * seq_len_new, n_kv, d_dim)[real_flat]
        v_flat = v_new.permute(0, 2, 1, 3).reshape(b_curr * seq_len_new, n_kv, d_dim)[real_flat]

        # Advanced indexing with broadcasting on the head dim. The whole
        # write happens in one scatter-style kernel per cache; no host sync.
        n_kv_idx = torch.arange(n_kv, device=k_new.device)
        self.k_cache[b_idx[:, None], n_kv_idx[None, :], s_idx[:, None]] = k_flat
        self.v_cache[b_idx[:, None], n_kv_idx[None, :], s_idx[:, None]] = v_flat

    # Return the updated cache for these indices.
    # We return the full pre-allocated buffer [B_curr, N_kv, max_seq_len, D]
    # to ensure compatibility with the global attention masks used in continuous batching.
    return self.k_cache[batch_indices], self.v_cache[batch_indices]

reset

reset()

Reset cache to empty state (does not deallocate memory).

源代码位于： src/llm/core/kv_cache.py

def reset(self) -> None:
    """Reset cache to empty state (does not deallocate memory)."""
    self._seq_len = 0

get_usable_length

get_usable_length(new_tokens)

Get the usable cache length after adding new tokens.

源代码位于： src/llm/core/kv_cache.py

def get_usable_length(self, new_tokens: int) -> int:
    """Get the usable cache length after adding new tokens."""
    return self._seq_len + new_tokens

from_model_config classmethod

from_model_config(max_batch_size, max_seq_len, num_layers, num_kv_heads, head_dim, device=None, dtype=None)

Create a list of KVCache objects, one per transformer layer.

参数：

名称	类型	描述	默认
max_batch_size	int	Maximum batch size.	必需
max_seq_len	int	Maximum sequence length.	必需
num_layers	int	Number of transformer layers.	必需
num_kv_heads	int	Number of KV heads per layer.	必需
head_dim	int	Dimension per head.	必需
device	device | str | None	Target device.	None
dtype	dtype | None	Data type.	None

返回：

类型	描述
list[KVCache]	List of KVCache objects, one for each layer.

源代码位于： src/llm/core/kv_cache.py

@classmethod
def from_model_config(
    cls,
    max_batch_size: int,
    max_seq_len: int,
    num_layers: int,
    num_kv_heads: int,
    head_dim: int,
    device: torch.device | str | None = None,
    dtype: torch.dtype | None = None,
) -> list[KVCache]:
    """Create a list of KVCache objects, one per transformer layer.

    Args:
        max_batch_size: Maximum batch size.
        max_seq_len: Maximum sequence length.
        num_layers: Number of transformer layers.
        num_kv_heads: Number of KV heads per layer.
        head_dim: Dimension per head.
        device: Target device.
        dtype: Data type.

    Returns:
        List of KVCache objects, one for each layer.
    """
    return [cls(max_batch_size, max_seq_len, num_kv_heads, head_dim, device, dtype) for _ in range(num_layers)]

create_decoder_kv_caches

create_decoder_kv_caches(model, batch_size)

Create per-layer KV caches sized for a DecoderModel.

源代码位于： src/llm/core/kv_cache.py

def create_decoder_kv_caches(model: Any, batch_size: int) -> list[KVCache]:
    """Create per-layer KV caches sized for a DecoderModel."""
    block = model.transformer_blocks[0]
    num_kv_heads = block.self_attn.num_kv_heads
    head_dim = block.self_attn.head_dim
    device = next(model.parameters()).device
    dtype = next(model.parameters()).dtype
    return KVCache.from_model_config(
        max_batch_size=batch_size,
        max_seq_len=model.max_seq_len,
        num_layers=len(model.transformer_blocks),
        num_kv_heads=num_kv_heads,
        head_dim=head_dim,
        device=device,
        dtype=dtype,
    )

reset_all_caches

reset_all_caches(caches)

Reset all caches in a list.

源代码位于： src/llm/core/kv_cache.py

def reset_all_caches(caches: list[KVCache]) -> None:
    """Reset all caches in a list."""
    for cache in caches:
        cache.reset()

Attention Implementations

mha

MultiHeadAttention

Bases: Module

源代码位于： src/llm/core/attn/mha.py

@register_attention("mha")
class MultiHeadAttention(nn.Module):
    # Standard MHA writes into the KV-cache pool during autoregressive decoding.
    # The continuous batching engine and training engine both depend on this.
    set_attention_kv_cache_capability("mha", supports=True)
    """
    Multi-Head Attention (MHA) mechanism.

    Integrates Layer Normalization and residual connection, supporting Pre-LN and Post-LN modes.

    Args:
        hidden_size (int): Total dimension of the model.
        num_heads (int): Number of attention heads. Must divide hidden_size. Defaults to 8.
        p (float): Dropout probability applied to attention weights and final output. Defaults to 0.1.
        bias (bool): Whether to use bias in the linear layers (QKV projection and output projection). Defaults to False.
        eps (float): Epsilon value for Layer Normalization. Defaults to 1e-5.
        norm_first (bool): Whether to use Pre-LN (True) or Post-LN (False) architecture. Defaults to True.
        is_causal (bool): Whether to apply causal masking by default (e.g., for decoders). Defaults to False.
        device (torch.device | str | None): Target device for model parameters. Defaults to None (inferred).
        dtype (torch.dtype | None): Target data type for model parameters. Defaults to None (inferred).

    Attributes:
        head_dim (int): Dimension of each attention head.
        qkv_proj (nn.Linear): Combined Q, K, V projection layer.
        out_proj (nn.Linear): Output projection layer.
        norm (nn.LayerNorm): Layer normalization module.
        dropout (nn.Dropout): Dropout layer applied after the output projection.
    """

    def __init__(
        self,
        hidden_size: int,
        num_heads: int = 8,
        p: float = 0.1,
        bias: bool = False,
        eps: float = 1e-5,
        norm_first: bool = True,
        is_causal: bool = False,
        include_norm_residual: bool = True,  # New parameter
        num_kv_heads: int | None = None,  # New: For GQA/MQA support
        window_size: int | None = None,  # Sliding window attention
        max_seq_len: int | None = None,  # RoPE max context (required if use_rope)
        use_rope: bool = False,  # Rotary position embedding (real Llama/Mistral)
        rope_theta: float = 10000.0,  # RoPE base frequency
        alibi: ALiBiPositionBias | None = None,  # Linear-bias positional encoding
        device: torch.device | str | None = None,
        dtype: torch.dtype | None = None,
    ):
        super().__init__()

        if hidden_size % num_heads != 0:
            raise ValueError(f"hidden_size ({hidden_size}) must be divisible by num_heads ({num_heads})")

        factory_kwargs = make_factory_kwargs(device, dtype)
        self.hidden_size = hidden_size
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
        self.head_dim = hidden_size // num_heads
        self.kv_dim = self.num_kv_heads * self.head_dim
        self.norm_first = norm_first  # Relevant only if include_norm_residual is True
        self.is_causal = is_causal
        self.p = p
        self.include_norm_residual = include_norm_residual
        self.window_size = window_size
        self.use_rope = use_rope
        self.alibi = alibi

        if use_rope:
            # RoPE rotates Q and K by head position (real Llama/Mistral inject
            # position here, not via additive embeddings). ``max_seq_len`` is
            # required so the cos/sin table is sized to the model's context
            # (RIL ISS-062 — core.rope had zero callers before this wiring).
            if max_seq_len is None:
                raise ValueError("use_rope=True requires max_seq_len for the RoPE cos/sin table")
            self.rope = RotaryPositionEmbedding(
                dim=self.head_dim,
                max_seq_len=max_seq_len,
                base=rope_theta,
                **factory_kwargs,
            )

        if self.num_heads % self.num_kv_heads != 0:
            raise ValueError(f"num_heads ({self.num_heads}) must be divisible by num_kv_heads ({self.num_kv_heads})")

        self.norm = None
        if self.include_norm_residual:
            self.norm = nn.LayerNorm(hidden_size, eps=eps, **factory_kwargs)

        self.qkv_dim = (self.num_heads + 2 * self.num_kv_heads) * self.head_dim
        self.qkv_proj = nn.Linear(hidden_size, self.qkv_dim, bias=bias, **factory_kwargs)
        self.out_proj = nn.Linear(hidden_size, hidden_size, bias=bias, **factory_kwargs)
        self.dropout = nn.Dropout(p)  # This is for the output projection

        self._init_weights()

    def _init_weights(self):
        """Initialize linear layer weights (Xavier uniform) and biases (zeros)."""
        for proj in [self.qkv_proj, self.out_proj]:
            nn.init.xavier_uniform_(proj.weight)
            if proj.bias is not None:
                nn.init.zeros_(proj.bias)

    def _rope_positions(
        self,
        batch_size: int,
        seq_len: int,
        start_pos: int | Tensor | None,
        device: torch.device,
    ) -> Tensor | None:
        """Compute per-token positions for RoPE.

        Three cases:

        * ``start_pos`` is a ``Tensor`` (the batch-serving path threads the
          per-row ``position_ids`` through as ``start_pos``) — return it
          unchanged so each row gets its own absolute positions.
        * ``start_pos`` is an ``int`` (KV-cache decode: the current chunk
          starts at cache length) — absolute positions are
          ``[start_pos, start_pos + seq_len)``, broadcast across the batch.
        * ``start_pos`` is ``None`` (pure prefill) — positions are
          ``[0, seq_len)`` (RoPE's internal default), return ``None``.

        Returns ``None`` when RoPE's internal ``0..seq_len`` default applies,
        otherwise a ``[B, S]`` long tensor of absolute positions.
        """
        if start_pos is None:
            return None
        if isinstance(start_pos, Tensor):
            # The batch-serving path threads a ``[B, S]`` ``position_ids``
            # tensor through as ``start_pos`` — return it unchanged. The
            # ``== 0`` early-return below is int-only; evaluating it on a
            # multi-element tensor raises a Boolean-ambiguity error, so the
            # tensor branch MUST come first (RIL ISS-112).
            return start_pos
        if start_pos == 0:
            return None
        base = int(start_pos)
        return torch.arange(base, base + seq_len, device=device, dtype=torch.long).expand(batch_size, -1)

    def forward(
        self,
        hidden_states: Tensor,
        attn_mask: Tensor | None = None,
        is_causal: bool | None = None,
        kv_cache: KVCache | None = None,
        use_cache: bool = False,
        batch_indices: Tensor | None = None,
        start_pos: int | Tensor | None = None,
        paged_kv_cache: PagedKVCache | None = None,
        layer_idx: int | None = None,
        prefix_kv: tuple[Tensor, Tensor] | None = None,
    ) -> Tensor | tuple[Tensor, tuple[Tensor, Tensor]]:
        """
        Forward pass.

        Args:
            hidden_states (Tensor): Input tensor of shape [B, S, H] (Batch, Sequence Length, Hidden Size).
            attn_mask (Tensor | None): Optional attention mask.
                - For F.scaled_dot_product_attention, expected to be a boolean tensor where `True` indicates masking.
                - Shape should be broadcastable to [B, N, S, S] (Batch, Num Heads, Seq Len, Seq Len).
                - E.g., Padding mask could be [B, 1, 1, S] or [B, 1, S, S].
            is_causal (bool | None): Whether to enforce causal masking for this forward pass.
                - If `None` (default), uses the default `self.is_causal` set during initialization.
                - If `True` or `False`, overrides the default setting.
            kv_cache (KVCache | None): Pre-allocated KV cache for efficient autoregressive generation.
                When provided, updates are done in-place without memory allocation.
            prefix_kv (tuple[Tensor, Tensor] | None): Optional prefix K/V to prepend to the
                projected K/V before the attention compute. Used by the Prefix Tuning slice
                (T2 PEFT) — see :class:`llm.core.prefix_tuning.PrefixTuningAttention`. Tensors
                must be shape ``[B, num_kv_heads, prefix_len, head_dim]``. Injected **after**
                the KV cache write so the cache only stores dynamic tokens.
            use_cache (bool): Whether to return the updated (key, value) pair.
            batch_indices (Tensor | None): Indices for specific KV cache slots [B]. Use with update_at_indices.
            start_pos (int | Tensor | None): Explicit write position for cache update. required if batch_indices is used.
            paged_kv_cache (PagedKVCache | None): Block-allocator KV cache. When set
                the linear ``kv_cache`` argument is ignored — the model writes K/V
                into the paged blocks and reads via ``paged_attention_forward``.
                ``batch_indices`` doubles as the per-row ``seq_id`` (the engine
                passes slot ids that we treat as ``PagedKVCache`` sequence ids).
            layer_idx (int | None): Index of this block in the decoder. Required
                when ``paged_kv_cache`` is set; used to slice the per-layer K/V
                tensor out of ``PagedKVCache.k_cache[layer_idx]``.

        Returns:
            Tensor or tuple[Tensor, tuple[Tensor, Tensor]]:
                - If use_cache=False: Output tensor of shape [B, S, H].
                - If use_cache=True: (Output tensor, (current_key, current_value))
        """
        batch_size, seq_len, _ = hidden_states.size()

        # Determine causality for this call
        use_causal = self.is_causal if is_causal is None else is_causal

        # Prepare input for QKV projection
        # If norm and residual are handled by this module, apply norm first (if pre-norm)
        if self.include_norm_residual and self.norm is not None:
            residual = hidden_states
            x_for_qkv = self.norm(hidden_states) if self.norm_first else hidden_states
        else:
            # If no norm/residual by this module, use hidden_states directly
            # No residual variable needed here if not added by this module
            x_for_qkv = hidden_states

        # 2. Project Q, K, V and reshape
        qkv = self.qkv_proj(x_for_qkv)  # [B, S, (N_q + 2*N_kv) * D]

        # Split Q, K, V
        q_size = self.num_heads * self.head_dim
        kv_size = self.num_kv_heads * self.head_dim

        q, k, v = torch.split(qkv, [q_size, kv_size, kv_size], dim=-1)

        # Reshape and transpose for attention calculation: [B, N, S, D]
        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = v.view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)

        # Rotary position embedding (real Llama/Mistral). Applies to Q/K
        # BEFORE they are written to the KV cache, so the cache stores the
        # rotated keys (matching HF Llama cache semantics — RIL ISS-062).
        if self.use_rope:
            positions = self._rope_positions(batch_size, seq_len, start_pos, device=q.device)
            q, k = self.rope(q, k, positions)

        # KV Cache handling
        if paged_kv_cache is not None:
            if self.alibi is not None:
                # ALiBi is only wired into the linear (non-paged) path in this
                # milestone; paged attention has no additive-bias channel, so
                # reject loudly instead of silently dropping the PE (RIL —
                # ALiBi milestone).
                raise NotImplementedError(
                    "ALiBi (use_alibi=True) is not supported with paged attention; "
                    "use attn_impl='mha' with the linear KV cache."
                )
            if prefix_kv is not None:
                # ``_forward_paged`` never threads ``prefix_kv`` through the
                # block allocator, so a PrefixTuningAttention-wrapped MHA under
                # paged/continuous-batching serving silently ignored the prefix
                # and ran as though the model was never prefix-tuned (silent
                # wrong answer). Reject loudly, matching FlashAttention's
                # existing refusal (RIL round-71 paged-prefix fix).
                raise NotImplementedError(
                    "prefix tuning (prefix_kv) is not supported with paged attention; "
                    "use the linear KV-cache path (kv_caches) instead."
                )
            return self._forward_paged(
                q=q,
                k=k,
                v=v,
                attn_mask=attn_mask,
                paged_kv_cache=paged_kv_cache,
                batch_indices=batch_indices,
                layer_idx=layer_idx,
                residual=residual if self.include_norm_residual and self.norm is not None else None,
            )

        has_past = False
        if kv_cache is not None:
            # Use efficient pre-allocated cache (in-place update)
            if batch_indices is not None:
                if start_pos is None:
                    raise ValueError("start_pos must be provided when using batch_indices for KV cache update.")
                k, v = kv_cache.update_at_indices(batch_indices, k, v, start_pos)
                # has_past logic for 'update_at_indices' scenario:
                # It implies we are manually managing positions, so usually we don't rely on global seq_len check?
                # SDPA 'is_causal' logic:
                # If we are in Decode (seq_len=1), is_causal=False usually (we attend to all past).
                # If we are in Prefill, is_causal=True.
                # Let's assume the caller sets is_causal correctly or we are managing 'has_past' effectively.
                # For 'update_at_indices', we likely have past data.
                has_past = True
            else:
                k, v = kv_cache.update(k, v)
                has_past = kv_cache.seq_len > seq_len

        if use_cache:
            current_kv = (k, v)

        # GQA: Repeat K, V if needed
        # Prefix injection happens BEFORE the GQA repeat so the prefix is
        # treated exactly like a regular token (repeated num_queries_per_kv
        # times across the query heads).
        if prefix_kv is not None:
            prefix_k, prefix_v = prefix_kv
            if prefix_k.shape != prefix_v.shape:
                raise ValueError(
                    f"prefix_k and prefix_v must share shape; got {tuple(prefix_k.shape)} vs {tuple(prefix_v.shape)}"
                )
            if prefix_k.shape[1] != self.num_kv_heads:
                raise ValueError(
                    f"prefix num_kv_heads ({prefix_k.shape[1]}) must match attention num_kv_heads ({self.num_kv_heads})"
                )
            if prefix_k.shape[3] != self.head_dim:
                raise ValueError(
                    f"prefix head_dim ({prefix_k.shape[3]}) must match attention head_dim ({self.head_dim})"
                )
            k = torch.cat([prefix_k, k], dim=2)
            v = torch.cat([prefix_v, v], dim=2)

        if self.num_kv_heads != self.num_heads:
            # k, v: [B, N_kv, S, D] -> [B, N_q, S, D]
            num_queries_per_kv = self.num_heads // self.num_kv_heads
            k = k.repeat_interleave(num_queries_per_kv, dim=1)
            v = v.repeat_interleave(num_queries_per_kv, dim=1)

        # 3. Attention computation
        # Use common SDPA wrapper to handle mask polarity and window size
        #
        # ALiBi bias (RIL — ALiBi milestone): a square [1, N, Sk, Sk] bias
        # reflects key column j at absolute position j. The queries sit at the
        # LAST ``Sq`` absolute positions (prefill with no cache: 0..Sq-1; with
        # KV-cache decode, the current tokens are at ``Sk-Sq..Sk-1``), and
        # ALiBi is translation-covariant (bias = slope * (key_pos - query_pos)),
        # so taking the last ``Sq`` rows is correct in all three cases.
        attn_bias = None
        if self.alibi is not None:
            seq_len_q = q.size(-2)
            seq_len_k = k.size(-2)
            attn_bias = self.alibi.get_bias(seq_len_k, device=q.device, dtype=q.dtype)[:, :, -seq_len_q:, :]

        attn_output = sdpa(
            query=q,
            key=k,
            value=v,
            attn_mask=attn_mask,
            attn_bias=attn_bias,
            dropout_p=self.p if self.training else 0.0,
            is_causal=use_causal if not has_past else False,
            scale=None,
            window_size=self.window_size,
            # The prefix was prepended to K/V above (k_len = prefix_len + q_len
            # in prefill), so an unshifted top-left causal mask hides every real
            # key. Shift the causal diagonal by the prefix length (round-71
            # prefix-causal fix).
            prefix_len=prefix_k.shape[2] if prefix_kv is not None else 0,
        )  # Output shape: [B, N, S, D]

        # 4. Combine head outputs
        # [B, N, S, D] -> [B, S, N, D] -> [B, S, H]
        attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_len, self.hidden_size)

        # 5. Output projection and dropout
        projected_output = self.dropout(self.out_proj(attn_output))

        if self.include_norm_residual and self.norm is not None:
            # 6. Residual connection
            output = residual + projected_output

            # 7. Layer Normalization (Post-LN mode)
            if not self.norm_first:
                output = self.norm(output)
        else:
            # No residual, no norm by this module
            output = projected_output

        if use_cache:
            return output, current_kv
        return output

    def _forward_paged(
        self,
        q: Tensor,
        k: Tensor,
        v: Tensor,
        attn_mask: Tensor | None,
        paged_kv_cache: PagedKVCache,
        batch_indices: Tensor | None,
        layer_idx: int | None,
        residual: Tensor | None,
    ) -> Tensor:
        """Run the attention computation through a :class:`PagedKVCache`.

        Per-row write: ``paged_kv_cache.update(seq_id, k_b.T, v_b.T)``
        appends the new tokens to that sequence's block table. The
        sequence id is taken from ``batch_indices`` (the engine passes
        slot ids; ``PagedKVCache`` treats them as sequence ids). After
        the writes we read the per-row block tables / seq lengths and
        call :func:`paged_attention_forward` to compute attention over
        the gathered context.

        Args:
            q: Projected query tensor ``[B, N_q, S, D]``.
            k: Projected key tensor ``[B, N_kv, S, D]``.
            v: Projected value tensor ``[B, N_kv, S, D]``.
            paged_kv_cache: The block-allocator cache (typed ``object``
                to avoid a circular import on ``core.paged_attention``).
            batch_indices: Slot ids per row ``[B]``; doubles as the
                ``seq_id`` for ``PagedKVCache.update``.
            layer_idx: Index of this block in the decoder; slices
                ``paged_kv_cache.k_cache[layer_idx]``.
            residual: Pre-norm residual tensor (``None`` when this
                block does not own the residual).

        Returns:
            Attention output ``[B, S, H]`` after output projection.
        """
        if layer_idx is None:
            raise ValueError(
                "layer_idx is required when paged_kv_cache is set; DecoderModel threads it through TransformerBlock."
            )
        if batch_indices is None:
            raise ValueError(
                "batch_indices is required when paged_kv_cache is set; the engine passes slot ids per row."
            )

        batch_size, _, seq_len, _ = q.shape

        # 1. Per-row write into the paged cache. ``PagedKVCache.update``
        #    expects ``[B, T, N_kv, D]`` (it transposes internally), so
        #    transpose our ``[B, N_kv, T, D]`` k/v to match.
        #
        #    Only the REAL tokens may be appended: continuous batching pads
        #    the batch to the longest prompt, and the padded rows' K/V is
        #    garbage (masked attention over pad tokens). The engine's
        #    attention mask marks padded query rows fully-masked, so the
        #    real length per row is the number of query rows whose first
        #    column is visible.
        seq_ids = batch_indices.tolist()
        lengths = None
        if attn_mask is not None:
            lengths = (~attn_mask[:, 0, :, 0]).sum(dim=-1)  # [B] real query rows
        for b, seq_id in enumerate(seq_ids):
            n = int(lengths[b]) if lengths is not None else k.shape[2]
            paged_kv_cache.update(
                seq_id=int(seq_id),
                k_new=k[b : b + 1, :, :n].transpose(1, 2),
                v_new=v[b : b + 1, :, :n].transpose(1, 2),
                # Scope the write to this layer's cache slice; layer 0 owns
                # block-table allocation/extension, later layers reuse it.
                layer_idx=layer_idx if layer_idx is not None else 0,
            )

        # 2. Build ``block_tables`` and ``seq_lens`` per row from the
        #    BlockManager's view of each sequence.
        block_size = paged_kv_cache.block_size
        max_blocks = max(
            (len(paged_kv_cache.get_block_table(int(sid))) for sid in seq_ids),
            default=1,
        )
        # Pad block-table columns to a single tensor shape.
        block_tables = torch.full((batch_size, max_blocks), -1, dtype=torch.long, device=q.device)
        seq_lens = torch.zeros(batch_size, dtype=torch.long, device=q.device)
        for b, seq_id in enumerate(seq_ids):
            table = paged_kv_cache.get_block_table(int(seq_id))
            block_tables[b, : len(table)] = torch.tensor(table, dtype=torch.long)
            seq_lens[b] = paged_kv_cache.block_manager.get_num_tokens(int(seq_id))

        # 3. Run the paged attention kernel over the per-layer slice.
        k_layer = paged_kv_cache.k_cache[layer_idx]
        v_layer = paged_kv_cache.v_cache[layer_idx]
        attn_output = paged_attention_forward(
            q=q,
            k_cache=k_layer,
            v_cache=v_layer,
            block_tables=block_tables,
            seq_lens=seq_lens,
            num_kv_heads=self.num_kv_heads,
            block_size=block_size,
            # Per-row real query-token counts so the kernel only applies the
            # causal overlay to prefill rows, not to decode rows whose single
            # query is left-padded into a larger batch-max row (RIL ISS-048).
            query_lens=lengths,
        )  # [B, N_q, S, D]

        # 4. Reshape and project — same post-processing as the linear path.
        attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_len, self.hidden_size)
        projected_output = self.dropout(self.out_proj(attn_output))

        if self.include_norm_residual and self.norm is not None and residual is not None:
            output = residual + projected_output
            if not self.norm_first:
                output = self.norm(output)
        else:
            output = projected_output
        return output

forward

forward(hidden_states, attn_mask=None, is_causal=None, kv_cache=None, use_cache=False, batch_indices=None, start_pos=None, paged_kv_cache=None, layer_idx=None, prefix_kv=None)

Forward pass.

参数：

名称	类型	描述	默认
hidden_states	Tensor	Input tensor of shape [B, S, H] (Batch, Sequence Length, Hidden Size).	必需
attn_mask	Tensor | None	Optional attention mask. - For F.scaled_dot_product_attention, expected to be a boolean tensor where True indicates masking. - Shape should be broadcastable to [B, N, S, S] (Batch, Num Heads, Seq Len, Seq Len). - E.g., Padding mask could be [B, 1, 1, S] or [B, 1, S, S].	None
is_causal	bool | None	Whether to enforce causal masking for this forward pass. - If None (default), uses the default self.is_causal set during initialization. - If True or False, overrides the default setting.	None
kv_cache	KVCache | None	Pre-allocated KV cache for efficient autoregressive generation. When provided, updates are done in-place without memory allocation.	None
prefix_kv	tuple[Tensor, Tensor] | None	Optional prefix K/V to prepend to the projected K/V before the attention compute. Used by the Prefix Tuning slice (T2 PEFT) — see :class:llm.core.prefix_tuning.PrefixTuningAttention. Tensors must be shape [B, num_kv_heads, prefix_len, head_dim]. Injected after the KV cache write so the cache only stores dynamic tokens.	None
use_cache	bool	Whether to return the updated (key, value) pair.	False
batch_indices	Tensor | None	Indices for specific KV cache slots [B]. Use with update_at_indices.	None
start_pos	int | Tensor | None	Explicit write position for cache update. required if batch_indices is used.	None
paged_kv_cache	PagedKVCache | None	Block-allocator KV cache. When set the linear kv_cache argument is ignored — the model writes K/V into the paged blocks and reads via paged_attention_forward. batch_indices doubles as the per-row seq_id (the engine passes slot ids that we treat as PagedKVCache sequence ids).	None
layer_idx	int | None	Index of this block in the decoder. Required when paged_kv_cache is set; used to slice the per-layer K/V tensor out of PagedKVCache.k_cache[layer_idx].	None

返回：

类型	描述
Tensor | tuple[Tensor, tuple[Tensor, Tensor]]	Tensor or tuple[Tensor, tuple[Tensor, Tensor]]: - If use_cache=False: Output tensor of shape [B, S, H]. - If use_cache=True: (Output tensor, (current_key, current_value))

源代码位于： src/llm/core/attn/mha.py

def forward(
    self,
    hidden_states: Tensor,
    attn_mask: Tensor | None = None,
    is_causal: bool | None = None,
    kv_cache: KVCache | None = None,
    use_cache: bool = False,
    batch_indices: Tensor | None = None,
    start_pos: int | Tensor | None = None,
    paged_kv_cache: PagedKVCache | None = None,
    layer_idx: int | None = None,
    prefix_kv: tuple[Tensor, Tensor] | None = None,
) -> Tensor | tuple[Tensor, tuple[Tensor, Tensor]]:
    """
    Forward pass.

    Args:
        hidden_states (Tensor): Input tensor of shape [B, S, H] (Batch, Sequence Length, Hidden Size).
        attn_mask (Tensor | None): Optional attention mask.
            - For F.scaled_dot_product_attention, expected to be a boolean tensor where `True` indicates masking.
            - Shape should be broadcastable to [B, N, S, S] (Batch, Num Heads, Seq Len, Seq Len).
            - E.g., Padding mask could be [B, 1, 1, S] or [B, 1, S, S].
        is_causal (bool | None): Whether to enforce causal masking for this forward pass.
            - If `None` (default), uses the default `self.is_causal` set during initialization.
            - If `True` or `False`, overrides the default setting.
        kv_cache (KVCache | None): Pre-allocated KV cache for efficient autoregressive generation.
            When provided, updates are done in-place without memory allocation.
        prefix_kv (tuple[Tensor, Tensor] | None): Optional prefix K/V to prepend to the
            projected K/V before the attention compute. Used by the Prefix Tuning slice
            (T2 PEFT) — see :class:`llm.core.prefix_tuning.PrefixTuningAttention`. Tensors
            must be shape ``[B, num_kv_heads, prefix_len, head_dim]``. Injected **after**
            the KV cache write so the cache only stores dynamic tokens.
        use_cache (bool): Whether to return the updated (key, value) pair.
        batch_indices (Tensor | None): Indices for specific KV cache slots [B]. Use with update_at_indices.
        start_pos (int | Tensor | None): Explicit write position for cache update. required if batch_indices is used.
        paged_kv_cache (PagedKVCache | None): Block-allocator KV cache. When set
            the linear ``kv_cache`` argument is ignored — the model writes K/V
            into the paged blocks and reads via ``paged_attention_forward``.
            ``batch_indices`` doubles as the per-row ``seq_id`` (the engine
            passes slot ids that we treat as ``PagedKVCache`` sequence ids).
        layer_idx (int | None): Index of this block in the decoder. Required
            when ``paged_kv_cache`` is set; used to slice the per-layer K/V
            tensor out of ``PagedKVCache.k_cache[layer_idx]``.

    Returns:
        Tensor or tuple[Tensor, tuple[Tensor, Tensor]]:
            - If use_cache=False: Output tensor of shape [B, S, H].
            - If use_cache=True: (Output tensor, (current_key, current_value))
    """
    batch_size, seq_len, _ = hidden_states.size()

    # Determine causality for this call
    use_causal = self.is_causal if is_causal is None else is_causal

    # Prepare input for QKV projection
    # If norm and residual are handled by this module, apply norm first (if pre-norm)
    if self.include_norm_residual and self.norm is not None:
        residual = hidden_states
        x_for_qkv = self.norm(hidden_states) if self.norm_first else hidden_states
    else:
        # If no norm/residual by this module, use hidden_states directly
        # No residual variable needed here if not added by this module
        x_for_qkv = hidden_states

    # 2. Project Q, K, V and reshape
    qkv = self.qkv_proj(x_for_qkv)  # [B, S, (N_q + 2*N_kv) * D]

    # Split Q, K, V
    q_size = self.num_heads * self.head_dim
    kv_size = self.num_kv_heads * self.head_dim

    q, k, v = torch.split(qkv, [q_size, kv_size, kv_size], dim=-1)

    # Reshape and transpose for attention calculation: [B, N, S, D]
    q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
    k = k.view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
    v = v.view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)

    # Rotary position embedding (real Llama/Mistral). Applies to Q/K
    # BEFORE they are written to the KV cache, so the cache stores the
    # rotated keys (matching HF Llama cache semantics — RIL ISS-062).
    if self.use_rope:
        positions = self._rope_positions(batch_size, seq_len, start_pos, device=q.device)
        q, k = self.rope(q, k, positions)

    # KV Cache handling
    if paged_kv_cache is not None:
        if self.alibi is not None:
            # ALiBi is only wired into the linear (non-paged) path in this
            # milestone; paged attention has no additive-bias channel, so
            # reject loudly instead of silently dropping the PE (RIL —
            # ALiBi milestone).
            raise NotImplementedError(
                "ALiBi (use_alibi=True) is not supported with paged attention; "
                "use attn_impl='mha' with the linear KV cache."
            )
        if prefix_kv is not None:
            # ``_forward_paged`` never threads ``prefix_kv`` through the
            # block allocator, so a PrefixTuningAttention-wrapped MHA under
            # paged/continuous-batching serving silently ignored the prefix
            # and ran as though the model was never prefix-tuned (silent
            # wrong answer). Reject loudly, matching FlashAttention's
            # existing refusal (RIL round-71 paged-prefix fix).
            raise NotImplementedError(
                "prefix tuning (prefix_kv) is not supported with paged attention; "
                "use the linear KV-cache path (kv_caches) instead."
            )
        return self._forward_paged(
            q=q,
            k=k,
            v=v,
            attn_mask=attn_mask,
            paged_kv_cache=paged_kv_cache,
            batch_indices=batch_indices,
            layer_idx=layer_idx,
            residual=residual if self.include_norm_residual and self.norm is not None else None,
        )

    has_past = False
    if kv_cache is not None:
        # Use efficient pre-allocated cache (in-place update)
        if batch_indices is not None:
            if start_pos is None:
                raise ValueError("start_pos must be provided when using batch_indices for KV cache update.")
            k, v = kv_cache.update_at_indices(batch_indices, k, v, start_pos)
            # has_past logic for 'update_at_indices' scenario:
            # It implies we are manually managing positions, so usually we don't rely on global seq_len check?
            # SDPA 'is_causal' logic:
            # If we are in Decode (seq_len=1), is_causal=False usually (we attend to all past).
            # If we are in Prefill, is_causal=True.
            # Let's assume the caller sets is_causal correctly or we are managing 'has_past' effectively.
            # For 'update_at_indices', we likely have past data.
            has_past = True
        else:
            k, v = kv_cache.update(k, v)
            has_past = kv_cache.seq_len > seq_len

    if use_cache:
        current_kv = (k, v)

    # GQA: Repeat K, V if needed
    # Prefix injection happens BEFORE the GQA repeat so the prefix is
    # treated exactly like a regular token (repeated num_queries_per_kv
    # times across the query heads).
    if prefix_kv is not None:
        prefix_k, prefix_v = prefix_kv
        if prefix_k.shape != prefix_v.shape:
            raise ValueError(
                f"prefix_k and prefix_v must share shape; got {tuple(prefix_k.shape)} vs {tuple(prefix_v.shape)}"
            )
        if prefix_k.shape[1] != self.num_kv_heads:
            raise ValueError(
                f"prefix num_kv_heads ({prefix_k.shape[1]}) must match attention num_kv_heads ({self.num_kv_heads})"
            )
        if prefix_k.shape[3] != self.head_dim:
            raise ValueError(
                f"prefix head_dim ({prefix_k.shape[3]}) must match attention head_dim ({self.head_dim})"
            )
        k = torch.cat([prefix_k, k], dim=2)
        v = torch.cat([prefix_v, v], dim=2)

    if self.num_kv_heads != self.num_heads:
        # k, v: [B, N_kv, S, D] -> [B, N_q, S, D]
        num_queries_per_kv = self.num_heads // self.num_kv_heads
        k = k.repeat_interleave(num_queries_per_kv, dim=1)
        v = v.repeat_interleave(num_queries_per_kv, dim=1)

    # 3. Attention computation
    # Use common SDPA wrapper to handle mask polarity and window size
    #
    # ALiBi bias (RIL — ALiBi milestone): a square [1, N, Sk, Sk] bias
    # reflects key column j at absolute position j. The queries sit at the
    # LAST ``Sq`` absolute positions (prefill with no cache: 0..Sq-1; with
    # KV-cache decode, the current tokens are at ``Sk-Sq..Sk-1``), and
    # ALiBi is translation-covariant (bias = slope * (key_pos - query_pos)),
    # so taking the last ``Sq`` rows is correct in all three cases.
    attn_bias = None
    if self.alibi is not None:
        seq_len_q = q.size(-2)
        seq_len_k = k.size(-2)
        attn_bias = self.alibi.get_bias(seq_len_k, device=q.device, dtype=q.dtype)[:, :, -seq_len_q:, :]

    attn_output = sdpa(
        query=q,
        key=k,
        value=v,
        attn_mask=attn_mask,
        attn_bias=attn_bias,
        dropout_p=self.p if self.training else 0.0,
        is_causal=use_causal if not has_past else False,
        scale=None,
        window_size=self.window_size,
        # The prefix was prepended to K/V above (k_len = prefix_len + q_len
        # in prefill), so an unshifted top-left causal mask hides every real
        # key. Shift the causal diagonal by the prefix length (round-71
        # prefix-causal fix).
        prefix_len=prefix_k.shape[2] if prefix_kv is not None else 0,
    )  # Output shape: [B, N, S, D]

    # 4. Combine head outputs
    # [B, N, S, D] -> [B, S, N, D] -> [B, S, H]
    attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_len, self.hidden_size)

    # 5. Output projection and dropout
    projected_output = self.dropout(self.out_proj(attn_output))

    if self.include_norm_residual and self.norm is not None:
        # 6. Residual connection
        output = residual + projected_output

        # 7. Layer Normalization (Post-LN mode)
        if not self.norm_first:
            output = self.norm(output)
    else:
        # No residual, no norm by this module
        output = projected_output

    if use_cache:
        return output, current_kv
    return output

sdpa

sdpa

sdpa(query, key, value, attn_mask=None, attn_bias=None, dropout_p=0.0, is_causal=False, scale=None, window_size=None, prefix_len=0)

Computes Scaled Dot-Product Attention using torch.nn.functional.scaled_dot_product_attention.

Acts as a compatibility wrapper for the codebase conventions: 1. Handles attn_mask where True indicates masking out (opposite to Torch SDPA). 2. Handles window_size by manually merging masks if necessary. 3. Folds a signed additive attn_bias (ALiBi) into the float mask.

参数：

名称	类型	描述	默认
query	Tensor	Shape (B, N, Sq, D).	必需
key	Tensor	Shape (B, N, Sk, D).	必需
value	Tensor	Shape (B, N, Sk, D).	必需
attn_mask	Tensor | None	Mask where True indicates elements to MASK OUT. Can be boolean or 0/1 float additive (legacy).	None
attn_bias	Tensor | None	Signed additive score bias (ALiBi), shape broadcastable to [B, N, Sq, Sk]. Summed into the mask (never through the 0/-inf masking channels).	None
dropout_p	float	Dropout probability.	0.0
is_causal	bool	Whether to apply causal masking.	False
scale	float | None	Scaling factor.	None
window_size	int | None	Sliding window size.	None
prefix_len	int	How many prefix K/V columns were prepended ahead of the real keys (Prefix Tuning). Shifts any causal diagonal by this amount so the real context stays visible (round-71 fix).	0

源代码位于： src/llm/core/attn/sdpa.py

def sdpa(
    query: Tensor,
    key: Tensor,
    value: Tensor,
    attn_mask: Tensor | None = None,
    attn_bias: Tensor | None = None,
    dropout_p: float = 0.0,
    is_causal: bool = False,
    scale: float | None = None,
    window_size: int | None = None,
    prefix_len: int = 0,
) -> Tensor:
    """
    Computes Scaled Dot-Product Attention using `torch.nn.functional.scaled_dot_product_attention`.

    Acts as a compatibility wrapper for the codebase conventions:
    1. Handles `attn_mask` where True indicates masking out (opposite to Torch SDPA).
    2. Handles `window_size` by manually merging masks if necessary.
    3. Folds a signed additive `attn_bias` (ALiBi) into the float mask.

    Args:
        query (Tensor): Shape (B, N, Sq, D).
        key (Tensor): Shape (B, N, Sk, D).
        value (Tensor): Shape (B, N, Sk, D).
        attn_mask (Tensor | None): Mask where True indicates elements to MASK OUT.
                                   Can be boolean or 0/1 float additive (legacy).
        attn_bias (Tensor | None): Signed additive score bias (ALiBi), shape
                                   broadcastable to [B, N, Sq, Sk]. Summed
                                   into the mask (never through the 0/-inf
                                   masking channels).
        dropout_p (float): Dropout probability.
        is_causal (bool): Whether to apply causal masking.
        scale (float | None): Scaling factor.
        window_size (int | None): Sliding window size.
        prefix_len (int): How many prefix K/V columns were prepended ahead of
            the real keys (Prefix Tuning). Shifts any causal diagonal by this
            amount so the real context stays visible (round-71 fix).
    """

    # 1. Handle Window Size and Mask Merging
    # If window_size is set, or if both attn_mask and is_causal are provided, we often need manual masking.
    has_window = window_size is not None and window_size > 0

    # Complex path: We need to construct a mask manually if:
    # - We have a window constraint (Torch SDPA doesn't support window_size directly yet for all backends/cases easily without mask)
    # - We have BOTH causal=True AND an attention mask (Torc SDPA generally prefers one or the other, or merged)
    # - causal=True with a prefix shift: torch's native is_causal mask is
    #   left-aligned (diagonal=1) and cannot express the prefix offset, so the
    #   shifted causal mask must be materialized here (round-71 prefix fix).
    if has_window or (is_causal and (attn_mask is not None or prefix_len != 0)):
        seq_len_q = query.size(-2)
        seq_len_k = key.size(-2)
        device = query.device

        # Start with efficient creation of causal mask if needed
        # We build a boolean mask where True = Mask Out (our convention)
        full_mask = None

        if is_causal:
            # True = Mask out (Upper triangle)
            # Shape: (Sq, Sk)
            #
            # ``prefix_len`` shifts the diagonal: when prefix K/V was prepended
            # ahead of the real keys (Prefix Tuning), key column ``j`` is at
            # absolute position ``j`` while query row ``i`` is at absolute
            # position ``prefix_len + i`` — so the causal bound is
            # ``j <= prefix_len + i``, i.e. ``diagonal = 1 + prefix_len``.
            # Without the shift the real (right-shifted) context keys were all
            # masked and the query attended only prefix keys (RIL round-71,
            # prefix-causal HIGH).
            full_mask = torch.triu(
                torch.ones((seq_len_q, seq_len_k), device=device, dtype=torch.bool),
                diagonal=1 + prefix_len,
            )

        if has_window:
            # The query block sits immediately after the cached keys, so its
            # rows carry absolute positions ``[seq_len_k - seq_len_q,
            # seq_len_k)``. During prefill (seq_len_k == seq_len_q) that is
            # ``[0, seq_len_k)``; during a KV-cache decode step
            # (seq_len_q == 1) it is ``[seq_len_k - 1]`` — the current
            # position. Using the *relative* row index (0..seq_len_q-1)
            # against absolute key columns attends the OLDEST window_size
            # keys at every decode step instead of the keys just before the
            # current position (RIL — decode window bug).
            row_offset = max(0, seq_len_k - seq_len_q)
            row_idx = (torch.arange(seq_len_q, device=device) + row_offset).unsqueeze(1)
            col_idx = torch.arange(seq_len_k, device=device).unsqueeze(0)
            # True = Mask out (distance > window)
            # Standard window attention: |i - j| > w
            # Note: For Causal Window, it's just i - j > w (past) ... but usually window is symmetric or causal.
            # Assuming standard generalized window constraint here.
            window_mask = torch.abs(row_idx - col_idx) > window_size
            full_mask = window_mask if full_mask is None else (full_mask | window_mask)

        if attn_mask is not None:
            # Normalize to the query device first: sparse/streaming masks are
            # built CPU-side (decoder.forward / serving), but the causal/window
            # ``full_mask`` below lives on ``device`` — OR-ing a CPU bool mask
            # into it raised on GPU machines (RIL ISS-300). No-op when already
            # on device, so this is cheap on the common path.
            attn_mask = attn_mask.to(device)
            # attn_mask: True = Mask out
            # We assume attn_mask is broadcastable.
            # If attn_mask is float/int 0/1. we should convert to bool for logical ops if we can,
            # but usually it's passed as bool in this codebase.
            # If it's float additive (-inf), this merging logic is trickier.
            # Assuming bool mask for complex merging.
            if attn_mask.dtype == torch.bool:
                full_mask = attn_mask if full_mask is None else (full_mask | attn_mask)
            elif attn_mask.dtype.is_floating_point:
                # Float additive mask (0 = keep, -inf = mask out, Torch SDPA
                # convention). It cannot be merged with the boolean
                # ``full_mask`` via ``|`` — additive and boolean masks are
                # different spaces. Convert the boolean part to an additive
                # mask (0 / -inf) and sum them, so the caller's additive mask
                # is NOT silently dropped on the window / causal+mask path
                # (RIL ISS-115).
                bool_part = full_mask
                # Normalize the float mask's shape to a key-additive layout
                # ([B, 1, 1, S_k]) so it plays well with the rank-2
                # causal/window ``full_mask``. reward_task / reward tests
                # pass a plain ``[B, S]`` float mask.
                float_mask = attn_mask
                if float_mask.ndim == 2:
                    float_mask = float_mask.unsqueeze(1).unsqueeze(1)
                if bool_part is not None:
                    additive_base = torch.where(
                        bool_part,
                        torch.tensor(float("-inf"), device=query.device, dtype=float_mask.dtype),
                        torch.zeros((), device=query.device, dtype=float_mask.dtype),
                    )
                    full_mask = additive_base + float_mask
                else:
                    full_mask = float_mask
            else:
                # Integer 0/1 mask (e.g. the long padding mask emitted by the
                # SFT/DPO/reward data pipelines, where 1 = real token and
                # 0 = pad). 1 is *keep*, so the mask-out predicate here is
                # ``== 0`` — NOT ``to(bool)`` (which would flip padding to
                # keep and *real* tokens to mask out). The data pipeline
                # emits ``[B, S]``; expand to ``[B, 1, S]`` so it broadcasts
                # against the ``[Sq, Sk]`` causal/window ``full_mask``.
                mask_out = attn_mask == 0
                if mask_out.ndim == 2:
                    mask_out = mask_out.unsqueeze(1)
                full_mask = mask_out if full_mask is None else (full_mask | mask_out)

        # Now we have a mask where True = Mask Out (bool) or 0/-inf additive
        # (float), matching the caller's convention.
        # F.sdpa expects True = Keep (for boolean masks); float masks are
        # passed through as additive. And we set is_causal=False because we
        # baked causal in.

        torch_mask = None
        if full_mask is not None:
            torch_mask = ~full_mask if full_mask.dtype == torch.bool else full_mask
            # F.sdpa expects a boolean mask of rank >= 3 (broadcastable to
            # [B, N, Sq, Sk]). The window/causal construction is rank-2
            # ([Sq, Sk]); a caller-supplied [B, 1, S] int mask merges to
            # rank-3 ([B, Sq, Sk]), which has no head dimension. Expand the
            # head axis so the merged mask actually applies per head instead
            # of raising "expanded size ... at non-singleton dimension 1".
            if torch_mask.ndim == 3 and query.ndim == 4:
                torch_mask = torch_mask.unsqueeze(1)

        if attn_bias is not None:
            torch_mask = _add_attn_bias(torch_mask, attn_bias, query)

        return functional.scaled_dot_product_attention(
            query,
            key,
            value,
            attn_mask=torch_mask,
            dropout_p=dropout_p,
            is_causal=False,
            scale=scale,
        )

    # 2. Fast Path: No complex masking conflict
    # We can rely on F.sdpa's native logic or simple inversion.

    torch_attn_mask = None
    if attn_mask is not None:
        # My convention: True = Mask Out
        # Torch convention: True = Keep
        if attn_mask.dtype == torch.bool:
            torch_attn_mask = ~attn_mask
        elif attn_mask.dtype.is_floating_point:
            # Float additive mask (0 = keep, -inf = mask out): pass through
            # unchanged — this is already the Torch SDPA convention. A
            # caller-supplied [B, S] float mask (reward_task / reward tests)
            # needs a broadcastable key axis for 4-D query, same as int.
            torch_attn_mask = attn_mask
            if torch_attn_mask.ndim == 2 and query.ndim == 4:
                torch_attn_mask = torch_attn_mask.unsqueeze(1).unsqueeze(1)
        else:
            # Integer 0/1 mask (e.g. a long padding mask from a data
            # pipeline): 1 = real = keep, 0 = pad. Torch wants True = Keep,
            # so the long mask maps 1:1 onto a bool mask with NO inversion
            # (1 -> True = keep); passing the raw long mask would be the
            # bitwise-not-free path but Torch SDPA rejects non-bool/float
            # dtypes. The data pipeline emits ``[B, S]``; F.sdpa rejects a
            # 2-D bool mask (and even ``[B, 1, S]``) for a 4-D query — it
            # needs at least ``[B, 1, 1, S]`` / ``[1, 1, S_q, S_k]``. Insert
            # a broadcastable head+key axis.
            torch_attn_mask = attn_mask.to(torch.bool)
            if torch_attn_mask.ndim == 2 and query.ndim == 4:
                torch_attn_mask = torch_attn_mask.unsqueeze(1).unsqueeze(1)

    if attn_bias is not None:
        if is_causal:
            # ``F.scaled_dot_product_attention`` rejects an explicit float
            # mask combined with ``is_causal=True``, and ALiBi forces a float
            # additive mask. Materialise the causal mask (upper triangle →
            # -inf) and fold ALiBi in; the mask now does the causality work.
            seq_len_q, seq_len_k = query.size(-2), key.size(-2)
            causal = torch.triu(
                torch.full(
                    (seq_len_q, seq_len_k),
                    float("-inf"),
                    device=query.device,
                    dtype=query.dtype,
                ),
                diagonal=1 + prefix_len,
            )  # [Sq, Sk]; lower triangle = 0 (keep); diagonal shifted by prefix
            torch_attn_mask = _add_attn_bias(torch_attn_mask, causal, query)
            torch_attn_mask = _add_attn_bias(torch_attn_mask, attn_bias.to(query.dtype), query)
            is_causal = False
        else:
            torch_attn_mask = _add_attn_bias(torch_attn_mask, attn_bias, query)

    return functional.scaled_dot_product_attention(
        query, key, value, attn_mask=torch_attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale
    )

mla

MultiLatentAttention

Bases: Module

源代码位于： src/llm/core/attn/mla.py

@register_attention("mla")
class MultiLatentAttention(nn.Module):
    # MLA caches the K, V from ``input_kv_proj`` into the standard
    # ``KVCache`` / ``PagedKVCache`` pool — same cache contract as MHA.
    # The architectural caveat is that this is the *placeholder* MLA
    # (learnable latent queries, uniform-mean output broadcast over the
    # sequence). Real DeepSeek-V2-style MLA with latent-compressed K, V
    # and decoupled RoPE is a separate, larger slice.
    set_attention_kv_cache_capability("mla", supports=True)
    """
    Multi-Latent Attention mechanism implementation.

    Similar to Multi-Head Attention but incorporates a set of learnable latent vectors
    that can capture different aspects of the input and enhance attention computations.

    Args:
        hidden_size: Hidden dimension size.
        num_heads: Number of attention heads. Defaults to 8.
        num_latents: Number of latent vectors. Defaults to 16.
        latent_dim: Size of each latent vector. If None, equals hidden_size. Defaults to None.
        dropout_p: Dropout probability. Defaults to 0.1.
        bias: Whether to use bias in the linear layers. Defaults to True.
        eps: Epsilon value for Layer Normalization. Defaults to 1e-5.
        norm_first: Whether to use Layer Normalization before attention. Defaults to True.
        is_causal: Whether to use causal attention. Defaults to False.
        device: Device for the model.
        dtype: Data type for the model parameters.

    Note:
        This is the **placeholder** MLA — the latent queries attend to
        the full ``input_kv_proj(x)`` (no latent-dim compression, no
        decoupled RoPE). The output is a uniform average over the
        ``num_latents`` latent outputs broadcast to every sequence
        position, so the architectural benefit of per-position KV cache
        is limited; the cache only saves the ``input_kv_proj`` cost on
        incremental decode. DeepSeek-V2-style MLA (latent-compressed
        K, V, decoupled RoPE) is a separate slice.
    """

    def __init__(
        self,
        hidden_size: int,
        num_heads: int = 8,
        p: float = 0.1,
        bias: bool = True,
        eps: float = 1e-5,
        norm_first: bool = True,
        is_causal: bool = False,
        include_norm_residual: bool = True,
        num_kv_heads: int | None = None,
        window_size: int | None = None,
        num_latents: int = 16,
        latent_dim: int | None = None,
        device: torch.device | None = None,
        dtype: torch.dtype | None = None,
        **_: object,
    ):
        super().__init__()

        if hidden_size % num_heads != 0:
            raise ValueError(f"hidden_size ({hidden_size}) must be divisible by num_heads ({num_heads})")

        factory_kwargs = make_factory_kwargs(device, dtype)
        self.hidden_size = hidden_size
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
        self.head_dim = hidden_size // num_heads
        self.scale = 1 / math.sqrt(self.head_dim)
        self.norm_first = norm_first
        self.num_latents = num_latents
        self.latent_dim = latent_dim if latent_dim is not None else hidden_size
        self.is_causal = is_causal
        self.dropout_p = p
        self.include_norm_residual = include_norm_residual

        # Layer Normalization - shared for all attention operations.
        # Only built when this module OWNS the norm+residual (the block passes
        # include_norm_residual=False, expecting a plain sublayer — MHA
        # mirrors this by leaving self.norm None). When False, MLA must NOT
        # apply an internal norm/residual; the block adds the residual (RIL
        # ISS-139, double-norm regression otherwise).
        self.norm = nn.LayerNorm(hidden_size, eps=eps, **factory_kwargs) if include_norm_residual else None

        # Learnable latent vectors - initialized directly with normal distribution
        self.latents = nn.Parameter(torch.randn(1, num_latents, self.latent_dim, **factory_kwargs) * 0.02)

        # Latent projections
        self.latent_q_proj = nn.Linear(self.latent_dim, hidden_size, bias=bias, **factory_kwargs)
        self.latent_v_proj = nn.Linear(hidden_size, self.latent_dim, bias=bias, **factory_kwargs)
        self.latent_output_proj = nn.Linear(self.latent_dim, hidden_size, bias=bias, **factory_kwargs)

        # Input projection for Key and Value
        self.input_kv_proj = nn.Linear(hidden_size, 2 * hidden_size, bias=bias, **factory_kwargs)

        # Output projection
        self.out_proj = nn.Linear(hidden_size, hidden_size, bias=bias, **factory_kwargs)

        # Single dropout layer
        self.dropout = nn.Dropout(p)

        # Initialize weights
        self._init_weights()

    def _init_weights(self):
        """Initialize module parameters with optimized scheme."""
        # Xavier/Glorot uniform initialization for all linear layers
        for module in [
            self.latent_q_proj,
            self.latent_v_proj,
            self.latent_output_proj,
            self.input_kv_proj,
            self.out_proj,
        ]:
            # Use improved initialization with gain based on activation
            gain = 1.0  # Linear activation gain
            fan_in, fan_out = module.weight.shape
            std = gain * math.sqrt(2.0 / (fan_in + fan_out))
            nn.init.trunc_normal_(module.weight, std=std)
            if module.bias is not None:
                nn.init.zeros_(module.bias)

    def _latent_attention(
        self,
        k: Tensor,
        v: Tensor,
        batch_size: int,
        attn_mask: Tensor | None = None,
        *,
        is_causal: bool,
    ) -> Tensor:
        """
        Process latent attention computation as a separate method for clarity.
        This function computes the attention between latent queries and input sequence.

        Args:
            k: Key tensor with shape [batch_size, num_heads, seq_len, head_dim]
            v: Value tensor with shape [batch_size, num_heads, seq_len, head_dim]
            batch_size: Batch size
            attn_mask: Optional attention mask.

                Incoming shape is the standard MHA convention
                ``[B, 1, S_q, S_k]`` where ``S_q`` is the new-token count
                and ``S_k`` is the cached context length. The latent
                attention's query axis is ``num_latents`` (latents are
                static parameters, not derived from the input), so we
                collapse ``S_q`` to the LAST position's mask (the
                canonical causal mask for the current generation step)
                and broadcast over ``num_latents``.

        Returns:
            Processed latent output with shape [batch_size, num_latents, hidden_size]
        """
        # Expand and project latent queries in one operation
        latent_q = self.latents.expand(batch_size, -1, -1)
        latent_q = self.latent_q_proj(latent_q.reshape(batch_size * self.num_latents, self.latent_dim))
        latent_q = latent_q.view(batch_size, self.num_latents, self.num_heads, self.head_dim)
        latent_q = latent_q.permute(0, 2, 1, 3)  # [batch_size, num_heads, num_latents, head_dim]

        # Reshape the MHA-style mask ``[B, 1, S_q, S_k]`` into the latent
        # attention's mask ``[B, 1, num_latents, S_k]``. The latent queries
        # are static parameters representing the "current token", so they all
        # share ONE key-visibility mask per row: the mask of that row's LAST
        # REAL query position.
        #
        # Under continuous batching S_q is the batch-max query length and
        # shorter rows (decode, or a shorter prefill) are right-padded with
        # all-masked rows. Collapsing to the batch-max last row
        # (``attn_mask[:, :, -1:, :]``) picked a pad row for those elements
        # — SDPA returned zeros for their latent attention (silent wrong
        # output modulo bias). A row is "real" iff it can see at least one
        # key (at least one False in True=masked convention); row 0 of a
        # decode row is always real.
        if attn_mask is not None:
            # Normalize the 2-D ``[B, S]`` padding mask emitted by the SFT /
            # DPO / Reward data pipelines (long 0/1, 1 = real token) into the
            # 4-D ``[B, 1, S_q, S_k]`` masked-out form the collapse below
            # expects — the same lowering ``sdpa()`` applies for MHA. The
            # latent queries are a single "current position" (``S_q == 1``),
            # so the key-only mask expands to ``[B, 1, 1, S_k]`` and the
            # ``flat = attn_mask[:, 0]`` slice still yields ``[B, S_q, S_k]``.
            # Without this, training a ``mla`` model crashed on the first
            # mask-consuming forward (no MLA engine e2e existed to catch it —
            # surfaced by the TASK-206 MLA TP engine test).
            if attn_mask.ndim == 2:
                if attn_mask.dtype != torch.bool:
                    # long/float 0/1: 1 = keep -> mask-out predicate is == 0
                    # (mirror of sdpa()'s integer-mask path).
                    attn_mask = attn_mask == 0
                attn_mask = attn_mask.unsqueeze(1).unsqueeze(1)  # [B, 1, 1, S_k]
            flat = attn_mask[:, 0]  # [B, S_q, S_k]
            row_has_visible = ~flat.all(dim=-1)  # [B, S_q]
            q_idx = torch.arange(flat.shape[1], device=flat.device)
            last_real = (q_idx.unsqueeze(0) * row_has_visible).max(dim=-1).values  # [B]
            view = flat[torch.arange(flat.shape[0], device=flat.device), last_real]  # [B, S_k]
            attn_mask = view.unsqueeze(1).unsqueeze(1).expand(-1, 1, self.num_latents, -1)

        # Compute attention with conditional dropout during training
        latent_output = sdpa(
            query=latent_q,
            key=k,
            value=v,
            attn_mask=attn_mask,
            dropout_p=self.dropout_p if self.training else 0.0,
            is_causal=is_causal,
            scale=self.scale,
        )  # [batch_size, num_heads, num_latents, head_dim]

        # Reshape for further processing
        latent_output = latent_output.permute(0, 2, 1, 3).reshape(batch_size, self.num_latents, self.hidden_size)

        # Transform through latent dimension
        latent_output = self.latent_v_proj(latent_output)
        latent_output = self.latent_output_proj(latent_output.reshape(batch_size * self.num_latents, self.latent_dim))

        return latent_output.reshape(batch_size, self.num_latents, self.hidden_size)

    def forward(
        self,
        hidden_states: Tensor,
        attn_mask: Tensor | None = None,
        is_causal: bool | None = None,
        kv_cache: KVCache | None = None,
        use_cache: bool = False,
        batch_indices: Tensor | None = None,
        start_pos: int | Tensor | None = None,
        paged_kv_cache: PagedKVCache | None = None,
        layer_idx: int | None = None,
        prefix_kv: tuple[Tensor, Tensor] | None = None,
    ) -> Tensor | tuple[Tensor, None]:
        """
        Optimized forward pass for the multi-latent attention mechanism.

        Args:
            hidden_states: Input tensor with shape [batch_size, seq_len, hidden_size].
            attn_mask: Optional mask tensor with shape [batch_size, 1, 1, seq_len].
                       1 indicates positions to attend to, 0 indicates positions to mask.
            kv_cache: Linear ``KVCache`` pool. Mutually exclusive with
                ``paged_kv_cache``.
            paged_kv_cache: Block-allocator ``PagedKVCache``. Mutually
                exclusive with ``kv_cache``.
            layer_idx: Required when ``paged_kv_cache`` is set; selects
                the per-layer K, V slice from the paged cache.
            prefix_kv: Optional ``(prefix_k, prefix_v)`` tuple for Prefix
                Tuning (Li & Liang 2021). Each tensor has shape
                ``[B, num_kv_heads, prefix_len, head_dim]``. The prefix
                is prepended to K and V before ``_latent_attention``
                runs — the latent queries attend to the prefix along
                with the cached context. ``attn_mask`` is extended by
                ``prefix_len`` ones on the S_k axis so the prefix is
                unconditionally visible (the prefix represents model
                parameters, not input tokens).

        Returns:
            Output tensor of shape ``[batch_size, seq_len, hidden_size]``.
            When ``use_cache=True``, returns ``(output, None)`` — the cache
            was updated in-place so there is nothing to return alongside
            the output (unlike MHA, which exposes the cached K, V).
        """
        if kv_cache is not None and paged_kv_cache is not None:
            raise ValueError("Pass either kv_cache or paged_kv_cache, not both.")

        use_causal = self.is_causal if is_causal is None else is_causal
        # Store residual connection
        residual = hidden_states

        # Get shape parameters once
        batch_size, seq_len = hidden_states.shape[:2]

        # Pre-LN norm only when this module owns the norm+residual (the block
        # passes include_norm_residual=False for a plain sublayer).
        if self.include_norm_residual and self.norm is not None and self.norm_first:
            hidden_states = self.norm(hidden_states)

        # Project input to key-value pairs in a single operation
        kv_proj = self.input_kv_proj(hidden_states)
        kv_proj = kv_proj.view(batch_size, seq_len, 2, self.num_heads, self.head_dim)
        kv_proj = kv_proj.permute(2, 0, 3, 1, 4)
        k, v = kv_proj[0], kv_proj[1]  # [batch_size, num_heads, seq_len, head_dim]

        # KV cache routing — same parallel-parameter pattern as MHA. The
        # latent attention then runs over the (possibly cached) K, V.
        if paged_kv_cache is not None:
            # ``target_seq_len`` aligns the per-row paged gather with the
            # mask's key-axis (the engine builds its mask against the
            # model's ``max_seq_len``). Without this hint the gather pads
            # only to the per-batch max, which can be smaller than the
            # mask's k-axis when the engine's running sequences are
            # short.
            target_seq_len = attn_mask.shape[-1] if attn_mask is not None else None
            # Per-row REAL query-token counts, same derivation as MHA's paged
            # write: the causal mask's column-0 visibility run. Prevents pad
            # (right-padded) positions from being appended into the paged
            # cache and inflating the block tables.
            lengths = (~attn_mask[:, 0, :, 0]).sum(dim=-1) if attn_mask is not None else None
            k, v = self._paged_kv_write(
                k=k,
                v=v,
                paged_kv_cache=paged_kv_cache,
                batch_indices=batch_indices,
                layer_idx=layer_idx,
                target_seq_len=target_seq_len,
                lengths=lengths,
            )
        elif kv_cache is not None:
            k, v = self._linear_kv_write(
                k=k,
                v=v,
                kv_cache=kv_cache,
                batch_indices=batch_indices,
                start_pos=start_pos,
            )

        # Prefix injection (Li & Liang 2021): prepend ``prefix_kv`` to
        # the (possibly cached) K, V so the latent queries attend to the
        # prefix in addition to the cached context. Done AFTER the
        # cache write (so the cache only stores dynamic tokens) and
        # BEFORE ``_latent_attention`` (so the latent attention sees
        # the extended sequence). MLA has no GQA (``num_heads ==
        # num_kv_heads``), so no repeat step is needed. The
        # ``attn_mask`` S_k axis is widened by ``prefix_len`` ones —
        # the prefix is unconditionally visible because it represents
        # model parameters, not input tokens.
        if prefix_kv is not None:
            prefix_k, prefix_v = prefix_kv
            if prefix_k.shape != prefix_v.shape:
                raise ValueError(
                    f"prefix_k and prefix_v must share shape; got {tuple(prefix_k.shape)} vs {tuple(prefix_v.shape)}"
                )
            if prefix_k.shape[1] != self.num_kv_heads:
                raise ValueError(
                    f"prefix num_kv_heads ({prefix_k.shape[1]}) must match attention num_kv_heads ({self.num_kv_heads})"
                )
            if prefix_k.shape[3] != self.head_dim:
                raise ValueError(
                    f"prefix head_dim ({prefix_k.shape[3]}) must match attention head_dim ({self.head_dim})"
                )
            if prefix_k.shape[0] != batch_size:
                raise ValueError(f"prefix batch ({prefix_k.shape[0]}) must match hidden_states batch ({batch_size})")
            target_dtype = k.dtype
            if prefix_k.dtype != target_dtype:
                prefix_k = prefix_k.to(target_dtype)
            if prefix_v.dtype != target_dtype:
                prefix_v = prefix_v.to(target_dtype)
            k = torch.cat([prefix_k, k], dim=2)
            v = torch.cat([prefix_v, v], dim=2)
            # Extend the mask's S_k axis by prefix_len ones. The
            # existing mask's S_k covers the cached/dynamic tokens; the
            # prefix segment is added at the front and is always
            # visible. Without this widening the latent attention's
            # mask reshape at ``_latent_attention`` would mismatch the
            # extended K shape.
            if attn_mask is not None:
                prefix_len = prefix_k.shape[2]
                ones_prefix = torch.ones(
                    attn_mask.shape[0],
                    attn_mask.shape[1],
                    attn_mask.shape[2],
                    prefix_len,
                    device=attn_mask.device,
                    dtype=attn_mask.dtype,
                )
                attn_mask = torch.cat([ones_prefix, attn_mask], dim=-1)

        # Process latent attention
        latent_output = self._latent_attention(k, v, batch_size, attn_mask, is_causal=use_causal)

        # Compute uniform weights across latents by default
        latent_weights = torch.ones(batch_size, 1, self.num_latents, device=latent_output.device) / self.num_latents

        # Apply weights to latent outputs
        output = torch.bmm(latent_weights, latent_output)

        # Expand to sequence length dimension efficiently
        output = output.expand(-1, seq_len, -1)

        # Apply output projection and dropout
        output = self.out_proj(output)
        output = self.dropout(output)

        # Residual + Post-LN only when this module owns the norm+residual (the
        # block passes include_norm_residual=False for a plain sublayer; adding
        # a residual here would double it, RIL ISS-139).
        if self.include_norm_residual and self.norm is not None:
            output = output + residual
            if not self.norm_first:
                output = self.norm(output)

        # Match the MHA contract:
        # - paged path returns the output tensor directly (the cache is
        #   mutated in place; the caller does not need a kv tuple);
        # - dense + use_cache returns ``(output, kv)``. For MLA the
        #   cached K, V are consumed internally by the latent attention,
        #   so we return ``None`` as the second element;
        # - no cache returns the output tensor.
        if paged_kv_cache is not None:
            return output
        if use_cache:
            return output, None
        return output

    def _linear_kv_write(
        self,
        k: Tensor,
        v: Tensor,
        kv_cache: KVCache,
        batch_indices: Tensor | None,
        start_pos: int | Tensor | None,
    ) -> tuple[Tensor, Tensor]:
        """Write the new K, V into the linear ``KVCache`` pool and return the cached view.

        For per-slot writes (``batch_indices`` is set) the cache contract
        is ``update_at_indices``; for a dense batch it's ``update``.
        """
        if batch_indices is not None:
            if start_pos is None:
                raise ValueError("start_pos must be provided when using batch_indices for KV cache update.")
            return kv_cache.update_at_indices(batch_indices, k, v, start_pos)
        return kv_cache.update(k, v)

    def _paged_kv_write(
        self,
        k: Tensor,
        v: Tensor,
        paged_kv_cache: PagedKVCache,
        batch_indices: Tensor | None,
        layer_idx: int | None,
        target_seq_len: int | None = None,
        lengths: Tensor | None = None,
    ) -> tuple[Tensor, Tensor]:
        """Write the new K, V into the paged cache and return the cached K, V slice.

        Per row: call ``paged_kv_cache.update(seq_id, k_b.T, v_b.T)``
        to append the new tokens. Then gather the per-row K, V via
        :meth:`PagedKVCache.get` so the latent attention runs over the
        full cached context (this is the cost the cache saves vs.
        recomputing ``input_kv_proj`` on every past token).

        Args:
            target_seq_len: Optional padding target for the returned K, V
                sequence axis. The latent attention's mask expects
                ``[B, 1, num_latents, target_seq_len]``; when set we pad
                the per-row K, V to this length with zeros (masked
                positions are out of range for the active sequence).
                When ``None``, we pad to the per-batch max seq length.

        Note:
            ``paged_attention_forward`` is *not* used here. That kernel
            returns the attended output — for MLA we need the raw K, V
            to feed the latent cross-attention block.
        """
        if layer_idx is None:
            raise ValueError(
                "layer_idx is required when paged_kv_cache is set; DecoderModel threads it through TransformerBlock."
            )
        if batch_indices is None:
            raise ValueError("batch_indices is required when paged_kv_cache is set.")

        # Per-row write into the paged cache. ``PagedKVCache.update``
        # expects ``[B, T, N_kv, D]`` (it transposes internally), so
        # transpose our ``[B, N_kv, T, D]`` k/v to match.
        #
        # Only each row's REAL tokens may be appended: continuous batching
        # right-pads the batch to the batch-max query length, and writing the
        # pad positions would (a) store garbage pad K/V that later extends the
        # block table and inflates ``get_num_tokens``, and (b) let short-lived
        # sequences swallow the whole block pool. ``lengths`` is the per-row
        # real query-token count derived from the causal mask's column-0
        # visibility run (same convention as MHA's paged write).
        seq_ids = batch_indices.tolist()
        if lengths is None:
            lengths = torch.tensor([k.shape[2]] * len(seq_ids), device=k.device)
        for b, seq_id in enumerate(seq_ids):
            n = int(lengths[b])
            paged_kv_cache.update(
                seq_id=int(seq_id),
                k_new=k[b : b + 1, :, :n].transpose(1, 2),
                v_new=v[b : b + 1, :, :n].transpose(1, 2),
                layer_idx=layer_idx,
            )

        # Gather the per-row K, V via the public ``PagedKVCache.get`` API.
        # The latent attention expects ``[B, N_heads, T_total, head_dim]``
        # per row, padded with zeros for shorter sequences.
        batch_size, num_heads, _, head_dim = k.shape
        per_row_seq_lens = [paged_kv_cache.block_manager.get_num_tokens(int(sid)) for sid in seq_ids]
        if target_seq_len is None:
            target_seq_len = max(per_row_seq_lens) if per_row_seq_lens else 1
            target_seq_len = max(target_seq_len, 1)

        k_gathered = torch.zeros(
            batch_size,
            num_heads,
            target_seq_len,
            head_dim,
            device=k.device,
            dtype=k.dtype,
        )
        v_gathered = torch.zeros_like(k_gathered)
        for b, seq_id in enumerate(seq_ids):
            seq_len = per_row_seq_lens[b]
            if seq_len == 0:
                continue
            # ``PagedKVCache.get`` returns ``[N_kv, num_tokens, D]``;
            # the MLA contract is ``[B, N_heads, T_total, D]`` — the
            # per-row head count equals ``self.num_heads`` (no GQA in
            # the placeholder MLA).
            k_row, v_row = paged_kv_cache.get(int(seq_id), 0, seq_len, layer_idx=layer_idx)
            k_gathered[b, :, :seq_len] = k_row
            v_gathered[b, :, :seq_len] = v_row

        return k_gathered, v_gathered

forward

forward(hidden_states, attn_mask=None, is_causal=None, kv_cache=None, use_cache=False, batch_indices=None, start_pos=None, paged_kv_cache=None, layer_idx=None, prefix_kv=None)

Optimized forward pass for the multi-latent attention mechanism.

参数：

名称	类型	描述	默认
hidden_states	Tensor	Input tensor with shape [batch_size, seq_len, hidden_size].	必需
attn_mask	Tensor | None	Optional mask tensor with shape [batch_size, 1, 1, seq_len]. 1 indicates positions to attend to, 0 indicates positions to mask.	None
kv_cache	KVCache | None	Linear KVCache pool. Mutually exclusive with paged_kv_cache.	None
paged_kv_cache	PagedKVCache | None	Block-allocator PagedKVCache. Mutually exclusive with kv_cache.	None
layer_idx	int | None	Required when paged_kv_cache is set; selects the per-layer K, V slice from the paged cache.	None
prefix_kv	tuple[Tensor, Tensor] | None	Optional (prefix_k, prefix_v) tuple for Prefix Tuning (Li & Liang 2021). Each tensor has shape [B, num_kv_heads, prefix_len, head_dim]. The prefix is prepended to K and V before _latent_attention runs — the latent queries attend to the prefix along with the cached context. attn_mask is extended by prefix_len ones on the S_k axis so the prefix is unconditionally visible (the prefix represents model parameters, not input tokens).	None

返回：

类型	描述
Tensor | tuple[Tensor, None]	Output tensor of shape [batch_size, seq_len, hidden_size].
Tensor | tuple[Tensor, None]	When use_cache=True, returns (output, None) — the cache
Tensor | tuple[Tensor, None]	was updated in-place so there is nothing to return alongside
Tensor | tuple[Tensor, None]	the output (unlike MHA, which exposes the cached K, V).

源代码位于： src/llm/core/attn/mla.py

def forward(
    self,
    hidden_states: Tensor,
    attn_mask: Tensor | None = None,
    is_causal: bool | None = None,
    kv_cache: KVCache | None = None,
    use_cache: bool = False,
    batch_indices: Tensor | None = None,
    start_pos: int | Tensor | None = None,
    paged_kv_cache: PagedKVCache | None = None,
    layer_idx: int | None = None,
    prefix_kv: tuple[Tensor, Tensor] | None = None,
) -> Tensor | tuple[Tensor, None]:
    """
    Optimized forward pass for the multi-latent attention mechanism.

    Args:
        hidden_states: Input tensor with shape [batch_size, seq_len, hidden_size].
        attn_mask: Optional mask tensor with shape [batch_size, 1, 1, seq_len].
                   1 indicates positions to attend to, 0 indicates positions to mask.
        kv_cache: Linear ``KVCache`` pool. Mutually exclusive with
            ``paged_kv_cache``.
        paged_kv_cache: Block-allocator ``PagedKVCache``. Mutually
            exclusive with ``kv_cache``.
        layer_idx: Required when ``paged_kv_cache`` is set; selects
            the per-layer K, V slice from the paged cache.
        prefix_kv: Optional ``(prefix_k, prefix_v)`` tuple for Prefix
            Tuning (Li & Liang 2021). Each tensor has shape
            ``[B, num_kv_heads, prefix_len, head_dim]``. The prefix
            is prepended to K and V before ``_latent_attention``
            runs — the latent queries attend to the prefix along
            with the cached context. ``attn_mask`` is extended by
            ``prefix_len`` ones on the S_k axis so the prefix is
            unconditionally visible (the prefix represents model
            parameters, not input tokens).

    Returns:
        Output tensor of shape ``[batch_size, seq_len, hidden_size]``.
        When ``use_cache=True``, returns ``(output, None)`` — the cache
        was updated in-place so there is nothing to return alongside
        the output (unlike MHA, which exposes the cached K, V).
    """
    if kv_cache is not None and paged_kv_cache is not None:
        raise ValueError("Pass either kv_cache or paged_kv_cache, not both.")

    use_causal = self.is_causal if is_causal is None else is_causal
    # Store residual connection
    residual = hidden_states

    # Get shape parameters once
    batch_size, seq_len = hidden_states.shape[:2]

    # Pre-LN norm only when this module owns the norm+residual (the block
    # passes include_norm_residual=False for a plain sublayer).
    if self.include_norm_residual and self.norm is not None and self.norm_first:
        hidden_states = self.norm(hidden_states)

    # Project input to key-value pairs in a single operation
    kv_proj = self.input_kv_proj(hidden_states)
    kv_proj = kv_proj.view(batch_size, seq_len, 2, self.num_heads, self.head_dim)
    kv_proj = kv_proj.permute(2, 0, 3, 1, 4)
    k, v = kv_proj[0], kv_proj[1]  # [batch_size, num_heads, seq_len, head_dim]

    # KV cache routing — same parallel-parameter pattern as MHA. The
    # latent attention then runs over the (possibly cached) K, V.
    if paged_kv_cache is not None:
        # ``target_seq_len`` aligns the per-row paged gather with the
        # mask's key-axis (the engine builds its mask against the
        # model's ``max_seq_len``). Without this hint the gather pads
        # only to the per-batch max, which can be smaller than the
        # mask's k-axis when the engine's running sequences are
        # short.
        target_seq_len = attn_mask.shape[-1] if attn_mask is not None else None
        # Per-row REAL query-token counts, same derivation as MHA's paged
        # write: the causal mask's column-0 visibility run. Prevents pad
        # (right-padded) positions from being appended into the paged
        # cache and inflating the block tables.
        lengths = (~attn_mask[:, 0, :, 0]).sum(dim=-1) if attn_mask is not None else None
        k, v = self._paged_kv_write(
            k=k,
            v=v,
            paged_kv_cache=paged_kv_cache,
            batch_indices=batch_indices,
            layer_idx=layer_idx,
            target_seq_len=target_seq_len,
            lengths=lengths,
        )
    elif kv_cache is not None:
        k, v = self._linear_kv_write(
            k=k,
            v=v,
            kv_cache=kv_cache,
            batch_indices=batch_indices,
            start_pos=start_pos,
        )

    # Prefix injection (Li & Liang 2021): prepend ``prefix_kv`` to
    # the (possibly cached) K, V so the latent queries attend to the
    # prefix in addition to the cached context. Done AFTER the
    # cache write (so the cache only stores dynamic tokens) and
    # BEFORE ``_latent_attention`` (so the latent attention sees
    # the extended sequence). MLA has no GQA (``num_heads ==
    # num_kv_heads``), so no repeat step is needed. The
    # ``attn_mask`` S_k axis is widened by ``prefix_len`` ones —
    # the prefix is unconditionally visible because it represents
    # model parameters, not input tokens.
    if prefix_kv is not None:
        prefix_k, prefix_v = prefix_kv
        if prefix_k.shape != prefix_v.shape:
            raise ValueError(
                f"prefix_k and prefix_v must share shape; got {tuple(prefix_k.shape)} vs {tuple(prefix_v.shape)}"
            )
        if prefix_k.shape[1] != self.num_kv_heads:
            raise ValueError(
                f"prefix num_kv_heads ({prefix_k.shape[1]}) must match attention num_kv_heads ({self.num_kv_heads})"
            )
        if prefix_k.shape[3] != self.head_dim:
            raise ValueError(
                f"prefix head_dim ({prefix_k.shape[3]}) must match attention head_dim ({self.head_dim})"
            )
        if prefix_k.shape[0] != batch_size:
            raise ValueError(f"prefix batch ({prefix_k.shape[0]}) must match hidden_states batch ({batch_size})")
        target_dtype = k.dtype
        if prefix_k.dtype != target_dtype:
            prefix_k = prefix_k.to(target_dtype)
        if prefix_v.dtype != target_dtype:
            prefix_v = prefix_v.to(target_dtype)
        k = torch.cat([prefix_k, k], dim=2)
        v = torch.cat([prefix_v, v], dim=2)
        # Extend the mask's S_k axis by prefix_len ones. The
        # existing mask's S_k covers the cached/dynamic tokens; the
        # prefix segment is added at the front and is always
        # visible. Without this widening the latent attention's
        # mask reshape at ``_latent_attention`` would mismatch the
        # extended K shape.
        if attn_mask is not None:
            prefix_len = prefix_k.shape[2]
            ones_prefix = torch.ones(
                attn_mask.shape[0],
                attn_mask.shape[1],
                attn_mask.shape[2],
                prefix_len,
                device=attn_mask.device,
                dtype=attn_mask.dtype,
            )
            attn_mask = torch.cat([ones_prefix, attn_mask], dim=-1)

    # Process latent attention
    latent_output = self._latent_attention(k, v, batch_size, attn_mask, is_causal=use_causal)

    # Compute uniform weights across latents by default
    latent_weights = torch.ones(batch_size, 1, self.num_latents, device=latent_output.device) / self.num_latents

    # Apply weights to latent outputs
    output = torch.bmm(latent_weights, latent_output)

    # Expand to sequence length dimension efficiently
    output = output.expand(-1, seq_len, -1)

    # Apply output projection and dropout
    output = self.out_proj(output)
    output = self.dropout(output)

    # Residual + Post-LN only when this module owns the norm+residual (the
    # block passes include_norm_residual=False for a plain sublayer; adding
    # a residual here would double it, RIL ISS-139).
    if self.include_norm_residual and self.norm is not None:
        output = output + residual
        if not self.norm_first:
            output = self.norm(output)

    # Match the MHA contract:
    # - paged path returns the output tensor directly (the cache is
    #   mutated in place; the caller does not need a kv tuple);
    # - dense + use_cache returns ``(output, kv)``. For MLA the
    #   cached K, V are consumed internally by the latent attention,
    #   so we return ``None`` as the second element;
    # - no cache returns the output tensor.
    if paged_kv_cache is not None:
        return output
    if use_cache:
        return output, None
    return output

flash_attn

Flash Attention 2 attention implementation.

Registered as attn_impl="flash_attn" through ATTENTION_REGISTRY. The flash-attn package is an optional dependency — importing this module never raises, but instantiating :class:FlashAttention does if the package is not installed. This mirrors the soft-dependency contract used elsewhere in the project (e.g. huggingface_hub in compat.hf_loader).

The class itself follows the same surface as :class:llm.core.attn.MultiHeadAttention so it can be substituted transparently: same projection layers, same KV-cache integration, same output projection. The only difference is the attention kernel: flash_attn.flash_attn_func instead of torch.nn.functional.scaled_dot_product_attention.

On supported hardware (Ampere/Hopper) this is meaningfully faster for training and (especially) long-context decode. On other devices the fallback path through PyTorch SDPA in MHA is the right choice.

FlashAttention

Bases: Module

源代码位于： src/llm/core/attn/flash_attn.py

@register_attention("flash_attn")
class FlashAttention(nn.Module):
    # KV-cache contract matches MHA: the wrapper writes K/V into the
    # shared ``KVCache`` pool during autoregressive decoding, so the
    # continuous batching engine and training engine work without
    # changes. ``ModelConfig.check_consistency`` consults this map.
    set_attention_kv_cache_capability("flash_attn", supports=True)
    """
    Flash Attention 2 wrapper exposing the standard MHA interface.

    Same projection layout as :class:`MultiHeadAttention` (combined
    QKV, output projection, optional pre-norm/residual). The attention
    kernel is ``flash_attn.flash_attn_func``; Q/K/V are reshaped from
    ``[B, N, S, D]`` to ``[B, S, H]`` for the call and the output is
    reshaped back.

    Args:
        hidden_size: Total dimension of the model.
        num_heads: Number of attention heads. Must divide ``hidden_size``.
        p: Dropout probability applied to attention weights and final
            output. (Forwarded to ``flash_attn_func``; ``flash-attn``
            applies dropout only during training.)
        bias: Whether to use bias in the QKV/output projections.
        eps: Epsilon for the optional pre-norm LayerNorm.
        norm_first: Whether to apply LayerNorm before (Pre-LN) or after
            (Post-LN) the attention block. Only used when
            ``include_norm_residual=True``.
        is_causal: Whether to apply causal masking by default.
        include_norm_residual: Whether this block owns its pre-norm and
            residual connection (same convention as ``mha``).
        num_kv_heads: For GQA/MQA. Defaults to ``num_heads``. Must
            divide ``num_heads``.
        device: Target device for parameters.
        dtype: Target dtype for parameters.

    Raises:
        ImportError: When constructed without ``flash-attn`` installed.
            Install with ``pip install flash-attn`` (builds a CUDA
            wheel; not available on plain PyPI for some platforms —
            see the project's ``[perf]`` extra).
    """

    def __init__(
        self,
        hidden_size: int,
        num_heads: int = 8,
        p: float = 0.1,
        bias: bool = False,
        eps: float = 1e-5,
        norm_first: bool = True,
        is_causal: bool = False,
        include_norm_residual: bool = True,
        num_kv_heads: int | None = None,
        device: torch.device | str | None = None,
        dtype: torch.dtype | None = None,
        # Same optional kwargs the TransformerBlock always threads (RIL
        # ISS-137): the block historically passed ``window_size=`` and (when
        # ``use_rope``) ``max_seq_len=``/``use_rope=True``/``rope_theta=``
        # UNCONDITIONALLY, so ``FlashAttention`` without these params raised
        # ``TypeError: got an unexpected keyword argument 'window_size'`` on
        # EVERY ``attn_impl='flash_attn'`` model build. Adopt the MHA
        # contract so the backend is actually constructible.
        window_size: int | None = None,  # Sliding window attention
        max_seq_len: int | None = None,  # RoPE max context (required if use_rope)
        use_rope: bool = False,  # Rotary position embedding (real Llama/Mistral)
        rope_theta: float = 10000.0,  # RoPE base frequency
    ):
        super().__init__()

        if not FLASH_ATTN_AVAILABLE:
            raise ImportError(
                "FlashAttention requires the optional 'flash-attn' package. "
                "Install it with `pip install 'llm[perf]'` "
                "(or `pip install flash-attn` directly on a CUDA host)."
            )

        if hidden_size % num_heads != 0:
            raise ValueError(f"hidden_size ({hidden_size}) must be divisible by num_heads ({num_heads})")

        factory_kwargs = make_factory_kwargs(device, dtype)
        self.hidden_size = hidden_size
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
        self.head_dim = hidden_size // num_heads
        self.kv_dim = self.num_kv_heads * self.head_dim
        self.norm_first = norm_first
        self.is_causal = is_causal
        self.p = p
        self.include_norm_residual = include_norm_residual
        self.window_size = window_size
        self.use_rope = use_rope

        if use_rope:
            # RoPE rotates Q/K by head position (real Llama/Mistral inject
            # position here, not via additive embeddings), identically to MHA
            # (RIL ISS-062). ``max_seq_len`` sizes the cos/sin table.
            if max_seq_len is None:
                raise ValueError("use_rope=True requires max_seq_len for the RoPE cos/sin table")
            self.rope = RotaryPositionEmbedding(
                dim=self.head_dim,
                max_seq_len=max_seq_len,
                base=rope_theta,
                **factory_kwargs,
            )

        if self.num_heads % self.num_kv_heads != 0:
            raise ValueError(f"num_heads ({self.num_heads}) must be divisible by num_kv_heads ({self.num_kv_heads})")

        self.norm = None
        if self.include_norm_residual:
            self.norm = nn.LayerNorm(hidden_size, eps=eps, **factory_kwargs)

        self.qkv_dim = (self.num_heads + 2 * self.num_kv_heads) * self.head_dim
        self.qkv_proj = nn.Linear(hidden_size, self.qkv_dim, bias=bias, **factory_kwargs)
        self.out_proj = nn.Linear(hidden_size, hidden_size, bias=bias, **factory_kwargs)
        self.dropout = nn.Dropout(p)

        self._init_weights()

    def _init_weights(self) -> None:
        """Xavier-uniform init for projections, zero bias."""
        for proj in (self.qkv_proj, self.out_proj):
            nn.init.xavier_uniform_(proj.weight)
            if proj.bias is not None:
                nn.init.zeros_(proj.bias)

    def _rope_positions(
        self,
        batch_size: int,
        seq_len: int,
        start_pos: int | Tensor | None,
        device: torch.device,
    ) -> Tensor | None:
        """Compute per-token positions for RoPE (same contract as MHA).

        * ``start_pos`` Tensor (batch-serving path): return unchanged.
        * ``start_pos`` int (KV-cache decode): ``[start_pos, start_pos+seq_len)``.
        * ``start_pos`` None (pure prefill): ``[0, seq_len)`` — return None so
          RoPE's internal default applies.
        """
        if start_pos is None:
            return None
        if isinstance(start_pos, Tensor):
            return start_pos
        if start_pos == 0:
            return None
        base = int(start_pos)
        return torch.arange(base, base + seq_len, device=device, dtype=torch.long).expand(batch_size, -1)

    def forward(
        self,
        hidden_states: Tensor,
        attn_mask: Tensor | None = None,
        is_causal: bool | None = None,
        kv_cache: KVCache | None = None,
        use_cache: bool = False,
        batch_indices: Tensor | None = None,
        start_pos: int | Tensor | None = None,
        paged_kv_cache: object | None = None,
        layer_idx: int | None = None,
        prefix_kv: tuple[Tensor, Tensor] | None = None,
    ) -> Tensor | tuple[Tensor, tuple[Tensor, Tensor]]:
        """Forward pass.

        See :meth:`MultiHeadAttention.forward` for the full argument
        contract — the signatures are intentionally identical so the
        continuous batching engine and training engine can pick
        ``flash_attn`` as a drop-in replacement.

        Note:
            ``attn_mask`` is **ignored** — ``flash_attn_func`` does not
            accept arbitrary masks. For non-causal masking or padded
            sequences use ``is_causal=False`` and pre-pad, or fall back
            to ``attn_impl="mha"``. Sliding-window on dense inputs is
            **supported** via ``window_size=`` (threaded by the decoder,
            RIL ISS-242); the still-future piece is combining a sliding
            window with a **padding** mask — that needs
            ``flash_attn_varlen_func`` (variable-length batches).
            ``paged_kv_cache`` is also rejected on this path — flash-attn
            does not expose a paged-attn kernel; use ``attn_impl="mha"``
            when serving with paged KV.

            ``prefix_kv`` is supported (Li & Liang 2021). The prefix
            K/V are concatenated to the projected K/V after the
            KV-cache write (so the cache stores only dynamic tokens)
            and before the GQA repeat (so the prefix is treated like a
            regular token and replicated to all query heads). Prefix
            dtype is auto-cast to the projected K/V dtype to satisfy
            ``flash_attn_func``'s fp16/bf16 requirement. The captured
            ``current_kv`` (when ``use_cache=True``) excludes the
            prefix — matching the MHA contract at
            ``MultiHeadAttention.forward``.
        """
        if paged_kv_cache is not None:
            raise NotImplementedError(
                "FlashAttention does not support paged_kv_cache; "
                "use attn_impl='mha' (which routes through paged_attention_forward)."
            )
        # Local import so an uninstalled ``flash-attn`` does not crash
        # the package at import time (we already gated on
        # ``FLASH_ATTN_AVAILABLE`` in ``__init__``).
        flash_attn_func = importlib.import_module("flash_attn").flash_attn_func  # optional CUDA-only dep

        batch_size, seq_len, _ = hidden_states.size()
        use_causal = self.is_causal if is_causal is None else is_causal

        if self.include_norm_residual and self.norm is not None:
            residual = hidden_states
            x_for_qkv = self.norm(hidden_states) if self.norm_first else hidden_states
        else:
            x_for_qkv = hidden_states

        # Project + split QKV (identical to MHA).
        qkv = self.qkv_proj(x_for_qkv)
        q_size = self.num_heads * self.head_dim
        kv_size = self.num_kv_heads * self.head_dim
        q, k, v = torch.split(qkv, [q_size, kv_size, kv_size], dim=-1)

        # Reshape: [B, S, N*D] -> [B, N, S, D]
        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = v.view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)

        if self.use_rope:
            # Apply RoPE to Q/K BEFORE the kernel (flash_attn_func performs no
            # positional encoding of its own). Identical to MHA's wiring — the
            # kernel then sees rotated Q/K and the KV cache stores the rotated
            # K, exactly like the MHA backend (RIL ISS-137).
            positions = self._rope_positions(batch_size, seq_len, start_pos, device=q.device)
            q, k = self.rope(q, k, positions)

        has_past = False
        if kv_cache is not None:
            if batch_indices is not None:
                if start_pos is None:
                    raise ValueError("start_pos must be provided when using batch_indices for KV cache update.")
                k, v = kv_cache.update_at_indices(batch_indices, k, v, start_pos)
                has_past = True
            else:
                k, v = kv_cache.update(k, v)
                has_past = kv_cache.seq_len > seq_len

        if use_cache:
            current_kv = (k, v)

        # Prefix injection (Li & Liang 2021): prepend ``prefix_kv`` to
        # the projected K/V so the new tokens attend to the prefix in
        # addition to the cached context. Done AFTER the KV-cache write
        # (so the cache only stores dynamic tokens) and BEFORE the GQA
        # repeat (so the prefix is treated like a regular token and
        # replicated to all query heads).
        if prefix_kv is not None:
            prefix_k, prefix_v = prefix_kv
            if prefix_k.shape != prefix_v.shape:
                raise ValueError(
                    f"prefix_k and prefix_v must share shape; got {tuple(prefix_k.shape)} vs {tuple(prefix_v.shape)}"
                )
            if prefix_k.shape[1] != self.num_kv_heads:
                raise ValueError(
                    f"prefix num_kv_heads ({prefix_k.shape[1]}) must match attention num_kv_heads ({self.num_kv_heads})"
                )
            if prefix_k.shape[3] != self.head_dim:
                raise ValueError(
                    f"prefix head_dim ({prefix_k.shape[3]}) must match attention head_dim ({self.head_dim})"
                )
            if prefix_k.shape[0] != batch_size:
                raise ValueError(f"prefix batch ({prefix_k.shape[0]}) must match hidden_states batch ({batch_size})")
            # flash_attn_func requires fp16/bf16. Cast the prefix to
            # the projected K/V dtype to satisfy the kernel contract.
            target_dtype = k.dtype
            if prefix_k.dtype != target_dtype:
                prefix_k = prefix_k.to(target_dtype)
            if prefix_v.dtype != target_dtype:
                prefix_v = prefix_v.to(target_dtype)
            k = torch.cat([prefix_k, k], dim=2)
            v = torch.cat([prefix_v, v], dim=2)

        # GQA: replicate K/V across query heads.
        if self.num_kv_heads != self.num_heads:
            num_queries_per_kv = self.num_heads // self.num_kv_heads
            k = k.repeat_interleave(num_queries_per_kv, dim=1)
            v = v.repeat_interleave(num_queries_per_kv, dim=1)

        # flash_attn_func expects [B, S, H] (heads merged into the last
        # axis). Transpose and make contiguous (the underlying CUDA
        # kernel does not accept strided inputs in this layout).
        q_bsh = q.transpose(1, 2).contiguous()
        k_bsh = k.transpose(1, 2).contiguous()
        v_bsh = v.transpose(1, 2).contiguous()

        # flash_attn_func requires fp16/bf16. We keep the call guarded
        # by the caller (MHA does the same — float dtype falls back via
        # the engine layer). Apply causal only when there is no past
        # context: with KV cache populated, the cache handles positions.
        # Sliding-window attention (``window_size``) is supported by the
        # kernel via a (left, right) tuple; unlimited defaults to (-1,-1).
        window_pair = (-1, -1) if self.window_size is None else (self.window_size, self.window_size)
        attn_output = flash_attn_func(
            q_bsh,
            k_bsh,
            v_bsh,
            dropout_p=self.p if self.training else 0.0,
            causal=use_causal if not has_past else False,
            window_size=window_pair,
        )

        # Output back to [B, N, S, D].
        attn_output = attn_output.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

        # [B, N, S, D] -> [B, S, H]
        attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_len, self.hidden_size)
        projected_output = self.dropout(self.out_proj(attn_output))

        if self.include_norm_residual and self.norm is not None:
            output = residual + projected_output
            if not self.norm_first:
                output = self.norm(output)
        else:
            output = projected_output

        if use_cache:
            return output, current_kv
        return output

forward

forward(hidden_states, attn_mask=None, is_causal=None, kv_cache=None, use_cache=False, batch_indices=None, start_pos=None, paged_kv_cache=None, layer_idx=None, prefix_kv=None)

Forward pass.

See :meth:MultiHeadAttention.forward for the full argument contract — the signatures are intentionally identical so the continuous batching engine and training engine can pick flash_attn as a drop-in replacement.

Note

attn_mask is ignored — flash_attn_func does not accept arbitrary masks. For non-causal masking or padded sequences use is_causal=False and pre-pad, or fall back to attn_impl="mha". Sliding-window on dense inputs is supported via window_size= (threaded by the decoder, RIL ISS-242); the still-future piece is combining a sliding window with a padding mask — that needs flash_attn_varlen_func (variable-length batches). paged_kv_cache is also rejected on this path — flash-attn does not expose a paged-attn kernel; use attn_impl="mha" when serving with paged KV.

prefix_kv is supported (Li & Liang 2021). The prefix K/V are concatenated to the projected K/V after the KV-cache write (so the cache stores only dynamic tokens) and before the GQA repeat (so the prefix is treated like a regular token and replicated to all query heads). Prefix dtype is auto-cast to the projected K/V dtype to satisfy flash_attn_func's fp16/bf16 requirement. The captured current_kv (when use_cache=True) excludes the prefix — matching the MHA contract at MultiHeadAttention.forward.

源代码位于： src/llm/core/attn/flash_attn.py

def forward(
    self,
    hidden_states: Tensor,
    attn_mask: Tensor | None = None,
    is_causal: bool | None = None,
    kv_cache: KVCache | None = None,
    use_cache: bool = False,
    batch_indices: Tensor | None = None,
    start_pos: int | Tensor | None = None,
    paged_kv_cache: object | None = None,
    layer_idx: int | None = None,
    prefix_kv: tuple[Tensor, Tensor] | None = None,
) -> Tensor | tuple[Tensor, tuple[Tensor, Tensor]]:
    """Forward pass.

    See :meth:`MultiHeadAttention.forward` for the full argument
    contract — the signatures are intentionally identical so the
    continuous batching engine and training engine can pick
    ``flash_attn`` as a drop-in replacement.

    Note:
        ``attn_mask`` is **ignored** — ``flash_attn_func`` does not
        accept arbitrary masks. For non-causal masking or padded
        sequences use ``is_causal=False`` and pre-pad, or fall back
        to ``attn_impl="mha"``. Sliding-window on dense inputs is
        **supported** via ``window_size=`` (threaded by the decoder,
        RIL ISS-242); the still-future piece is combining a sliding
        window with a **padding** mask — that needs
        ``flash_attn_varlen_func`` (variable-length batches).
        ``paged_kv_cache`` is also rejected on this path — flash-attn
        does not expose a paged-attn kernel; use ``attn_impl="mha"``
        when serving with paged KV.

        ``prefix_kv`` is supported (Li & Liang 2021). The prefix
        K/V are concatenated to the projected K/V after the
        KV-cache write (so the cache stores only dynamic tokens)
        and before the GQA repeat (so the prefix is treated like a
        regular token and replicated to all query heads). Prefix
        dtype is auto-cast to the projected K/V dtype to satisfy
        ``flash_attn_func``'s fp16/bf16 requirement. The captured
        ``current_kv`` (when ``use_cache=True``) excludes the
        prefix — matching the MHA contract at
        ``MultiHeadAttention.forward``.
    """
    if paged_kv_cache is not None:
        raise NotImplementedError(
            "FlashAttention does not support paged_kv_cache; "
            "use attn_impl='mha' (which routes through paged_attention_forward)."
        )
    # Local import so an uninstalled ``flash-attn`` does not crash
    # the package at import time (we already gated on
    # ``FLASH_ATTN_AVAILABLE`` in ``__init__``).
    flash_attn_func = importlib.import_module("flash_attn").flash_attn_func  # optional CUDA-only dep

    batch_size, seq_len, _ = hidden_states.size()
    use_causal = self.is_causal if is_causal is None else is_causal

    if self.include_norm_residual and self.norm is not None:
        residual = hidden_states
        x_for_qkv = self.norm(hidden_states) if self.norm_first else hidden_states
    else:
        x_for_qkv = hidden_states

    # Project + split QKV (identical to MHA).
    qkv = self.qkv_proj(x_for_qkv)
    q_size = self.num_heads * self.head_dim
    kv_size = self.num_kv_heads * self.head_dim
    q, k, v = torch.split(qkv, [q_size, kv_size, kv_size], dim=-1)

    # Reshape: [B, S, N*D] -> [B, N, S, D]
    q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
    k = k.view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
    v = v.view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)

    if self.use_rope:
        # Apply RoPE to Q/K BEFORE the kernel (flash_attn_func performs no
        # positional encoding of its own). Identical to MHA's wiring — the
        # kernel then sees rotated Q/K and the KV cache stores the rotated
        # K, exactly like the MHA backend (RIL ISS-137).
        positions = self._rope_positions(batch_size, seq_len, start_pos, device=q.device)
        q, k = self.rope(q, k, positions)

    has_past = False
    if kv_cache is not None:
        if batch_indices is not None:
            if start_pos is None:
                raise ValueError("start_pos must be provided when using batch_indices for KV cache update.")
            k, v = kv_cache.update_at_indices(batch_indices, k, v, start_pos)
            has_past = True
        else:
            k, v = kv_cache.update(k, v)
            has_past = kv_cache.seq_len > seq_len

    if use_cache:
        current_kv = (k, v)

    # Prefix injection (Li & Liang 2021): prepend ``prefix_kv`` to
    # the projected K/V so the new tokens attend to the prefix in
    # addition to the cached context. Done AFTER the KV-cache write
    # (so the cache only stores dynamic tokens) and BEFORE the GQA
    # repeat (so the prefix is treated like a regular token and
    # replicated to all query heads).
    if prefix_kv is not None:
        prefix_k, prefix_v = prefix_kv
        if prefix_k.shape != prefix_v.shape:
            raise ValueError(
                f"prefix_k and prefix_v must share shape; got {tuple(prefix_k.shape)} vs {tuple(prefix_v.shape)}"
            )
        if prefix_k.shape[1] != self.num_kv_heads:
            raise ValueError(
                f"prefix num_kv_heads ({prefix_k.shape[1]}) must match attention num_kv_heads ({self.num_kv_heads})"
            )
        if prefix_k.shape[3] != self.head_dim:
            raise ValueError(
                f"prefix head_dim ({prefix_k.shape[3]}) must match attention head_dim ({self.head_dim})"
            )
        if prefix_k.shape[0] != batch_size:
            raise ValueError(f"prefix batch ({prefix_k.shape[0]}) must match hidden_states batch ({batch_size})")
        # flash_attn_func requires fp16/bf16. Cast the prefix to
        # the projected K/V dtype to satisfy the kernel contract.
        target_dtype = k.dtype
        if prefix_k.dtype != target_dtype:
            prefix_k = prefix_k.to(target_dtype)
        if prefix_v.dtype != target_dtype:
            prefix_v = prefix_v.to(target_dtype)
        k = torch.cat([prefix_k, k], dim=2)
        v = torch.cat([prefix_v, v], dim=2)

    # GQA: replicate K/V across query heads.
    if self.num_kv_heads != self.num_heads:
        num_queries_per_kv = self.num_heads // self.num_kv_heads
        k = k.repeat_interleave(num_queries_per_kv, dim=1)
        v = v.repeat_interleave(num_queries_per_kv, dim=1)

    # flash_attn_func expects [B, S, H] (heads merged into the last
    # axis). Transpose and make contiguous (the underlying CUDA
    # kernel does not accept strided inputs in this layout).
    q_bsh = q.transpose(1, 2).contiguous()
    k_bsh = k.transpose(1, 2).contiguous()
    v_bsh = v.transpose(1, 2).contiguous()

    # flash_attn_func requires fp16/bf16. We keep the call guarded
    # by the caller (MHA does the same — float dtype falls back via
    # the engine layer). Apply causal only when there is no past
    # context: with KV cache populated, the cache handles positions.
    # Sliding-window attention (``window_size``) is supported by the
    # kernel via a (left, right) tuple; unlimited defaults to (-1,-1).
    window_pair = (-1, -1) if self.window_size is None else (self.window_size, self.window_size)
    attn_output = flash_attn_func(
        q_bsh,
        k_bsh,
        v_bsh,
        dropout_p=self.p if self.training else 0.0,
        causal=use_causal if not has_past else False,
        window_size=window_pair,
    )

    # Output back to [B, N, S, D].
    attn_output = attn_output.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

    # [B, N, S, D] -> [B, S, H]
    attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_len, self.hidden_size)
    projected_output = self.dropout(self.out_proj(attn_output))

    if self.include_norm_residual and self.norm is not None:
        output = residual + projected_output
        if not self.norm_first:
            output = self.norm(output)
    else:
        output = projected_output

    if use_cache:
        return output, current_kv
    return output

base

Attention-layer protocols shared by all attn/ implementations.

Defines the optional extension points adapters (e.g. Prefix Tuning) rely on. Implementations only need to honor the protocols they genuinely support — anything else raises NotImplementedError at adapter construction time so the failure mode is loud, not silent.

PrefixCapableAttention

Bases: Protocol

Attention impls that accept an optional prefix_kv to prepend.

The prefix K/V are concatenated to the projected K/V along the sequence dimension before the attention compute and after any KV-cache write — so the cache only holds the dynamically- generated tokens and the prefix is recomputed (or folded to a static buffer) on every forward.

Implementations that don't support prefix tuning (MLA, Flash Attention fused kernels, etc.) should leave prefix_kv as an accepted-but-ignored arg or raise NotImplementedError. The PrefixTuningAttention wrapper guards construction with an explicit isinstance check so the failure is at adapter build time, not at first forward.

prefix_kv is a (prefix_k, prefix_v) tuple of shape [B, num_kv_heads, prefix_len, head_dim], or None to skip prefix injection (the default).

源代码位于： src/llm/core/attn/base.py

@runtime_checkable
class PrefixCapableAttention(Protocol):
    """Attention impls that accept an optional ``prefix_kv`` to prepend.

    The prefix K/V are concatenated to the projected K/V along the
    sequence dimension **before** the attention compute and **after**
    any KV-cache write — so the cache only holds the dynamically-
    generated tokens and the prefix is recomputed (or folded to a
    static buffer) on every forward.

    Implementations that don't support prefix tuning (MLA, Flash
    Attention fused kernels, etc.) should leave ``prefix_kv`` as an
    accepted-but-ignored arg or raise ``NotImplementedError``. The
    ``PrefixTuningAttention`` wrapper guards construction with an
    explicit ``isinstance`` check so the failure is at adapter build
    time, not at first forward.

    ``prefix_kv`` is a ``(prefix_k, prefix_v)`` tuple of shape
    ``[B, num_kv_heads, prefix_len, head_dim]``, or ``None`` to skip
    prefix injection (the default).
    """

    num_kv_heads: int
    head_dim: int

    def __call__(
        self,
        hidden_states: Tensor,
        prefix_kv: tuple[Tensor, Tensor] | None = None,
        **kwargs,
    ) -> Tensor: ...
    def forward(
        self,
        hidden_states: Tensor,
        prefix_kv: tuple[Tensor, Tensor] | None = None,
        **kwargs,
    ) -> Tensor: ...
    def parameters(self, recurse: bool = True): ...

Paged Attention

Block-allocator KV cache for serving (see ADR-004):

paged_kv_cache

Paged KV Cache for memory-efficient inference.

PrefixCache

Cache for storing prefix KV blocks (block_ids only).

源代码位于： src/llm/core/paged_attention/paged_kv_cache.py

class PrefixCache:
    """Cache for storing prefix KV blocks (block_ids only)."""

    def __init__(self, max_prefixes: int = 10):
        self.max_prefixes = max_prefixes
        self.cache: OrderedDict[str, list[int]] = OrderedDict()

    def add(self, prefix_hash: str, block_ids: list[int]) -> None:
        """Add prefix blocks to cache."""
        if len(self.cache) >= self.max_prefixes:
            self.cache.popitem(last=False)

        self.cache[prefix_hash] = block_ids
        self.cache.move_to_end(prefix_hash)

    def get(self, prefix_hash: str) -> list[int] | None:
        """Get cached block IDs for prefix."""
        if prefix_hash in self.cache:
            self.cache.move_to_end(prefix_hash)
            return self.cache[prefix_hash]
        return None

    def remove(self, prefix_hash: str) -> None:
        """Drop a prefix entry (no-op if absent).

        Used when the sequence that owned the cached blocks is freed — the
        stored block IDs become dangling and must not be replayed (RIL
        ISS-071).  Because LRU eviction may already have dropped the entry,
        absence is not an error.
        """
        self.cache.pop(prefix_hash, None)

add

add(prefix_hash, block_ids)

Add prefix blocks to cache.

源代码位于： src/llm/core/paged_attention/paged_kv_cache.py

def add(self, prefix_hash: str, block_ids: list[int]) -> None:
    """Add prefix blocks to cache."""
    if len(self.cache) >= self.max_prefixes:
        self.cache.popitem(last=False)

    self.cache[prefix_hash] = block_ids
    self.cache.move_to_end(prefix_hash)

get

get(prefix_hash)

Get cached block IDs for prefix.

源代码位于： src/llm/core/paged_attention/paged_kv_cache.py

def get(self, prefix_hash: str) -> list[int] | None:
    """Get cached block IDs for prefix."""
    if prefix_hash in self.cache:
        self.cache.move_to_end(prefix_hash)
        return self.cache[prefix_hash]
    return None

remove

remove(prefix_hash)

Drop a prefix entry (no-op if absent).

Used when the sequence that owned the cached blocks is freed — the stored block IDs become dangling and must not be replayed (RIL ISS-071). Because LRU eviction may already have dropped the entry, absence is not an error.

源代码位于： src/llm/core/paged_attention/paged_kv_cache.py

def remove(self, prefix_hash: str) -> None:
    """Drop a prefix entry (no-op if absent).

    Used when the sequence that owned the cached blocks is freed — the
    stored block IDs become dangling and must not be replayed (RIL
    ISS-071).  Because LRU eviction may already have dropped the entry,
    absence is not an error.
    """
    self.cache.pop(prefix_hash, None)

PagedKVCache

Block-level KV cache for paged attention.

源代码位于： src/llm/core/paged_attention/paged_kv_cache.py

class PagedKVCache:
    """Block-level KV cache for paged attention."""

    def __init__(
        self,
        num_layers: int,
        num_kv_heads: int,
        head_dim: int,
        num_blocks: int,
        block_size: int = 16,
        device: str = "cuda",
        dtype: torch.dtype = torch.float16,
        enable_prefix_cache: bool = False,
        max_prefixes: int = 10,
    ):
        self.num_layers = num_layers
        self.num_kv_heads = num_kv_heads
        self.head_dim = head_dim
        self.block_size = block_size
        self.num_blocks = num_blocks

        self.k_cache = torch.zeros(
            num_layers, num_blocks, num_kv_heads, block_size, head_dim, device=device, dtype=dtype
        )
        self.v_cache = torch.zeros_like(self.k_cache)

        self.block_manager = BlockManager(num_blocks, block_size, num_layers)

        self.enable_prefix_cache = enable_prefix_cache
        self.prefix_cache = PrefixCache(max_prefixes) if enable_prefix_cache else None
        self._seq_to_hash: dict[int, str] = {}
        # Hash -> the sequence that currently OWNS the cached entry. ``free``
        # may only remove an entry whose current owner is the freed sequence:
        # a stale ``_seq_to_hash`` leftover from a sequence whose entry was
        # LRU-evicted (or overwritten by another sequence registering the same
        # prompt) must not steal a live entry registered by someone else.
        self._hash_to_owner: dict[str, int] = {}

    def _hash_tokens(self, tokens: list[int]) -> str:
        """Generate hash for token list.

        Uses ``array.array('i', ...)`` so token ids outside ``[0, 256)``
        (the norm for BPE/SentencePiece vocabularies) are handled correctly.
        """
        return hashlib.sha256(array.array("i", tokens).tobytes()).hexdigest()

    def add_prefix(self, seq_id: int, prefix_tokens: list[int], block_ids: list[int]) -> None:
        """Add prefix blocks to cache, owned by ``seq_id``.

        A sequence registers at most ONE live prefix entry: :meth:`free`
        drops the entry its owner registered (and only that entry), so a
        re-registration under a different hash evicts the previous one here
        to keep that bookkeeping exact (RIL TASK-065).
        """
        if not self.enable_prefix_cache or self.prefix_cache is None:
            return
        prefix_hash = self._hash_tokens(prefix_tokens)
        prev_hash = self._seq_to_hash.pop(seq_id, None)
        if prev_hash is not None and prev_hash != prefix_hash:
            self.prefix_cache.remove(prev_hash)
        # Snapshot, don't alias: the caller (the engine) passes the sequence's
        # LIVE block table, and the prefix must describe only the PREFIX's
        # blocks. Storing the list by reference would let the owner's later
        # ``extend_sequence`` grow the cached entry past the prompt into its
        # decode blocks; a subsequent hit would then fork the owner's WHOLE
        # decode table, pinning those blocks and forcing the owner into a
        # copy-on-write spiral on every decode step (RIL TASK-065 follow-up).
        self.prefix_cache.add(prefix_hash, list(block_ids))
        self._seq_to_hash[seq_id] = prefix_hash
        # The LAST registrar of a hash is the entry's owner (a fresh prefill
        # that re-encountered an LRU-evicted prompt re-registers its own
        # blocks). ``free`` consults this to avoid cross-sequence eviction.
        self._hash_to_owner[prefix_hash] = seq_id
        # Prune owner bookkeeping for hashes no longer cached (LRU eviction
        # drops the entry but not this map); keeps it bounded by max_prefixes.
        for stale in [k for k in self._hash_to_owner if k not in self.prefix_cache.cache]:
            del self._hash_to_owner[stale]

    def try_get_prefix_blocks(self, prefix_tokens: list[int]) -> list[int] | None:
        """Try to get cached prefix blocks."""
        if not self.enable_prefix_cache or self.prefix_cache is None:
            return None
        prefix_hash = self._hash_tokens(prefix_tokens)
        return self.prefix_cache.get(prefix_hash)

    def stage_prefix(self, seq_id: int, prefix_block_ids: list[int], num_prefix_tokens: int) -> list[int]:
        """Start a new sequence sharing a cached prefix (prefix replay).

        Creates the sequence in the block manager with the cached blocks
        forked into its table (shared, refcounted — no K/V is copied). The
        sequence is recorded as already owning ``num_prefix_tokens`` tokens,
        so a subsequent :meth:`update` appends at the right offset and copies
        on write before a token lands inside a still-shared block.

        Args:
            seq_id: Sequence identifier (slot id on the engine path).
            prefix_block_ids: Block ids from :meth:`try_get_prefix_blocks`.
            num_prefix_tokens: Number of already-owned prefix tokens.

        Returns:
            The staged block table.
        """
        if not self.enable_prefix_cache:
            raise RuntimeError("stage_prefix requires enable_prefix_cache=True")
        return self.block_manager.allocate_sequence_shared_prefix(seq_id, prefix_block_ids, num_prefix_tokens)

    def update(self, seq_id: int, k_new: Tensor, v_new: Tensor, layer_idx: int = 0) -> list[int]:
        """Append new tokens to sequence.

        For a brand-new sequence this allocates fresh blocks; for an
        existing sequence it extends the block table only if the new
        tokens cross a block boundary.

        ``layer_idx`` scopes the write to a single transformer layer: the
        decoder calls :meth:`update` once per layer with that layer's own K/V,
        and each call must write **only** its own slice of ``k_cache`` /
        ``v_cache`` (which are ``[num_layers, ...]``).  Only the layer that
        calls first (``layer_idx == 0``) allocates / extends the block table
        and advances the sequence's token count; the remaining layers reuse
        the same block table and write at the same token offsets without
        re-advancing.  (Default 0 keeps the single-layer contract intact.)

        Args:
            seq_id: Sequence identifier.
            k_new: [batch, tokens, num_kv_heads, head_dim]
            v_new: [batch, tokens, num_kv_heads, head_dim]
            layer_idx: Which layer's cache slice to write.

        Returns:
            List of physical block IDs allocated for this sequence
            (initial allocation) or the current full block table
            (extension).
        """
        num_new_tokens = k_new.shape[1]
        k_transposed = k_new.transpose(1, 2)
        v_transposed = v_new.transpose(1, 2)

        if seq_id in self.block_manager.sequences:
            if layer_idx == 0:
                # First layer of this step extends / accounts the sequence.
                existing_num_tokens = self.block_manager.get_num_tokens(seq_id)
                block_table = self.block_manager.extend_sequence(seq_id, num_new_tokens)
                start_token_offset = existing_num_tokens
            else:
                # Later layers already saw layer 0 extend the table this
                # step; they write the same new tokens without re-advancing
                # the token count (which lives on the shared block manager).
                block_table = self.get_block_table(seq_id)
                start_token_offset = self.block_manager.get_num_tokens(seq_id) - num_new_tokens
        else:
            if layer_idx != 0:
                raise RuntimeError(
                    f"Sequence {seq_id} has no block table yet; layer 0 must update() before layer {layer_idx}."
                )
            if not self.block_manager.can_allocate_sequence(num_new_tokens):
                raise RuntimeError("No free blocks available for new sequence")
            # Re-fetch the LIVE table (``allocate_sequence`` returns a copy)
            # so a layer-0 COW remap below propagates to ``get_block_table``
            # callers instead of vanishing into the throwaway copy.
            self.block_manager.allocate_sequence(seq_id, num_new_tokens)
            block_table = self.block_manager.get_block_table(seq_id)
            start_token_offset = 0

        # Write the new tokens into the (possibly extended) block table.
        # Each new token goes into the block whose relative index matches
        # ``(start_token_offset + i) // block_size``.  The leading index is
        # ``layer_idx`` (NOT ``:``) — ``:`` would broadcast that layer's K/V
        # into every layer's slice, so a multi-layer decoder attends over the
        # wrong K/V for every layer but the one that wrote last.
        #
        # Layer 0 owns the block-table remapping: when a staged-prefix
        # sequence (or any shared-block holder) writes into a block that is
        # still referenced by the prefix-cache owner, it must copy-on-write
        # first — writing in place would corrupt the cached K/V the owner
        # (and the model's attention) still reads (RIL TASK-065). Later
        # layers read the remapped table fresh from the manager, so they
        # land in the private block without re-running the COW.
        for i in range(num_new_tokens):
            global_pos = start_token_offset + i
            block_idx = global_pos // self.block_size
            in_block_offset = global_pos % self.block_size
            block_id = block_table[block_idx]
            if layer_idx == 0:
                block_id = self._copy_on_write_if_shared(block_id, block_table, block_idx)
            self.k_cache[layer_idx, block_id, :, in_block_offset, :] = k_transposed[:, :, i, :]
            self.v_cache[layer_idx, block_id, :, in_block_offset, :] = v_transposed[:, :, i, :]

        return block_table

    def _copy_on_write_if_shared(self, block_id: int, block_table: list[int], block_idx: int) -> int:
        """Private-copy ``block_id`` before an in-place write if it is shared.

        A block is shared when another sequence is still reading it — the
        prefix-cache owner whose blocks a staged sequence references, or a
        sibling sequence forked from the same prefix. Overwriting it in place
        would leak the write into the other sequence's context. The fresh
        block receives the FULL multi-layer content (all transformer layers
        share one physical block) before the write, and the sequence's block
        table is remapped; the original block is left byte-identical for its
        other readers.

        Returns the id to write into (unchanged when not shared).
        """
        if not self.block_manager.is_block_shared(block_id):
            return block_id
        old_id = block_id
        new_id = self.block_manager.cow_block(old_id)
        # Copy the preserved prefix content from the shared block before this
        # layer's write overwrites its slice; later layers read the same
        # logical block and must see the full history (TASK-065).
        with torch.no_grad():
            self.k_cache[:, new_id, :, :, :].copy_(self.k_cache[:, old_id, :, :, :])
            self.v_cache[:, new_id, :, :, :].copy_(self.v_cache[:, old_id, :, :, :])
        block_table[block_idx] = new_id
        return new_id

    def get_block_table(self, seq_id: int) -> list[int]:
        """Get block IDs for a sequence."""
        return self.block_manager.get_block_table(seq_id)

    def get(self, seq_id: int, start_idx: int, end_idx: int, layer_idx: int = 0) -> tuple[Tensor, Tensor]:
        """Get KV cache slice for a sequence range.

        Args:
            seq_id: Sequence identifier.
            start_idx: Starting token index (inclusive).
            end_idx: Ending token index (exclusive).
            layer_idx: Which layer's cache slice to read (default 0 keeps
                the single-layer contract; multi-layer readers must pass the
                layer whose KV they're attending over).

        Raises:
            ValueError: If ``start_idx`` or ``end_idx`` are out of bounds
                or ``start_idx >= end_idx``.
        """
        block_table = self.get_block_table(seq_id)

        num_tokens = self.block_manager.get_num_tokens(seq_id)
        if start_idx < 0 or end_idx > num_tokens:
            raise ValueError(
                f"Index range [{start_idx}:{end_idx}] out of bounds for sequence {seq_id} with {num_tokens} tokens"
            )
        if start_idx >= end_idx:
            raise ValueError(f"start_idx ({start_idx}) must be less than end_idx ({end_idx})")

        k_seq = []
        v_seq = []

        start_block = start_idx // self.block_size
        end_block = (end_idx - 1) // self.block_size + 1

        for block_id in block_table[start_block:end_block]:
            # Index on the layer axis ``layer_idx``, not ``:`` — the cache is
            # ``[num_layers, num_blocks, ...]`` and each layer read/writes only
            # its own slice.  After the layer slice each block is
            # ``[num_kv_heads, block_size, head_dim]`` so blocks concatenate
            # along dim 1 (the token axis).
            k_seq.append(self.k_cache[layer_idx, block_id, :, : self.block_size, :])
            v_seq.append(self.v_cache[layer_idx, block_id, :, : self.block_size, :])

        k_full = torch.cat(k_seq, dim=1)
        v_full = torch.cat(v_seq, dim=1)

        start_offset = start_idx % self.block_size
        num_tokens = end_idx - start_idx

        # ``k_full`` is ``[num_kv_heads, num_tokens, head_dim]`` (the layer
        # axis was consumed by the ``layer_idx`` slice above).
        return k_full[:, start_offset : start_offset + num_tokens, :], v_full[
            :, start_offset : start_offset + num_tokens, :
        ]

    def free(self, seq_id: int):
        """Free blocks when sequence completes.

        The sequence's OWN prefix-cache entry (the one it registered via
        :meth:`add_prefix`) is dropped BEFORE its blocks are freed: the entry
        stores this sequence's physical block IDs, and once
        ``free_sequence`` returns them to the allocator a later request may
        be handed the same blocks. Leaving it in place would replay
        another/in-flight sequence's newly-written K/V as a cached prefix —
        use-after-free of the KV blocks (RIL ISS-071).

        Only the entry the sequence itself registered is removed (via
        ``_seq_to_hash``, which :meth:`add_prefix` keeps exact by evicting a
        stale prior entry on re-registration). A sequence that merely
        REPLAYED a prefix via :meth:`stage_prefix` shared the *owner's*
        blocks, which remain live (and pristine — every write into a shared
        block copy-on-writes) for as long as the owner holds them, so its
        free must not evict an entry it does not own; doing so would make the
        first replay destroy the very cache entry that served it (RIL
        TASK-065).
        """
        if self.prefix_cache is not None:
            prefix_hash = self._seq_to_hash.pop(seq_id, None)
            # Only the CURRENT owner of the entry may remove it. A stale
            # ``_seq_to_hash`` mapping from a sequence whose entry was
            # LRU-evicted, or superseded by another sequence's registration of
            # the same prompt, must not evict a live entry owned by another
            # still-running sequence (cross-owner theft).
            if prefix_hash is not None and self._hash_to_owner.get(prefix_hash) == seq_id:
                self.prefix_cache.remove(prefix_hash)
                self._hash_to_owner.pop(prefix_hash, None)
        self.block_manager.free_sequence(seq_id)

add_prefix

add_prefix(seq_id, prefix_tokens, block_ids)

Add prefix blocks to cache, owned by seq_id.

A sequence registers at most ONE live prefix entry: :meth:free drops the entry its owner registered (and only that entry), so a re-registration under a different hash evicts the previous one here to keep that bookkeeping exact (RIL TASK-065).

源代码位于： src/llm/core/paged_attention/paged_kv_cache.py

def add_prefix(self, seq_id: int, prefix_tokens: list[int], block_ids: list[int]) -> None:
    """Add prefix blocks to cache, owned by ``seq_id``.

    A sequence registers at most ONE live prefix entry: :meth:`free`
    drops the entry its owner registered (and only that entry), so a
    re-registration under a different hash evicts the previous one here
    to keep that bookkeeping exact (RIL TASK-065).
    """
    if not self.enable_prefix_cache or self.prefix_cache is None:
        return
    prefix_hash = self._hash_tokens(prefix_tokens)
    prev_hash = self._seq_to_hash.pop(seq_id, None)
    if prev_hash is not None and prev_hash != prefix_hash:
        self.prefix_cache.remove(prev_hash)
    # Snapshot, don't alias: the caller (the engine) passes the sequence's
    # LIVE block table, and the prefix must describe only the PREFIX's
    # blocks. Storing the list by reference would let the owner's later
    # ``extend_sequence`` grow the cached entry past the prompt into its
    # decode blocks; a subsequent hit would then fork the owner's WHOLE
    # decode table, pinning those blocks and forcing the owner into a
    # copy-on-write spiral on every decode step (RIL TASK-065 follow-up).
    self.prefix_cache.add(prefix_hash, list(block_ids))
    self._seq_to_hash[seq_id] = prefix_hash
    # The LAST registrar of a hash is the entry's owner (a fresh prefill
    # that re-encountered an LRU-evicted prompt re-registers its own
    # blocks). ``free`` consults this to avoid cross-sequence eviction.
    self._hash_to_owner[prefix_hash] = seq_id
    # Prune owner bookkeeping for hashes no longer cached (LRU eviction
    # drops the entry but not this map); keeps it bounded by max_prefixes.
    for stale in [k for k in self._hash_to_owner if k not in self.prefix_cache.cache]:
        del self._hash_to_owner[stale]

try_get_prefix_blocks

try_get_prefix_blocks(prefix_tokens)

Try to get cached prefix blocks.

源代码位于： src/llm/core/paged_attention/paged_kv_cache.py

def try_get_prefix_blocks(self, prefix_tokens: list[int]) -> list[int] | None:
    """Try to get cached prefix blocks."""
    if not self.enable_prefix_cache or self.prefix_cache is None:
        return None
    prefix_hash = self._hash_tokens(prefix_tokens)
    return self.prefix_cache.get(prefix_hash)

stage_prefix

stage_prefix(seq_id, prefix_block_ids, num_prefix_tokens)

Start a new sequence sharing a cached prefix (prefix replay).

Creates the sequence in the block manager with the cached blocks forked into its table (shared, refcounted — no K/V is copied). The sequence is recorded as already owning num_prefix_tokens tokens, so a subsequent :meth:update appends at the right offset and copies on write before a token lands inside a still-shared block.

参数：

名称	类型	描述	默认
seq_id	int	Sequence identifier (slot id on the engine path).	必需
prefix_block_ids	list[int]	Block ids from :meth:try_get_prefix_blocks.	必需
num_prefix_tokens	int	Number of already-owned prefix tokens.	必需

返回：

类型	描述
list[int]	The staged block table.

源代码位于： src/llm/core/paged_attention/paged_kv_cache.py

def stage_prefix(self, seq_id: int, prefix_block_ids: list[int], num_prefix_tokens: int) -> list[int]:
    """Start a new sequence sharing a cached prefix (prefix replay).

    Creates the sequence in the block manager with the cached blocks
    forked into its table (shared, refcounted — no K/V is copied). The
    sequence is recorded as already owning ``num_prefix_tokens`` tokens,
    so a subsequent :meth:`update` appends at the right offset and copies
    on write before a token lands inside a still-shared block.

    Args:
        seq_id: Sequence identifier (slot id on the engine path).
        prefix_block_ids: Block ids from :meth:`try_get_prefix_blocks`.
        num_prefix_tokens: Number of already-owned prefix tokens.

    Returns:
        The staged block table.
    """
    if not self.enable_prefix_cache:
        raise RuntimeError("stage_prefix requires enable_prefix_cache=True")
    return self.block_manager.allocate_sequence_shared_prefix(seq_id, prefix_block_ids, num_prefix_tokens)

update

update(seq_id, k_new, v_new, layer_idx=0)

Append new tokens to sequence.

For a brand-new sequence this allocates fresh blocks; for an existing sequence it extends the block table only if the new tokens cross a block boundary.

layer_idx scopes the write to a single transformer layer: the decoder calls :meth:update once per layer with that layer's own K/V, and each call must write only its own slice of k_cache / v_cache (which are [num_layers, ...]). Only the layer that calls first (layer_idx == 0) allocates / extends the block table and advances the sequence's token count; the remaining layers reuse the same block table and write at the same token offsets without re-advancing. (Default 0 keeps the single-layer contract intact.)

参数：

名称	类型	描述	默认
seq_id	int	Sequence identifier.	必需
k_new	Tensor	[batch, tokens, num_kv_heads, head_dim]	必需
v_new	Tensor	[batch, tokens, num_kv_heads, head_dim]	必需
layer_idx	int	Which layer's cache slice to write.	0

返回：

类型	描述
list[int]	List of physical block IDs allocated for this sequence
list[int]	(initial allocation) or the current full block table
list[int]	(extension).

源代码位于： src/llm/core/paged_attention/paged_kv_cache.py

def update(self, seq_id: int, k_new: Tensor, v_new: Tensor, layer_idx: int = 0) -> list[int]:
    """Append new tokens to sequence.

    For a brand-new sequence this allocates fresh blocks; for an
    existing sequence it extends the block table only if the new
    tokens cross a block boundary.

    ``layer_idx`` scopes the write to a single transformer layer: the
    decoder calls :meth:`update` once per layer with that layer's own K/V,
    and each call must write **only** its own slice of ``k_cache`` /
    ``v_cache`` (which are ``[num_layers, ...]``).  Only the layer that
    calls first (``layer_idx == 0``) allocates / extends the block table
    and advances the sequence's token count; the remaining layers reuse
    the same block table and write at the same token offsets without
    re-advancing.  (Default 0 keeps the single-layer contract intact.)

    Args:
        seq_id: Sequence identifier.
        k_new: [batch, tokens, num_kv_heads, head_dim]
        v_new: [batch, tokens, num_kv_heads, head_dim]
        layer_idx: Which layer's cache slice to write.

    Returns:
        List of physical block IDs allocated for this sequence
        (initial allocation) or the current full block table
        (extension).
    """
    num_new_tokens = k_new.shape[1]
    k_transposed = k_new.transpose(1, 2)
    v_transposed = v_new.transpose(1, 2)

    if seq_id in self.block_manager.sequences:
        if layer_idx == 0:
            # First layer of this step extends / accounts the sequence.
            existing_num_tokens = self.block_manager.get_num_tokens(seq_id)
            block_table = self.block_manager.extend_sequence(seq_id, num_new_tokens)
            start_token_offset = existing_num_tokens
        else:
            # Later layers already saw layer 0 extend the table this
            # step; they write the same new tokens without re-advancing
            # the token count (which lives on the shared block manager).
            block_table = self.get_block_table(seq_id)
            start_token_offset = self.block_manager.get_num_tokens(seq_id) - num_new_tokens
    else:
        if layer_idx != 0:
            raise RuntimeError(
                f"Sequence {seq_id} has no block table yet; layer 0 must update() before layer {layer_idx}."
            )
        if not self.block_manager.can_allocate_sequence(num_new_tokens):
            raise RuntimeError("No free blocks available for new sequence")
        # Re-fetch the LIVE table (``allocate_sequence`` returns a copy)
        # so a layer-0 COW remap below propagates to ``get_block_table``
        # callers instead of vanishing into the throwaway copy.
        self.block_manager.allocate_sequence(seq_id, num_new_tokens)
        block_table = self.block_manager.get_block_table(seq_id)
        start_token_offset = 0

    # Write the new tokens into the (possibly extended) block table.
    # Each new token goes into the block whose relative index matches
    # ``(start_token_offset + i) // block_size``.  The leading index is
    # ``layer_idx`` (NOT ``:``) — ``:`` would broadcast that layer's K/V
    # into every layer's slice, so a multi-layer decoder attends over the
    # wrong K/V for every layer but the one that wrote last.
    #
    # Layer 0 owns the block-table remapping: when a staged-prefix
    # sequence (or any shared-block holder) writes into a block that is
    # still referenced by the prefix-cache owner, it must copy-on-write
    # first — writing in place would corrupt the cached K/V the owner
    # (and the model's attention) still reads (RIL TASK-065). Later
    # layers read the remapped table fresh from the manager, so they
    # land in the private block without re-running the COW.
    for i in range(num_new_tokens):
        global_pos = start_token_offset + i
        block_idx = global_pos // self.block_size
        in_block_offset = global_pos % self.block_size
        block_id = block_table[block_idx]
        if layer_idx == 0:
            block_id = self._copy_on_write_if_shared(block_id, block_table, block_idx)
        self.k_cache[layer_idx, block_id, :, in_block_offset, :] = k_transposed[:, :, i, :]
        self.v_cache[layer_idx, block_id, :, in_block_offset, :] = v_transposed[:, :, i, :]

    return block_table

get_block_table

get_block_table(seq_id)

Get block IDs for a sequence.

源代码位于： src/llm/core/paged_attention/paged_kv_cache.py

def get_block_table(self, seq_id: int) -> list[int]:
    """Get block IDs for a sequence."""
    return self.block_manager.get_block_table(seq_id)

get

get(seq_id, start_idx, end_idx, layer_idx=0)

Get KV cache slice for a sequence range.

参数：

名称	类型	描述	默认
seq_id	int	Sequence identifier.	必需
start_idx	int	Starting token index (inclusive).	必需
end_idx	int	Ending token index (exclusive).	必需
layer_idx	int	Which layer's cache slice to read (default 0 keeps the single-layer contract; multi-layer readers must pass the layer whose KV they're attending over).	0

引发：

类型	描述
ValueError	If start_idx or end_idx are out of bounds or start_idx >= end_idx.

源代码位于： src/llm/core/paged_attention/paged_kv_cache.py

def get(self, seq_id: int, start_idx: int, end_idx: int, layer_idx: int = 0) -> tuple[Tensor, Tensor]:
    """Get KV cache slice for a sequence range.

    Args:
        seq_id: Sequence identifier.
        start_idx: Starting token index (inclusive).
        end_idx: Ending token index (exclusive).
        layer_idx: Which layer's cache slice to read (default 0 keeps
            the single-layer contract; multi-layer readers must pass the
            layer whose KV they're attending over).

    Raises:
        ValueError: If ``start_idx`` or ``end_idx`` are out of bounds
            or ``start_idx >= end_idx``.
    """
    block_table = self.get_block_table(seq_id)

    num_tokens = self.block_manager.get_num_tokens(seq_id)
    if start_idx < 0 or end_idx > num_tokens:
        raise ValueError(
            f"Index range [{start_idx}:{end_idx}] out of bounds for sequence {seq_id} with {num_tokens} tokens"
        )
    if start_idx >= end_idx:
        raise ValueError(f"start_idx ({start_idx}) must be less than end_idx ({end_idx})")

    k_seq = []
    v_seq = []

    start_block = start_idx // self.block_size
    end_block = (end_idx - 1) // self.block_size + 1

    for block_id in block_table[start_block:end_block]:
        # Index on the layer axis ``layer_idx``, not ``:`` — the cache is
        # ``[num_layers, num_blocks, ...]`` and each layer read/writes only
        # its own slice.  After the layer slice each block is
        # ``[num_kv_heads, block_size, head_dim]`` so blocks concatenate
        # along dim 1 (the token axis).
        k_seq.append(self.k_cache[layer_idx, block_id, :, : self.block_size, :])
        v_seq.append(self.v_cache[layer_idx, block_id, :, : self.block_size, :])

    k_full = torch.cat(k_seq, dim=1)
    v_full = torch.cat(v_seq, dim=1)

    start_offset = start_idx % self.block_size
    num_tokens = end_idx - start_idx

    # ``k_full`` is ``[num_kv_heads, num_tokens, head_dim]`` (the layer
    # axis was consumed by the ``layer_idx`` slice above).
    return k_full[:, start_offset : start_offset + num_tokens, :], v_full[
        :, start_offset : start_offset + num_tokens, :
    ]

free

free(seq_id)

Free blocks when sequence completes.

The sequence's OWN prefix-cache entry (the one it registered via :meth:add_prefix) is dropped BEFORE its blocks are freed: the entry stores this sequence's physical block IDs, and once free_sequence returns them to the allocator a later request may be handed the same blocks. Leaving it in place would replay another/in-flight sequence's newly-written K/V as a cached prefix — use-after-free of the KV blocks (RIL ISS-071).

Only the entry the sequence itself registered is removed (via _seq_to_hash, which :meth:add_prefix keeps exact by evicting a stale prior entry on re-registration). A sequence that merely REPLAYED a prefix via :meth:stage_prefix shared the owner's blocks, which remain live (and pristine — every write into a shared block copy-on-writes) for as long as the owner holds them, so its free must not evict an entry it does not own; doing so would make the first replay destroy the very cache entry that served it (RIL TASK-065).

源代码位于： src/llm/core/paged_attention/paged_kv_cache.py

def free(self, seq_id: int):
    """Free blocks when sequence completes.

    The sequence's OWN prefix-cache entry (the one it registered via
    :meth:`add_prefix`) is dropped BEFORE its blocks are freed: the entry
    stores this sequence's physical block IDs, and once
    ``free_sequence`` returns them to the allocator a later request may
    be handed the same blocks. Leaving it in place would replay
    another/in-flight sequence's newly-written K/V as a cached prefix —
    use-after-free of the KV blocks (RIL ISS-071).

    Only the entry the sequence itself registered is removed (via
    ``_seq_to_hash``, which :meth:`add_prefix` keeps exact by evicting a
    stale prior entry on re-registration). A sequence that merely
    REPLAYED a prefix via :meth:`stage_prefix` shared the *owner's*
    blocks, which remain live (and pristine — every write into a shared
    block copy-on-writes) for as long as the owner holds them, so its
    free must not evict an entry it does not own; doing so would make the
    first replay destroy the very cache entry that served it (RIL
    TASK-065).
    """
    if self.prefix_cache is not None:
        prefix_hash = self._seq_to_hash.pop(seq_id, None)
        # Only the CURRENT owner of the entry may remove it. A stale
        # ``_seq_to_hash`` mapping from a sequence whose entry was
        # LRU-evicted, or superseded by another sequence's registration of
        # the same prompt, must not evict a live entry owned by another
        # still-running sequence (cross-owner theft).
        if prefix_hash is not None and self._hash_to_owner.get(prefix_hash) == seq_id:
            self.prefix_cache.remove(prefix_hash)
            self._hash_to_owner.pop(prefix_hash, None)
    self.block_manager.free_sequence(seq_id)

attention

Paged Attention forward implementation.

paged_attention_forward

paged_attention_forward(q, k_cache, v_cache, block_tables, seq_lens, num_kv_heads, block_size=16, query_lens=None)

Paged attention forward pass.

Supports both prefill (q with multiple query tokens) and decode (q with a single query token per sequence). Each row of q attends to its full cached context (the first seq_lens[b] tokens of sequence b), regardless of how many query tokens that row carries. The original Python-fallback kernel gathered the whole k/v slice per sequence regardless of S_q; the multi-token generalisation just lets the matmul produce S_q outputs instead of one.

参数：

名称	类型	描述	默认
q	Tensor	Query tensor [batch, num_heads, query_len, head_dim].	必需
k_cache	Tensor	KV cache tensor. Either the per-layer slice [num_blocks, num_kv_heads, block_size, head_dim] (the production path — caller slices PagedKVCache.k_cache[layer_idx] before passing) or the full PagedKVCache shape [num_layers, num_blocks, num_kv_heads, block_size, head_dim] (legacy / direct tests). When 5-D, the layer axis is collapsed by taking index 0; for a multi-layer model the caller must slice.	必需
v_cache	Tensor	Same shape as k_cache.	必需
block_tables	Tensor	[batch, max_blocks] physical block IDs per sequence.	必需
seq_lens	Tensor	Current sequence lengths [batch].	必需
num_kv_heads	int	Number of KV heads.	必需
block_size	int	Tokens per block.	16
query_lens	Tensor | None	Optional per-row number of REAL query tokens [batch]. In continuous batching, q is left-padded to the batch-max query length so a decode row's single query lives at local index 0, and the causal overlay (which is only meaningful for prefill rows with multiple query tokens) must not mask a decode row's full context (RIL ISS-048). When None (direct tests / deprecated callers), the overlay applies to every row whenever query_len > 1, mirroring the historical behaviour.	None

返回：

类型	描述
Tensor	Attention output tensor [batch, num_heads, query_len, head_dim].

源代码位于： src/llm/core/paged_attention/attention.py

def paged_attention_forward(
    q: Tensor,
    k_cache: Tensor,
    v_cache: Tensor,
    block_tables: Tensor,
    seq_lens: Tensor,
    num_kv_heads: int,
    block_size: int = 16,
    query_lens: Tensor | None = None,
) -> Tensor:
    """Paged attention forward pass.

    Supports both prefill (``q`` with multiple query tokens) and decode
    (``q`` with a single query token per sequence). Each row of ``q``
    attends to its full cached context (the first ``seq_lens[b]`` tokens
    of sequence ``b``), regardless of how many query tokens that row
    carries. The original Python-fallback kernel gathered the whole k/v
    slice per sequence regardless of ``S_q``; the multi-token generalisation
    just lets the matmul produce ``S_q`` outputs instead of one.

    Args:
        q: Query tensor [batch, num_heads, query_len, head_dim].
        k_cache: KV cache tensor. Either the per-layer slice
            ``[num_blocks, num_kv_heads, block_size, head_dim]`` (the
            production path — caller slices ``PagedKVCache.k_cache[layer_idx]``
            before passing) or the full ``PagedKVCache`` shape
            ``[num_layers, num_blocks, num_kv_heads, block_size, head_dim]``
            (legacy / direct tests). When 5-D, the layer axis is collapsed
            by taking index 0; for a multi-layer model the caller must slice.
        v_cache: Same shape as ``k_cache``.
        block_tables: [batch, max_blocks] physical block IDs per sequence.
        seq_lens: Current sequence lengths [batch].
        num_kv_heads: Number of KV heads.
        block_size: Tokens per block.
        query_lens: Optional per-row number of REAL query tokens [batch].
            In continuous batching, ``q`` is left-padded to the batch-max
            query length so a decode row's single query lives at local index
            0, and the causal overlay (which is only meaningful for prefill
            rows with multiple query tokens) must not mask a decode row's
            full context (RIL ISS-048). When ``None`` (direct tests /
            deprecated callers), the overlay applies to every row whenever
            ``query_len > 1``, mirroring the historical behaviour.

    Returns:
        Attention output tensor [batch, num_heads, query_len, head_dim].
    """
    batch_size, num_heads, _query_len, head_dim = q.shape
    if k_cache.ndim == 5:
        # Legacy / direct-test path: full ``PagedKVCache`` buffer.
        # The caller should have sliced already for a multi-layer model;
        # we fall back to layer 0 to match the historical behaviour.
        k_cache = k_cache[0]
        v_cache = v_cache[0]
    num_blocks = k_cache.shape[0]

    max_seq_len = seq_lens.max().item()

    k_gathered = []
    v_gathered = []

    for b in range(batch_size):
        seq_len = seq_lens[b].item()
        num_blocks_needed = (seq_len + block_size - 1) // block_size

        seq_block_ids = block_tables[b, :num_blocks_needed].tolist()

        k_seq = []
        v_seq = []
        for i, block_id in enumerate(seq_block_ids):
            if block_id < 0 or block_id >= num_blocks:
                continue
            start = 0
            end = seq_len - (num_blocks_needed - 1) * block_size if i == num_blocks_needed - 1 else block_size

            k_seq.append(k_cache[block_id, :, start:end, :])
            v_seq.append(v_cache[block_id, :, start:end, :])

        if k_seq:
            k_full = torch.cat(k_seq, dim=1)
            v_full = torch.cat(v_seq, dim=1)
        else:
            # No valid blocks for this sequence (e.g. seq_len == 0 or all
            # block IDs out of range): skip to the padding step below by
            # starting from an empty tensor that pad_len will extend.
            k_full = k_cache.new_zeros(num_kv_heads, 0, head_dim)
            v_full = v_cache.new_zeros(num_kv_heads, 0, head_dim)

        if k_full.shape[1] < max_seq_len:
            pad_len = int(max_seq_len - k_full.shape[1])
            k_full = torch.cat([k_full, k_full.new_zeros(num_kv_heads, pad_len, head_dim)], dim=1)
            v_full = torch.cat([v_full, v_full.new_zeros(num_kv_heads, pad_len, head_dim)], dim=1)

        k_gathered.append(k_full)
        v_gathered.append(v_full)

    k_full = torch.stack(k_gathered, dim=0).to(q.device)
    v_full = torch.stack(v_gathered, dim=0).to(q.device)

    if num_kv_heads != num_heads:
        repeat_factor = num_heads // num_kv_heads
        k_full = k_full.repeat_interleave(repeat_factor, dim=1)
        v_full = v_full.repeat_interleave(repeat_factor, dim=1)

    scale = head_dim**-0.5
    attn_weights = torch.matmul(q, k_full.transpose(-2, -1)) * scale

    # ``k_full`` is padded with zeros to the batch-max context length, so
    # columns beyond each row's real ``seq_lens[b]`` must be masked or they
    # would participate in the softmax and distort the attention weights.
    col_idx = torch.arange(max_seq_len, device=q.device).reshape(1, 1, 1, -1)
    attn_mask = col_idx >= seq_lens.reshape(-1, 1, 1, 1)  # [B, 1, 1, max_seq_len]
    if _query_len > 1:
        # Causality between query positions (query row ``s`` attends to keys
        # ``0..s``) is only meaningful for prefill rows that carry multiple
        # real query tokens. In continuous batching ``q`` is left-padded to
        # the batch-max query length, so a decode row's single real query sits
        # at local index 0 with padding after it; applying the overlay to that
        # row masks every key > 0, collapsing decode attention onto key 0 and
        # producing garbage output (RIL ISS-048). Gate the overlay per-row on
        # having more than one real query token.
        if query_lens is not None:
            prefill_rows = query_lens.reshape(-1, 1, 1, 1) > 1  # [B, 1, 1, 1]
            row_idx = torch.arange(_query_len, device=q.device).reshape(1, 1, -1, 1)
            causal = col_idx > row_idx  # [1, 1, query_len, max_seq_len]
            attn_mask = attn_mask | (causal & prefill_rows)
        else:
            # Legacy / direct-test callers that don't pass query_lens: apply
            # the overlay to every row (historical behaviour).
            row_idx = torch.arange(_query_len, device=q.device).reshape(1, 1, -1, 1)
            attn_mask = attn_mask | (col_idx > row_idx)
    attn_weights = attn_weights.masked_fill(attn_mask, float("-inf"))

    attn_weights = torch.softmax(attn_weights, dim=-1)

    output = torch.matmul(attn_weights, v_full)

    return output

block_allocator

Block Allocator for Paged Attention.

Manages allocation and deallocation of physical memory blocks.

BlockAllocator

Allocator for physical memory blocks.

Uses a free-list approach to efficiently manage block allocation and deallocation. Supports reference counting for copy-on-write.

源代码位于： src/llm/core/paged_attention/block_allocator.py

class BlockAllocator:
    """
    Allocator for physical memory blocks.

    Uses a free-list approach to efficiently manage block allocation
    and deallocation. Supports reference counting for copy-on-write.
    """

    def __init__(self, num_blocks: int, block_size: int = 16):
        """
        Initialize the block allocator.

        Args:
            num_blocks: Total number of physical blocks available.
            block_size: Number of tokens per block.
        """
        self.num_blocks = num_blocks
        self.block_size = block_size

        # Free list of block indices
        self.free_blocks: deque[int] = deque(range(num_blocks))

        # Reference counts for copy-on-write
        self.ref_counts: dict[int, int] = {}

    @property
    def num_free_blocks(self) -> int:
        """Number of available blocks."""
        return len(self.free_blocks)

    @property
    def num_allocated_blocks(self) -> int:
        """Number of allocated blocks."""
        return self.num_blocks - self.num_free_blocks

    def can_allocate(self, num_blocks: int = 1) -> bool:
        """Check if the requested number of blocks can be allocated."""
        return len(self.free_blocks) >= num_blocks

    def allocate(self) -> int:
        """
        Allocate a single block.

        Returns:
            Block index.

        Raises:
            RuntimeError: If no free blocks are available.
        """
        if not self.free_blocks:
            raise RuntimeError("No free blocks available")

        block_id = self.free_blocks.popleft()
        self.ref_counts[block_id] = 1
        return block_id

    def allocate_n(self, n: int) -> list[int]:
        """
        Allocate multiple blocks.

        Args:
            n: Number of blocks to allocate.

        Returns:
            List of block indices.

        Raises:
            RuntimeError: If not enough free blocks are available.
        """
        if len(self.free_blocks) < n:
            raise RuntimeError(f"Not enough free blocks: requested {n}, available {len(self.free_blocks)}")

        return [self.allocate() for _ in range(n)]

    def free(self, block_id: int) -> None:
        """
        Free a block (decrement reference count).

        The block is only returned to the free list when its
        reference count reaches zero.

        Args:
            block_id: Block index to free.
        """
        if block_id not in self.ref_counts:
            raise ValueError(f"Block {block_id} is not allocated")

        self.ref_counts[block_id] -= 1

        if self.ref_counts[block_id] == 0:
            del self.ref_counts[block_id]
            self.free_blocks.append(block_id)

    def free_all(self, block_ids: list[int]) -> None:
        """Free multiple blocks."""
        for block_id in block_ids:
            self.free(block_id)

    def fork(self, block_id: int) -> int:
        """
        Fork a block for copy-on-write.

        Increments the reference count instead of copying.

        Args:
            block_id: Block to fork.

        Returns:
            Same block_id (sharing the physical block).
        """
        if block_id not in self.ref_counts:
            raise ValueError(f"Block {block_id} is not allocated")

        self.ref_counts[block_id] += 1
        return block_id

    def get_ref_count(self, block_id: int) -> int:
        """Get the reference count of a block."""
        return self.ref_counts.get(block_id, 0)

    def is_shared(self, block_id: int) -> bool:
        """Check if a block is shared (ref_count > 1)."""
        return self.get_ref_count(block_id) > 1

    def copy_on_write(self, block_id: int) -> int:
        """
        Perform copy-on-write if block is shared.

        If the block has ref_count > 1, allocates a new block
        and decrements the old block's ref_count.

        Args:
            block_id: Block that may need copying.

        Returns:
            New block_id (may be same if not shared).
        """
        if not self.is_shared(block_id):
            return block_id

        # Allocate new block
        new_block_id = self.allocate()

        # Decrement old block's ref_count
        self.ref_counts[block_id] -= 1
        if self.ref_counts[block_id] == 0:
            del self.ref_counts[block_id]
            self.free_blocks.append(block_id)

        return new_block_id

    def reset(self) -> None:
        """Reset allocator to initial state."""
        self.free_blocks = deque(range(self.num_blocks))
        self.ref_counts.clear()

    def __repr__(self) -> str:
        return (
            f"BlockAllocator(num_blocks={self.num_blocks}, "
            f"free={self.num_free_blocks}, "
            f"allocated={self.num_allocated_blocks})"
        )

num_free_blocks property

num_free_blocks

Number of available blocks.

num_allocated_blocks property

num_allocated_blocks

Number of allocated blocks.

can_allocate

can_allocate(num_blocks=1)

Check if the requested number of blocks can be allocated.

源代码位于： src/llm/core/paged_attention/block_allocator.py

def can_allocate(self, num_blocks: int = 1) -> bool:
    """Check if the requested number of blocks can be allocated."""
    return len(self.free_blocks) >= num_blocks

allocate

allocate()

Allocate a single block.

返回：

类型	描述
int	Block index.

引发：

类型	描述
RuntimeError	If no free blocks are available.

源代码位于： src/llm/core/paged_attention/block_allocator.py

def allocate(self) -> int:
    """
    Allocate a single block.

    Returns:
        Block index.

    Raises:
        RuntimeError: If no free blocks are available.
    """
    if not self.free_blocks:
        raise RuntimeError("No free blocks available")

    block_id = self.free_blocks.popleft()
    self.ref_counts[block_id] = 1
    return block_id

allocate_n

allocate_n(n)

Allocate multiple blocks.

参数：

名称	类型	描述	默认
n	int	Number of blocks to allocate.	必需

返回：

类型	描述
list[int]	List of block indices.

引发：

类型	描述
RuntimeError	If not enough free blocks are available.

源代码位于： src/llm/core/paged_attention/block_allocator.py

def allocate_n(self, n: int) -> list[int]:
    """
    Allocate multiple blocks.

    Args:
        n: Number of blocks to allocate.

    Returns:
        List of block indices.

    Raises:
        RuntimeError: If not enough free blocks are available.
    """
    if len(self.free_blocks) < n:
        raise RuntimeError(f"Not enough free blocks: requested {n}, available {len(self.free_blocks)}")

    return [self.allocate() for _ in range(n)]

free

free(block_id)

Free a block (decrement reference count).

The block is only returned to the free list when its reference count reaches zero.

参数：

名称	类型	描述	默认
block_id	int	Block index to free.	必需

源代码位于： src/llm/core/paged_attention/block_allocator.py

def free(self, block_id: int) -> None:
    """
    Free a block (decrement reference count).

    The block is only returned to the free list when its
    reference count reaches zero.

    Args:
        block_id: Block index to free.
    """
    if block_id not in self.ref_counts:
        raise ValueError(f"Block {block_id} is not allocated")

    self.ref_counts[block_id] -= 1

    if self.ref_counts[block_id] == 0:
        del self.ref_counts[block_id]
        self.free_blocks.append(block_id)

free_all

free_all(block_ids)

Free multiple blocks.

源代码位于： src/llm/core/paged_attention/block_allocator.py

def free_all(self, block_ids: list[int]) -> None:
    """Free multiple blocks."""
    for block_id in block_ids:
        self.free(block_id)

fork

fork(block_id)

Fork a block for copy-on-write.

Increments the reference count instead of copying.

参数：

名称	类型	描述	默认
block_id	int	Block to fork.	必需

返回：

类型	描述
int	Same block_id (sharing the physical block).

源代码位于： src/llm/core/paged_attention/block_allocator.py

def fork(self, block_id: int) -> int:
    """
    Fork a block for copy-on-write.

    Increments the reference count instead of copying.

    Args:
        block_id: Block to fork.

    Returns:
        Same block_id (sharing the physical block).
    """
    if block_id not in self.ref_counts:
        raise ValueError(f"Block {block_id} is not allocated")

    self.ref_counts[block_id] += 1
    return block_id

get_ref_count

get_ref_count(block_id)

Get the reference count of a block.

源代码位于： src/llm/core/paged_attention/block_allocator.py

def get_ref_count(self, block_id: int) -> int:
    """Get the reference count of a block."""
    return self.ref_counts.get(block_id, 0)

is_shared

is_shared(block_id)

Check if a block is shared (ref_count > 1).

源代码位于： src/llm/core/paged_attention/block_allocator.py

def is_shared(self, block_id: int) -> bool:
    """Check if a block is shared (ref_count > 1)."""
    return self.get_ref_count(block_id) > 1

copy_on_write

copy_on_write(block_id)

Perform copy-on-write if block is shared.

If the block has ref_count > 1, allocates a new block and decrements the old block's ref_count.

参数：

名称	类型	描述	默认
block_id	int	Block that may need copying.	必需

返回：

类型	描述
int	New block_id (may be same if not shared).

源代码位于： src/llm/core/paged_attention/block_allocator.py

def copy_on_write(self, block_id: int) -> int:
    """
    Perform copy-on-write if block is shared.

    If the block has ref_count > 1, allocates a new block
    and decrements the old block's ref_count.

    Args:
        block_id: Block that may need copying.

    Returns:
        New block_id (may be same if not shared).
    """
    if not self.is_shared(block_id):
        return block_id

    # Allocate new block
    new_block_id = self.allocate()

    # Decrement old block's ref_count
    self.ref_counts[block_id] -= 1
    if self.ref_counts[block_id] == 0:
        del self.ref_counts[block_id]
        self.free_blocks.append(block_id)

    return new_block_id

reset

reset()

Reset allocator to initial state.

源代码位于： src/llm/core/paged_attention/block_allocator.py

def reset(self) -> None:
    """Reset allocator to initial state."""
    self.free_blocks = deque(range(self.num_blocks))
    self.ref_counts.clear()

block_manager

Block Manager for Paged Attention.

Manages logical-to-physical block mapping for sequences.

SequenceBlockInfo dataclass

Block information for a single sequence.

源代码位于： src/llm/core/paged_attention/block_manager.py

@dataclass
class SequenceBlockInfo:
    """Block information for a single sequence."""

    seq_id: int
    block_table: list[int] = field(default_factory=list)  # List of physical block IDs
    num_tokens: int = 0  # Total tokens in this sequence

BlockManager

Manages block allocation for multiple sequences.

Maintains a mapping from sequence IDs to their block tables, handling allocation, extension, and freeing of blocks.

源代码位于： src/llm/core/paged_attention/block_manager.py

class BlockManager:
    """
    Manages block allocation for multiple sequences.

    Maintains a mapping from sequence IDs to their block tables,
    handling allocation, extension, and freeing of blocks.
    """

    def __init__(
        self,
        num_blocks: int,
        block_size: int = 16,
        num_layers: int = 1,
    ):
        """
        Initialize the block manager.

        Args:
            num_blocks: Total physical blocks per layer.
            block_size: Tokens per block.
            num_layers: Number of transformer layers.
        """
        self.block_size = block_size
        self.num_layers = num_layers

        # One allocator per layer for independent management
        self.allocators = [BlockAllocator(num_blocks, block_size) for _ in range(num_layers)]

        # Sequence block tables: seq_id -> SequenceBlockInfo
        self.sequences: dict[int, SequenceBlockInfo] = {}

    @property
    def num_free_blocks(self) -> int:
        """Minimum free blocks across all layers."""
        return min(alloc.num_free_blocks for alloc in self.allocators)

    def can_allocate_sequence(self, num_tokens: int) -> bool:
        """
        Check if a new sequence with given tokens can be allocated.

        Args:
            num_tokens: Number of tokens in the sequence.

        Returns:
            True if allocation is possible.
        """
        num_blocks_needed = self._tokens_to_blocks(num_tokens)
        return all(alloc.can_allocate(num_blocks_needed) for alloc in self.allocators)

    def allocate_sequence(self, seq_id: int, num_tokens: int) -> list[int]:
        """
        Allocate blocks for a new sequence.

        Args:
            seq_id: Unique sequence identifier.
            num_tokens: Initial number of tokens.

        Returns:
            List of physical block IDs for the first layer
            (all layers get the same logical structure).

        Raises:
            RuntimeError: If allocation fails.
            ValueError: If sequence already exists.
        """
        if seq_id in self.sequences:
            raise ValueError(f"Sequence {seq_id} already exists")

        num_blocks = self._tokens_to_blocks(num_tokens)

        # Allocate blocks for each layer
        block_tables_per_layer = []
        try:
            for allocator in self.allocators:
                blocks = allocator.allocate_n(num_blocks)
                block_tables_per_layer.append(blocks)
        except RuntimeError:
            # Rollback on failure
            for i, blocks in enumerate(block_tables_per_layer):
                self.allocators[i].free_all(blocks)
            raise

        # Store first layer's block table (representative)
        first_layer_blocks = block_tables_per_layer[0]
        self.sequences[seq_id] = SequenceBlockInfo(
            seq_id=seq_id,
            block_table=first_layer_blocks.copy(),
            num_tokens=num_tokens,
        )

        return first_layer_blocks

    def extend_sequence(self, seq_id: int, num_new_tokens: int) -> list[int]:
        """
        Extend a sequence with additional tokens.

        Allocates new blocks if needed.

        Args:
            seq_id: Sequence to extend.
            num_new_tokens: Number of new tokens to add.

        Returns:
            Updated block table for the sequence.
        """
        if seq_id not in self.sequences:
            raise ValueError(f"Sequence {seq_id} does not exist")

        info = self.sequences[seq_id]
        old_num_blocks = len(info.block_table)
        new_total_tokens = info.num_tokens + num_new_tokens
        new_num_blocks = self._tokens_to_blocks(new_total_tokens)

        blocks_to_add = new_num_blocks - old_num_blocks

        if blocks_to_add > 0:
            # Allocate additional blocks from each layer's allocator.
            # All allocators start from ``deque(range(num_blocks))`` and
            # allocate in the same order, so the block IDs are identical
            # across layers. We only extend the tracked block table (layer
            # 0's) since ``free_sequence`` frees the same IDs on every
            # allocator and they are guaranteed to match.
            for i, allocator in enumerate(self.allocators):
                new_blocks = allocator.allocate_n(blocks_to_add)
                if i == 0:
                    info.block_table.extend(new_blocks)

        info.num_tokens = new_total_tokens
        return info.block_table

    def free_sequence(self, seq_id: int) -> None:
        """
        Free all blocks associated with a sequence.

        Args:
            seq_id: Sequence to free.
        """
        if seq_id not in self.sequences:
            return  # Already freed or never existed

        info = self.sequences.pop(seq_id)

        # Free blocks in all layers
        for allocator in self.allocators:
            allocator.free_all(info.block_table)

    def fork_sequence(self, src_seq_id: int, dst_seq_id: int) -> list[int]:
        """
        Fork a sequence using copy-on-write.

        Creates a new sequence that shares blocks with the source
        until either is modified.

        Args:
            src_seq_id: Source sequence to fork from.
            dst_seq_id: New sequence ID.

        Returns:
            Block table for the new sequence.
        """
        if src_seq_id not in self.sequences:
            raise ValueError(f"Source sequence {src_seq_id} does not exist")
        if dst_seq_id in self.sequences:
            raise ValueError(f"Destination sequence {dst_seq_id} already exists")

        src_info = self.sequences[src_seq_id]

        # Fork blocks (increment ref counts)
        for allocator in self.allocators:
            for block_id in src_info.block_table:
                allocator.fork(block_id)

        # Create new sequence with shared blocks
        self.sequences[dst_seq_id] = SequenceBlockInfo(
            seq_id=dst_seq_id,
            block_table=src_info.block_table.copy(),
            num_tokens=src_info.num_tokens,
        )

        return src_info.block_table.copy()

    def allocate_sequence_shared_prefix(self, seq_id: int, prefix_block_ids: list[int], num_tokens: int) -> list[int]:
        """Create a new sequence whose initial blocks are a SHARED prefix.

        Used for paged prefix replay: the cached prefix blocks belong to
        another (still-live) sequence, but a new request with the same prompt
        may reference them for reads without recomputing their K/V. Each
        block table entry is forked (refcount +1 on every layer allocator)
        — nothing is copied; the sequence must copy-on-write before any write
        lands inside a shared block (:meth:`cow_block`), or the prefix cache
        owner's K/V would be corrupted.

        Args:
            seq_id: Unique sequence identifier.
            prefix_block_ids: Physical block ids (layer 0's view; ids are
                identical across layers) of the cached prefix.
            num_tokens: Number of prefix tokens this sequence already owns
                (the staged prefix length).

        Returns:
            The sequence's block table (the shared prefix blocks).

        Raises:
            ValueError: If the sequence already exists, the prefix is empty,
                or ``num_tokens`` spans more blocks than the cached table.
        """
        if seq_id in self.sequences:
            raise ValueError(f"Sequence {seq_id} already exists")
        if not prefix_block_ids:
            raise ValueError("prefix_block_ids must be non-empty")
        if self._tokens_to_blocks(num_tokens) > len(prefix_block_ids):
            raise ValueError(
                f"{num_tokens} prefix tokens need {self._tokens_to_blocks(num_tokens)} blocks "
                f"but the cached table has {len(prefix_block_ids)}"
            )

        for allocator in self.allocators:
            for block_id in prefix_block_ids:
                allocator.fork(block_id)

        self.sequences[seq_id] = SequenceBlockInfo(
            seq_id=seq_id,
            block_table=list(prefix_block_ids),
            num_tokens=num_tokens,
        )
        return list(prefix_block_ids)

    def is_block_shared(self, block_id: int) -> bool:
        """Whether a block is referenced by more than one sequence.

        Blocks are forked/freed uniformly across the per-layer allocators, so
        the layer-0 refcount is authoritative.
        """
        return self.allocators[0].is_shared(block_id)

    def cow_block(self, block_id: int) -> int:
        """Copy-on-write a logical block across every layer allocator.

        Allocates one fresh physical block per layer and decrements the
        shared block's refcount on each. Because all allocators allocate in
        lock-step (identical call order), every layer returns the same new id,
        preserving the cross-layer block-id invariant. The caller must copy
        the block's data into the new id before writing into it.

        Returns:
            The new private block id (or ``block_id`` itself when it was not
            shared, in which case nothing is allocated or decremented).
        """
        new_id = self.allocators[0].copy_on_write(block_id)
        for allocator in self.allocators[1:]:
            other_id = allocator.copy_on_write(block_id)
            if other_id != new_id:
                raise RuntimeError(
                    f"allocators diverged on COW of block {block_id}: "
                    f"layer 0 produced {new_id}, another produced {other_id}"
                )
        return new_id

    def get_block_table(self, seq_id: int) -> list[int]:
        """Get the block table for a sequence."""
        if seq_id not in self.sequences:
            raise ValueError(f"Sequence {seq_id} does not exist")
        return self.sequences[seq_id].block_table

    def get_num_tokens(self, seq_id: int) -> int:
        """Get the number of tokens in a sequence."""
        if seq_id not in self.sequences:
            raise ValueError(f"Sequence {seq_id} does not exist")
        return self.sequences[seq_id].num_tokens

    def _tokens_to_blocks(self, num_tokens: int) -> int:
        """Calculate number of blocks needed for given tokens."""
        if num_tokens <= 0:
            return 0
        return (num_tokens + self.block_size - 1) // self.block_size

    def get_all_sequence_ids(self) -> list[int]:
        """Get all active sequence IDs."""
        return list(self.sequences.keys())

    def reset(self) -> None:
        """Reset manager to initial state."""
        self.sequences.clear()
        for allocator in self.allocators:
            allocator.reset()

    def __repr__(self) -> str:
        return (
            f"BlockManager(block_size={self.block_size}, "
            f"num_layers={self.num_layers}, "
            f"free_blocks={self.num_free_blocks}, "
            f"sequences={len(self.sequences)})"
        )

num_free_blocks property

num_free_blocks

Minimum free blocks across all layers.

can_allocate_sequence

can_allocate_sequence(num_tokens)

Check if a new sequence with given tokens can be allocated.

参数：

名称	类型	描述	默认
num_tokens	int	Number of tokens in the sequence.	必需

返回：

类型	描述
bool	True if allocation is possible.

源代码位于： src/llm/core/paged_attention/block_manager.py

def can_allocate_sequence(self, num_tokens: int) -> bool:
    """
    Check if a new sequence with given tokens can be allocated.

    Args:
        num_tokens: Number of tokens in the sequence.

    Returns:
        True if allocation is possible.
    """
    num_blocks_needed = self._tokens_to_blocks(num_tokens)
    return all(alloc.can_allocate(num_blocks_needed) for alloc in self.allocators)

allocate_sequence

allocate_sequence(seq_id, num_tokens)

Allocate blocks for a new sequence.

参数：

名称	类型	描述	默认
seq_id	int	Unique sequence identifier.	必需
num_tokens	int	Initial number of tokens.	必需

返回：

类型	描述
list[int]	List of physical block IDs for the first layer
list[int]	(all layers get the same logical structure).

引发：

类型	描述
RuntimeError	If allocation fails.
ValueError	If sequence already exists.

源代码位于： src/llm/core/paged_attention/block_manager.py

def allocate_sequence(self, seq_id: int, num_tokens: int) -> list[int]:
    """
    Allocate blocks for a new sequence.

    Args:
        seq_id: Unique sequence identifier.
        num_tokens: Initial number of tokens.

    Returns:
        List of physical block IDs for the first layer
        (all layers get the same logical structure).

    Raises:
        RuntimeError: If allocation fails.
        ValueError: If sequence already exists.
    """
    if seq_id in self.sequences:
        raise ValueError(f"Sequence {seq_id} already exists")

    num_blocks = self._tokens_to_blocks(num_tokens)

    # Allocate blocks for each layer
    block_tables_per_layer = []
    try:
        for allocator in self.allocators:
            blocks = allocator.allocate_n(num_blocks)
            block_tables_per_layer.append(blocks)
    except RuntimeError:
        # Rollback on failure
        for i, blocks in enumerate(block_tables_per_layer):
            self.allocators[i].free_all(blocks)
        raise

    # Store first layer's block table (representative)
    first_layer_blocks = block_tables_per_layer[0]
    self.sequences[seq_id] = SequenceBlockInfo(
        seq_id=seq_id,
        block_table=first_layer_blocks.copy(),
        num_tokens=num_tokens,
    )

    return first_layer_blocks

extend_sequence

extend_sequence(seq_id, num_new_tokens)

Extend a sequence with additional tokens.

Allocates new blocks if needed.

参数：

名称	类型	描述	默认
seq_id	int	Sequence to extend.	必需
num_new_tokens	int	Number of new tokens to add.	必需

返回：

类型	描述
list[int]	Updated block table for the sequence.

源代码位于： src/llm/core/paged_attention/block_manager.py

def extend_sequence(self, seq_id: int, num_new_tokens: int) -> list[int]:
    """
    Extend a sequence with additional tokens.

    Allocates new blocks if needed.

    Args:
        seq_id: Sequence to extend.
        num_new_tokens: Number of new tokens to add.

    Returns:
        Updated block table for the sequence.
    """
    if seq_id not in self.sequences:
        raise ValueError(f"Sequence {seq_id} does not exist")

    info = self.sequences[seq_id]
    old_num_blocks = len(info.block_table)
    new_total_tokens = info.num_tokens + num_new_tokens
    new_num_blocks = self._tokens_to_blocks(new_total_tokens)

    blocks_to_add = new_num_blocks - old_num_blocks

    if blocks_to_add > 0:
        # Allocate additional blocks from each layer's allocator.
        # All allocators start from ``deque(range(num_blocks))`` and
        # allocate in the same order, so the block IDs are identical
        # across layers. We only extend the tracked block table (layer
        # 0's) since ``free_sequence`` frees the same IDs on every
        # allocator and they are guaranteed to match.
        for i, allocator in enumerate(self.allocators):
            new_blocks = allocator.allocate_n(blocks_to_add)
            if i == 0:
                info.block_table.extend(new_blocks)

    info.num_tokens = new_total_tokens
    return info.block_table

free_sequence

free_sequence(seq_id)

Free all blocks associated with a sequence.

参数：

名称	类型	描述	默认
seq_id	int	Sequence to free.	必需

源代码位于： src/llm/core/paged_attention/block_manager.py

def free_sequence(self, seq_id: int) -> None:
    """
    Free all blocks associated with a sequence.

    Args:
        seq_id: Sequence to free.
    """
    if seq_id not in self.sequences:
        return  # Already freed or never existed

    info = self.sequences.pop(seq_id)

    # Free blocks in all layers
    for allocator in self.allocators:
        allocator.free_all(info.block_table)

fork_sequence

fork_sequence(src_seq_id, dst_seq_id)

Fork a sequence using copy-on-write.

Creates a new sequence that shares blocks with the source until either is modified.

参数：

名称	类型	描述	默认
src_seq_id	int	Source sequence to fork from.	必需
dst_seq_id	int	New sequence ID.	必需

返回：

类型	描述
list[int]	Block table for the new sequence.

源代码位于： src/llm/core/paged_attention/block_manager.py

def fork_sequence(self, src_seq_id: int, dst_seq_id: int) -> list[int]:
    """
    Fork a sequence using copy-on-write.

    Creates a new sequence that shares blocks with the source
    until either is modified.

    Args:
        src_seq_id: Source sequence to fork from.
        dst_seq_id: New sequence ID.

    Returns:
        Block table for the new sequence.
    """
    if src_seq_id not in self.sequences:
        raise ValueError(f"Source sequence {src_seq_id} does not exist")
    if dst_seq_id in self.sequences:
        raise ValueError(f"Destination sequence {dst_seq_id} already exists")

    src_info = self.sequences[src_seq_id]

    # Fork blocks (increment ref counts)
    for allocator in self.allocators:
        for block_id in src_info.block_table:
            allocator.fork(block_id)

    # Create new sequence with shared blocks
    self.sequences[dst_seq_id] = SequenceBlockInfo(
        seq_id=dst_seq_id,
        block_table=src_info.block_table.copy(),
        num_tokens=src_info.num_tokens,
    )

    return src_info.block_table.copy()

allocate_sequence_shared_prefix

allocate_sequence_shared_prefix(seq_id, prefix_block_ids, num_tokens)

Create a new sequence whose initial blocks are a SHARED prefix.

Used for paged prefix replay: the cached prefix blocks belong to another (still-live) sequence, but a new request with the same prompt may reference them for reads without recomputing their K/V. Each block table entry is forked (refcount +1 on every layer allocator) — nothing is copied; the sequence must copy-on-write before any write lands inside a shared block (:meth:cow_block), or the prefix cache owner's K/V would be corrupted.

参数：

名称	类型	描述	默认
seq_id	int	Unique sequence identifier.	必需
prefix_block_ids	list[int]	Physical block ids (layer 0's view; ids are identical across layers) of the cached prefix.	必需
num_tokens	int	Number of prefix tokens this sequence already owns (the staged prefix length).	必需

返回：

类型	描述
list[int]	The sequence's block table (the shared prefix blocks).

引发：

类型	描述
ValueError	If the sequence already exists, the prefix is empty, or num_tokens spans more blocks than the cached table.

源代码位于： src/llm/core/paged_attention/block_manager.py

def allocate_sequence_shared_prefix(self, seq_id: int, prefix_block_ids: list[int], num_tokens: int) -> list[int]:
    """Create a new sequence whose initial blocks are a SHARED prefix.

    Used for paged prefix replay: the cached prefix blocks belong to
    another (still-live) sequence, but a new request with the same prompt
    may reference them for reads without recomputing their K/V. Each
    block table entry is forked (refcount +1 on every layer allocator)
    — nothing is copied; the sequence must copy-on-write before any write
    lands inside a shared block (:meth:`cow_block`), or the prefix cache
    owner's K/V would be corrupted.

    Args:
        seq_id: Unique sequence identifier.
        prefix_block_ids: Physical block ids (layer 0's view; ids are
            identical across layers) of the cached prefix.
        num_tokens: Number of prefix tokens this sequence already owns
            (the staged prefix length).

    Returns:
        The sequence's block table (the shared prefix blocks).

    Raises:
        ValueError: If the sequence already exists, the prefix is empty,
            or ``num_tokens`` spans more blocks than the cached table.
    """
    if seq_id in self.sequences:
        raise ValueError(f"Sequence {seq_id} already exists")
    if not prefix_block_ids:
        raise ValueError("prefix_block_ids must be non-empty")
    if self._tokens_to_blocks(num_tokens) > len(prefix_block_ids):
        raise ValueError(
            f"{num_tokens} prefix tokens need {self._tokens_to_blocks(num_tokens)} blocks "
            f"but the cached table has {len(prefix_block_ids)}"
        )

    for allocator in self.allocators:
        for block_id in prefix_block_ids:
            allocator.fork(block_id)

    self.sequences[seq_id] = SequenceBlockInfo(
        seq_id=seq_id,
        block_table=list(prefix_block_ids),
        num_tokens=num_tokens,
    )
    return list(prefix_block_ids)

is_block_shared

is_block_shared(block_id)

Whether a block is referenced by more than one sequence.

Blocks are forked/freed uniformly across the per-layer allocators, so the layer-0 refcount is authoritative.

源代码位于： src/llm/core/paged_attention/block_manager.py

def is_block_shared(self, block_id: int) -> bool:
    """Whether a block is referenced by more than one sequence.

    Blocks are forked/freed uniformly across the per-layer allocators, so
    the layer-0 refcount is authoritative.
    """
    return self.allocators[0].is_shared(block_id)

cow_block

cow_block(block_id)

Copy-on-write a logical block across every layer allocator.

Allocates one fresh physical block per layer and decrements the shared block's refcount on each. Because all allocators allocate in lock-step (identical call order), every layer returns the same new id, preserving the cross-layer block-id invariant. The caller must copy the block's data into the new id before writing into it.

返回：

类型	描述
int	The new private block id (or block_id itself when it was not
int	shared, in which case nothing is allocated or decremented).

源代码位于： src/llm/core/paged_attention/block_manager.py

def cow_block(self, block_id: int) -> int:
    """Copy-on-write a logical block across every layer allocator.

    Allocates one fresh physical block per layer and decrements the
    shared block's refcount on each. Because all allocators allocate in
    lock-step (identical call order), every layer returns the same new id,
    preserving the cross-layer block-id invariant. The caller must copy
    the block's data into the new id before writing into it.

    Returns:
        The new private block id (or ``block_id`` itself when it was not
        shared, in which case nothing is allocated or decremented).
    """
    new_id = self.allocators[0].copy_on_write(block_id)
    for allocator in self.allocators[1:]:
        other_id = allocator.copy_on_write(block_id)
        if other_id != new_id:
            raise RuntimeError(
                f"allocators diverged on COW of block {block_id}: "
                f"layer 0 produced {new_id}, another produced {other_id}"
            )
    return new_id

get_block_table

get_block_table(seq_id)

Get the block table for a sequence.

源代码位于： src/llm/core/paged_attention/block_manager.py

def get_block_table(self, seq_id: int) -> list[int]:
    """Get the block table for a sequence."""
    if seq_id not in self.sequences:
        raise ValueError(f"Sequence {seq_id} does not exist")
    return self.sequences[seq_id].block_table

get_num_tokens

get_num_tokens(seq_id)

Get the number of tokens in a sequence.

源代码位于： src/llm/core/paged_attention/block_manager.py

def get_num_tokens(self, seq_id: int) -> int:
    """Get the number of tokens in a sequence."""
    if seq_id not in self.sequences:
        raise ValueError(f"Sequence {seq_id} does not exist")
    return self.sequences[seq_id].num_tokens

get_all_sequence_ids

get_all_sequence_ids()

Get all active sequence IDs.

源代码位于： src/llm/core/paged_attention/block_manager.py

def get_all_sequence_ids(self) -> list[int]:
    """Get all active sequence IDs."""
    return list(self.sequences.keys())

reset

reset()

Reset manager to initial state.

源代码位于： src/llm/core/paged_attention/block_manager.py

def reset(self) -> None:
    """Reset manager to initial state."""
    self.sequences.clear()
    for allocator in self.allocators:
        allocator.reset()

MLP Variants

mlp

MLP

Bases: Module

Multi-Layer Perceptron (MLP).

Transformer-style Feed-Forward Networks with flexible normalization.

Args:
    hidden_size (int): Dimensionality of inputs and outputs.
    intermediate_size (int, optional): Dimensionality of the inner layer. Defaults to 4 * hidden_size.
    activation (str or nn.Module): Activation name or module. Defaults to "gelu".
    dropout_p (float): Dropout probability. Defaults to 0.1.
    norm_first (bool): Whether to apply normalization before or after (pre-LN vs post-LN). Defaults to True.
    norm_type (Type[nn.Module] or nn.Module): Normalization layer type or instance. Defaults to nn.LayerNorm.
    norm_eps (float): Epsilon for normalization layers. Defaults to 1e-5.
    bias (bool): Whether to include bias terms in Linear layers. Defaults to True.
    device (torch.device, optional): Device for parameters. Defaults to None.
    dtype (torch.dtype, optional): Dtype for parameters. Defaults to None.

源代码位于： src/llm/core/mlp.py

@register_mlp("mlp")
class MLP(nn.Module):
    """
        Multi-Layer Perceptron (MLP).
    Transformer-style Feed-Forward Networks with flexible normalization.

        Args:
            hidden_size (int): Dimensionality of inputs and outputs.
            intermediate_size (int, optional): Dimensionality of the inner layer. Defaults to 4 * hidden_size.
            activation (str or nn.Module): Activation name or module. Defaults to "gelu".
            dropout_p (float): Dropout probability. Defaults to 0.1.
            norm_first (bool): Whether to apply normalization before or after (pre-LN vs post-LN). Defaults to True.
            norm_type (Type[nn.Module] or nn.Module): Normalization layer type or instance. Defaults to nn.LayerNorm.
            norm_eps (float): Epsilon for normalization layers. Defaults to 1e-5.
            bias (bool): Whether to include bias terms in Linear layers. Defaults to True.
            device (torch.device, optional): Device for parameters. Defaults to None.
            dtype (torch.dtype, optional): Dtype for parameters. Defaults to None.
    """

    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int | None = None,
        activation: str | nn.Module = "gelu",
        dropout_p: float = 0.1,
        norm_first: bool = True,
        norm_type: type[nn.Module] | nn.Module = nn.LayerNorm,
        norm_eps: float = 1e-5,
        bias: bool = True,
        use_glu: bool = False,  # New parameter for SwiGLU/GLU support
        include_norm_residual: bool = True,  # New parameter
        device: torch.device | None = None,
        dtype: torch.dtype | None = None,
    ):
        super().__init__()
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size or (4 * hidden_size)
        self.norm_first = norm_first
        self.use_glu = use_glu
        self.include_norm_residual = include_norm_residual

        self.norm = None
        if self.include_norm_residual:
            # Create normalization layer based on provided type or instance
            if isinstance(norm_type, type):
                # Specific handling for known norm types with eps
                if norm_type == nn.LayerNorm or norm_type == nn.RMSNorm:
                    self.norm = norm_type(hidden_size, eps=norm_eps, device=device, dtype=dtype)
                else:
                    # For other norm types that might have different init parameters
                    # This might need adjustment if they don't follow `norm(hidden_size, **kwargs)`
                    try:
                        self.norm = norm_type(hidden_size, device=device, dtype=dtype)
                    except TypeError:  # Fallback if eps is not accepted but common for other norms
                        self.norm = norm_type(hidden_size, eps=norm_eps, device=device, dtype=dtype)

            else:
                # If an instance is provided, use it directly and ensure it's on the
                # requested device/dtype (if provided). This prevents mismatch when
                # callers pass a pre-created normalization layer that defaults to CPU.
                self.norm = norm_type
                # Move provided norm instance to the same device/dtype as other params
                if device is not None or dtype is not None:
                    # .to() is in-place for nn.Module, but assign back in case it returns a new object
                    try:
                        moved_norm = self.norm.to(device=device, dtype=dtype)
                        # If .to() returned a new module, keep that reference
                        self.norm = moved_norm
                    except RuntimeError, TypeError, AttributeError:
                        logger.warning(
                            "Failed to move custom norm module to %s/%s — caller must place it",
                            device,
                            dtype,
                        )

        factory_kwargs = make_factory_kwargs(device, dtype)
        self.fc1 = nn.Linear(hidden_size, self.intermediate_size, bias=bias, **factory_kwargs)
        if self.use_glu:
            self.gate_proj = nn.Linear(hidden_size, self.intermediate_size, bias=bias, **factory_kwargs)
        self.fc2 = nn.Linear(self.intermediate_size, hidden_size, bias=bias, **factory_kwargs)

        # Determine activation module and name
        if isinstance(activation, str):
            self.activation_name = activation.lower()
            self.activation = get_activation_layer(self.activation_name)()
        else:
            self.activation = activation
            self.activation_name = activation.__class__.__name__.lower()

        self.dropout = nn.Dropout(dropout_p)

        # Initialize weights
        self._init_weights()

    def _init_weights(self):
        """Initializes the weights of the MLP based on the activation."""
        act = self.activation_name
        neg_slope = getattr(self.activation, "negative_slope", 0.0)

        # Dynamic weight initialization
        if act in ("relu", "leaky_relu"):
            # He initialization for ReLU variants
            nn.init.kaiming_uniform_(self.fc1.weight, a=neg_slope, nonlinearity=act)
            if self.use_glu:
                nn.init.kaiming_uniform_(self.gate_proj.weight, a=neg_slope, nonlinearity=act)
            nn.init.kaiming_uniform_(self.fc2.weight, a=neg_slope, nonlinearity=act)
        elif act in ("gelu", "silu", "swish"):
            # Truncated normal for smoother activations
            std1 = 1.0 / math.sqrt(self.hidden_size)
            std2 = 1.0 / math.sqrt(self.intermediate_size)
            try:
                nn.init.trunc_normal_(self.fc1.weight, std=std1)
                if self.use_glu:
                    nn.init.trunc_normal_(self.gate_proj.weight, std=std1)
                nn.init.trunc_normal_(self.fc2.weight, std=std2)
            except AttributeError:
                nn.init.normal_(self.fc1.weight, mean=0.0, std=std1)
                if self.use_glu:
                    nn.init.normal_(self.gate_proj.weight, mean=0.0, std=std1)
                nn.init.normal_(self.fc2.weight, mean=0.0, std=std2)
        else:
            # Default Xavier/Glorot
            nn.init.xavier_uniform_(self.fc1.weight)
            nn.init.xavier_uniform_(self.fc2.weight)

        # Zero out biases for stable training
        if self.fc1.bias is not None:
            nn.init.zeros_(self.fc1.bias)
        if self.use_glu and self.gate_proj.bias is not None:
            nn.init.zeros_(self.gate_proj.bias)
        if self.fc2.bias is not None:
            nn.init.zeros_(self.fc2.bias)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """
        Forward pass with optional pre-LayerNorm and residual connection.

        Args:
            hidden_states (torch.Tensor): Input tensor of shape [..., hidden_size].

        Returns:
            torch.Tensor: Output tensor with same shape as input.
        """
        if self.include_norm_residual:
            residual = hidden_states

            # Apply normalization first if using pre-norm
            # Ensure self.norm exists before calling it
            x = self.norm(hidden_states) if self.norm_first and self.norm else hidden_states

            # MLP computation (common for both pre-norm and post-norm)
            if self.use_glu:
                # GLU logic
                x_fc1 = self.fc1(x)
                x_fc1 = self.activation(x_fc1)
                x_gate = self.gate_proj(x)
                x_mlp = x_fc1 * x_gate
            else:
                x_mlp = self.fc1(x)
                x_mlp = self.activation(x_mlp)
            x_mlp = self.dropout(x_mlp)
            x_mlp = self.fc2(x_mlp)

            # Add residual connection
            x_mlp = residual + x_mlp

            # Apply normalization after if using post-norm
            # Ensure self.norm exists before calling it
            output = x_mlp if self.norm_first or not self.norm else self.norm(x_mlp)
            return output
        else:
            # No internal norm or residual connection
            x = hidden_states  # Direct input to MLP
            if self.use_glu:
                # GLU logic: (x * activation(gate(x)))
                # For SwiGLU, activation is SiLU
                x_fc1 = self.fc1(x)
                x_fc1 = self.activation(x_fc1)
                x_gate = self.gate_proj(x)
                x = x_fc1 * x_gate
            else:
                x = self.fc1(x)
                x = self.activation(x)
            x = self.dropout(x)
            x = self.fc2(x)
            return x

forward

forward(hidden_states)

Forward pass with optional pre-LayerNorm and residual connection.

参数：

名称	类型	描述	默认
hidden_states	Tensor	Input tensor of shape [..., hidden_size].	必需

返回：

类型	描述
Tensor	torch.Tensor: Output tensor with same shape as input.

源代码位于： src/llm/core/mlp.py

def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    """
    Forward pass with optional pre-LayerNorm and residual connection.

    Args:
        hidden_states (torch.Tensor): Input tensor of shape [..., hidden_size].

    Returns:
        torch.Tensor: Output tensor with same shape as input.
    """
    if self.include_norm_residual:
        residual = hidden_states

        # Apply normalization first if using pre-norm
        # Ensure self.norm exists before calling it
        x = self.norm(hidden_states) if self.norm_first and self.norm else hidden_states

        # MLP computation (common for both pre-norm and post-norm)
        if self.use_glu:
            # GLU logic
            x_fc1 = self.fc1(x)
            x_fc1 = self.activation(x_fc1)
            x_gate = self.gate_proj(x)
            x_mlp = x_fc1 * x_gate
        else:
            x_mlp = self.fc1(x)
            x_mlp = self.activation(x_mlp)
        x_mlp = self.dropout(x_mlp)
        x_mlp = self.fc2(x_mlp)

        # Add residual connection
        x_mlp = residual + x_mlp

        # Apply normalization after if using post-norm
        # Ensure self.norm exists before calling it
        output = x_mlp if self.norm_first or not self.norm else self.norm(x_mlp)
        return output
    else:
        # No internal norm or residual connection
        x = hidden_states  # Direct input to MLP
        if self.use_glu:
            # GLU logic: (x * activation(gate(x)))
            # For SwiGLU, activation is SiLU
            x_fc1 = self.fc1(x)
            x_fc1 = self.activation(x_fc1)
            x_gate = self.gate_proj(x)
            x = x_fc1 * x_gate
        else:
            x = self.fc1(x)
            x = self.activation(x)
        x = self.dropout(x)
        x = self.fc2(x)
        return x

moe

MoE

Bases: Module

Mixture of Experts (MoE) Layer.

参数：

名称	类型	描述	默认
hidden_size	int	The dimensionality of the input and output.	必需
num_experts	int	The total number of experts.	必需
top_k	int	The number of top experts to select for each token.	必需
intermediate_size	int	The intermediate size for each expert's MLP. Defaults to 4 * hidden_size.	None
activation	str or Module	Activation function for experts. Defaults to "gelu".	'gelu'
dropout_p	float	Dropout probability for experts. Defaults to 0.1.	0.1
norm_type	Type[Module] or Module	Normalization layer type for experts. Defaults to nn.LayerNorm.	LayerNorm
norm_eps	float	Epsilon for normalization layers in experts. Defaults to 1e-5.	1e-05
bias	bool	Whether to include bias terms in Linear layers of experts. Defaults to True.	True
device	device	Device for parameters. Defaults to None.	None
dtype	dtype	Dtype for parameters. Defaults to None.	None

源代码位于： src/llm/core/moe/moe.py

@register_mlp("moe")
class MoE(nn.Module):
    """
    Mixture of Experts (MoE) Layer.

    Args:
        hidden_size (int): The dimensionality of the input and output.
        num_experts (int): The total number of experts.
        top_k (int): The number of top experts to select for each token.
        intermediate_size (int, optional): The intermediate size for each expert's MLP.
                                           Defaults to 4 * hidden_size.
        activation (str or nn.Module): Activation function for experts. Defaults to "gelu".
        dropout_p (float): Dropout probability for experts. Defaults to 0.1.
        norm_type (Type[nn.Module] or nn.Module): Normalization layer type for experts. Defaults to nn.LayerNorm.
        norm_eps (float): Epsilon for normalization layers in experts. Defaults to 1e-5.
        bias (bool): Whether to include bias terms in Linear layers of experts. Defaults to True.
        device (torch.device, optional): Device for parameters. Defaults to None.
        dtype (torch.dtype, optional): Dtype for parameters. Defaults to None.
    """

    def __init__(
        self,
        hidden_size: int,
        num_experts: int,
        top_k: int,
        intermediate_size: int | None = None,
        activation: str | nn.Module = "gelu",
        dropout_p: float = 0.1,
        norm_type: type[nn.Module] | nn.Module = nn.LayerNorm,
        norm_eps: float = 1e-5,
        bias: bool = True,
        device: torch.device | None = None,
        dtype: torch.dtype | None = None,
    ):
        super().__init__()
        if top_k > num_experts:
            raise ValueError(f"top_k ({top_k}) cannot be greater than num_experts ({num_experts})")
        if top_k == 0:
            raise ValueError("top_k must be at least 1")

        self.hidden_size = hidden_size
        self.num_experts = num_experts
        self.top_k = top_k

        factory_kwargs = make_factory_kwargs(device, dtype)

        # Gating network (router)
        self.gate = nn.Linear(hidden_size, num_experts, bias=False, **factory_kwargs)

        # Experts (always using MLP)
        self.experts = nn.ModuleList(
            [
                MLP(
                    hidden_size=hidden_size,
                    intermediate_size=intermediate_size,
                    activation=activation,
                    dropout_p=dropout_p,
                    norm_first=False,  # MoE typically handles norm/residual externally
                    norm_type=norm_type,
                    norm_eps=norm_eps,
                    bias=bias,
                    include_norm_residual=False,  # Experts are simple MLPs, norm/residual handled by TransformerBlock
                    **factory_kwargs,
                )
                for _ in range(num_experts)
            ]
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass of the MoE layer.

        Args:
            x (torch.Tensor): Input tensor of shape [..., hidden_size].

        Returns:
            torch.Tensor: Output tensor with same shape as input.
        """
        original_shape = x.shape
        x = x.view(-1, self.hidden_size)  # Flatten to [batch_size * seq_len, hidden_size]

        # 1. Gating network to get expert scores
        gate_logits = self.gate(x)  # [batch_size * seq_len, num_experts]

        # 2. Select top-k experts
        # top_k_logits: [batch_size * seq_len, top_k]
        # top_k_indices: [batch_size * seq_len, top_k]
        top_k_logits, top_k_indices = torch.topk(gate_logits, self.top_k, dim=-1)

        # 3. Apply softmax to get weights for selected experts
        # expert_weights: [batch_size * seq_len, top_k]
        expert_weights = functional.softmax(top_k_logits, dim=-1, dtype=x.dtype)

        # Initialize output tensor
        output = torch.zeros_like(x)  # [batch_size * seq_len, hidden_size]

        # Create a list of lists, where each inner list contains the indices
        # of tokens routed to that expert.
        expert_inputs: list[list[torch.Tensor]] = [[] for _ in range(self.num_experts)]
        expert_weights_per_token: list[list[torch.Tensor]] = [[] for _ in range(self.num_experts)]
        expert_original_indices: list[list[int]] = [[] for _ in range(self.num_experts)]

        for i in range(x.size(0)):  # Iterate over each token
            for k_idx in range(self.top_k):
                expert_idx = int(top_k_indices[i, k_idx].item())
                expert_inputs[expert_idx].append(x[i])
                expert_weights_per_token[expert_idx].append(expert_weights[i, k_idx])
                expert_original_indices[expert_idx].append(i)

        # Process each expert
        for i, expert in enumerate(self.experts):
            if expert_inputs[i]:
                expert_input_batch = torch.stack(expert_inputs[i])
                expert_output_batch = expert(expert_input_batch)
                expert_weights_batch = torch.stack(expert_weights_per_token[i]).unsqueeze(-1)
                # Weighted sum and scatter
                weighted_expert_output = expert_output_batch * expert_weights_batch
                output.index_add_(0, torch.tensor(expert_original_indices[i], device=x.device), weighted_expert_output)

        return output.view(original_shape)

forward

forward(x)

Forward pass of the MoE layer.

参数：

名称	类型	描述	默认
x	Tensor	Input tensor of shape [..., hidden_size].	必需

返回：

类型	描述
Tensor	torch.Tensor: Output tensor with same shape as input.

源代码位于： src/llm/core/moe/moe.py

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """
    Forward pass of the MoE layer.

    Args:
        x (torch.Tensor): Input tensor of shape [..., hidden_size].

    Returns:
        torch.Tensor: Output tensor with same shape as input.
    """
    original_shape = x.shape
    x = x.view(-1, self.hidden_size)  # Flatten to [batch_size * seq_len, hidden_size]

    # 1. Gating network to get expert scores
    gate_logits = self.gate(x)  # [batch_size * seq_len, num_experts]

    # 2. Select top-k experts
    # top_k_logits: [batch_size * seq_len, top_k]
    # top_k_indices: [batch_size * seq_len, top_k]
    top_k_logits, top_k_indices = torch.topk(gate_logits, self.top_k, dim=-1)

    # 3. Apply softmax to get weights for selected experts
    # expert_weights: [batch_size * seq_len, top_k]
    expert_weights = functional.softmax(top_k_logits, dim=-1, dtype=x.dtype)

    # Initialize output tensor
    output = torch.zeros_like(x)  # [batch_size * seq_len, hidden_size]

    # Create a list of lists, where each inner list contains the indices
    # of tokens routed to that expert.
    expert_inputs: list[list[torch.Tensor]] = [[] for _ in range(self.num_experts)]
    expert_weights_per_token: list[list[torch.Tensor]] = [[] for _ in range(self.num_experts)]
    expert_original_indices: list[list[int]] = [[] for _ in range(self.num_experts)]

    for i in range(x.size(0)):  # Iterate over each token
        for k_idx in range(self.top_k):
            expert_idx = int(top_k_indices[i, k_idx].item())
            expert_inputs[expert_idx].append(x[i])
            expert_weights_per_token[expert_idx].append(expert_weights[i, k_idx])
            expert_original_indices[expert_idx].append(i)

    # Process each expert
    for i, expert in enumerate(self.experts):
        if expert_inputs[i]:
            expert_input_batch = torch.stack(expert_inputs[i])
            expert_output_batch = expert(expert_input_batch)
            expert_weights_batch = torch.stack(expert_weights_per_token[i]).unsqueeze(-1)
            # Weighted sum and scatter
            weighted_expert_output = expert_output_batch * expert_weights_batch
            output.index_add_(0, torch.tensor(expert_original_indices[i], device=x.device), weighted_expert_output)

    return output.view(original_shape)

Transformer Block

transformer_block

TransformerBlock

Bases: Module

A single Transformer block, comprising a Multi-Head Attention (MHA) layer and a Multi-Layer Perceptron (MLP) layer, with normalization and residual connections.

The block can be configured for Pre-LN (Layer Normalization before sublayer) or Post-LN (Layer Normalization after sublayer and residual connection).

源代码位于： src/llm/core/transformer_block.py

class TransformerBlock(nn.Module):
    """
    A single Transformer block, comprising a Multi-Head Attention (MHA) layer
    and a Multi-Layer Perceptron (MLP) layer, with normalization and residual connections.

    The block can be configured for Pre-LN (Layer Normalization before sublayer)
    or Post-LN (Layer Normalization after sublayer and residual connection).
    """

    def __init__(
        self,
        hidden_size: int,
        num_heads: int,
        intermediate_size: int | None = None,
        attn_dropout_p: float = 0.1,
        mlp_dropout_p: float = 0.1,
        mlp_activation: str | nn.Module = "gelu",
        norm_eps: float = 1e-5,
        norm_first: bool = True,
        is_causal: bool = False,  # Default causality for MHA within this block
        qkv_bias: bool = True,  # Bias for QKV projections in MHA
        mlp_bias: bool = True,  # Bias for Linear layers in MLP
        num_experts: int = 0,
        top_k: int = 0,
        num_kv_heads: int | None = None,  # For GQA support
        use_glu: bool = False,  # New: For SwiGLU support
        norm_type: type[nn.Module] | nn.Module | Callable[..., nn.Module] = nn.LayerNorm,
        window_size: int | None = None,  # Sliding window attention
        max_seq_len: int | None = None,  # RoPE context (required if use_rope)
        use_rope: bool = False,  # Rotary position embedding (real Llama/Mistral)
        rope_theta: float = 10000.0,  # RoPE base frequency
        alibi: ALiBiPositionBias | None = None,  # Linear-bias PE (mha backend)
        device: torch.device | str | None = None,
        dtype: torch.dtype | None = None,
        # Registry keys
        attn_impl: str = "mha",
        mlp_impl: str = "mlp",
    ):
        """
        Initializes the TransformerBlock.
        """
        super().__init__()
        factory_kwargs = make_factory_kwargs(device, dtype)

        from llm.core.registry import ATTENTION_REGISTRY, MLP_REGISTRY

        self.norm_first = norm_first
        self.hidden_size = hidden_size

        # Initialize Norms. ``norm_type`` must be a factory callable
        # ``(**kwargs) -> nn.Module`` — typically one of the entries in
        # ``NORM_REGISTRY``. Already-instantiated ``nn.Module`` instances are
        # rejected because the previous isinstance(type) branch that
        # deep-copied them was a code smell (Finding C).
        if isinstance(norm_type, nn.Module):
            raise TypeError(
                "norm_type must be a factory callable (e.g. from NORM_REGISTRY), "
                "not an already-constructed nn.Module. Pass norm_impl='rms_norm' "
                "or 'layer_norm' to DecoderModel instead of pre-constructed norm modules."
            )
        if not callable(norm_type):
            raise TypeError(f"norm_type must be a callable factory, got {type(norm_type).__name__}.")
        self.norm1 = norm_type(hidden_size, eps=norm_eps, **factory_kwargs)
        self.norm2 = norm_type(hidden_size, eps=norm_eps, **factory_kwargs)

        # Initialize Attention via Registry
        attn_cls = ATTENTION_REGISTRY.get(attn_impl)
        attention_kwargs: dict = dict(
            hidden_size=hidden_size,
            num_heads=num_heads,
            p=attn_dropout_p,
            bias=qkv_bias,
            is_causal=is_causal,
            include_norm_residual=False,
            eps=norm_eps,
            norm_first=False,
            num_kv_heads=num_kv_heads,
            window_size=window_size,
            **factory_kwargs,
        )
        if use_rope:
            # Only the MHA backend accepts RoPE wiring today; other backends
            # (flash_attn, MLA) do not declare these kwargs. Thread them only
            # when requested so the default path's call shape is unchanged.
            attention_kwargs.update(max_seq_len=max_seq_len, use_rope=True, rope_theta=rope_theta)
        if alibi is not None:
            # ALiBi is wired into the mha backend only (RIL — ALiBi milestone);
            # DecoderModel rejects alibi with any other attn_impl before this
            # point, so the kwargs legitimately only reach MHA.
            attention_kwargs.update(alibi=alibi)
        self.self_attn = attn_cls(**attention_kwargs)

        # Initialize MLP via Registry
        if intermediate_size is None:
            intermediate_size = 4 * hidden_size

        mlp_cls = MLP_REGISTRY.get(mlp_impl)

        # Prepare kwargs for MLP/MoE
        # Note: Different implementations might need different kwargs.
        # Ideally we pass a config object, but here we pass common args.
        # MoE needs num_experts and top_k, MLP doesn't.
        # We pass them as **kwargs, assuming constructors handle extra args or we filter.
        # But our classes strictly define __init__.
        # So we construct specific kwargs map.

        common_mlp_kwargs = {
            "hidden_size": hidden_size,
            "intermediate_size": intermediate_size,
            "activation": mlp_activation,
            "dropout_p": mlp_dropout_p,
            "bias": mlp_bias,
            "norm_eps": norm_eps,
            **factory_kwargs,
        }

        if mlp_impl == "moe":
            # Add MoE specific args
            if num_experts <= 0 or top_k <= 0:
                raise ValueError("num_experts and top_k must be positive for MoE.")
            common_mlp_kwargs["num_experts"] = num_experts
            common_mlp_kwargs["top_k"] = top_k
        else:
            # Add MLP specific args (standard MLP doesn't need num_experts)
            common_mlp_kwargs["include_norm_residual"] = False
            common_mlp_kwargs["use_glu"] = use_glu
            common_mlp_kwargs["norm_first"] = False

        self.mlp = mlp_cls(**common_mlp_kwargs)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attn_mask: torch.Tensor | None = None,
        is_causal: bool | None = None,
        kv_cache: KVCache | None = None,
        use_cache: bool = False,
        batch_indices: torch.Tensor | None = None,
        start_pos: int | torch.Tensor | None = None,
        paged_kv_cache: object | None = None,
        layer_idx: int | None = None,
    ) -> torch.Tensor | tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]:
        """
        Forward pass of the Transformer block.

        Args:
            hidden_states (torch.Tensor): Input tensor of shape [B, S, H].
            attn_mask (torch.Tensor, optional): Attention mask for MHA.
            is_causal (bool, optional): Overrides the default MHA causality for this pass.
                                        If None, MHA's default `is_causal` is used.
            kv_cache (KVCache | None): Pre-allocated KV cache for efficient autoregressive generation.
            use_cache (bool): Whether to return the updated (key, value) pair.
            batch_indices(torch.Tensor | None): Cache update indices.
            start_pos (int | torch.Tensor | None): Cache update position.
            paged_kv_cache (object | None): Block-allocator KV cache; ignored
                when ``None``. When set, ``kv_cache`` is unused and ``layer_idx``
                must point at this block's index in the decoder.
            layer_idx (int | None): Index of this block in the decoder; required
                when ``paged_kv_cache`` is set.

        Returns:
            torch.Tensor or tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
                - If use_cache=False: Output tensor of shape [B, S, H].
                - If use_cache=True: (Output tensor, (current_key, current_value))
        """
        # Determine causality for the MHA call
        # If is_causal is provided as an argument, it overrides the MHA's default.
        # Otherwise, MHA uses its own self.is_causal.
        # The MHA forward method handles this logic if is_causal=None is passed.

        residual = hidden_states

        # 1. Multi-Head Attention Sublayer
        if self.norm_first:
            hidden_states = self.norm1(hidden_states)

        attn_outputs = self.self_attn(
            hidden_states,
            attn_mask=attn_mask,
            is_causal=is_causal,
            kv_cache=kv_cache,
            use_cache=use_cache,
            batch_indices=batch_indices,
            start_pos=start_pos,
            paged_kv_cache=paged_kv_cache,
            layer_idx=layer_idx,
        )

        if paged_kv_cache is not None:
            # Paged path returns the output directly (no separate kv
            # tuple to surface — the cache is mutated in place).
            attn_output = attn_outputs
            current_kv = None
        elif use_cache:
            attn_output, current_kv = attn_outputs
        else:
            attn_output = attn_outputs

        # Apply residual connection
        # Pre-LN MHA: output = residual + Attention(Norm(x))
        # Post-LN MHA: output = Norm(residual + Attention(x))
        if self.norm_first:
            hidden_states = residual + attn_output
            residual = hidden_states  # Update residual for next block
        else:
            hidden_states = self.norm1(residual + attn_output)
            residual = hidden_states  # Update residual for next block

        # 2. MLP Sublayer
        if self.norm_first:
            hidden_states = self.norm2(hidden_states)

        mlp_output = self.mlp(hidden_states)

        # Apply residual connection
        output = residual + mlp_output if self.norm_first else self.norm2(residual + mlp_output)

        if paged_kv_cache is not None:
            # The paged cache is mutated in place; there is no per-block
            # KV tuple to surface to the caller.
            return output
        if use_cache:
            # ``current_kv`` may be None for backends that manage their own
            # cache (MLA); the decoder ignores the second tuple element.
            return output, current_kv
        return output

forward

forward(hidden_states, attn_mask=None, is_causal=None, kv_cache=None, use_cache=False, batch_indices=None, start_pos=None, paged_kv_cache=None, layer_idx=None)

Forward pass of the Transformer block.

参数：

名称	类型	描述	默认
hidden_states	Tensor	Input tensor of shape [B, S, H].	必需
attn_mask	Tensor	Attention mask for MHA.	None
is_causal	bool	Overrides the default MHA causality for this pass. If None, MHA's default is_causal is used.	None
kv_cache	KVCache | None	Pre-allocated KV cache for efficient autoregressive generation.	None
use_cache	bool	Whether to return the updated (key, value) pair.	False
batch_indices	Tensor | None	Cache update indices.	None
start_pos	int | Tensor | None	Cache update position.	None
paged_kv_cache	object | None	Block-allocator KV cache; ignored when None. When set, kv_cache is unused and layer_idx must point at this block's index in the decoder.	None
layer_idx	int | None	Index of this block in the decoder; required when paged_kv_cache is set.	None

返回：

类型	描述
Tensor | tuple[Tensor, tuple[Tensor, Tensor] | None]	torch.Tensor or tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: - If use_cache=False: Output tensor of shape [B, S, H]. - If use_cache=True: (Output tensor, (current_key, current_value))

源代码位于： src/llm/core/transformer_block.py

def forward(
    self,
    hidden_states: torch.Tensor,
    attn_mask: torch.Tensor | None = None,
    is_causal: bool | None = None,
    kv_cache: KVCache | None = None,
    use_cache: bool = False,
    batch_indices: torch.Tensor | None = None,
    start_pos: int | torch.Tensor | None = None,
    paged_kv_cache: object | None = None,
    layer_idx: int | None = None,
) -> torch.Tensor | tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]:
    """
    Forward pass of the Transformer block.

    Args:
        hidden_states (torch.Tensor): Input tensor of shape [B, S, H].
        attn_mask (torch.Tensor, optional): Attention mask for MHA.
        is_causal (bool, optional): Overrides the default MHA causality for this pass.
                                    If None, MHA's default `is_causal` is used.
        kv_cache (KVCache | None): Pre-allocated KV cache for efficient autoregressive generation.
        use_cache (bool): Whether to return the updated (key, value) pair.
        batch_indices(torch.Tensor | None): Cache update indices.
        start_pos (int | torch.Tensor | None): Cache update position.
        paged_kv_cache (object | None): Block-allocator KV cache; ignored
            when ``None``. When set, ``kv_cache`` is unused and ``layer_idx``
            must point at this block's index in the decoder.
        layer_idx (int | None): Index of this block in the decoder; required
            when ``paged_kv_cache`` is set.

    Returns:
        torch.Tensor or tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
            - If use_cache=False: Output tensor of shape [B, S, H].
            - If use_cache=True: (Output tensor, (current_key, current_value))
    """
    # Determine causality for the MHA call
    # If is_causal is provided as an argument, it overrides the MHA's default.
    # Otherwise, MHA uses its own self.is_causal.
    # The MHA forward method handles this logic if is_causal=None is passed.

    residual = hidden_states

    # 1. Multi-Head Attention Sublayer
    if self.norm_first:
        hidden_states = self.norm1(hidden_states)

    attn_outputs = self.self_attn(
        hidden_states,
        attn_mask=attn_mask,
        is_causal=is_causal,
        kv_cache=kv_cache,
        use_cache=use_cache,
        batch_indices=batch_indices,
        start_pos=start_pos,
        paged_kv_cache=paged_kv_cache,
        layer_idx=layer_idx,
    )

    if paged_kv_cache is not None:
        # Paged path returns the output directly (no separate kv
        # tuple to surface — the cache is mutated in place).
        attn_output = attn_outputs
        current_kv = None
    elif use_cache:
        attn_output, current_kv = attn_outputs
    else:
        attn_output = attn_outputs

    # Apply residual connection
    # Pre-LN MHA: output = residual + Attention(Norm(x))
    # Post-LN MHA: output = Norm(residual + Attention(x))
    if self.norm_first:
        hidden_states = residual + attn_output
        residual = hidden_states  # Update residual for next block
    else:
        hidden_states = self.norm1(residual + attn_output)
        residual = hidden_states  # Update residual for next block

    # 2. MLP Sublayer
    if self.norm_first:
        hidden_states = self.norm2(hidden_states)

    mlp_output = self.mlp(hidden_states)

    # Apply residual connection
    output = residual + mlp_output if self.norm_first else self.norm2(residual + mlp_output)

    if paged_kv_cache is not None:
        # The paged cache is mutated in place; there is no per-block
        # KV tuple to surface to the caller.
        return output
    if use_cache:
        # ``current_kv`` may be None for backends that manage their own
        # cache (MLA); the decoder ignores the second tuple element.
        return output, current_kv
    return output

Normalization Layers

rms_norm

RMSNorm

Bases: Module

Root Mean Square Normalization (RMSNorm) 实现.

RMSNorm 是 LayerNorm 的一种简化形式. 它仅使用激活值的均方根统计量进行归一化, 而不进行中心化 (减去均值). 通常只包含一个可学习的缩放参数 (gamma/weight), 不包含可学习的偏置参数 (beta).

参考文献: Zhang, Biao, and Rico Sennrich. "Root mean square layer normalization." Advances in Neural Information Processing Systems 32 (2019). 论文链接: https://arxiv.org/abs/1910.07467

数学公式 (对于特征向量 x): RMS(x) = sqrt( (1/H) * Σ(x_i²) + ε ) (H = 归一化维度的大小) x_normalized = x / RMS(x) output = gamma * x_normalized

参数

normalized_shape (int 或 list/tuple of ints): 需要进行归一化的输入张量的结尾维度形状. 与 LayerNorm 中的定义相同. eps (float): 加在均方根计算中的小常数, 防止除零错误并提高数值稳定性. 默认为 1e-6 (常见于 RMSNorm 实现). elementwise_affine (bool): 如果为 True, 则此模块包含可学习的缩放参数 gamma (gamma/weight), 形状与 normalized_shape 相同. gamma 初始化为 1. 注意: RMSNorm 通常不使用偏置项 (beta). 默认为 True.

源代码位于： src/llm/core/rms_norm.py

class RMSNorm(nn.Module):
    """
    Root Mean Square Normalization (RMSNorm) 实现.

    RMSNorm 是 LayerNorm 的一种简化形式. 它仅使用激活值的均方根统计量进行归一化,
    而不进行中心化 (减去均值). 通常只包含一个可学习的缩放参数 (gamma/weight),
    不包含可学习的偏置参数 (beta).

    参考文献: Zhang, Biao, and Rico Sennrich. "Root mean square layer normalization." Advances in Neural Information Processing Systems 32 (2019).
    论文链接: https://arxiv.org/abs/1910.07467

    数学公式 (对于特征向量 x):
        RMS(x) = sqrt( (1/H) * Σ(x_i²) + ε )  (H = 归一化维度的大小)
        x_normalized = x / RMS(x)
        output = gamma * x_normalized

    参数:
        normalized_shape (int 或 list/tuple of ints):
            需要进行归一化的输入张量的结尾维度形状.
            与 LayerNorm 中的定义相同.
        eps (float):
            加在均方根计算中的小常数, 防止除零错误并提高数值稳定性.
            默认为 1e-6 (常见于 RMSNorm 实现).
        elementwise_affine (bool):
            如果为 True, 则此模块包含可学习的缩放参数 gamma (gamma/weight),
            形状与 `normalized_shape` 相同. gamma 初始化为 1.
            注意: RMSNorm 通常不使用偏置项 (beta).
            默认为 True.
    """

    def __init__(
        self,
        normalized_shape: int | list[int] | tuple[int, ...],
        eps: float = 1e-6,  # 注意: RMSNorm 文献中 eps 常为 1e-6
        elementwise_affine: bool = True,
        device: torch.device | str | None = None,
        dtype: torch.dtype | None = None,
    ):
        super().__init__()

        # 存储配置
        if isinstance(normalized_shape, int):
            self.normalized_shape: tuple[int, ...] = (normalized_shape,)
        else:
            self.normalized_shape = tuple(normalized_shape)

        self.eps = eps
        self.elementwise_affine = elementwise_affine

        factory_kwargs = make_factory_kwargs(device, dtype)
        if self.elementwise_affine:
            # 初始化可学习的缩放参数 gamma
            self.weight = nn.Parameter(torch.ones(self.normalized_shape, **factory_kwargs))  # gamma
        else:
            # 如果无仿射变换, 将 weight 注册为 None
            self.register_parameter("weight", None)

        # RMSNorm 通常没有偏置项 (beta)
        # self.register_parameter("bias", None) # 显式表明无偏置

    def _compute_rms(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """计算 RMS 的辅助函数"""
        # 确定需要归一化的维度
        num_normalized_dims = len(self.normalized_shape)
        dims_to_normalize = tuple(range(hidden_states.ndim - num_normalized_dims, hidden_states.ndim))

        # 计算均方值 (不减去均值的方差)
        # mean(x^2)
        mean_square = torch.mean(hidden_states.pow(2), dim=dims_to_normalize, keepdim=True)

        # 计算 RMS: sqrt(mean(x^2) + eps)
        rms = torch.sqrt(mean_square + self.eps)
        return rms

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """
        前向传播函数

        参数:
            hidden_states: 输入张量, 其尾部维度应与 `normalized_shape` 匹配.
               形状例如: [batch_size, ..., *normalized_shape]

        返回:
            归一化后的张量, 形状与输入 hidden_states 相同.
        """
        # 1. 计算 RMS
        rms = self._compute_rms(hidden_states)

        # 2. 归一化: x / RMS(x)
        x_normalized = hidden_states / rms

        # 3. 应用缩放因子 (gamma/weight) (如果启用)
        if self.elementwise_affine:
            return self.weight * x_normalized
        else:
            return x_normalized

    def extra_repr(self) -> str:
        # 自定义打印模块信息时的显示内容
        return f"{self.normalized_shape}, eps={self.eps}, elementwise_affine={self.elementwise_affine}"

forward

forward(hidden_states)

前向传播函数

参数

hidden_states: 输入张量, 其尾部维度应与 normalized_shape 匹配. 形状例如: [batch_size, ..., *normalized_shape]

返回

归一化后的张量, 形状与输入 hidden_states 相同.

源代码位于： src/llm/core/rms_norm.py

def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    """
    前向传播函数

    参数:
        hidden_states: 输入张量, 其尾部维度应与 `normalized_shape` 匹配.
           形状例如: [batch_size, ..., *normalized_shape]

    返回:
        归一化后的张量, 形状与输入 hidden_states 相同.
    """
    # 1. 计算 RMS
    rms = self._compute_rms(hidden_states)

    # 2. 归一化: x / RMS(x)
    x_normalized = hidden_states / rms

    # 3. 应用缩放因子 (gamma/weight) (如果启用)
    if self.elementwise_affine:
        return self.weight * x_normalized
    else:
        return x_normalized

rms_norm_numpy

rms_norm_numpy(x, gamma=None, eps=1e-06)

RMS Normalization 的 NumPy 实现 (简化版)

注意: 此版本为了简洁, 固定在最后一个轴 (axis=-1) 上进行归一化. 主要用于帮助理解 RMSNorm 的核心计算步骤. 不包含偏置项.

参数

x: 输入 NumPy 数组, 形状例如 [batch_size, ..., feature_dim] gamma: 缩放参数 (如果提供), 形状应为 [feature_dim] eps: 防止除零错误的小常数

返回

归一化后的 NumPy 数组, 形状与输入 x 相同.

源代码位于： src/llm/core/rms_norm.py

def rms_norm_numpy(x: np.ndarray, gamma: np.ndarray | None = None, eps: float = 1e-6) -> np.ndarray:
    """
    RMS Normalization 的 NumPy 实现 (简化版)

    注意: 此版本为了简洁, *固定*在最后一个轴 (axis=-1) 上进行归一化.
    主要用于帮助理解 RMSNorm 的核心计算步骤. 不包含偏置项.

    参数:
        x: 输入 NumPy 数组, 形状例如 [batch_size, ..., feature_dim]
        gamma: 缩放参数 (如果提供), 形状应为 [feature_dim]
        eps: 防止除零错误的小常数

    返回:
        归一化后的 NumPy 数组, 形状与输入 x 相同.
    """
    # 1. 沿最后一个轴计算均方值
    mean_square = np.mean(np.square(x), axis=-1, keepdims=True)

    # 2. 计算 RMS
    rms = np.sqrt(mean_square + eps)

    # 3. 归一化
    x_normalized = x / rms

    # 4. 应用缩放因子 gamma
    if gamma is not None:
        return gamma * x_normalized
    else:
        return x_normalized

layer_norm

LayerNorm

Bases: Module

自定义 Layer Normalization 实现 (更接近 PyTorch 内置版本)

Layer Normalization 通过对单个样本内的特征维度进行归一化来稳定训练过程. 与 Batch Normalization 不同, LayerNorm 对每个样本独立操作, 不依赖于 Batch Size, 因此特别适用于序列模型 (RNN, Transformer).

数学公式 (对于一个样本 x 中的一个元素 x_i): μ = (1/H) * Σ(x_i) (在归一化维度 H 上求均值) σ² = (1/H) * Σ((x_i - μ)²) (在归一化维度 H 上求方差) x_normalized = (x - μ) / sqrt(σ² + ε) output = gamma * x_normalized + beta

参数

normalized_shape (int 或 list/tuple of ints): 需要进行归一化的输入张量的结尾维度形状. 例如, 如果输入形状是 (N, C, H, W) 且希望对最后两个维度 (H, W) 进行归一化, 则 normalized_shape 应为 (H, W) 或 [H, W]. 如果只对最后一个维度进行归一化, 可以传入一个整数, 如 W. eps (float): 加在分母中的小常数, 防止除零错误并提高数值稳定性. 默认为 1e-5. elementwise_affine (bool): 如果为 True, 则此模块包含可学习的仿射参数 gamma (weight) 和 beta (bias), 形状与 normalized_shape 相同. gamma 初始化为 1, beta 初始化为 0. 默认为 True.

源代码位于： src/llm/core/layer_norm.py

class LayerNorm(nn.Module):
    """
    自定义 Layer Normalization 实现 (更接近 PyTorch 内置版本)

    Layer Normalization 通过对单个样本内的特征维度进行归一化来稳定训练过程.
    与 Batch Normalization 不同, LayerNorm 对每个样本独立操作, 不依赖于 Batch Size,
    因此特别适用于序列模型 (RNN, Transformer).

    数学公式 (对于一个样本 x 中的一个元素 x_i):
        μ = (1/H) * Σ(x_i)  (在归一化维度 H 上求均值)
        σ² = (1/H) * Σ((x_i - μ)²) (在归一化维度 H 上求方差)
        x_normalized = (x - μ) / sqrt(σ² + ε)
        output = gamma * x_normalized + beta

    参数:
        normalized_shape (int 或 list/tuple of ints):
            需要进行归一化的输入张量的结尾维度形状.
            例如, 如果输入形状是 (N, C, H, W) 且希望对最后两个维度 (H, W) 进行归一化,
            则 normalized_shape 应为 (H, W) 或 `[H, W]`.
            如果只对最后一个维度进行归一化, 可以传入一个整数, 如 `W`.
        eps (float):
            加在分母中的小常数, 防止除零错误并提高数值稳定性. 默认为 1e-5.
        elementwise_affine (bool):
            如果为 True, 则此模块包含可学习的仿射参数 gamma (weight) 和 beta (bias),
            形状与 `normalized_shape` 相同. gamma 初始化为 1, beta 初始化为 0.
            默认为 True.
    """

    def __init__(
        self,
        normalized_shape: int | list[int] | tuple[int, ...],
        eps: float = 1e-5,
        elementwise_affine: bool = True,
    ):
        super().__init__()

        # 将 normalized_shape 统一处理为 tuple
        if isinstance(normalized_shape, int):
            self.normalized_shape: tuple[int, ...] = (normalized_shape,)
        else:
            self.normalized_shape = tuple(normalized_shape)

        self.eps = eps
        self.elementwise_affine = elementwise_affine

        if self.elementwise_affine:
            # 初始化可学习的缩放参数 gamma 和偏移参数 beta
            # 形状与需要归一化的维度一致
            self.weight = nn.Parameter(torch.ones(self.normalized_shape))  # gamma
            self.bias = nn.Parameter(torch.zeros(self.normalized_shape))  # beta
        else:
            # 如果不使用仿射变换, 则注册为 None, 这是标准做法
            self.register_parameter("weight", None)
            self.register_parameter("bias", None)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """
        前向传播函数

        参数:
            hidden_states: 输入张量, 其尾部维度应与 `normalized_shape` 匹配.
               例如, 形状可以是 [batch_size, ..., *normalized_shape]

        返回:
            归一化后的张量, 形状与输入 hidden_states 相同.
        """
        # 1. 确定需要计算均值和方差的维度
        # normalized_shape 定义了最后几个维度, 我们需要在这些维度上计算统计量
        # 例如, 如果 hidden_states.shape = (N, C, H, W) 且 normalized_shape = (H, W)
        # 则 dims_to_normalize = (-2, -1)
        num_normalized_dims = len(self.normalized_shape)
        dims_to_normalize = tuple(range(hidden_states.ndim - num_normalized_dims, hidden_states.ndim))

        # 2. 计算均值 (μ) 和方差 (σ²)
        # 在指定的维度上计算, 并保持维度以便广播
        # 注意: 计算方差时使用 unbiased=False, 与 PyTorch 官方实现一致
        mean = torch.mean(hidden_states, dim=dims_to_normalize, keepdim=True)
        # var = torch.var(hidden_states, dim=dims_to_normalize, unbiased=False, keepdim=True) # 简洁写法
        # 或者, 使用定义式计算方差(对初学者更清晰):
        var = ((hidden_states - mean) ** 2).mean(dim=dims_to_normalize, keepdim=True)

        # 3. 归一化 (x_normalized)
        # (x - μ) / sqrt(σ² + ε)
        x_normalized = (hidden_states - mean) / torch.sqrt(var + self.eps)

        # 4. 应用仿射变换 (gamma * x_normalized + beta)
        if self.elementwise_affine:
            # self.weight (gamma) 和 self.bias (beta) 的形状是 normalized_shape
            # PyTorch 的广播机制会自动将它们应用到 x_normalized 的对应维度上
            return self.weight * x_normalized + self.bias
        else:
            return x_normalized

    def extra_repr(self) -> str:
        # 自定义打印模块信息时的显示内容, 使其更像官方版本
        return f"{self.normalized_shape}, eps={self.eps}, elementwise_affine={self.elementwise_affine}"

forward

forward(hidden_states)

前向传播函数

参数

hidden_states: 输入张量, 其尾部维度应与 normalized_shape 匹配. 例如, 形状可以是 [batch_size, ..., *normalized_shape]

返回

归一化后的张量, 形状与输入 hidden_states 相同.

源代码位于： src/llm/core/layer_norm.py

def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    """
    前向传播函数

    参数:
        hidden_states: 输入张量, 其尾部维度应与 `normalized_shape` 匹配.
           例如, 形状可以是 [batch_size, ..., *normalized_shape]

    返回:
        归一化后的张量, 形状与输入 hidden_states 相同.
    """
    # 1. 确定需要计算均值和方差的维度
    # normalized_shape 定义了最后几个维度, 我们需要在这些维度上计算统计量
    # 例如, 如果 hidden_states.shape = (N, C, H, W) 且 normalized_shape = (H, W)
    # 则 dims_to_normalize = (-2, -1)
    num_normalized_dims = len(self.normalized_shape)
    dims_to_normalize = tuple(range(hidden_states.ndim - num_normalized_dims, hidden_states.ndim))

    # 2. 计算均值 (μ) 和方差 (σ²)
    # 在指定的维度上计算, 并保持维度以便广播
    # 注意: 计算方差时使用 unbiased=False, 与 PyTorch 官方实现一致
    mean = torch.mean(hidden_states, dim=dims_to_normalize, keepdim=True)
    # var = torch.var(hidden_states, dim=dims_to_normalize, unbiased=False, keepdim=True) # 简洁写法
    # 或者, 使用定义式计算方差(对初学者更清晰):
    var = ((hidden_states - mean) ** 2).mean(dim=dims_to_normalize, keepdim=True)

    # 3. 归一化 (x_normalized)
    # (x - μ) / sqrt(σ² + ε)
    x_normalized = (hidden_states - mean) / torch.sqrt(var + self.eps)

    # 4. 应用仿射变换 (gamma * x_normalized + beta)
    if self.elementwise_affine:
        # self.weight (gamma) 和 self.bias (beta) 的形状是 normalized_shape
        # PyTorch 的广播机制会自动将它们应用到 x_normalized 的对应维度上
        return self.weight * x_normalized + self.bias
    else:
        return x_normalized

layer_norm_numpy

layer_norm_numpy(x, gamma=None, beta=None, eps=1e-05)

Layer Normalization 的 NumPy 实现 (简化版)

注意: 此版本为了简洁, 固定在最后一个轴 (axis=-1) 上进行归一化. 主要用于帮助理解 LayerNorm 的核心计算步骤.

参数

x: 输入 NumPy 数组, 形状例如 [batch_size, ..., feature_dim] gamma: 缩放参数 (如果提供), 形状应为 [feature_dim] beta: 偏移参数 (如果提供), 形状应为 [feature_dim] eps: 防止除零错误的小常数

返回

归一化后的 NumPy 数组, 形状与输入 x 相同.

源代码位于： src/llm/core/layer_norm.py

def layer_norm_numpy(
    x: np.ndarray, gamma: np.ndarray | None = None, beta: np.ndarray | None = None, eps: float = 1e-5
) -> np.ndarray:
    """
    Layer Normalization 的 NumPy 实现 (简化版)

    注意: 此版本为了简洁, *固定*在最后一个轴 (axis=-1) 上进行归一化.
    主要用于帮助理解 LayerNorm 的核心计算步骤.

    参数:
        x: 输入 NumPy 数组, 形状例如 [batch_size, ..., feature_dim]
        gamma: 缩放参数 (如果提供), 形状应为 [feature_dim]
        beta: 偏移参数 (如果提供), 形状应为 [feature_dim]
        eps: 防止除零错误的小常数

    返回:
        归一化后的 NumPy 数组, 形状与输入 x 相同.
    """
    # 1. 在最后一个轴上计算均值 μ
    mean = np.mean(x, axis=-1, keepdims=True)

    # 2. 在最后一个轴上计算方差 σ²
    var = np.mean((x - mean) ** 2, axis=-1, keepdims=True)

    # 3. 归一化 x_normalized
    x_normalized = (x - mean) / np.sqrt(var + eps)

    # 4. 应用仿射变换
    if gamma is not None and beta is not None:
        return gamma * x_normalized + beta
    else:
        return x_normalized

Embeddings and Positional Encoding

embedding

EmbeddingLayer

Bases: Module

Combines token embeddings with positional encodings.

The layer first embeds input token IDs into dense vectors, then scales these embeddings by the square root of the hidden size, and finally adds positional encodings.

源代码位于： src/llm/core/embedding.py

class EmbeddingLayer(nn.Module):
    """
    Combines token embeddings with positional encodings.

    The layer first embeds input token IDs into dense vectors, then scales these
    embeddings by the square root of the hidden size, and finally adds
    positional encodings.
    """

    def __init__(
        self,
        vocab_size: int,
        hidden_size: int,
        max_seq_len: int = 512,
        pos_encoding_learned: bool = False,
        dropout_p: float = 0.1,
        padding_idx: int | None = None,
        use_rope: bool = False,
        device: torch.device | str | None = None,
        dtype: torch.dtype | None = None,
    ):
        """
        Initializes the EmbeddingLayer.

        Args:
            vocab_size (int): The size of the vocabulary.
            hidden_size (int): The embedding dimension.
            max_seq_len (int, default=512): Maximum sequence length for positional encoding.
            pos_encoding_learned (bool, default=False): If True, use learned positional embeddings.
                                                       If False, use sinusoidal.
            dropout_p (float, default=0.1): Dropout probability for positional encoding.
            padding_idx (int, optional, default=None): If specified, the entries at `padding_idx`
                                                       in `token_embeddings` do not contribute to
                                                       the gradient; furthermore, the embedding vector
                                                       for `padding_idx` is initialized to all zeros.
            use_rope (bool, default=False): If True, no additive positional encoding is applied —
                position comes from RoPE inside attention (real Llama/Mistral). The
                ``positional_encoding`` module is still constructed (learned=False) so
                attribute reads by hf_publisher keep working, but its output is skipped.
            device (torch.device | str | None, default=None): Target device for the layers.
            dtype (torch.dtype | None, default=None): Target data type for the layers.
        """
        factory_kwargs = make_factory_kwargs(device, dtype)
        super().__init__()

        self.hidden_size = hidden_size
        self.padding_idx = padding_idx
        self.use_rope = use_rope

        self.token_embeddings = nn.Embedding(
            num_embeddings=vocab_size, embedding_dim=hidden_size, padding_idx=padding_idx, **factory_kwargs
        )

        # PositionalEncoding should be initialized with factory_kwargs
        # This assumes PositionalEncoding's __init__ is modified to accept and use **factory_kwargs
        # for its internal nn.Embedding if learned=True.
        self.positional_encoding = PositionalEncoding(
            hidden_size=hidden_size,
            max_seq_len=max_seq_len,
            dropout_p=dropout_p,
            learned=pos_encoding_learned,
            **factory_kwargs,  # Pass factory_kwargs here
        )
        # No explicit .to(device, dtype) for self.positional_encoding is needed here
        # if PositionalEncoding correctly uses factory_kwargs for its parameters (learned case)
        # and its buffers are handled by the parent module's .to() method (sinusoidal case).

    def forward(
        self, input_ids: torch.Tensor, start_pos: int = 0, position_ids: torch.Tensor | None = None
    ) -> torch.Tensor:
        """
        Forward pass of the EmbeddingLayer.

        Args:
            input_ids (torch.Tensor): Tensor of token IDs of shape (batch_size, seq_len).
            start_pos (int): Initial position index for the sequence.
            position_ids (torch.Tensor, optional): Explicit position IDs.

        Returns:
            torch.Tensor: Tensor of embeddings with positional encodings,
                          of shape (batch_size, seq_len, hidden_size).
        """
        token_embs = self.token_embeddings(input_ids)
        scaled_embs = token_embs * math.sqrt(self.hidden_size)
        if self.use_rope:
            # RoPE models inject position inside attention; adding the
            # sinusoidal/learned positional encoding here would double-count
            # position information (real Llama/Mistral semantics — RIL ISS-062).
            return scaled_embs
        output_embs = self.positional_encoding(scaled_embs, start_pos=start_pos, position_ids=position_ids)
        return output_embs

forward

forward(input_ids, start_pos=0, position_ids=None)

Forward pass of the EmbeddingLayer.

参数：

名称	类型	描述	默认
input_ids	Tensor	Tensor of token IDs of shape (batch_size, seq_len).	必需
start_pos	int	Initial position index for the sequence.	0
position_ids	Tensor	Explicit position IDs.	None

返回：

类型	描述
Tensor	torch.Tensor: Tensor of embeddings with positional encodings, of shape (batch_size, seq_len, hidden_size).

源代码位于： src/llm/core/embedding.py

def forward(
    self, input_ids: torch.Tensor, start_pos: int = 0, position_ids: torch.Tensor | None = None
) -> torch.Tensor:
    """
    Forward pass of the EmbeddingLayer.

    Args:
        input_ids (torch.Tensor): Tensor of token IDs of shape (batch_size, seq_len).
        start_pos (int): Initial position index for the sequence.
        position_ids (torch.Tensor, optional): Explicit position IDs.

    Returns:
        torch.Tensor: Tensor of embeddings with positional encodings,
                      of shape (batch_size, seq_len, hidden_size).
    """
    token_embs = self.token_embeddings(input_ids)
    scaled_embs = token_embs * math.sqrt(self.hidden_size)
    if self.use_rope:
        # RoPE models inject position inside attention; adding the
        # sinusoidal/learned positional encoding here would double-count
        # position information (real Llama/Mistral semantics — RIL ISS-062).
        return scaled_embs
    output_embs = self.positional_encoding(scaled_embs, start_pos=start_pos, position_ids=position_ids)
    return output_embs

positional_encoding

PositionalEncoding

Bases: Module

Implements Positional Encoding for Transformer models.

Supports both sinusoidal (fixed) and learned positional embeddings. Positional information is added to the input embeddings to provide order awareness.

参数：

名称	类型	描述	默认
hidden_size	int	The dimension of the model's hidden states.	必需
max_seq_len	int	Maximum sequence length supported. Defaults to 512.	512
dropout_p	float	Dropout probability applied after adding positional signal. Defaults to 0.1.	0.1
learned	bool	If True, uses learned embeddings; otherwise sinusoidal. Defaults to False.	False
device	device | None	Device to place parameters on.	None
dtype	dtype | None	Data type for parameters.	None

源代码位于： src/llm/core/positional_encoding.py

class PositionalEncoding(nn.Module):
    """
    Implements Positional Encoding for Transformer models.

    Supports both sinusoidal (fixed) and learned positional embeddings.
    Positional information is added to the input embeddings to provide order awareness.

    Args:
        hidden_size (int): The dimension of the model's hidden states.
        max_seq_len (int): Maximum sequence length supported. Defaults to 512.
        dropout_p (float): Dropout probability applied after adding positional signal. Defaults to 0.1.
        learned (bool): If True, uses learned embeddings; otherwise sinusoidal. Defaults to False.
        device (torch.device | None): Device to place parameters on.
        dtype (torch.dtype | None): Data type for parameters.
    """

    def __init__(
        self,
        hidden_size: int,
        max_seq_len: int = 512,
        dropout_p: float = 0.1,
        learned: bool = False,
        device: torch.device | str | None = None,
        dtype: torch.dtype | None = None,
    ):
        super().__init__()
        self.hidden_size = hidden_size
        self.max_seq_len = max_seq_len
        self.dropout_p = dropout_p
        self.learned = learned
        self.dropout = nn.Dropout(p=dropout_p)
        self.pos_embedding: nn.Embedding
        self.pe: torch.Tensor

        factory_kwargs = make_factory_kwargs(device, dtype)

        if self.learned:
            self.pos_embedding = nn.Embedding(max_seq_len, hidden_size, **factory_kwargs)
            # Optional: Initialize weights, though default nn.Embedding initialization is often fine.
            # self.pos_embedding.weight.data.normal_(mean=0.0, std=0.02)
        else:
            # Create a buffer for sinusoidal positional encodings
            # pe shape: (1, max_seq_len, hidden_size)
            pe = torch.zeros(1, max_seq_len, hidden_size, **factory_kwargs)

            # position shape: (max_seq_len, 1)
            # Calculations for sinusoidal encoding should ideally be done in float32 for precision,
            # then cast to the target dtype if needed. However, for simplicity and directness,
            # we can try to use the target device and a float dtype for calculation tensors.
            # If factory_kwargs['dtype'] is a float type, use it, else default to torch.float for calculations.
            calc_dtype = dtype if dtype is not None and dtype.is_floating_point else torch.float

            position = torch.arange(0, max_seq_len, device=device, dtype=calc_dtype).unsqueeze(1)

            # div_term shape: (hidden_size / 2)
            div_term_base = torch.arange(0, hidden_size, 2, device=device, dtype=calc_dtype)
            div_term = torch.exp(div_term_base * (-math.log(10000.0) / hidden_size))

            # Apply sin to even indices in the hidden_size dimension.
            # ``0::2`` selects ``ceil(hidden_size/2)`` columns, so the full
            # ``div_term`` (length ``ceil(hidden_size/2)``) covers them.
            pe[0, :, 0::2] = torch.sin(position * div_term)
            # Apply cos to odd indices. ``1::2`` selects ``floor(h/2)``
            # columns, so we use only the first ``hidden_size // 2``
            # frequencies. For even ``hidden_size`` this slice is a no-op
            # (``div_term[:h/2] == div_term``); for odd it prevents a
            # shape-mismatch crash that otherwise broke construction.
            pe[0, :, 1::2] = torch.cos(position * div_term[: hidden_size // 2])

            # Non-persistent: the sine table is a pure deterministic function
            # of (max_seq_len, hidden_size), regenerated every time the module
            # is constructed (models rebuild from config on load, so a
            # persisted copy is dead weight). Storing it persistent bloated
            # every checkpoint and exported artifact (ONNX/TorchScript/GGUF
            # embed it as a constant) and got Q4_0-quantized by the GGUF
            # exporter under a name llama.cpp doesn't recognize (RIL ISS-146).
            self.register_buffer("pe", pe, persistent=False)

    def forward(self, x: torch.Tensor, start_pos: int = 0, position_ids: torch.Tensor | None = None) -> torch.Tensor:
        """
        Args:
            x: Tensor, shape [batch_size, seq_len, hidden_size]
            start_pos: Initial position index for the sequence (used if position_ids is None).
            position_ids: Optional Tensor of shape [batch_size, seq_len] containing explicit position indices.
        """
        seq_len = x.size(1)

        # Validation logic (only check start_pos if position_ids not provided)
        if position_ids is None and start_pos + seq_len > self.max_seq_len:
            raise ValueError(
                f"Sequence endpoint {start_pos + seq_len} exceeds maximum sequence length {self.max_seq_len}"
            )
        if position_ids is not None and position_ids.numel():
            # Explicit position ids must stay inside the table. An out-of-range
            # id would index the embedding/PE rows past the buffer; on CUDA
            # that surfaces as a device-side assert that poisons the whole
            # process (every in-flight request), rather than a Python error.
            min_id, max_id = int(position_ids.min().item()), int(position_ids.max().item())
            if min_id < 0 or max_id >= self.max_seq_len:
                raise ValueError(
                    f"position_ids out of range [{min_id}, {max_id}]: must lie within "
                    f"[0, {self.max_seq_len}) (the embedding table has only "
                    f"{self.max_seq_len} positions)"
                )

        if self.learned:
            if position_ids is None:
                # Create position IDs [start_pos, ..., start_pos + seq_len - 1]
                # Broadcast across batch
                pos_ids = torch.arange(start_pos, start_pos + seq_len, dtype=torch.long, device=x.device).unsqueeze(0)
            else:
                pos_ids = position_ids

            pos_enc = self.pos_embedding(pos_ids)
            x = x + pos_enc
        else:
            # self.pe is [1, max_seq_len, hidden_size]
            if position_ids is None:
                x = x + self.pe[:, start_pos : start_pos + seq_len, :]
            else:
                # Gather positional encodings based on position_ids
                # position_ids shape: [B, S]
                # self.pe shape: [1, MaxLen, H]
                # We want [B, S, H]

                # Expand PE to match batch size? No need if we index properly.
                # self.pe[0] is [MaxLen, H]
                # We gather rows specified by position_ids
                # F.embedding can do this if we treat pe[0] as weight matrix?
                # Or advanced indexing: self.pe[0, position_ids] -> [B, S, H]
                pos_enc = self.pe[0, position_ids]
                x = x + pos_enc

        return self.dropout(x)

forward

forward(x, start_pos=0, position_ids=None)

参数：

名称	类型	描述	默认
x	Tensor	Tensor, shape [batch_size, seq_len, hidden_size]	必需
start_pos	int	Initial position index for the sequence (used if position_ids is None).	0
position_ids	Tensor | None	Optional Tensor of shape [batch_size, seq_len] containing explicit position indices.	None

源代码位于： src/llm/core/positional_encoding.py

def forward(self, x: torch.Tensor, start_pos: int = 0, position_ids: torch.Tensor | None = None) -> torch.Tensor:
    """
    Args:
        x: Tensor, shape [batch_size, seq_len, hidden_size]
        start_pos: Initial position index for the sequence (used if position_ids is None).
        position_ids: Optional Tensor of shape [batch_size, seq_len] containing explicit position indices.
    """
    seq_len = x.size(1)

    # Validation logic (only check start_pos if position_ids not provided)
    if position_ids is None and start_pos + seq_len > self.max_seq_len:
        raise ValueError(
            f"Sequence endpoint {start_pos + seq_len} exceeds maximum sequence length {self.max_seq_len}"
        )
    if position_ids is not None and position_ids.numel():
        # Explicit position ids must stay inside the table. An out-of-range
        # id would index the embedding/PE rows past the buffer; on CUDA
        # that surfaces as a device-side assert that poisons the whole
        # process (every in-flight request), rather than a Python error.
        min_id, max_id = int(position_ids.min().item()), int(position_ids.max().item())
        if min_id < 0 or max_id >= self.max_seq_len:
            raise ValueError(
                f"position_ids out of range [{min_id}, {max_id}]: must lie within "
                f"[0, {self.max_seq_len}) (the embedding table has only "
                f"{self.max_seq_len} positions)"
            )

    if self.learned:
        if position_ids is None:
            # Create position IDs [start_pos, ..., start_pos + seq_len - 1]
            # Broadcast across batch
            pos_ids = torch.arange(start_pos, start_pos + seq_len, dtype=torch.long, device=x.device).unsqueeze(0)
        else:
            pos_ids = position_ids

        pos_enc = self.pos_embedding(pos_ids)
        x = x + pos_enc
    else:
        # self.pe is [1, max_seq_len, hidden_size]
        if position_ids is None:
            x = x + self.pe[:, start_pos : start_pos + seq_len, :]
        else:
            # Gather positional encodings based on position_ids
            # position_ids shape: [B, S]
            # self.pe shape: [1, MaxLen, H]
            # We want [B, S, H]

            # Expand PE to match batch size? No need if we index properly.
            # self.pe[0] is [MaxLen, H]
            # We gather rows specified by position_ids
            # F.embedding can do this if we treat pe[0] as weight matrix?
            # Or advanced indexing: self.pe[0, position_ids] -> [B, S, H]
            pos_enc = self.pe[0, position_ids]
            x = x + pos_enc

    return self.dropout(x)

rope

Rotary Position Embedding (RoPE) Module.

Implements RoPE with scaling support for extended context lengths. Supports linear scaling, dynamic scaling, and NTK-aware scaling.

Reference: https://arxiv.org/abs/2104.09864

RotaryPositionEmbedding

Bases: Module

Rotary Position Embedding (RoPE).

Encodes position information by rotating query and key vectors.

参数：

名称	类型	描述	默认
dim	int	Dimension of the embedding (typically head_dim)	必需
max_seq_len	int	Maximum sequence length for precomputed embeddings	2048
base	float	Base for computing rotation frequencies (default: 10000)	10000.0
scaling_type	str | None	Type of RoPE scaling: None, 'linear', 'dynamic', 'ntk'	None
scaling_factor	float	Scaling factor for extended context (default: 1.0)	1.0
device	device | str | None	Device for the embedding	None
dtype	dtype | None	Data type for the embedding	None

源代码位于： src/llm/core/rope.py

class RotaryPositionEmbedding(nn.Module):
    """
    Rotary Position Embedding (RoPE).

    Encodes position information by rotating query and key vectors.

    Args:
        dim: Dimension of the embedding (typically head_dim)
        max_seq_len: Maximum sequence length for precomputed embeddings
        base: Base for computing rotation frequencies (default: 10000)
        scaling_type: Type of RoPE scaling: None, 'linear', 'dynamic', 'ntk'
        scaling_factor: Scaling factor for extended context (default: 1.0)
        device: Device for the embedding
        dtype: Data type for the embedding
    """

    def __init__(
        self,
        dim: int,
        max_seq_len: int = 2048,
        base: float = 10000.0,
        scaling_type: str | None = None,
        scaling_factor: float = 1.0,
        device: torch.device | str | None = None,
        dtype: torch.dtype | None = None,
    ):
        super().__init__()
        self.dim = dim
        self.max_seq_len = max_seq_len
        self.base = base
        self.scaling_type = scaling_type
        self.scaling_factor = scaling_factor

        if dim % 2 != 0:
            raise ValueError(
                f"RoPE requires an even head_dim, got dim={dim}. "
                "rotate_half pairs each dimension with its counterpart at "
                "dim/2, which is undefined when dim is odd."
            )

        # Compute inverse frequencies
        inv_freq = self._compute_inv_freq(device, dtype)
        self.inv_freq: torch.Tensor
        self.register_buffer("inv_freq", inv_freq, persistent=False)

        # Precompute cos and sin for efficiency
        self._seq_len_cached = 0
        self._cos_cached: torch.Tensor | None = None
        self._sin_cached: torch.Tensor | None = None

    def _compute_inv_freq(
        self,
        device: torch.device | str | None = None,
        dtype: torch.dtype | None = None,
    ) -> torch.Tensor:
        """Compute inverse frequencies with optional NTK scaling."""
        calc_dtype = dtype if dtype is not None and dtype.is_floating_point else torch.float32

        if self.scaling_type == "ntk" and self.scaling_factor > 1.0:
            # NTK-aware scaling: adjust base instead of positions
            base = self.base * (self.scaling_factor ** (self.dim / (self.dim - 2)))
        else:
            base = self.base

        inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2, device=device, dtype=calc_dtype) / self.dim))
        return inv_freq

    def _update_cos_sin_cache(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> None:
        """Update cached cos/sin values if sequence length changed."""
        if seq_len > self._seq_len_cached or self._cos_cached is None:
            self._seq_len_cached = max(seq_len, self.max_seq_len)

            # Position indices
            t = torch.arange(self._seq_len_cached, device=device, dtype=self.inv_freq.dtype)

            # Apply linear or dynamic scaling to positions
            if self.scaling_type == "linear" and self.scaling_factor > 1.0:
                t = t / self.scaling_factor
            elif self.scaling_type == "dynamic" and seq_len > self.max_seq_len:
                # Dynamic scaling: scale positions based on current sequence length
                scale = seq_len / self.max_seq_len
                t = t / scale

            # Compute frequencies: [seq_len, dim/2]
            freqs = torch.outer(t, self.inv_freq)

            # Compute cos and sin: [seq_len, dim]
            emb = torch.cat((freqs, freqs), dim=-1)
            self._cos_cached = emb.cos().to(dtype)
            self._sin_cached = emb.sin().to(dtype)

    def forward(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        position_ids: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Apply rotary position embedding to query and key tensors.

        Args:
            q: Query tensor of shape [batch, heads, seq_len, head_dim]
            k: Key tensor of shape [batch, heads, seq_len, head_dim]
            position_ids: Optional position indices [batch, seq_len]

        Returns:
            Tuple of (rotated_q, rotated_k)
        """
        seq_len = q.size(2)
        # Size the cos/sin cache to cover the positions we will index, not
        # just the (relative) sequence length. ``_update_cos_sin_cache`` only
        # grows the table to ``max(seq_len, max_seq_len)``, but when explicit
        # ``position_ids`` are threaded (KV-cache decode at absolute
        # position >= max_seq_len, batch-serving position_ids) the table is
        # indexed by ABSOLUTE position — an absolute position beyond
        # ``max_seq_len`` used to raise a raw ``IndexError: index N out of
        # bounds`` (RIL ISS-141; the serving tier guards over-context prompts
        # up front, training/dense-cache RoPE callers did not).
        cover_len = int(position_ids.max().item()) + 1 if position_ids is not None else seq_len
        self._update_cos_sin_cache(max(cover_len, seq_len), q.device, q.dtype)

        assert self._cos_cached is not None  # noqa: S101
        assert self._sin_cached is not None  # noqa: S101

        if position_ids is None:
            cos = self._cos_cached[:seq_len]
            sin = self._sin_cached[:seq_len]
        else:
            cos = self._cos_cached[position_ids]
            sin = self._sin_cached[position_ids]

        # Reshape for broadcasting: [1, 1, seq_len, dim] or [batch, 1, seq_len, dim]
        if position_ids is None:
            cos = cos.unsqueeze(0).unsqueeze(0)
            sin = sin.unsqueeze(0).unsqueeze(0)
        else:
            cos = cos.unsqueeze(1)
            sin = sin.unsqueeze(1)

        q_embed = apply_rotary_pos_emb(q, cos, sin)
        k_embed = apply_rotary_pos_emb(k, cos, sin)

        return q_embed, k_embed

    def extra_repr(self) -> str:
        scaling_info = (
            f", scaling_type={self.scaling_type}, scaling_factor={self.scaling_factor}" if self.scaling_type else ""
        )
        return f"dim={self.dim}, max_seq_len={self.max_seq_len}, base={self.base}{scaling_info}"

forward

forward(q, k, position_ids=None)

Apply rotary position embedding to query and key tensors.

参数：

名称	类型	描述	默认
q	Tensor	Query tensor of shape [batch, heads, seq_len, head_dim]	必需
k	Tensor	Key tensor of shape [batch, heads, seq_len, head_dim]	必需
position_ids	Tensor | None	Optional position indices [batch, seq_len]	None

返回：

类型	描述
tuple[Tensor, Tensor]	Tuple of (rotated_q, rotated_k)

源代码位于： src/llm/core/rope.py

def forward(
    self,
    q: torch.Tensor,
    k: torch.Tensor,
    position_ids: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Apply rotary position embedding to query and key tensors.

    Args:
        q: Query tensor of shape [batch, heads, seq_len, head_dim]
        k: Key tensor of shape [batch, heads, seq_len, head_dim]
        position_ids: Optional position indices [batch, seq_len]

    Returns:
        Tuple of (rotated_q, rotated_k)
    """
    seq_len = q.size(2)
    # Size the cos/sin cache to cover the positions we will index, not
    # just the (relative) sequence length. ``_update_cos_sin_cache`` only
    # grows the table to ``max(seq_len, max_seq_len)``, but when explicit
    # ``position_ids`` are threaded (KV-cache decode at absolute
    # position >= max_seq_len, batch-serving position_ids) the table is
    # indexed by ABSOLUTE position — an absolute position beyond
    # ``max_seq_len`` used to raise a raw ``IndexError: index N out of
    # bounds`` (RIL ISS-141; the serving tier guards over-context prompts
    # up front, training/dense-cache RoPE callers did not).
    cover_len = int(position_ids.max().item()) + 1 if position_ids is not None else seq_len
    self._update_cos_sin_cache(max(cover_len, seq_len), q.device, q.dtype)

    assert self._cos_cached is not None  # noqa: S101
    assert self._sin_cached is not None  # noqa: S101

    if position_ids is None:
        cos = self._cos_cached[:seq_len]
        sin = self._sin_cached[:seq_len]
    else:
        cos = self._cos_cached[position_ids]
        sin = self._sin_cached[position_ids]

    # Reshape for broadcasting: [1, 1, seq_len, dim] or [batch, 1, seq_len, dim]
    if position_ids is None:
        cos = cos.unsqueeze(0).unsqueeze(0)
        sin = sin.unsqueeze(0).unsqueeze(0)
    else:
        cos = cos.unsqueeze(1)
        sin = sin.unsqueeze(1)

    q_embed = apply_rotary_pos_emb(q, cos, sin)
    k_embed = apply_rotary_pos_emb(k, cos, sin)

    return q_embed, k_embed

rotate_half

rotate_half(x)

Rotate half the hidden dims of the input.

源代码位于： src/llm/core/rope.py

def rotate_half(x: torch.Tensor) -> torch.Tensor:
    """Rotate half the hidden dims of the input."""
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)

apply_rotary_pos_emb

apply_rotary_pos_emb(x, cos, sin)

Apply rotary position embedding to a tensor.

参数：

名称	类型	描述	默认
x	Tensor	Input tensor of shape [..., seq_len, dim]	必需
cos	Tensor	Cosine embeddings	必需
sin	Tensor	Sine embeddings	必需

返回：

类型	描述
Tensor	Rotated tensor of the same shape

源代码位于： src/llm/core/rope.py

def apply_rotary_pos_emb(
    x: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
) -> torch.Tensor:
    """
    Apply rotary position embedding to a tensor.

    Args:
        x: Input tensor of shape [..., seq_len, dim]
        cos: Cosine embeddings
        sin: Sine embeddings

    Returns:
        Rotated tensor of the same shape
    """
    return (x * cos) + (rotate_half(x) * sin)

get_rope_scaling_factor

get_rope_scaling_factor(seq_len, max_trained_len, scaling_type='linear')

Compute appropriate RoPE scaling factor.

参数：

名称	类型	描述	默认
seq_len	int	Current sequence length	必需
max_trained_len	int	Maximum length the model was trained on	必需
scaling_type	str	Type of scaling ('linear' or 'dynamic')	'linear'

返回：

类型	描述
float	Scaling factor

源代码位于： src/llm/core/rope.py

def get_rope_scaling_factor(
    seq_len: int,
    max_trained_len: int,
    scaling_type: str = "linear",
) -> float:
    """
    Compute appropriate RoPE scaling factor.

    Args:
        seq_len: Current sequence length
        max_trained_len: Maximum length the model was trained on
        scaling_type: Type of scaling ('linear' or 'dynamic')

    Returns:
        Scaling factor
    """
    if seq_len <= max_trained_len:
        return 1.0

    if scaling_type == "linear":
        return seq_len / max_trained_len
    elif scaling_type == "dynamic":
        return math.sqrt(seq_len / max_trained_len)
    else:
        return 1.0

alibi

ALiBi (Attention with Linear Biases) Module.

Implements ALiBi, a simple position encoding method that adds linear biases to attention scores based on token distance.

Reference: https://arxiv.org/abs/2108.12409

ALiBiPositionBias

Bases: Module

ALiBi Position Bias Module.

Generates position-dependent biases to be added to attention scores.

参数：

名称	类型	描述	默认
num_heads	int	Number of attention heads	必需
max_seq_len	int	Maximum sequence length for cached bias	2048
device	device | str | None	Target device	None
dtype	dtype | None	Target data type	None

源代码位于： src/llm/core/alibi.py

class ALiBiPositionBias(nn.Module):
    """
    ALiBi Position Bias Module.

    Generates position-dependent biases to be added to attention scores.

    Args:
        num_heads: Number of attention heads
        max_seq_len: Maximum sequence length for cached bias
        device: Target device
        dtype: Target data type
    """

    def __init__(
        self,
        num_heads: int,
        max_seq_len: int = 2048,
        device: torch.device | str | None = None,
        dtype: torch.dtype | None = None,
    ):
        super().__init__()
        self.num_heads = num_heads
        self.max_seq_len = max_seq_len

        # Register slopes as buffer (not learned)
        slopes = get_alibi_slopes(num_heads)
        self.slopes: torch.Tensor
        self.register_buffer("slopes", slopes, persistent=False)

        # Cache for bias matrix
        self._cached_bias: torch.Tensor | None = None
        self._cached_seq_len: int = 0

    def _update_cache(
        self, seq_len: int, device: torch.device | str | None = None, dtype: torch.dtype | None = None
    ) -> None:
        """Update cached bias if sequence length changed."""
        if seq_len > self._cached_seq_len or self._cached_bias is None:
            self._cached_seq_len = max(seq_len, self.max_seq_len)
            self._cached_bias = build_alibi_bias(
                self.num_heads,
                self._cached_seq_len,
                device=device,
                dtype=dtype,
            )

    def forward(self, attention_scores: torch.Tensor) -> torch.Tensor:
        """
        Add ALiBi bias to attention scores.

        Args:
            attention_scores: Attention scores of shape [batch, heads, seq_len, seq_len]

        Returns:
            Attention scores with ALiBi bias added
        """
        seq_len = attention_scores.size(-1)
        self._update_cache(seq_len, attention_scores.device, attention_scores.dtype)
        assert self._cached_bias is not None  # noqa: S101

        # Extract relevant portion of cached bias
        bias = self._cached_bias[:, :, :seq_len, :seq_len]
        return attention_scores + bias

    def get_bias(
        self,
        seq_len: int,
        device: torch.device | str | None = None,
        dtype: torch.dtype | None = None,
    ) -> torch.Tensor:
        """
        Get ALiBi bias matrix for given sequence length.

        Args:
            seq_len: Sequence length
            device: Target device
            dtype: Target data type

        Returns:
            Bias tensor of shape [1, num_heads, seq_len, seq_len]
        """
        device = device if device is not None else self.slopes.device
        dtype = dtype if dtype is not None else self.slopes.dtype
        self._update_cache(seq_len, device, dtype)
        assert self._cached_bias is not None  # noqa: S101
        return self._cached_bias[:, :, :seq_len, :seq_len]

    def extra_repr(self) -> str:
        return f"num_heads={self.num_heads}, max_seq_len={self.max_seq_len}"

forward

forward(attention_scores)

Add ALiBi bias to attention scores.

参数：

名称	类型	描述	默认
attention_scores	Tensor	Attention scores of shape [batch, heads, seq_len, seq_len]	必需

返回：

类型	描述
Tensor	Attention scores with ALiBi bias added

源代码位于： src/llm/core/alibi.py

def forward(self, attention_scores: torch.Tensor) -> torch.Tensor:
    """
    Add ALiBi bias to attention scores.

    Args:
        attention_scores: Attention scores of shape [batch, heads, seq_len, seq_len]

    Returns:
        Attention scores with ALiBi bias added
    """
    seq_len = attention_scores.size(-1)
    self._update_cache(seq_len, attention_scores.device, attention_scores.dtype)
    assert self._cached_bias is not None  # noqa: S101

    # Extract relevant portion of cached bias
    bias = self._cached_bias[:, :, :seq_len, :seq_len]
    return attention_scores + bias

get_bias

get_bias(seq_len, device=None, dtype=None)

Get ALiBi bias matrix for given sequence length.

参数：

名称	类型	描述	默认
seq_len	int	Sequence length	必需
device	device | str | None	Target device	None
dtype	dtype | None	Target data type	None

返回：

类型	描述
Tensor	Bias tensor of shape [1, num_heads, seq_len, seq_len]

源代码位于： src/llm/core/alibi.py

def get_bias(
    self,
    seq_len: int,
    device: torch.device | str | None = None,
    dtype: torch.dtype | None = None,
) -> torch.Tensor:
    """
    Get ALiBi bias matrix for given sequence length.

    Args:
        seq_len: Sequence length
        device: Target device
        dtype: Target data type

    Returns:
        Bias tensor of shape [1, num_heads, seq_len, seq_len]
    """
    device = device if device is not None else self.slopes.device
    dtype = dtype if dtype is not None else self.slopes.dtype
    self._update_cache(seq_len, device, dtype)
    assert self._cached_bias is not None  # noqa: S101
    return self._cached_bias[:, :, :seq_len, :seq_len]

get_alibi_slopes

get_alibi_slopes(num_heads)

Compute ALiBi slopes for each attention head.

Slopes form a geometric sequence: 2^(-8/n), 2^(-16/n), ..., 2^(-8) where n is the number of heads.

参数：

名称	类型	描述	默认
num_heads	int	Number of attention heads	必需

返回：

类型	描述
Tensor	Tensor of shape [num_heads] containing slopes

源代码位于： src/llm/core/alibi.py

def get_alibi_slopes(num_heads: int) -> torch.Tensor:
    """
    Compute ALiBi slopes for each attention head.

    Slopes form a geometric sequence: 2^(-8/n), 2^(-16/n), ..., 2^(-8)
    where n is the number of heads.

    Args:
        num_heads: Number of attention heads

    Returns:
        Tensor of shape [num_heads] containing slopes
    """
    # Get closest power of 2 for consistent slopes
    closest_power_of_2 = 2 ** math.floor(math.log2(num_heads))

    # Compute base for the geometric sequence
    base = 2 ** (-(2 ** -(math.log2(closest_power_of_2) - 3)))

    # Generate slopes for powers of 2
    powers = torch.arange(1, closest_power_of_2 + 1)
    slopes = torch.pow(base, powers)

    if closest_power_of_2 != num_heads:
        # If num_heads is not a power of 2, compute extra slopes
        extra_base = 2 ** (-(2 ** -(math.log2(2 * closest_power_of_2) - 3)))
        num_remaining = num_heads - closest_power_of_2
        extra_powers = torch.arange(1, 2 * num_remaining + 1, 2)
        extra_slopes = torch.pow(extra_base, extra_powers)
        slopes = torch.cat([slopes, extra_slopes])

    return slopes

build_alibi_bias

build_alibi_bias(num_heads, seq_len, device=None, dtype=None)

Build the ALiBi bias matrix.

参数：

名称	类型	描述	默认
num_heads	int	Number of attention heads	必需
seq_len	int	Sequence length	必需
device	device | str | None	Target device	None
dtype	dtype | None	Target data type	None

返回：

类型	描述
Tensor	Tensor of shape [1, num_heads, seq_len, seq_len]

源代码位于： src/llm/core/alibi.py

def build_alibi_bias(
    num_heads: int,
    seq_len: int,
    device: torch.device | str | None = None,
    dtype: torch.dtype | None = None,
) -> torch.Tensor:
    """
    Build the ALiBi bias matrix.

    Args:
        num_heads: Number of attention heads
        seq_len: Sequence length
        device: Target device
        dtype: Target data type

    Returns:
        Tensor of shape [1, num_heads, seq_len, seq_len]
    """
    # Get slopes for each head
    slopes = get_alibi_slopes(num_heads)
    slopes = slopes.to(device=device, dtype=dtype)

    # Create distance matrix: position i attending to position j
    # Distance = j - i (positive for future, negative for past)
    positions = torch.arange(seq_len, device=device, dtype=dtype)
    distance = positions.unsqueeze(0) - positions.unsqueeze(1)  # [seq_len, seq_len]

    # Compute bias: -slope * |distance| (penalize far tokens)
    # For causal attention, we only care about past positions (distance <= 0)
    # Use distance directly (negative for past) as bias
    bias = distance.unsqueeze(0) * slopes.unsqueeze(1).unsqueeze(2)  # [num_heads, seq_len, seq_len]

    return bias.unsqueeze(0)  # [1, num_heads, seq_len, seq_len]

PEFT Helpers

The unified PEFT_REGISTRY lives in llm.core.peft; these are the standalone helper modules for the individual methods.

lora

LoRA (Low-Rank Adaptation) Module.

Implements parameter-efficient fine-tuning by adding trainable low-rank matrices to frozen linear layers.

Reference: https://arxiv.org/abs/2106.09685

LoRALinear

Bases: Module

LoRA-adapted linear layer.

Wraps a frozen nn.Linear and adds trainable low-rank matrices A and B. Output: base_output + (input @ A) @ B * scaling

Snapshot of scaling, set by

:func:disable_lora, cleared by :func:enable_lora.

参数：

名称	类型	描述	默认
base_layer	Linear	The original nn.Linear layer to adapt (will be frozen)	必需
rank	int	Rank of the low-rank matrices (default: 8)	8
alpha	float	Scaling factor (default: 16.0)	16.0
dropout	float	Dropout probability for LoRA path (default: 0.0)	0.0

源代码位于： src/llm/core/lora.py

class LoRALinear(nn.Module):
    """
    LoRA-adapted linear layer.

    Wraps a frozen nn.Linear and adds trainable low-rank matrices A and B.
    Output: base_output + (input @ A) @ B * scaling

    _original_scaling: Snapshot of ``scaling``, set by
        :func:`disable_lora`, cleared by :func:`enable_lora`.

    Args:
        base_layer: The original nn.Linear layer to adapt (will be frozen)
        rank: Rank of the low-rank matrices (default: 8)
        alpha: Scaling factor (default: 16.0)
        dropout: Dropout probability for LoRA path (default: 0.0)
    """

    _original_scaling: float | None
    # Set only by ``merge_weights`` (and read/cleared by ``unmerge_weights``);
    # declared here so static analysis sees a float, not an inferred union.
    _merged_scaling: float | None = None

    def __init__(
        self,
        base_layer: nn.Linear,
        rank: int = 8,
        alpha: float = 16.0,
        dropout: float = 0.0,
    ):
        super().__init__()
        self.base_layer = base_layer
        self.rank = rank
        self.alpha = alpha
        self.scaling = alpha / rank

        in_features = base_layer.in_features
        out_features = base_layer.out_features
        device = base_layer.weight.device
        dtype = base_layer.weight.dtype

        # LoRA matrices (same device/dtype as base)
        self.lora_A = nn.Parameter(torch.empty(in_features, rank, device=device, dtype=dtype))
        self.lora_B = nn.Parameter(torch.empty(rank, out_features, device=device, dtype=dtype))

        # Dropout for regularization
        self.lora_dropout = nn.Dropout(p=dropout) if dropout > 0 else nn.Identity()

        # Initialize
        self._init_lora_weights()

        # Freeze base layer
        self.base_layer.weight.requires_grad = False
        if self.base_layer.bias is not None:
            self.base_layer.bias.requires_grad = False

    def _init_lora_weights(self) -> None:
        """Initialize LoRA weights: A with Kaiming, B with zeros."""
        init_lora_weights(self.lora_A, self.lora_B)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass: frozen base + LoRA adaptation."""
        base_output = self.base_layer(x)
        lora_output = self.lora_dropout(x) @ self.lora_A @ self.lora_B
        return base_output + lora_output * self.scaling

    def merge_weights(self) -> None:
        """Merge LoRA weights into the base layer for efficient inference.

        Folds ``ΔW · scaling`` into ``base_layer.weight`` and then
        disables the adapter path (``scaling`` → 0) so that
        :meth:`forward` continues to produce the **same** output as
        before merging — the wrapper becomes a pure pass-through to the
        already-merged base.  Without disabling the path the adapter
        would be double-counted on every forward (a roughly 2x-strength
        model at serve time).

        :meth:`unmerge_weights` restores both the original base weight
        and this scaling snapshot.

        Idempotent: a second call while already merged is a no-op. Re-running
        the fold would re-store ``_merged_scaling = self.scaling`` where
        ``scaling`` is already 0 (post first merge), so the re-fold no-ops
        AND the snapshot becomes 0 — the later ``unmerge_weights`` then
        restores scaling=0 and leaves the base permanently folded at 2x
        adapter strength (RIL ISS-159).
        """
        if self._merged_scaling is not None:
            return
        with torch.no_grad():
            self._merged_scaling = self.scaling
            delta_w = (self.lora_A @ self.lora_B) * self.scaling
            self.base_layer.weight.add_(delta_w.T)
            self.scaling = 0.0

    def unmerge_weights(self) -> None:
        """Unmerge LoRA weights from the base layer.

        Restores the pre-merge base weight and re-enables the adapter
        path.  No-op if :meth:`merge_weights` was never called.
        """
        with torch.no_grad():
            merged_scaling = self._merged_scaling
            if merged_scaling is None:
                return
            delta_w = (self.lora_A @ self.lora_B) * merged_scaling
            self.base_layer.weight.sub_(delta_w.T)
            self.scaling = merged_scaling
            self._merged_scaling = None

    @property
    def trainable_parameters(self) -> int:
        """Number of trainable LoRA parameters."""
        return self.lora_A.numel() + self.lora_B.numel()

    def extra_repr(self) -> str:
        return f"rank={self.rank}, alpha={self.alpha}, scaling={self.scaling:.4f}"

trainable_parameters property

trainable_parameters

Number of trainable LoRA parameters.

forward

forward(x)

Forward pass: frozen base + LoRA adaptation.

源代码位于： src/llm/core/lora.py

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Forward pass: frozen base + LoRA adaptation."""
    base_output = self.base_layer(x)
    lora_output = self.lora_dropout(x) @ self.lora_A @ self.lora_B
    return base_output + lora_output * self.scaling

merge_weights

merge_weights()

Merge LoRA weights into the base layer for efficient inference.

Folds ΔW · scaling into base_layer.weight and then disables the adapter path (scaling → 0) so that :meth:forward continues to produce the same output as before merging — the wrapper becomes a pure pass-through to the already-merged base. Without disabling the path the adapter would be double-counted on every forward (a roughly 2x-strength model at serve time).

:meth:unmerge_weights restores both the original base weight and this scaling snapshot.

Idempotent: a second call while already merged is a no-op. Re-running the fold would re-store _merged_scaling = self.scaling where scaling is already 0 (post first merge), so the re-fold no-ops AND the snapshot becomes 0 — the later unmerge_weights then restores scaling=0 and leaves the base permanently folded at 2x adapter strength (RIL ISS-159).

源代码位于： src/llm/core/lora.py

def merge_weights(self) -> None:
    """Merge LoRA weights into the base layer for efficient inference.

    Folds ``ΔW · scaling`` into ``base_layer.weight`` and then
    disables the adapter path (``scaling`` → 0) so that
    :meth:`forward` continues to produce the **same** output as
    before merging — the wrapper becomes a pure pass-through to the
    already-merged base.  Without disabling the path the adapter
    would be double-counted on every forward (a roughly 2x-strength
    model at serve time).

    :meth:`unmerge_weights` restores both the original base weight
    and this scaling snapshot.

    Idempotent: a second call while already merged is a no-op. Re-running
    the fold would re-store ``_merged_scaling = self.scaling`` where
    ``scaling`` is already 0 (post first merge), so the re-fold no-ops
    AND the snapshot becomes 0 — the later ``unmerge_weights`` then
    restores scaling=0 and leaves the base permanently folded at 2x
    adapter strength (RIL ISS-159).
    """
    if self._merged_scaling is not None:
        return
    with torch.no_grad():
        self._merged_scaling = self.scaling
        delta_w = (self.lora_A @ self.lora_B) * self.scaling
        self.base_layer.weight.add_(delta_w.T)
        self.scaling = 0.0

unmerge_weights

unmerge_weights()

Unmerge LoRA weights from the base layer.

Restores the pre-merge base weight and re-enables the adapter path. No-op if :meth:merge_weights was never called.

源代码位于： src/llm/core/lora.py

def unmerge_weights(self) -> None:
    """Unmerge LoRA weights from the base layer.

    Restores the pre-merge base weight and re-enables the adapter
    path.  No-op if :meth:`merge_weights` was never called.
    """
    with torch.no_grad():
        merged_scaling = self._merged_scaling
        if merged_scaling is None:
            return
        delta_w = (self.lora_A @ self.lora_B) * merged_scaling
        self.base_layer.weight.sub_(delta_w.T)
        self.scaling = merged_scaling
        self._merged_scaling = None

apply_lora

apply_lora(model, rank=8, alpha=16.0, dropout=0.0, target_modules=None)

Apply LoRA to specified linear layers in a model.

参数：

名称	类型	描述	默认
model	Module	The model to adapt	必需
rank	int	LoRA rank	8
alpha	float	LoRA alpha (scaling = alpha / rank)	16.0
dropout	float	Dropout probability for LoRA path	0.0
target_modules	list[str] | None	List of module name patterns to target. If None, targets all nn.Linear layers.	None

返回：

类型	描述
Module	The model with LoRA applied (modified in-place)

源代码位于： src/llm/core/lora.py

def apply_lora(
    model: nn.Module,
    rank: int = 8,
    alpha: float = 16.0,
    dropout: float = 0.0,
    target_modules: list[str] | None = None,
) -> nn.Module:
    """
    Apply LoRA to specified linear layers in a model.

    Args:
        model: The model to adapt
        rank: LoRA rank
        alpha: LoRA alpha (scaling = alpha / rank)
        dropout: Dropout probability for LoRA path
        target_modules: List of module name patterns to target.
                        If None, targets all nn.Linear layers.

    Returns:
        The model with LoRA applied (modified in-place)
    """
    if target_modules is None:
        target_modules = []

    def should_apply(name: str) -> bool:
        if not target_modules:
            return True
        return any(pattern in name for pattern in target_modules)

    # Collect modules to replace
    replacements = []
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear) and should_apply(name):
            replacements.append((name, module))

    # Apply replacements
    for name, module in replacements:
        lora_layer = LoRALinear(module, rank=rank, alpha=alpha, dropout=dropout)

        # Navigate to parent and replace
        parts = name.split(".")
        parent = model
        for part in parts[:-1]:
            parent = getattr(parent, part)
        setattr(parent, parts[-1], lora_layer)

    return model

merge_lora

merge_lora(model)

Merge all LoRA weights into base layers for efficient inference.

参数：

名称	类型	描述	默认
model	Module	Model with LoRA layers	必需

返回：

类型	描述
Module	The model with merged weights (modified in-place)

源代码位于： src/llm/core/lora.py

def merge_lora(model: nn.Module) -> nn.Module:
    """
    Merge all LoRA weights into base layers for efficient inference.

    Args:
        model: Model with LoRA layers

    Returns:
        The model with merged weights (modified in-place)
    """
    for module in model.modules():
        if isinstance(module, LoRALinear):
            module.merge_weights()
    return model

unmerge_lora

unmerge_lora(model)

Unmerge all LoRA weights from base layers.

参数：

名称	类型	描述	默认
model	Module	Model with merged LoRA layers	必需

返回：

类型	描述
Module	The model with unmerged weights (modified in-place)

源代码位于： src/llm/core/lora.py

def unmerge_lora(model: nn.Module) -> nn.Module:
    """
    Unmerge all LoRA weights from base layers.

    Args:
        model: Model with merged LoRA layers

    Returns:
        The model with unmerged weights (modified in-place)
    """
    for module in model.modules():
        if isinstance(module, LoRALinear):
            module.unmerge_weights()
    return model

get_lora_parameters

get_lora_parameters(model)

Get only LoRA parameters for optimizer.

参数：

名称	类型	描述	默认
model	Module	Model with LoRA layers	必需

产生：

类型	描述
Parameter	LoRA parameters (lora_A and lora_B)

源代码位于： src/llm/core/lora.py

def get_lora_parameters(model: nn.Module) -> Iterator[nn.Parameter]:
    """
    Get only LoRA parameters for optimizer.

    Args:
        model: Model with LoRA layers

    Yields:
        LoRA parameters (lora_A and lora_B)
    """
    for module in model.modules():
        if isinstance(module, LoRALinear):
            yield module.lora_A
            yield module.lora_B

count_lora_parameters

count_lora_parameters(model)

Count trainable and total parameters in a model with LoRA.

参数：

名称	类型	描述	默认
model	Module	Model with LoRA layers	必需

返回：

类型	描述
tuple[int, int]	Tuple of (trainable_params, total_params)

源代码位于： src/llm/core/lora.py

def count_lora_parameters(model: nn.Module) -> tuple[int, int]:
    """
    Count trainable and total parameters in a model with LoRA.

    Args:
        model: Model with LoRA layers

    Returns:
        Tuple of (trainable_params, total_params)
    """
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    return trainable, total

disable_lora

disable_lora(model)

Disable LoRA by setting scaling to 0.

源代码位于： src/llm/core/lora.py

def disable_lora(model: nn.Module) -> None:
    """Disable LoRA by setting scaling to 0."""
    for module in model.modules():
        if isinstance(module, LoRALinear):
            module._original_scaling = module.scaling
            module.scaling = 0.0

enable_lora

enable_lora(model)

Re-enable LoRA after disabling.

源代码位于： src/llm/core/lora.py

def enable_lora(model: nn.Module) -> None:
    """Re-enable LoRA after disabling."""
    for module in model.modules():
        orig = getattr(module, "_original_scaling", None)
        if isinstance(module, LoRALinear) and orig is not None:
            module.scaling = orig

qlora

QLoRA (Quantized LoRA) Module.

Implements memory-efficient fine-tuning by combining: 1. 4-bit quantization of base model weights (NF4 format) 2. Low-rank adaptation (LoRA) in full precision

Reference: https://arxiv.org/abs/2305.14314

QLoRALinear

Bases: Module

QLoRA-adapted linear layer.

Combines 4-bit quantized base weights with full-precision LoRA adapters. Memory usage: ~4x reduction for base weights.

参数：

名称	类型	描述	默认
base_layer	Linear	Original nn.Linear layer to adapt (will be quantized)	必需
rank	int	LoRA rank (default: 8)	8
alpha	float	LoRA alpha (default: 16.0)	16.0
dropout	float	Dropout for LoRA path (default: 0.0)	0.0
block_size	int	Quantization block size (default: 64)	64

源代码位于： src/llm/core/qlora.py

class QLoRALinear(nn.Module):
    """
    QLoRA-adapted linear layer.

    Combines 4-bit quantized base weights with full-precision LoRA adapters.
    Memory usage: ~4x reduction for base weights.

    Args:
        base_layer: Original nn.Linear layer to adapt (will be quantized)
        rank: LoRA rank (default: 8)
        alpha: LoRA alpha (default: 16.0)
        dropout: Dropout for LoRA path (default: 0.0)
        block_size: Quantization block size (default: 64)
    """

    def __init__(
        self,
        base_layer: nn.Linear,
        rank: int = 8,
        alpha: float = 16.0,
        dropout: float = 0.0,
        block_size: int = 64,
    ):
        super().__init__()

        self.in_features = base_layer.in_features
        self.out_features = base_layer.out_features
        self.rank = rank
        self.alpha = alpha
        self.scaling = alpha / rank
        self.block_size = block_size

        # Store original dtype for dequantization
        self.compute_dtype = base_layer.weight.dtype
        device = base_layer.weight.device

        # Quantize base weights to 4-bit NF4
        with torch.no_grad():
            indices, scales = quantize_nf4(base_layer.weight.detach(), block_size)
            self.register_buffer("weight_indices", indices)
            self.register_buffer("weight_scales", scales)
            self.original_shape = base_layer.weight.shape

        # Handle bias
        if base_layer.bias is not None:
            self.register_buffer("bias", base_layer.bias.detach().clone())
        else:
            self.bias = None

        # LoRA adapters in full precision
        self.lora_A = nn.Parameter(torch.empty(self.in_features, rank, device=device, dtype=self.compute_dtype))
        self.lora_B = nn.Parameter(torch.empty(rank, self.out_features, device=device, dtype=self.compute_dtype))

        # Dropout for regularization
        self.lora_dropout = nn.Dropout(p=dropout) if dropout > 0 else nn.Identity()

        # Initialize LoRA weights
        self._init_lora_weights()

    def _init_lora_weights(self) -> None:
        """Initialize LoRA weights following standard LoRA initialization."""
        init_lora_weights(self.lora_A, self.lora_B)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass with dequantized base + LoRA adaptation."""
        # Dequantize base weights on-the-fly
        indices = self.weight_indices
        scales = self.weight_scales
        if not isinstance(indices, torch.Tensor) or not isinstance(scales, torch.Tensor):
            raise RuntimeError("NF4 quantized buffers were not initialized")
        weight = dequantize_nf4(
            indices,
            scales,
            self.original_shape,
            self.block_size,
            self.compute_dtype,
        )

        # Base layer output
        base_output = functional.linear(x, weight, self.bias)

        # LoRA output
        lora_output = self.lora_dropout(x) @ self.lora_A @ self.lora_B

        return base_output + lora_output * self.scaling

    @property
    def trainable_parameters(self) -> int:
        """Number of trainable LoRA parameters."""
        return self.lora_A.numel() + self.lora_B.numel()

    def extra_repr(self) -> str:
        return (
            f"in={self.in_features}, out={self.out_features}, "
            f"rank={self.rank}, alpha={self.alpha}, block_size={self.block_size}"
        )

trainable_parameters property

trainable_parameters

Number of trainable LoRA parameters.

forward

forward(x)

Forward pass with dequantized base + LoRA adaptation.

源代码位于： src/llm/core/qlora.py

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Forward pass with dequantized base + LoRA adaptation."""
    # Dequantize base weights on-the-fly
    indices = self.weight_indices
    scales = self.weight_scales
    if not isinstance(indices, torch.Tensor) or not isinstance(scales, torch.Tensor):
        raise RuntimeError("NF4 quantized buffers were not initialized")
    weight = dequantize_nf4(
        indices,
        scales,
        self.original_shape,
        self.block_size,
        self.compute_dtype,
    )

    # Base layer output
    base_output = functional.linear(x, weight, self.bias)

    # LoRA output
    lora_output = self.lora_dropout(x) @ self.lora_A @ self.lora_B

    return base_output + lora_output * self.scaling

quantize_nf4

quantize_nf4(weight, block_size=64)

Quantize weights to 4-bit NF4 format with block-wise scaling.

参数：

名称	类型	描述	默认
weight	Tensor	Weight tensor to quantize [out_features, in_features]	必需
block_size	int	Number of elements per quantization block	64

返回：

类型	描述
Tensor	Tuple of (quantized_indices, scales)
Tensor	quantized_indices: uint8 tensor with 4-bit indices packed
tuple[Tensor, Tensor]	scales: Absmax scale per block

源代码位于： src/llm/core/qlora.py

def quantize_nf4(weight: torch.Tensor, block_size: int = 64) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Quantize weights to 4-bit NF4 format with block-wise scaling.

    Args:
        weight: Weight tensor to quantize [out_features, in_features]
        block_size: Number of elements per quantization block

    Returns:
        Tuple of (quantized_indices, scales)
        - quantized_indices: uint8 tensor with 4-bit indices packed
        - scales: Absmax scale per block
    """
    weight_flat = weight.flatten().float()
    n_elements = weight_flat.numel()

    # Pad to multiple of block_size
    n_blocks = (n_elements + block_size - 1) // block_size
    padded_size = n_blocks * block_size
    if padded_size > n_elements:
        weight_flat = functional.pad(weight_flat, (0, padded_size - n_elements))

    # Reshape to blocks
    weight_blocks = weight_flat.view(n_blocks, block_size)

    # Compute absmax scale per block
    scales = weight_blocks.abs().max(dim=1, keepdim=True).values.clamp(min=1e-8)

    # Normalize to [-1, 1]
    normalized = weight_blocks / scales

    # Quantize to nearest NF4 level
    nf4 = NF4_LEVELS.to(weight.device)
    distances = (normalized.unsqueeze(-1) - nf4.view(1, 1, -1)).abs()
    indices = distances.argmin(dim=-1).to(torch.uint8)

    # Store original shape info
    return indices.flatten()[:n_elements].view(weight.shape), scales.flatten()[:n_blocks]

dequantize_nf4

dequantize_nf4(indices, scales, original_shape, block_size=64, dtype=torch.float16)

Dequantize 4-bit NF4 indices back to floating point.

参数：

名称	类型	描述	默认
indices	Tensor	Quantized indices (uint8)	必需
scales	Tensor	Block-wise scales	必需
original_shape	tuple[int, ...]	Original weight shape	必需
block_size	int	Block size used during quantization	64
dtype	dtype	Target dtype for dequantized weights	float16

返回：

类型	描述
Tensor	Dequantized weight tensor

源代码位于： src/llm/core/qlora.py

def dequantize_nf4(
    indices: torch.Tensor,
    scales: torch.Tensor,
    original_shape: tuple[int, ...],
    block_size: int = 64,
    dtype: torch.dtype = torch.float16,
) -> torch.Tensor:
    """
    Dequantize 4-bit NF4 indices back to floating point.

    Args:
        indices: Quantized indices (uint8)
        scales: Block-wise scales
        original_shape: Original weight shape
        block_size: Block size used during quantization
        dtype: Target dtype for dequantized weights

    Returns:
        Dequantized weight tensor
    """
    nf4 = NF4_LEVELS.to(indices.device)
    weight_flat = nf4[indices.flatten().long()]

    n_elements = weight_flat.numel()
    n_blocks = scales.numel()

    # Pad to match blocks
    padded_size = n_blocks * block_size
    if padded_size > n_elements:
        weight_flat = functional.pad(weight_flat, (0, padded_size - n_elements))

    # Apply scales
    weight_blocks = weight_flat.view(n_blocks, block_size)
    weight_scaled = weight_blocks * scales.view(-1, 1)

    return weight_scaled.flatten()[:n_elements].view(original_shape).to(dtype)

apply_qlora

apply_qlora(model, rank=8, alpha=16.0, dropout=0.0, block_size=64, target_modules=None)

Apply QLoRA to specified linear layers in a model.

参数：

名称	类型	描述	默认
model	Module	The model to adapt	必需
rank	int	LoRA rank	8
alpha	float	LoRA alpha (scaling = alpha / rank)	16.0
dropout	float	Dropout probability for LoRA path	0.0
block_size	int	Quantization block size	64
target_modules	list[str] | None	List of module name patterns to target. If None, targets all nn.Linear layers.	None

返回：

类型	描述
Module	The model with QLoRA applied (modified in-place)

源代码位于： src/llm/core/qlora.py

def apply_qlora(
    model: nn.Module,
    rank: int = 8,
    alpha: float = 16.0,
    dropout: float = 0.0,
    block_size: int = 64,
    target_modules: list[str] | None = None,
) -> nn.Module:
    """
    Apply QLoRA to specified linear layers in a model.

    Args:
        model: The model to adapt
        rank: LoRA rank
        alpha: LoRA alpha (scaling = alpha / rank)
        dropout: Dropout probability for LoRA path
        block_size: Quantization block size
        target_modules: List of module name patterns to target.
                        If None, targets all nn.Linear layers.

    Returns:
        The model with QLoRA applied (modified in-place)
    """
    if target_modules is None:
        target_modules = []

    def should_apply(name: str) -> bool:
        if not target_modules:
            return True
        return any(pattern in name for pattern in target_modules)

    # Collect modules to replace
    replacements = []
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear) and should_apply(name):
            replacements.append((name, module))

    # Apply replacements
    for name, module in replacements:
        qlora_layer = QLoRALinear(module, rank=rank, alpha=alpha, dropout=dropout, block_size=block_size)

        # Navigate to parent and replace
        parts = name.split(".")
        parent = model
        for part in parts[:-1]:
            parent = getattr(parent, part)
        setattr(parent, parts[-1], qlora_layer)

    return model

get_qlora_parameters

get_qlora_parameters(model)

Get only QLoRA trainable parameters for optimizer.

源代码位于： src/llm/core/qlora.py

def get_qlora_parameters(model: nn.Module):
    """Get only QLoRA trainable parameters for optimizer."""
    for module in model.modules():
        if isinstance(module, QLoRALinear):
            yield module.lora_A
            yield module.lora_B

adalora

AdaLoRA (Adaptive LoRA) module.

SVD-form parameter-efficient fine-tuning. The increment matrix is parameterized as ΔW = P · diag(λ · mask) · Q where P, Q are orthonormalized on every forward (QR decomposition) and λ is a learnable diagonal. A buffer mask lets the future pruning slice zero out low-importance components without a public-API break.

Reference: Zhang et al., 2023 — Adaptive Budget Allocation for Parameter-Efficient Fine-Tuning, arXiv:2303.10512.

AdaLoRALinear

Bases: Module

AdaLoRA-adapted linear layer with SVD-form parameterization.

The increment matrix is parameterized as ΔW = P · diag(λ · mask) · Q where:

• P: (out_features, init_rank) — left singular vectors, trainable.
• Q: (init_rank, in_features) — right singular vectors, trainable.
• λ: (init_rank,) — singular values, trainable.
• mask: (init_rank,) — binary mask registered as a buffer (not a Parameter) so the future pruning slice can write to it without breaking autograd.

Forward pass:

1. Orthonormalize P via QR → P̃.
2. Orthonormalize Q via QR on Qᵀ → Q̃ (rows of Q̃ are orthonormal — i.e. Q̃ Q̃ᵀ = I).
3. Compute ΔW = P̃ · diag(λ · mask) · Q̃ of shape (out_features, in_features).
4. Return base(x) + scaling · x · ΔWᵀ.

At initialization λ = 0 so ΔW = 0 exactly — the layer behaves identically to the base layer, matching LoRA's zero-initialized-B invariant.

参数：

名称	类型	描述	默认
base_layer	Linear	The original nn.Linear to adapt (will be frozen).	必需
init_rank	int	Initial rank budget (upper bound on the number of singular components). Must satisfy init_rank ≤ min(in_features, out_features) so QR can preserve the full column / row count.	12
target_rank	int | None	Final target rank after pruning. Stored on the layer so the future pruning slice can consult it. This foundation slice does not prune — the mask starts all-ones.	None
alpha	float	Scaling factor. Forward scales ΔW by alpha / init_rank (the same convention LoRA uses for alpha / rank).	32.0
dropout	float	Dropout probability for the LoRA path. 0.0 (default) keeps the layer fully deterministic.	0.0
orth_reg_weight	float	Default weight for the orthogonality regularization. Stored on the layer so trainers can read it without hard-coding; the loss itself is opt-in (call :meth:orth_reg_loss and add to the total loss).	0.5

源代码位于： src/llm/core/adalora.py

class AdaLoRALinear(nn.Module):
    """AdaLoRA-adapted linear layer with SVD-form parameterization.

    The increment matrix is parameterized as ``ΔW = P · diag(λ · mask)
    · Q`` where:

    - ``P: (out_features, init_rank)`` — left singular vectors, trainable.
    - ``Q: (init_rank, in_features)`` — right singular vectors, trainable.
    - ``λ: (init_rank,)`` — singular values, trainable.
    - ``mask: (init_rank,)`` — binary mask registered as a buffer (not
      a Parameter) so the future pruning slice can write to it
      without breaking autograd.

    Forward pass:

    1. Orthonormalize P via QR → ``P̃``.
    2. Orthonormalize Q via QR on ``Qᵀ`` → ``Q̃`` (rows of Q̃ are
       orthonormal — i.e. ``Q̃ Q̃ᵀ = I``).
    3. Compute ``ΔW = P̃ · diag(λ · mask) · Q̃`` of shape
       ``(out_features, in_features)``.
    4. Return ``base(x) + scaling · x · ΔWᵀ``.

    At initialization ``λ = 0`` so ``ΔW = 0`` exactly — the layer
    behaves identically to the base layer, matching LoRA's
    zero-initialized-B invariant.

    Args:
        base_layer: The original ``nn.Linear`` to adapt (will be frozen).
        init_rank: Initial rank budget (upper bound on the number of
            singular components). Must satisfy ``init_rank ≤
            min(in_features, out_features)`` so QR can preserve the
            full column / row count.
        target_rank: Final target rank after pruning. Stored on the
            layer so the future pruning slice can consult it. This
            foundation slice does **not** prune — the mask starts
            all-ones.
        alpha: Scaling factor. Forward scales ``ΔW`` by
            ``alpha / init_rank`` (the same convention LoRA uses for
            ``alpha / rank``).
        dropout: Dropout probability for the LoRA path. ``0.0`` (default)
            keeps the layer fully deterministic.
        orth_reg_weight: Default weight for the orthogonality
            regularization. Stored on the layer so trainers can read
            it without hard-coding; the loss itself is opt-in (call
            :meth:`orth_reg_loss` and add to the total loss).
    """

    _original_scaling: float | None
    # Set only by ``merge_weights`` (and read/cleared by ``unmerge_weights``);
    # declared here so static analysis sees a float, not an inferred union.
    _merged_scaling: float | None = None

    def __init__(
        self,
        base_layer: nn.Linear,
        init_rank: int = 12,
        target_rank: int | None = None,
        alpha: float = 32.0,
        dropout: float = 0.0,
        orth_reg_weight: float = 0.5,
    ):
        super().__init__()
        in_features = base_layer.in_features
        out_features = base_layer.out_features

        if init_rank <= 0:
            raise ValueError(f"init_rank must be positive, got {init_rank}")
        if init_rank > min(in_features, out_features):
            raise ValueError(
                f"init_rank ({init_rank}) must be ≤ "
                f"min(in_features={in_features}, out_features={out_features}) "
                "so QR decomposition can preserve the full column/row count."
            )
        if target_rank is not None:
            if target_rank <= 0:
                raise ValueError(f"target_rank must be positive, got {target_rank}")
            if target_rank > init_rank:
                raise ValueError(f"target_rank ({target_rank}) must be ≤ init_rank ({init_rank})")

        self.base_layer = base_layer
        self.init_rank = init_rank
        self.target_rank = target_rank if target_rank is not None else init_rank // 2
        self.alpha = alpha
        self.scaling = alpha / init_rank
        self.orth_reg_weight = orth_reg_weight

        device = base_layer.weight.device
        dtype = base_layer.weight.dtype

        # SVD-form parameters. P and Q are overwritten by their
        # orthonormalized versions on every forward, so the persistent
        # values stored in the state-dict are best understood as
        # "raw" coefficients — they're not used directly in forward.
        self.lora_P = nn.Parameter(torch.empty(out_features, init_rank, device=device, dtype=dtype))
        self.lora_Q = nn.Parameter(torch.empty(init_rank, in_features, device=device, dtype=dtype))
        self.lora_lambda = nn.Parameter(torch.empty(init_rank, device=device, dtype=dtype))

        # Pruning mask. Registered as a buffer (not Parameter) because
        # it is not trained by gradient descent — the future pruning
        # slice updates it based on importance scores. ``persistent=True``
        # so it travels through state-dict and checkpoint load/save.
        self.register_buffer(
            "mask",
            torch.ones(init_rank, device=device, dtype=dtype),
            persistent=True,
        )

        self.lora_dropout = nn.Dropout(p=dropout) if dropout > 0 else nn.Identity()

        # Init per the paper: P and Q random Gaussian, Λ zero.
        # Resulting ΔW = P · 0 · Q = 0, matching LoRA's zero-B invariant.
        nn.init.normal_(self.lora_P, mean=0.0, std=0.02)
        nn.init.normal_(self.lora_Q, mean=0.0, std=0.02)
        nn.init.zeros_(self.lora_lambda)

        # Freeze base layer.
        self.base_layer.weight.requires_grad = False
        if self.base_layer.bias is not None:
            self.base_layer.bias.requires_grad = False

    @property
    def effective_rank(self) -> int:
        """Number of currently active (masked-in) components.

        Reads through ``self.mask`` so the future pruning slice can
        drive this property by zeroing mask entries — no other
        bookkeeping needed.
        """
        mask = cast(torch.Tensor, self.mask)
        return int(mask.sum().item())

    @property
    def trainable_parameters(self) -> int:
        """Number of trainable AdaLoRA parameters (P + Q + λ)."""
        return self.lora_P.numel() + self.lora_Q.numel() + self.lora_lambda.numel()

    def _orthonormalized_P(self) -> torch.Tensor:  # noqa: N802
        """Orthonormal columns of P via QR decomposition.

        ``P`` has shape ``(out_features, init_rank)``. With
        ``init_rank ≤ out_features`` (validated in ``__init__``) the
        reduced-QR returns Q of shape ``(out_features, init_rank)``
        with orthonormal columns.
        """
        return _orthonormalize(self.lora_P)

    def _orthonormalized_Q(self) -> torch.Tensor:  # noqa: N802
        """Orthonormal rows of Q via QR on the transpose.

        We want ``Q · Qᵀ = I_{init_rank}``. Computing QR on ``Qᵀ``
        orthonormalizes its columns, which are the rows of Q in
        transposed form. We then transpose back to recover a
        ``(init_rank, in_features)`` tensor whose rows are orthonormal.
        """
        # ``Qᵀ`` has shape ``(in_features, init_rank)``; with
        # ``init_rank ≤ in_features`` (validated) reduced-QR returns
        # Q of shape ``(in_features, init_rank)`` whose columns (i.e.
        # the rows of the transposed-back tensor) are orthonormal.
        Q_ortho_T = _orthonormalize(self.lora_Q.T)  # noqa: N806
        return Q_ortho_T.T

    def _effective_increment(self) -> torch.Tensor:
        """Compute ``ΔW = P̃ · diag(λ · mask) · Q̃``.

        Returns:
            Tensor of shape ``(out_features, in_features)``.
        """
        P_ortho = self._orthonormalized_P()  # noqa: N806
        Q_ortho = self._orthonormalized_Q()  # noqa: N806
        # Broadcast (out, rank) * (rank,) → (out, rank), then matmul
        # with Q_ortho (rank, in) → (out, in).
        mask = cast(torch.Tensor, self.mask)
        scaled = self.lora_lambda * mask
        return (P_ortho * scaled) @ Q_ortho

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass: ``base(x) + scaling · x · ΔWᵀ``."""
        base_output = self.base_layer(x)
        if self.scaling == 0.0:
            return base_output
        delta_w = self._effective_increment()
        lora_output = self.lora_dropout(x) @ delta_w.T
        return base_output + lora_output * self.scaling

    def orth_reg_loss(self) -> torch.Tensor:
        """Orthogonality regularization loss.

        Per the AdaLoRA paper (arXiv:2303.10512), the loss penalises
        deviation from orthonormality of the *learned parameters*::

            ||P^T P - I||_F^2 + ||Q Q^T - I||_F^2

        Returns a scalar (sum of both terms). Trainers add this to the
        task loss (scaled by :attr:`orth_reg_weight`) to keep P and Q
        well-conditioned during training.

        The loss is computed over ``lora_P`` / ``lora_Q`` themselves —
        the matrices the optimizer updates — NOT over the QR factors used
        in forward (RIL ISS-157). Those factors are orthonormal by
        construction, so penalising them returns ~float noise no matter
        how degenerate the underlying P/Q are: a regularizer that cannot
        distinguish a rank-1 P from a perfectly orthonormal one is dead
        weight. Penalising the raw parameters gives gradient descent a
        real signal to keep them well-conditioned.
        """
        P = self.lora_P  # noqa: N806
        Q = self.lora_Q  # noqa: N806
        identity = torch.eye(self.init_rank, device=P.device, dtype=P.dtype)
        # P is (out, rank) -> PᵀP is (rank, rank). Q is (rank, in) ->
        # Q Qᵀ is (rank, rank).
        loss_p = torch.linalg.norm(P.T @ P - identity, ord="fro") ** 2
        loss_q = torch.linalg.norm(Q @ Q.T - identity, ord="fro") ** 2
        return loss_p + loss_q

    def merge_weights(self) -> None:
        """Merge ``ΔW`` into the base layer for efficient inference.

        Folds ``ΔW · scaling`` into ``base_layer.weight`` and then
        disables the adapter path (``scaling`` → 0) so the layer's
        forward — which early-returns the base when
        ``self.scaling == 0.0`` — stays equivalent to the folded base
        instead of double-counting the adapter at serve time. The
        companion :meth:`unmerge_weights` restores both the original
        base weight and this scaling snapshot for continued training.

        Idempotent: a second call while already merged is a no-op (mirrors
        the LoRA fix — see :meth:`llm.core.lora.LoRALinear.merge_weights`).
        Re-running the fold would re-store ``_merged_scaling =
        self.scaling`` where ``scaling`` is already 0 (post first merge),
        so the re-fold no-ops AND the snapshot becomes 0 — the later
        ``unmerge_weights`` then restores scaling=0 and leaves the base
        permanently folded at 2x adapter strength (RIL ISS-159).
        """
        if self._merged_scaling is not None:
            return
        with torch.no_grad():
            self._merged_scaling = self.scaling
            delta_w = self._effective_increment()
            self.base_layer.weight.add_(delta_w * self.scaling)
            self.scaling = 0.0

    def unmerge_weights(self) -> None:
        """Unmerge ``ΔW`` from the base layer.

        Restores the pre-merge base weight and re-enables the adapter
        path. No-op if :meth:`merge_weights` was never called.
        """
        with torch.no_grad():
            merged_scaling = self._merged_scaling
            if merged_scaling is None:
                return
            delta_w = self._effective_increment()
            self.base_layer.weight.sub_(delta_w * merged_scaling)
            self.scaling = merged_scaling
            self._merged_scaling = None

    def compute_importance_scores(self, gradient_ema: torch.Tensor | None = None) -> torch.Tensor:
        """Per-component importance scores, one per singular value.

        Per the AdaLoRA paper (Algorithm 1, line 7), the combined
        importance score is ``|λ_i| · |∂L/∂λ_i|``. Trainers compute
        the EMA of ``|∂L/∂λ_i|`` themselves (the optimizer owns
        gradient statistics), then pass it here.

        Args:
            gradient_ema: Optional EMA tensor of shape
                ``(init_rank,)`` holding ``|∂L/∂λ_i|`` averages.
                ``None`` falls back to magnitude-only scoring
                ``(|λ_i|)``, which is enough to rank components when
                the trainer does not track gradients.

        Returns:
            Tensor of shape ``(init_rank,)`` with one score per
            component. Components with higher scores carry more of
            the model's expressive capacity and should be kept under
            pruning.
        """
        magnitude = self.lora_lambda.abs()
        if gradient_ema is None:
            return magnitude
        if gradient_ema.shape != magnitude.shape:
            raise ValueError(
                f"gradient_ema shape {tuple(gradient_ema.shape)} does not "
                f"match lora_lambda shape {tuple(magnitude.shape)}"
            )
        return magnitude * gradient_ema.abs()

    def prune_to_rank(
        self,
        target_rank: int,
        scores: torch.Tensor | None = None,
    ) -> None:
        """Zero out mask entries for the lowest-importance components.

        Mutates :attr:`mask` in-place so that exactly ``target_rank``
        entries remain ``1.0`` and the rest are ``0.0``. The kept
        entries are the ``target_rank`` components with **highest**
        importance score (see :meth:`compute_importance_scores`).

        Args:
            target_rank: Number of components to keep. Must satisfy
                ``0 ≤ target_rank ≤ self.effective_rank`` — un-pruning
                is not supported (the dropped λ entries have been
                overwritten by the optimizer and cannot be recovered
                in-place).
            scores: Optional pre-computed importance scores of shape
                ``(init_rank,)``. When ``None``, falls back to
                :meth:`compute_importance_scores` with default
                magnitude-only scoring. Pass scores explicitly when
                wiring gradient-EMA scoring through a trainer.

        Raises:
            ValueError: if ``target_rank`` is out of range or larger
                than the current ``effective_rank``.
        """
        if target_rank < 0:
            raise ValueError(f"target_rank must be ≥ 0, got {target_rank}")
        if target_rank > self.effective_rank:
            raise ValueError(
                f"target_rank ({target_rank}) exceeds effective_rank "
                f"({self.effective_rank}); un-pruning is not supported"
            )
        if target_rank == self.effective_rank:
            # Already at the requested rank (or below); nothing to do.
            return

        if scores is None:
            scores = self.compute_importance_scores()
        # `topk` with largest=True returns the indices of the
        # highest-scoring components. Build a fresh mask from those
        # indices — this is robust to repeated calls (idempotent:
        # pruning to k twice yields the same mask, modulo ties).
        _, keep_indices = torch.topk(scores, k=target_rank, largest=True)
        mask = cast(torch.Tensor, self.mask)
        new_mask = torch.zeros_like(mask)
        new_mask[keep_indices] = 1.0
        mask.copy_(new_mask)

    def update_budget(
        self,
        current_step: int,
        tinit: int,
        tfinal: int,
    ) -> int:
        """Return the rank budget for the current training step.

        Linear schedule from ``init_rank`` at ``tinit`` to
        ``target_rank`` at ``tfinal``. Useful for periodic pruning
        during fine-tuning: train at full rank through warmup, then
        gradually reallocate the budget down to ``target_rank``.

        Args:
            current_step: The current training step (≥ 0).
            tinit: Step at and before which the budget is held at
                ``init_rank``. The first pruning-eligible step is
                ``tinit + 1``.
            tfinal: Step at and after which the budget is held at
                ``target_rank``. Must be strictly greater than
                ``tinit``.

        Returns:
            Integer rank budget to use for this step. Round to
            ``int`` to keep the mask-integer contract.

        Raises:
            ValueError: if ``current_step < 0`` or ``tinit >= tfinal``.
        """
        if current_step < 0:
            raise ValueError(f"current_step must be ≥ 0, got {current_step}")
        if tinit >= tfinal:
            raise ValueError(f"tinit ({tinit}) must be strictly less than tfinal ({tfinal})")
        if current_step <= tinit:
            return self.init_rank
        if current_step >= tfinal:
            return self.target_rank
        progress = (current_step - tinit) / (tfinal - tinit)
        return round(self.init_rank - progress * (self.init_rank - self.target_rank))

    def extra_repr(self) -> str:
        return (
            f"init_rank={self.init_rank}, target_rank={self.target_rank}, "
            f"alpha={self.alpha}, scaling={self.scaling:.4f}, "
            f"effective_rank={self.effective_rank}"
        )

effective_rank property

effective_rank

Number of currently active (masked-in) components.

Reads through self.mask so the future pruning slice can drive this property by zeroing mask entries — no other bookkeeping needed.

trainable_parameters property

trainable_parameters

Number of trainable AdaLoRA parameters (P + Q + λ).

forward

forward(x)

Forward pass: base(x) + scaling · x · ΔWᵀ.

源代码位于： src/llm/core/adalora.py

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Forward pass: ``base(x) + scaling · x · ΔWᵀ``."""
    base_output = self.base_layer(x)
    if self.scaling == 0.0:
        return base_output
    delta_w = self._effective_increment()
    lora_output = self.lora_dropout(x) @ delta_w.T
    return base_output + lora_output * self.scaling

orth_reg_loss

orth_reg_loss()

Orthogonality regularization loss.

Per the AdaLoRA paper (arXiv:2303.10512), the loss penalises deviation from orthonormality of the learned parameters::

||P^T P - I||_F^2 + ||Q Q^T - I||_F^2

Returns a scalar (sum of both terms). Trainers add this to the task loss (scaled by :attr:orth_reg_weight) to keep P and Q well-conditioned during training.

The loss is computed over lora_P / lora_Q themselves — the matrices the optimizer updates — NOT over the QR factors used in forward (RIL ISS-157). Those factors are orthonormal by construction, so penalising them returns ~float noise no matter how degenerate the underlying P/Q are: a regularizer that cannot distinguish a rank-1 P from a perfectly orthonormal one is dead weight. Penalising the raw parameters gives gradient descent a real signal to keep them well-conditioned.

源代码位于： src/llm/core/adalora.py

def orth_reg_loss(self) -> torch.Tensor:
    """Orthogonality regularization loss.

    Per the AdaLoRA paper (arXiv:2303.10512), the loss penalises
    deviation from orthonormality of the *learned parameters*::

        ||P^T P - I||_F^2 + ||Q Q^T - I||_F^2

    Returns a scalar (sum of both terms). Trainers add this to the
    task loss (scaled by :attr:`orth_reg_weight`) to keep P and Q
    well-conditioned during training.

    The loss is computed over ``lora_P`` / ``lora_Q`` themselves —
    the matrices the optimizer updates — NOT over the QR factors used
    in forward (RIL ISS-157). Those factors are orthonormal by
    construction, so penalising them returns ~float noise no matter
    how degenerate the underlying P/Q are: a regularizer that cannot
    distinguish a rank-1 P from a perfectly orthonormal one is dead
    weight. Penalising the raw parameters gives gradient descent a
    real signal to keep them well-conditioned.
    """
    P = self.lora_P  # noqa: N806
    Q = self.lora_Q  # noqa: N806
    identity = torch.eye(self.init_rank, device=P.device, dtype=P.dtype)
    # P is (out, rank) -> PᵀP is (rank, rank). Q is (rank, in) ->
    # Q Qᵀ is (rank, rank).
    loss_p = torch.linalg.norm(P.T @ P - identity, ord="fro") ** 2
    loss_q = torch.linalg.norm(Q @ Q.T - identity, ord="fro") ** 2
    return loss_p + loss_q

merge_weights

merge_weights()

Merge ΔW into the base layer for efficient inference.

Folds ΔW · scaling into base_layer.weight and then disables the adapter path (scaling → 0) so the layer's forward — which early-returns the base when self.scaling == 0.0 — stays equivalent to the folded base instead of double-counting the adapter at serve time. The companion :meth:unmerge_weights restores both the original base weight and this scaling snapshot for continued training.

Idempotent: a second call while already merged is a no-op (mirrors the LoRA fix — see :meth:llm.core.lora.LoRALinear.merge_weights). Re-running the fold would re-store _merged_scaling = self.scaling where scaling is already 0 (post first merge), so the re-fold no-ops AND the snapshot becomes 0 — the later unmerge_weights then restores scaling=0 and leaves the base permanently folded at 2x adapter strength (RIL ISS-159).

源代码位于： src/llm/core/adalora.py

def merge_weights(self) -> None:
    """Merge ``ΔW`` into the base layer for efficient inference.

    Folds ``ΔW · scaling`` into ``base_layer.weight`` and then
    disables the adapter path (``scaling`` → 0) so the layer's
    forward — which early-returns the base when
    ``self.scaling == 0.0`` — stays equivalent to the folded base
    instead of double-counting the adapter at serve time. The
    companion :meth:`unmerge_weights` restores both the original
    base weight and this scaling snapshot for continued training.

    Idempotent: a second call while already merged is a no-op (mirrors
    the LoRA fix — see :meth:`llm.core.lora.LoRALinear.merge_weights`).
    Re-running the fold would re-store ``_merged_scaling =
    self.scaling`` where ``scaling`` is already 0 (post first merge),
    so the re-fold no-ops AND the snapshot becomes 0 — the later
    ``unmerge_weights`` then restores scaling=0 and leaves the base
    permanently folded at 2x adapter strength (RIL ISS-159).
    """
    if self._merged_scaling is not None:
        return
    with torch.no_grad():
        self._merged_scaling = self.scaling
        delta_w = self._effective_increment()
        self.base_layer.weight.add_(delta_w * self.scaling)
        self.scaling = 0.0

unmerge_weights

unmerge_weights()

Unmerge ΔW from the base layer.

Restores the pre-merge base weight and re-enables the adapter path. No-op if :meth:merge_weights was never called.

源代码位于： src/llm/core/adalora.py

def unmerge_weights(self) -> None:
    """Unmerge ``ΔW`` from the base layer.

    Restores the pre-merge base weight and re-enables the adapter
    path. No-op if :meth:`merge_weights` was never called.
    """
    with torch.no_grad():
        merged_scaling = self._merged_scaling
        if merged_scaling is None:
            return
        delta_w = self._effective_increment()
        self.base_layer.weight.sub_(delta_w * merged_scaling)
        self.scaling = merged_scaling
        self._merged_scaling = None

compute_importance_scores

compute_importance_scores(gradient_ema=None)

Per-component importance scores, one per singular value.

Per the AdaLoRA paper (Algorithm 1, line 7), the combined importance score is |λ_i| · |∂L/∂λ_i|. Trainers compute the EMA of |∂L/∂λ_i| themselves (the optimizer owns gradient statistics), then pass it here.

参数：

名称	类型	描述	默认
gradient_ema	Tensor | None	Optional EMA tensor of shape (init_rank,) holding |∂L/∂λ_i| averages. None falls back to magnitude-only scoring (|λ_i|), which is enough to rank components when the trainer does not track gradients.	None

返回：

类型	描述
Tensor	Tensor of shape (init_rank,) with one score per
Tensor	component. Components with higher scores carry more of
Tensor	the model's expressive capacity and should be kept under
Tensor	pruning.

源代码位于： src/llm/core/adalora.py

def compute_importance_scores(self, gradient_ema: torch.Tensor | None = None) -> torch.Tensor:
    """Per-component importance scores, one per singular value.

    Per the AdaLoRA paper (Algorithm 1, line 7), the combined
    importance score is ``|λ_i| · |∂L/∂λ_i|``. Trainers compute
    the EMA of ``|∂L/∂λ_i|`` themselves (the optimizer owns
    gradient statistics), then pass it here.

    Args:
        gradient_ema: Optional EMA tensor of shape
            ``(init_rank,)`` holding ``|∂L/∂λ_i|`` averages.
            ``None`` falls back to magnitude-only scoring
            ``(|λ_i|)``, which is enough to rank components when
            the trainer does not track gradients.

    Returns:
        Tensor of shape ``(init_rank,)`` with one score per
        component. Components with higher scores carry more of
        the model's expressive capacity and should be kept under
        pruning.
    """
    magnitude = self.lora_lambda.abs()
    if gradient_ema is None:
        return magnitude
    if gradient_ema.shape != magnitude.shape:
        raise ValueError(
            f"gradient_ema shape {tuple(gradient_ema.shape)} does not "
            f"match lora_lambda shape {tuple(magnitude.shape)}"
        )
    return magnitude * gradient_ema.abs()

prune_to_rank

prune_to_rank(target_rank, scores=None)

Zero out mask entries for the lowest-importance components.

Mutates :attr:mask in-place so that exactly target_rank entries remain 1.0 and the rest are 0.0. The kept entries are the target_rank components with highest importance score (see :meth:compute_importance_scores).

参数：

名称	类型	描述	默认
target_rank	int	Number of components to keep. Must satisfy 0 ≤ target_rank ≤ self.effective_rank — un-pruning is not supported (the dropped λ entries have been overwritten by the optimizer and cannot be recovered in-place).	必需
scores	Tensor | None	Optional pre-computed importance scores of shape (init_rank,). When None, falls back to :meth:compute_importance_scores with default magnitude-only scoring. Pass scores explicitly when wiring gradient-EMA scoring through a trainer.	None

引发：

类型	描述
ValueError	if target_rank is out of range or larger than the current effective_rank.

源代码位于： src/llm/core/adalora.py

def prune_to_rank(
    self,
    target_rank: int,
    scores: torch.Tensor | None = None,
) -> None:
    """Zero out mask entries for the lowest-importance components.

    Mutates :attr:`mask` in-place so that exactly ``target_rank``
    entries remain ``1.0`` and the rest are ``0.0``. The kept
    entries are the ``target_rank`` components with **highest**
    importance score (see :meth:`compute_importance_scores`).

    Args:
        target_rank: Number of components to keep. Must satisfy
            ``0 ≤ target_rank ≤ self.effective_rank`` — un-pruning
            is not supported (the dropped λ entries have been
            overwritten by the optimizer and cannot be recovered
            in-place).
        scores: Optional pre-computed importance scores of shape
            ``(init_rank,)``. When ``None``, falls back to
            :meth:`compute_importance_scores` with default
            magnitude-only scoring. Pass scores explicitly when
            wiring gradient-EMA scoring through a trainer.

    Raises:
        ValueError: if ``target_rank`` is out of range or larger
            than the current ``effective_rank``.
    """
    if target_rank < 0:
        raise ValueError(f"target_rank must be ≥ 0, got {target_rank}")
    if target_rank > self.effective_rank:
        raise ValueError(
            f"target_rank ({target_rank}) exceeds effective_rank "
            f"({self.effective_rank}); un-pruning is not supported"
        )
    if target_rank == self.effective_rank:
        # Already at the requested rank (or below); nothing to do.
        return

    if scores is None:
        scores = self.compute_importance_scores()
    # `topk` with largest=True returns the indices of the
    # highest-scoring components. Build a fresh mask from those
    # indices — this is robust to repeated calls (idempotent:
    # pruning to k twice yields the same mask, modulo ties).
    _, keep_indices = torch.topk(scores, k=target_rank, largest=True)
    mask = cast(torch.Tensor, self.mask)
    new_mask = torch.zeros_like(mask)
    new_mask[keep_indices] = 1.0
    mask.copy_(new_mask)

update_budget

update_budget(current_step, tinit, tfinal)

Return the rank budget for the current training step.

Linear schedule from init_rank at tinit to target_rank at tfinal. Useful for periodic pruning during fine-tuning: train at full rank through warmup, then gradually reallocate the budget down to target_rank.

参数：

名称	类型	描述	默认
current_step	int	The current training step (≥ 0).	必需
tinit	int	Step at and before which the budget is held at init_rank. The first pruning-eligible step is tinit + 1.	必需
tfinal	int	Step at and after which the budget is held at target_rank. Must be strictly greater than tinit.	必需

返回：

类型	描述
int	Integer rank budget to use for this step. Round to
int	int to keep the mask-integer contract.

引发：

类型	描述
ValueError	if current_step < 0 or tinit >= tfinal.

源代码位于： src/llm/core/adalora.py

def update_budget(
    self,
    current_step: int,
    tinit: int,
    tfinal: int,
) -> int:
    """Return the rank budget for the current training step.

    Linear schedule from ``init_rank`` at ``tinit`` to
    ``target_rank`` at ``tfinal``. Useful for periodic pruning
    during fine-tuning: train at full rank through warmup, then
    gradually reallocate the budget down to ``target_rank``.

    Args:
        current_step: The current training step (≥ 0).
        tinit: Step at and before which the budget is held at
            ``init_rank``. The first pruning-eligible step is
            ``tinit + 1``.
        tfinal: Step at and after which the budget is held at
            ``target_rank``. Must be strictly greater than
            ``tinit``.

    Returns:
        Integer rank budget to use for this step. Round to
        ``int`` to keep the mask-integer contract.

    Raises:
        ValueError: if ``current_step < 0`` or ``tinit >= tfinal``.
    """
    if current_step < 0:
        raise ValueError(f"current_step must be ≥ 0, got {current_step}")
    if tinit >= tfinal:
        raise ValueError(f"tinit ({tinit}) must be strictly less than tfinal ({tfinal})")
    if current_step <= tinit:
        return self.init_rank
    if current_step >= tfinal:
        return self.target_rank
    progress = (current_step - tinit) / (tfinal - tinit)
    return round(self.init_rank - progress * (self.init_rank - self.target_rank))

AdaLoRAGradientEMA

Per-layer EMA of |∂L/∂λ| for AdaLoRA's importance scoring.

Implements the EMA half of AdaLoRA Algorithm 1 (Zhang et al. 2023, page 4)::

I_avg_i <- alpha * I_avg_i + (1 - alpha) * |dL/dlambda_i|

The tracker is constructed against a model that already has AdaLoRALinear layers in place. After each backward pass, the trainer calls :meth:update to fold the current gradient into the EMA. The result is consumed by :func:prune_adalora (via :meth:as_dict) to weight components by their combined |λ| · |∂L/∂λ| score during pruning.

State is checkpointable through :meth:state_dict / :meth:load_state_dict, both keyed by the layer's qualified name ("0", "layer1.attn", ...). id(layer) is not stable across pickle roundtrips, so the checkpoint key must be a structural path.

参数：

名称	类型	描述	默认
model	Module	The model whose AdaLoRALinear layers will be tracked. Walked via :meth:nn.Module.named_modules, so the DDP/FSDP unwrap path is just model.modules().	必需
alpha	float	EMA smoothing factor (the weight on the previous EMA). alpha=0.95 matches the paper's recommendation. Must satisfy 0 < alpha < 1.	0.95

引发：

类型	描述
ValueError	if alpha is outside (0, 1).

源代码位于： src/llm/core/adalora.py

class AdaLoRAGradientEMA:
    """Per-layer EMA of ``|∂L/∂λ|`` for AdaLoRA's importance scoring.

    Implements the EMA half of AdaLoRA Algorithm 1
    (Zhang et al. 2023, page 4)::

        I_avg_i <- alpha * I_avg_i + (1 - alpha) * |dL/dlambda_i|

    The tracker is constructed against a model that already has
    ``AdaLoRALinear`` layers in place. After each backward pass, the
    trainer calls :meth:`update` to fold the current gradient into the
    EMA. The result is consumed by :func:`prune_adalora` (via
    :meth:`as_dict`) to weight components by their combined
    ``|λ| · |∂L/∂λ|`` score during pruning.

    State is checkpointable through :meth:`state_dict` /
    :meth:`load_state_dict`, both keyed by the layer's **qualified
    name** (``"0"``, ``"layer1.attn"``, ...). ``id(layer)`` is *not*
    stable across pickle roundtrips, so the checkpoint key must be a
    structural path.

    Args:
        model: The model whose ``AdaLoRALinear`` layers will be
            tracked. Walked via :meth:`nn.Module.named_modules`, so the
            DDP/FSDP unwrap path is just ``model.modules()``.
        alpha: EMA smoothing factor (the weight on the *previous* EMA).
            ``alpha=0.95`` matches the paper's recommendation. Must
            satisfy ``0 < alpha < 1``.

    Raises:
        ValueError: if ``alpha`` is outside ``(0, 1)``.
    """

    def __init__(self, model: nn.Module, alpha: float = 0.95):
        if not (0.0 < alpha < 1.0):
            raise ValueError(f"alpha must be in (0, 1), got {alpha}")
        self.alpha = alpha
        self._current_step_weight = 1.0 - alpha

        # Walk named_modules so we capture the qualified name once at
        # construction. The same walk also picks up the AdaLoRALinear
        # in case the user is passing a sub-module (e.g. just the
        # decoder body); every AdaLoRALinear reachable from ``model``
        # gets its own EMA tensor.
        self._layers: dict[str, AdaLoRALinear] = {}
        self._emas: dict[str, torch.Tensor] = {}
        for name, module in model.named_modules():
            if isinstance(module, AdaLoRALinear):
                self._layers[name] = module
                # Match the layer's dtype and device so a subsequent
                # ``ema + grad_abs`` does not silently promote / move.
                params = list(module.parameters())
                ref = params[0] if params else module.mask
                self._emas[name] = torch.zeros(module.init_rank, dtype=ref.dtype, device=ref.device)

    def update(self) -> None:
        """Fold ``|∂L/∂λ|`` from each layer's last backward into the EMA.

        Layers whose ``lora_lambda.grad`` is ``None`` (frozen, or no
        path through them this step) are left untouched — only the
        components training actually drove get smoothed in.
        """
        for name, layer in self._layers.items():
            grad = layer.lora_lambda.grad
            if grad is None:
                continue
            ema = self._emas[name]
            ema.mul_(self.alpha).add_(grad.abs(), alpha=self._current_step_weight)

    def as_dict(self) -> dict[int, torch.Tensor]:
        """Return ``{id(layer): ema_tensor}`` for :func:`prune_adalora`.

        The trainer passes this directly as ``gradient_emas=`` to
        :func:`prune_adalora`. ``id()`` is fine here because the
        tracker and the prune call share the same Python process; it
        is **only** the checkpoint roundtrip that needs a stable key.
        """
        return {id(layer): self._emas[name] for name, layer in self._layers.items()}

    def state_dict(self) -> dict[str, torch.Tensor]:
        """Return a serializable snapshot keyed by qualified name.

        Tensors are detached so a subsequent ``load_state_dict`` on a
        fresh tracker doesn't keep a reference to the autograd graph.
        """
        return {name: ema.detach().clone() for name, ema in self._emas.items()}

    def load_state_dict(self, state: dict[str, torch.Tensor] | None) -> None:
        """Restore EMA tensors from a :meth:`state_dict` snapshot.

        Unknown keys (e.g. from a stale checkpoint where the model has
        since been pruned of a layer) are silently ignored. ``None``
        is a no-op so a fresh checkpoint doesn't crash the trainer.
        """
        if not state:
            return
        for name, tensor in state.items():
            target = self._emas.get(name)
            if target is None:
                # Stale key — layer was renamed or removed. Skip.
                continue
            # Move + cast defensively so a checkpoint saved on a
            # different device/dtype still loads without surprises.
            target.copy_(tensor.to(device=target.device, dtype=target.dtype))

update

update()

Fold |∂L/∂λ| from each layer's last backward into the EMA.

Layers whose lora_lambda.grad is None (frozen, or no path through them this step) are left untouched — only the components training actually drove get smoothed in.

源代码位于： src/llm/core/adalora.py

def update(self) -> None:
    """Fold ``|∂L/∂λ|`` from each layer's last backward into the EMA.

    Layers whose ``lora_lambda.grad`` is ``None`` (frozen, or no
    path through them this step) are left untouched — only the
    components training actually drove get smoothed in.
    """
    for name, layer in self._layers.items():
        grad = layer.lora_lambda.grad
        if grad is None:
            continue
        ema = self._emas[name]
        ema.mul_(self.alpha).add_(grad.abs(), alpha=self._current_step_weight)

as_dict

as_dict()

Return {id(layer): ema_tensor} for :func:prune_adalora.

The trainer passes this directly as gradient_emas= to :func:prune_adalora. id() is fine here because the tracker and the prune call share the same Python process; it is only the checkpoint roundtrip that needs a stable key.

源代码位于： src/llm/core/adalora.py

def as_dict(self) -> dict[int, torch.Tensor]:
    """Return ``{id(layer): ema_tensor}`` for :func:`prune_adalora`.

    The trainer passes this directly as ``gradient_emas=`` to
    :func:`prune_adalora`. ``id()`` is fine here because the
    tracker and the prune call share the same Python process; it
    is **only** the checkpoint roundtrip that needs a stable key.
    """
    return {id(layer): self._emas[name] for name, layer in self._layers.items()}

state_dict

state_dict()

Return a serializable snapshot keyed by qualified name.

Tensors are detached so a subsequent load_state_dict on a fresh tracker doesn't keep a reference to the autograd graph.

源代码位于： src/llm/core/adalora.py

def state_dict(self) -> dict[str, torch.Tensor]:
    """Return a serializable snapshot keyed by qualified name.

    Tensors are detached so a subsequent ``load_state_dict`` on a
    fresh tracker doesn't keep a reference to the autograd graph.
    """
    return {name: ema.detach().clone() for name, ema in self._emas.items()}

load_state_dict

load_state_dict(state)

Restore EMA tensors from a :meth:state_dict snapshot.

Unknown keys (e.g. from a stale checkpoint where the model has since been pruned of a layer) are silently ignored. None is a no-op so a fresh checkpoint doesn't crash the trainer.

源代码位于： src/llm/core/adalora.py

def load_state_dict(self, state: dict[str, torch.Tensor] | None) -> None:
    """Restore EMA tensors from a :meth:`state_dict` snapshot.

    Unknown keys (e.g. from a stale checkpoint where the model has
    since been pruned of a layer) are silently ignored. ``None``
    is a no-op so a fresh checkpoint doesn't crash the trainer.
    """
    if not state:
        return
    for name, tensor in state.items():
        target = self._emas.get(name)
        if target is None:
            # Stale key — layer was renamed or removed. Skip.
            continue
        # Move + cast defensively so a checkpoint saved on a
        # different device/dtype still loads without surprises.
        target.copy_(tensor.to(device=target.device, dtype=target.dtype))

apply_adalora

apply_adalora(model, init_rank=12, target_rank=None, alpha=32.0, dropout=0.0, target_modules=None, orth_reg_weight=0.5)

Apply AdaLoRA to specified linear layers in a model.

Mirrors :func:llm.core.lora.apply_lora so swapping LoRA → AdaLoRA in user code is a one-import change. target_rank is stored on each layer but does not trigger pruning in this foundation slice — that lands in the follow-up.

参数：

名称	类型	描述	默认
model	Module	The model to adapt. Modified in-place.	必需
init_rank	int	Initial rank budget.	12
target_rank	int | None	Final target rank after pruning (deferred).	None
alpha	float	Scaling factor (alpha / init_rank).	32.0
dropout	float	Dropout probability for the AdaLoRA path.	0.0
target_modules	list[str] | None	List of module-name substring patterns. If None, every nn.Linear is replaced.	None
orth_reg_weight	float	Default weight for orthogonality regularization.	0.5

返回：

类型	描述
Module	The same model, modified in-place.

源代码位于： src/llm/core/adalora.py

def apply_adalora(
    model: nn.Module,
    init_rank: int = 12,
    target_rank: int | None = None,
    alpha: float = 32.0,
    dropout: float = 0.0,
    target_modules: list[str] | None = None,
    orth_reg_weight: float = 0.5,
) -> nn.Module:
    """Apply AdaLoRA to specified linear layers in a model.

    Mirrors :func:`llm.core.lora.apply_lora` so swapping LoRA → AdaLoRA
    in user code is a one-import change. ``target_rank`` is stored on
    each layer but does **not** trigger pruning in this foundation
    slice — that lands in the follow-up.

    Args:
        model: The model to adapt. Modified in-place.
        init_rank: Initial rank budget.
        target_rank: Final target rank after pruning (deferred).
        alpha: Scaling factor (``alpha / init_rank``).
        dropout: Dropout probability for the AdaLoRA path.
        target_modules: List of module-name substring patterns. If
            ``None``, every ``nn.Linear`` is replaced.
        orth_reg_weight: Default weight for orthogonality regularization.

    Returns:
        The same model, modified in-place.
    """
    if target_modules is None:
        target_modules = []

    def should_apply(name: str) -> bool:
        if not target_modules:
            return True
        return any(pattern in name for pattern in target_modules)

    replacements: list[tuple[str, nn.Linear]] = []
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear) and should_apply(name):
            replacements.append((name, module))

    for name, module in replacements:
        adalora_layer = AdaLoRALinear(
            module,
            init_rank=init_rank,
            target_rank=target_rank,
            alpha=alpha,
            dropout=dropout,
            orth_reg_weight=orth_reg_weight,
        )
        parts = name.split(".")
        parent = model
        for part in parts[:-1]:
            parent = getattr(parent, part)
        setattr(parent, parts[-1], adalora_layer)

    return model

merge_adalora

merge_adalora(model)

Merge all AdaLoRA deltas into the corresponding base layers.

源代码位于： src/llm/core/adalora.py

def merge_adalora(model: nn.Module) -> nn.Module:
    """Merge all AdaLoRA deltas into the corresponding base layers."""
    for module in model.modules():
        if isinstance(module, AdaLoRALinear):
            module.merge_weights()
    return model

unmerge_adalora

unmerge_adalora(model)

Unmerge all AdaLoRA deltas from the corresponding base layers.

源代码位于： src/llm/core/adalora.py

def unmerge_adalora(model: nn.Module) -> nn.Module:
    """Unmerge all AdaLoRA deltas from the corresponding base layers."""
    for module in model.modules():
        if isinstance(module, AdaLoRALinear):
            module.unmerge_weights()
    return model

get_adalora_parameters

get_adalora_parameters(model)

Yield every AdaLoRA P, Q, and λ parameter in the model.

Trainers pass this to the optimizer so only the AdaLoRA path is updated — the base weights stay frozen.

源代码位于： src/llm/core/adalora.py

def get_adalora_parameters(model: nn.Module) -> Iterator[nn.Parameter]:
    """Yield every AdaLoRA ``P``, ``Q``, and ``λ`` parameter in the model.

    Trainers pass this to the optimizer so only the AdaLoRA path is
    updated — the base weights stay frozen.
    """
    for module in model.modules():
        if isinstance(module, AdaLoRALinear):
            yield module.lora_P
            yield module.lora_Q
            yield module.lora_lambda

count_adalora_parameters

count_adalora_parameters(model)

Return (trainable_params, total_params) for a model with AdaLoRA.

Trainable count follows whatever requires_grad is set on every parameter — the base layer weights are frozen at construction, so trainable_params will equal the sum of all AdaLoRA trainable_parameters plus any other unfrozen parameters the caller has chosen to add (e.g. norms).

源代码位于： src/llm/core/adalora.py

def count_adalora_parameters(model: nn.Module) -> tuple[int, int]:
    """Return ``(trainable_params, total_params)`` for a model with AdaLoRA.

    Trainable count follows whatever ``requires_grad`` is set on every
    parameter — the base layer weights are frozen at construction, so
    ``trainable_params`` will equal the sum of all AdaLoRA
    ``trainable_parameters`` plus any other unfrozen parameters the
    caller has chosen to add (e.g. norms).
    """
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    return trainable, total

disable_adalora

disable_adalora(model)

Disable AdaLoRA by setting scaling to 0 (model falls back to base).

源代码位于： src/llm/core/adalora.py

def disable_adalora(model: nn.Module) -> None:
    """Disable AdaLoRA by setting scaling to 0 (model falls back to base)."""
    for module in model.modules():
        if isinstance(module, AdaLoRALinear):
            module._original_scaling = module.scaling  # type: ignore[attr-defined]
            module.scaling = 0.0

enable_adalora

enable_adalora(model)

Re-enable AdaLoRA after :func:disable_adalora.

源代码位于： src/llm/core/adalora.py

def enable_adalora(model: nn.Module) -> None:
    """Re-enable AdaLoRA after :func:`disable_adalora`."""
    for module in model.modules():
        if isinstance(module, AdaLoRALinear):
            original = getattr(module, "_original_scaling", None)
            if original is not None:
                module.scaling = original

prune_adalora

prune_adalora(model, target_rank=None, schedule=None, current_step=None, gradient_emas=None)

Walk every AdaLoRALinear in model and prune to a target rank.

Two calling modes:

1. Explicit target rank::

   prune_adalora(model, target_rank=8)

Every AdaLoRALinear layer is pruned to target_rank (subject to the per-layer effective_rank upper bound).

1. Budget schedule::

   prune_adalora(model, schedule=(tinit, tfinal), current_step=step)

Each layer is pruned to the rank that layer.update_budget(step, tinit, tfinal) returns — i.e. the budget is re-evaluated per call, so a training loop can invoke this helper periodically and the rank shrinks over time.

参数：

名称	类型	描述	默认
model	Module	Model containing one or more AdaLoRALinear layers.	必需
target_rank	int | None	Explicit rank to prune every layer to. Mutually exclusive with schedule.	None
schedule	tuple[int, int] | None	(tinit, tfinal) tuple for a linear rank-budget schedule. Must be paired with current_step.	None
current_step	int | None	Current training step, used only when schedule is given.	None
gradient_emas	dict[int, Tensor] | None	Optional dict mapping id(layer) to a gradient-EMA tensor of shape (init_rank,) for that layer's λ. Passed through to :meth:AdaLoRALinear.compute_importance_scores so the trainer can supply |∂L/∂λ_i| averages.	None

引发：

类型	描述
ValueError	if neither target_rank nor schedule is provided, or if schedule is given without current_step.

源代码位于： src/llm/core/adalora.py

def prune_adalora(
    model: nn.Module,
    target_rank: int | None = None,
    schedule: tuple[int, int] | None = None,
    current_step: int | None = None,
    gradient_emas: dict[int, torch.Tensor] | None = None,
) -> None:
    """Walk every ``AdaLoRALinear`` in ``model`` and prune to a target rank.

    Two calling modes:

    1. **Explicit target rank**::

        prune_adalora(model, target_rank=8)

       Every AdaLoRALinear layer is pruned to ``target_rank``
       (subject to the per-layer ``effective_rank`` upper bound).

    2. **Budget schedule**::

        prune_adalora(model, schedule=(tinit, tfinal), current_step=step)

       Each layer is pruned to the rank that
       ``layer.update_budget(step, tinit, tfinal)`` returns — i.e.
       the budget is re-evaluated per call, so a training loop can
       invoke this helper periodically and the rank shrinks over time.

    Args:
        model: Model containing one or more ``AdaLoRALinear`` layers.
        target_rank: Explicit rank to prune every layer to. Mutually
            exclusive with ``schedule``.
        schedule: ``(tinit, tfinal)`` tuple for a linear rank-budget
            schedule. Must be paired with ``current_step``.
        current_step: Current training step, used only when
            ``schedule`` is given.
        gradient_emas: Optional dict mapping ``id(layer)`` to a
            gradient-EMA tensor of shape ``(init_rank,)`` for that
            layer's λ. Passed through to
            :meth:`AdaLoRALinear.compute_importance_scores` so the
            trainer can supply ``|∂L/∂λ_i|`` averages.

    Raises:
        ValueError: if neither ``target_rank`` nor ``schedule`` is
            provided, or if ``schedule`` is given without
            ``current_step``.
    """
    if (target_rank is None) == (schedule is None):
        raise ValueError(
            "prune_adalora requires exactly one of target_rank or "
            "schedule=(tinit, tfinal); got "
            f"target_rank={target_rank!r}, schedule={schedule!r}"
        )

    layers = [m for m in model.modules() if isinstance(m, AdaLoRALinear)]
    if not layers:
        return

    if target_rank is not None:
        for layer in layers:
            rank = min(target_rank, layer.effective_rank)
            scores = (
                layer.compute_importance_scores(gradient_emas.get(id(layer))) if gradient_emas is not None else None
            )
            layer.prune_to_rank(rank, scores=scores)
        return

    # Schedule branch. The mutually-exclusive check at the top of
    # this function guarantees ``schedule`` is not None here.
    assert schedule is not None  # noqa: S101
    if current_step is None:
        raise ValueError("prune_adalora with schedule=(tinit, tfinal) requires current_step to be provided")
    tinit, tfinal = schedule
    for layer in layers:
        rank = layer.update_budget(current_step, tinit, tfinal)
        scores = layer.compute_importance_scores(gradient_emas.get(id(layer))) if gradient_emas is not None else None
        layer.prune_to_rank(rank, scores=scores)

bitfit

BitFit (Bias-Term Fine-Tuning).

The simplest parameter-efficient fine-tuning method: train only the bias parameters, freeze everything else. Unlike LoRA / AdaLoRA / IA³, BitFit adds no new parameters and wraps no modules — it just toggles requires_grad on every bias in the model.

Per the paper (Ben-Zaken et al. 2021), all bias parameters are trainable: attention Q/K/V/O projection biases, FFN intermediate / output projection biases, and LayerNorm / RMSNorm biases. The user can opt to filter by target_modules (substring match on the parameter's qualified name) to bias-select a subset — e.g. only attention biases — but the default is to train every bias.

The helper API mirrors LoRA / AdaLoRA / IA³ so swapping PEFT methods in user code is a one-import change:

from llm.core.bitfit import apply_bitfit, get_bitfit_parameters

apply_bitfit(model)
optimizer = torch.optim.AdamW(get_bitfit_parameters(model), lr=1e-3)

Per-model cost: O(num_biases) trainable params — typically <0.1% of total parameters. BitFit is the lightest possible PEFT method: no math, no wrappers, no scheduler.

Reference: Ben-Zaken et al., 2021 — BitFit: Simple Parameter-efficient Fine-tuning for Transformer-based Masked Language-models.

apply_bitfit

apply_bitfit(model, target_modules=None)

Freeze every parameter, then enable gradients on every bias.

参数：

名称	类型	描述	默认
model	Module	The model to adapt (modified in-place).	必需
target_modules	list[str] | None	Optional list of module-name substring patterns. A bias is trainable only if its qualified name (e.g. "layers.0.attn.q_proj.bias") contains at least one of the patterns. None (default) → every bias is trainable.	None

返回：

类型	描述
Module	The model with BitFit applied (modified in-place).

Note

BitFit saves the original requires_grad state on the model under _bitfit_original_requires_grad so :func:unapply_bitfit can restore it. Calling :func:apply_bitfit twice without an intervening :func:unapply_bitfit re-saves on the second call (so the snapshot always reflects the pre-BitFit state).

源代码位于： src/llm/core/bitfit.py

def apply_bitfit(
    model: nn.Module,
    target_modules: list[str] | None = None,
) -> nn.Module:
    """Freeze every parameter, then enable gradients on every bias.

    Args:
        model: The model to adapt (modified in-place).
        target_modules: Optional list of module-name substring patterns.
            A bias is trainable only if its qualified name (e.g.
            ``"layers.0.attn.q_proj.bias"``) contains at least one of
            the patterns. ``None`` (default) → every bias is trainable.

    Returns:
        The model with BitFit applied (modified in-place).

    Note:
        BitFit saves the original ``requires_grad`` state on the model
        under ``_bitfit_original_requires_grad`` so :func:`unapply_bitfit`
        can restore it. Calling :func:`apply_bitfit` twice without an
        intervening :func:`unapply_bitfit` re-saves on the second call
        (so the snapshot always reflects the pre-BitFit state).
    """
    if target_modules is None:
        target_modules = []

    # Snapshot the original requires_grad state so unapply_bitfit can
    # restore it. Save BEFORE toggling anything — that way repeated
    # calls of apply_bitfit converge to the same final state.
    object.__setattr__(
        model,
        "_bitfit_original_requires_grad",
        {name: p.requires_grad for name, p in model.named_parameters()},
    )

    # Freeze every parameter.
    for p in model.parameters():
        p.requires_grad = False

    # Enable gradients on biases whose qualified name matches.
    # We check for ``.bias`` SUFFIX (not just substring) — substring
    # would falsely match module names like ``fc_with_bias`` whose
    # ``weight`` parameter is NOT a bias. The qualified name of a
    # bias parameter always ends in ``.bias`` (or equals ``bias`` at
    # the top level).
    for name, p in model.named_parameters():
        if not (name == "bias" or name.endswith(".bias")):
            continue
        if target_modules and not any(pattern in name for pattern in target_modules):
            continue
        p.requires_grad = True

    return model

unapply_bitfit

unapply_bitfit(model)

Restore the pre-BitFit requires_grad state.

Reverses :func:apply_bitfit — every parameter is set back to whatever its requires_grad was before :func:apply_bitfit was called. No-op if :func:apply_bitfit was never called.

源代码位于： src/llm/core/bitfit.py

def unapply_bitfit(model: nn.Module) -> nn.Module:
    """Restore the pre-BitFit ``requires_grad`` state.

    Reverses :func:`apply_bitfit` — every parameter is set back to
    whatever its ``requires_grad`` was before :func:`apply_bitfit` was
    called. No-op if :func:`apply_bitfit` was never called.
    """
    snapshot = getattr(model, "_bitfit_original_requires_grad", None)
    if snapshot is None:
        return model

    for name, p in model.named_parameters():
        if name in snapshot:
            p.requires_grad = snapshot[name]
    del model._bitfit_original_requires_grad
    return model

get_bitfit_parameters

get_bitfit_parameters(model)

Yield every trainable bias parameter.

Use this to wire the optimizer

torch.optim.AdamW(get_bitfit_parameters(model), lr=...)

After :func:apply_bitfit, exactly the bias parameters are trainable, so this is equivalent to iter(model.parameters()) filtered by p.requires_grad — but we re-check the .bias suffix explicitly so the helper remains correct if a user manually enables a non-bias parameter after applying BitFit (the helper still yields only biases, not the non-bias ones).

Substring matching is intentionally avoided: fc_with_bias.weight contains the substring bias but is not a bias parameter.

源代码位于： src/llm/core/bitfit.py

def get_bitfit_parameters(model: nn.Module) -> Iterator[nn.Parameter]:
    """Yield every trainable bias parameter.

    Use this to wire the optimizer:
        ``torch.optim.AdamW(get_bitfit_parameters(model), lr=...)``

    After :func:`apply_bitfit`, exactly the bias parameters are
    trainable, so this is equivalent to ``iter(model.parameters())``
    filtered by ``p.requires_grad`` — but we re-check the ``.bias``
    suffix explicitly so the helper remains correct if a user
    manually enables a non-bias parameter after applying BitFit
    (the helper still yields only biases, not the non-bias ones).

    Substring matching is intentionally avoided: ``fc_with_bias.weight``
    contains the substring ``bias`` but is not a bias parameter.
    """
    for name, p in model.named_parameters():
        if not (name == "bias" or name.endswith(".bias")):
            continue
        if p.requires_grad:
            yield p

count_bitfit_parameters

count_bitfit_parameters(model)

Count trainable vs. total parameters in a BitFit-adapted model.

返回：

类型	描述
int	(trainable_params, total_params) — the BitFit contribution
int	is trainable_params, dominated by the frozen base weights.

Before :func:apply_bitfit is called, the helper reports whatever requires_grad state the model was constructed with (typically all-trainable, since :class:nn.Linear defaults to True).

源代码位于： src/llm/core/bitfit.py

def count_bitfit_parameters(model: nn.Module) -> tuple[int, int]:
    """Count trainable vs. total parameters in a BitFit-adapted model.

    Returns:
        ``(trainable_params, total_params)`` — the BitFit contribution
        is ``trainable_params``, dominated by the frozen base weights.

    Before :func:`apply_bitfit` is called, the helper reports whatever
    ``requires_grad`` state the model was constructed with (typically
    all-trainable, since :class:`nn.Linear` defaults to ``True``).
    """
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    return trainable, total

is_bitfit_applied

is_bitfit_applied(model)

Return True if :func:apply_bitfit was called on this model and not yet reversed by :func:unapply_bitfit.

Useful for checkpoint validation: if a checkpoint claims to be BitFit-adapted, the snapshot attribute should be present (or absent, if the user already called unapply_bitfit).

源代码位于： src/llm/core/bitfit.py

def is_bitfit_applied(model: nn.Module) -> bool:
    """Return ``True`` if :func:`apply_bitfit` was called on this model
    and not yet reversed by :func:`unapply_bitfit`.

    Useful for checkpoint validation: if a checkpoint claims to be
    BitFit-adapted, the snapshot attribute should be present (or
    absent, if the user already called ``unapply_bitfit``).
    """
    return hasattr(model, "_bitfit_original_requires_grad")

ia3

IA³ (Infused Adapter by Inhibiting and Amplifying Inner Activations).

A parameter-efficient fine-tuning method that wraps a frozen nn.Linear with a single trainable vector that multiplicatively scales the output. IA³ is the multiplicative counterpart to LoRA's additive design — instead of adding ΔW · x it scales the existing output W · x element-wise on the output dimension.

Per-layer cost: out_features trainable parameters (vs. LoRA's rank * (in_features + out_features) and AdaLoRA's init_rank * (in_features + out_features)). At out_features=4096 and rank=8, in_features=4096 that is ~4k vs. ~65k parameters per adapted linear — typically two orders of magnitude fewer than LoRA.

Forward

y = (W · x + b) * l where l is a learned vector of shape (out_features,) broadcast across batch and sequence dims.

Initialization

l ← ones so the wrapped layer starts as the identity transform — the model's existing forward is preserved at step 1, which avoids the chicken-and-egg problem that a zero-init would cause.

Merge for inference

W ← W * l[None, :] and b ← b * l. The learned vector folds into the base weight and can be discarded — no extra params at serve time, no extra matmul. Symmetric unmerge_weights reverses the merge so the same model can be checkpointed pre- and post-merge.

Reference: Liu et al., 2022 — Few-Shot Parameter-Efficient Fine-Tuning is Better and Cheaper than In-Context Learning, arXiv:2205.05638 (aka "T-Few"). The paper applies IA³ to attention K/V/output and FFN intermediate projections — the helper API mirrors LoRA so swapping apply_lora for apply_ia3 is a one-import change.

IA3Linear

Bases: Module

Wrap a frozen nn.Linear with a trainable multiplicative scale.

参数：

名称	类型	描述	默认
base_layer	Linear	The original nn.Linear to adapt (frozen at construction). The base layer's weight is not merged on init — the wrapper just multiplies the output of the base layer by the learned ia3_l vector.	必需
init_scale	float	Initial value of the multiplicative scale. Defaults to 1.0 so the wrapper starts as the identity transform — the model behaves identically to the base at step 1.	1.0

源代码位于： src/llm/core/ia3.py

class IA3Linear(nn.Module):
    """Wrap a frozen ``nn.Linear`` with a trainable multiplicative scale.

    Args:
        base_layer: The original ``nn.Linear`` to adapt (frozen at
            construction). The base layer's weight is **not** merged
            on init — the wrapper just multiplies the output of the
            base layer by the learned ``ia3_l`` vector.
        init_scale: Initial value of the multiplicative scale. Defaults
            to ``1.0`` so the wrapper starts as the identity transform
            — the model behaves identically to the base at step 1.
    """

    _original_ia3_l: torch.Tensor | None
    # Set only by ``merge_weights`` (and read/cleared by ``unmerge_weights``);
    # declared here so static analysis sees a Tensor, not an inferred union.
    _merged_ia3_l: torch.Tensor | None = None

    def __init__(
        self,
        base_layer: nn.Linear,
        init_scale: float = 1.0,
    ):
        super().__init__()
        self.base_layer = base_layer

        out_features = base_layer.out_features
        device = base_layer.weight.device
        dtype = base_layer.weight.dtype

        # One learned multiplier per output channel. Broadcasts over
        # batch and sequence dims at forward time.
        self.ia3_l = nn.Parameter(torch.full((out_features,), init_scale, device=device, dtype=dtype))

        # Freeze the base layer so only ``ia3_l`` is trainable.
        self.base_layer.weight.requires_grad = False
        if self.base_layer.bias is not None:
            self.base_layer.bias.requires_grad = False

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass: frozen base output, multiplicatively scaled."""
        base_output = self.base_layer(x)
        # ``ia3_l`` has shape (out_features,). The base output is
        # (..., out_features) — broadcasting handles the rest.
        return base_output * self.ia3_l

    def merge_weights(self) -> None:
        """Merge the multiplicative scale into the base weight for inference.

        After this call, ``self.base_layer.weight`` already contains
        the scale (``W * l[None, :]``) and ``self.ia3_l`` is set to
        ones so the wrapper is the identity on top of the already-scaled
        base. ``unmerge_weights`` reverses the operation, restoring
        both the original ``ia3_l`` snapshot and the pre-merge base
        weight — useful for checkpoint roundtrip where the same model
        needs to be saved both pre- and post-merge.
        """
        # Idempotent: a second call while already merged is a no-op (mirrors
        # the LoRA fix — see LoRALinear.merge_weights). Re-running would
        # re-snapshot ``_merged_ia3_l`` as ones (the active scale
        # post-first-merge) and multiply by ones (no-op), so the snapshot
        # becomes ones; the later ``unmerge_weights`` then divides by ones
        # and leaves the base permanently scaled (RIL ISS-159).
        if self._merged_ia3_l is not None:
            return
        with torch.no_grad():
            # Save the original scale so ``unmerge_weights`` can
            # restore it after dividing the base weight back out.
            self._merged_ia3_l = self.ia3_l.detach().clone()
            # Fold the scale into the base weight.
            # ``weight`` has shape ``(out_features, in_features)``; we
            # scale each output-channel row by the matching ``ia3_l``
            # entry, which means multiplying on dim 0 with a
            # ``(out_features, 1)`` broadcast.
            self.base_layer.weight.mul_(self.ia3_l.unsqueeze(1))
            if self.base_layer.bias is not None:
                self.base_layer.bias.mul_(self.ia3_l)
            # Zero-out the active scale — the wrapper is now identity
            # on top of the already-folded weight, so forward still
            # produces the same output as pre-merge.
            self.ia3_l.fill_(1.0)

    def unmerge_weights(self) -> None:
        """Reverse :meth:`merge_weights` — restores both the saved
        ``ia3_l`` snapshot and the pre-merge base weight. No-op if
        :meth:`merge_weights` was never called.
        """
        with torch.no_grad():
            merged_ia3_l = self._merged_ia3_l
            if merged_ia3_l is None:
                return
            # Restore the active scale to its pre-merge value, then
            # divide it back out of the base weight.
            self.ia3_l.copy_(merged_ia3_l)
            self._merged_ia3_l = None
            self.base_layer.weight.div_(self.ia3_l.unsqueeze(1))
            if self.base_layer.bias is not None:
                self.base_layer.bias.div_(self.ia3_l)

    @property
    def trainable_parameters(self) -> int:
        """Number of trainable IA³ parameters (just ``ia3_l.numel()``)."""
        return self.ia3_l.numel()

    def extra_repr(self) -> str:
        return (
            f"in_features={self.base_layer.in_features}, "
            f"out_features={self.base_layer.out_features}, "
            f"trainable={self.ia3_l.numel()}"
        )

trainable_parameters property

trainable_parameters

Number of trainable IA³ parameters (just ia3_l.numel()).

forward

forward(x)

Forward pass: frozen base output, multiplicatively scaled.

源代码位于： src/llm/core/ia3.py

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Forward pass: frozen base output, multiplicatively scaled."""
    base_output = self.base_layer(x)
    # ``ia3_l`` has shape (out_features,). The base output is
    # (..., out_features) — broadcasting handles the rest.
    return base_output * self.ia3_l

merge_weights

merge_weights()

Merge the multiplicative scale into the base weight for inference.

After this call, self.base_layer.weight already contains the scale (W * l[None, :]) and self.ia3_l is set to ones so the wrapper is the identity on top of the already-scaled base. unmerge_weights reverses the operation, restoring both the original ia3_l snapshot and the pre-merge base weight — useful for checkpoint roundtrip where the same model needs to be saved both pre- and post-merge.

源代码位于： src/llm/core/ia3.py

def merge_weights(self) -> None:
    """Merge the multiplicative scale into the base weight for inference.

    After this call, ``self.base_layer.weight`` already contains
    the scale (``W * l[None, :]``) and ``self.ia3_l`` is set to
    ones so the wrapper is the identity on top of the already-scaled
    base. ``unmerge_weights`` reverses the operation, restoring
    both the original ``ia3_l`` snapshot and the pre-merge base
    weight — useful for checkpoint roundtrip where the same model
    needs to be saved both pre- and post-merge.
    """
    # Idempotent: a second call while already merged is a no-op (mirrors
    # the LoRA fix — see LoRALinear.merge_weights). Re-running would
    # re-snapshot ``_merged_ia3_l`` as ones (the active scale
    # post-first-merge) and multiply by ones (no-op), so the snapshot
    # becomes ones; the later ``unmerge_weights`` then divides by ones
    # and leaves the base permanently scaled (RIL ISS-159).
    if self._merged_ia3_l is not None:
        return
    with torch.no_grad():
        # Save the original scale so ``unmerge_weights`` can
        # restore it after dividing the base weight back out.
        self._merged_ia3_l = self.ia3_l.detach().clone()
        # Fold the scale into the base weight.
        # ``weight`` has shape ``(out_features, in_features)``; we
        # scale each output-channel row by the matching ``ia3_l``
        # entry, which means multiplying on dim 0 with a
        # ``(out_features, 1)`` broadcast.
        self.base_layer.weight.mul_(self.ia3_l.unsqueeze(1))
        if self.base_layer.bias is not None:
            self.base_layer.bias.mul_(self.ia3_l)
        # Zero-out the active scale — the wrapper is now identity
        # on top of the already-folded weight, so forward still
        # produces the same output as pre-merge.
        self.ia3_l.fill_(1.0)

unmerge_weights

unmerge_weights()

Reverse :meth:merge_weights — restores both the saved ia3_l snapshot and the pre-merge base weight. No-op if :meth:merge_weights was never called.

源代码位于： src/llm/core/ia3.py

def unmerge_weights(self) -> None:
    """Reverse :meth:`merge_weights` — restores both the saved
    ``ia3_l`` snapshot and the pre-merge base weight. No-op if
    :meth:`merge_weights` was never called.
    """
    with torch.no_grad():
        merged_ia3_l = self._merged_ia3_l
        if merged_ia3_l is None:
            return
        # Restore the active scale to its pre-merge value, then
        # divide it back out of the base weight.
        self.ia3_l.copy_(merged_ia3_l)
        self._merged_ia3_l = None
        self.base_layer.weight.div_(self.ia3_l.unsqueeze(1))
        if self.base_layer.bias is not None:
            self.base_layer.bias.div_(self.ia3_l)

apply_ia3

apply_ia3(model, init_scale=1.0, target_modules=None)

Apply IA³ to specified linear layers in a model.

参数：

名称	类型	描述	默认
model	Module	The model to adapt (modified in-place).	必需
init_scale	float	Initial value of the multiplicative scale (passed through to :class:IA3Linear).	1.0
target_modules	list[str] | None	List of module-name substring patterns. If None (default), every nn.Linear is wrapped — same default as apply_lora. Pass e.g. ["q_proj", "v_proj"] to wrap only attention projections.	None

返回：

类型	描述
Module	The model with IA³ applied (modified in-place).

源代码位于： src/llm/core/ia3.py

def apply_ia3(
    model: nn.Module,
    init_scale: float = 1.0,
    target_modules: list[str] | None = None,
) -> nn.Module:
    """Apply IA³ to specified linear layers in a model.

    Args:
        model: The model to adapt (modified in-place).
        init_scale: Initial value of the multiplicative scale (passed
            through to :class:`IA3Linear`).
        target_modules: List of module-name substring patterns. If
            ``None`` (default), every ``nn.Linear`` is wrapped — same
            default as ``apply_lora``. Pass e.g. ``["q_proj", "v_proj"]``
            to wrap only attention projections.

    Returns:
        The model with IA³ applied (modified in-place).
    """
    if target_modules is None:
        target_modules = []

    def should_apply(name: str) -> bool:
        if not target_modules:
            return True
        return any(pattern in name for pattern in target_modules)

    replacements = []
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear) and should_apply(name):
            replacements.append((name, module))

    for name, module in replacements:
        ia3_layer = IA3Linear(module, init_scale=init_scale)
        parts = name.split(".")
        parent = model
        for part in parts[:-1]:
            parent = getattr(parent, part)
        setattr(parent, parts[-1], ia3_layer)

    return model

merge_ia3

merge_ia3(model)

Merge every IA³ scale into the wrapped base weight.

After this, the model is identical to the original at inference but no longer has trainable IA³ params. Reversible via :func:unmerge_ia3.

源代码位于： src/llm/core/ia3.py

def merge_ia3(model: nn.Module) -> nn.Module:
    """Merge every IA³ scale into the wrapped base weight.

    After this, the model is identical to the original at inference
    but no longer has trainable IA³ params. Reversible via
    :func:`unmerge_ia3`.
    """
    for module in model.modules():
        if isinstance(module, IA3Linear):
            module.merge_weights()
    return model

unmerge_ia3

unmerge_ia3(model)

Reverse :func:merge_ia3 — restores the trained ia3_l as the active scale. Useful for checkpoint roundtrip.

源代码位于： src/llm/core/ia3.py

def unmerge_ia3(model: nn.Module) -> nn.Module:
    """Reverse :func:`merge_ia3` — restores the trained ``ia3_l`` as the
    active scale. Useful for checkpoint roundtrip.
    """
    for module in model.modules():
        if isinstance(module, IA3Linear):
            module.unmerge_weights()
    return model

get_ia3_parameters

get_ia3_parameters(model)

Yield every IA³ trainable parameter — one ia3_l per wrapper.

Use this to wire the optimizer

torch.optim.Adam(get_ia3_parameters(model), lr=...)

源代码位于： src/llm/core/ia3.py

def get_ia3_parameters(model: nn.Module) -> Iterator[nn.Parameter]:
    """Yield every IA³ trainable parameter — one ``ia3_l`` per wrapper.

    Use this to wire the optimizer:
        ``torch.optim.Adam(get_ia3_parameters(model), lr=...)``
    """
    for module in model.modules():
        if isinstance(module, IA3Linear):
            yield module.ia3_l

count_ia3_parameters

count_ia3_parameters(model)

Count trainable vs. total parameters in an IA³-adapted model.

返回：

类型	描述
int	(trainable_params, total_params) — the IA³ contribution
int	is trainable_params, dominated by the frozen base weights.

源代码位于： src/llm/core/ia3.py

def count_ia3_parameters(model: nn.Module) -> tuple[int, int]:
    """Count trainable vs. total parameters in an IA³-adapted model.

    Returns:
        ``(trainable_params, total_params)`` — the IA³ contribution
        is ``trainable_params``, dominated by the frozen base weights.
    """
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    return trainable, total

disable_ia3

disable_ia3(model)

Disable IA³ at inference — sets every ia3_l to all-ones so the wrapper is the identity transform. Use this when you want to evaluate the base model behaviour without un-wrapping.

源代码位于： src/llm/core/ia3.py

def disable_ia3(model: nn.Module) -> None:
    """Disable IA³ at inference — sets every ``ia3_l`` to all-ones so
    the wrapper is the identity transform. Use this when you want to
    evaluate the base model behaviour without un-wrapping.
    """
    for module in model.modules():
        if isinstance(module, IA3Linear):
            module._original_ia3_l = module.ia3_l.detach().clone()
            with torch.no_grad():
                module.ia3_l.fill_(1.0)

enable_ia3

enable_ia3(model)

Re-enable IA³ after :func:disable_ia3 — restores the saved ia3_l snapshot. No-op if disable_ia3 was never called.

源代码位于： src/llm/core/ia3.py

def enable_ia3(model: nn.Module) -> None:
    """Re-enable IA³ after :func:`disable_ia3` — restores the saved
    ``ia3_l`` snapshot. No-op if ``disable_ia3`` was never called.
    """
    for module in model.modules():
        orig = getattr(module, "_original_ia3_l", None)
        if isinstance(module, IA3Linear) and orig is not None:
            with torch.no_grad():
                module.ia3_l.copy_(orig)
            del module._original_ia3_l

adapter

Adapter Layers (Houlsby et al. 2019).

Parameter-efficient fine-tuning via bottleneck modules inserted into transformer blocks. The adapter is a small feed-forward block with a residual connection - the base Linear is frozen, and only the adapter parameters train.

Per the original paper, adapters are inserted:

``x → Linear → activation → Linear → + residual → output``

i.e. a down-projection (to a small bottleneck dim), a non-linearity, and an up-projection back to the hidden dim, summed with the base output. The up-projection is zero-initialized so the adapter is the identity transform at step 1 - no chicken-and-egg training stall.

Per-layer trainable cost: hidden_size x bottleneck_dim + bottleneck_dim x hidden_size + bottleneck_dim + hidden_size (down weight + up weight + down bias + up bias). At hidden_size=4096, bottleneck_dim=64 that's 2 x 4096 x 64 + 64 + 4096 ≈ 528k params per adapted Linear - vs. LoRA's rank x (in + out) = 8 x (4096 + 4096) ≈ 65k and IA³'s 4096. Adapters are usually bigger than LoRA / IA³ but smaller than full fine-tuning.

Forward

y = base_linear(x) + up(activation(down(base_linear(x))))

Initialization

• down: Kaiming uniform (standard Linear init).
• up: zeros - so the adapter contributes 0 to the output at step 1, and the wrapper is the identity on top of the base.
• activation: nn.ReLU (the original paper uses ReLU; later work uses GELU or Tanh - picked here to match Houlsby 2019).

The helper API (apply_adapter / merge_adapter / unmerge_adapter / get_adapter_parameters / count_adapter_parameters / disable_adapter / enable_adapter) mirrors LoRA / IA³ / BitFit so swapping PEFT methods in user code is a one-import change. Note that merge_adapter is a near no-op for adapters - the up-projection being zero means the adapter contributes nothing, so merging the adapter into the base would just add zeros. The function is kept for API parity.

Reference: Houlsby et al., 2019 - Parameter-Efficient Transfer Learning for NLP, arXiv:1902.00751. The bottleneck-only-after-FFN variant (Pfeiffer et al. 2020) and the Compacter / MAD-X decompositions are deliberate follow-ups.

AdapterLinear

Bases: Module

Wrap a frozen nn.Linear with a bottleneck adapter on the output.

参数：

名称	类型	描述	默认
base_layer	Linear	The original nn.Linear to adapt (frozen at construction).	必需
bottleneck_dim	int	Width of the adapter's hidden dim. Smaller values reduce trainable parameters; typical values are 8-256 depending on the hidden size.	必需
activation	type[Module]	Non-linearity class (defaults to nn.ReLU to match Houlsby 2019).	ReLU

属性：

名称	类型	描述
_original_up_weight	Tensor | None	Snapshot of up-projection weight, set by :func:disable_adapter, cleared by :func:enable_adapter.
_original_up_bias	Tensor | None	Snapshot of up-projection bias, set/cleared alongside _original_up_weight.

源代码位于： src/llm/core/adapter.py

class AdapterLinear(nn.Module):
    """Wrap a frozen ``nn.Linear`` with a bottleneck adapter on the output.

    Args:
        base_layer: The original ``nn.Linear`` to adapt (frozen at
            construction).
        bottleneck_dim: Width of the adapter's hidden dim. Smaller
            values reduce trainable parameters; typical values are
            8-256 depending on the hidden size.
        activation: Non-linearity class (defaults to ``nn.ReLU`` to
            match Houlsby 2019).

    Attributes:
        _original_up_weight: Snapshot of up-projection weight, set by
            :func:`disable_adapter`, cleared by :func:`enable_adapter`.
        _original_up_bias: Snapshot of up-projection bias, set/cleared
            alongside ``_original_up_weight``.
    """

    _original_up_weight: torch.Tensor | None
    _original_up_bias: torch.Tensor | None

    def __init__(
        self,
        base_layer: nn.Linear,
        bottleneck_dim: int,
        activation: type[nn.Module] = nn.ReLU,
    ):
        super().__init__()
        if bottleneck_dim <= 0:
            raise ValueError(f"bottleneck_dim must be positive, got {bottleneck_dim}")
        self.base_layer = base_layer
        self.bottleneck_dim = bottleneck_dim

        hidden = base_layer.out_features
        device = base_layer.weight.device
        dtype = base_layer.weight.dtype

        # Down-project: hidden → bottleneck.
        self.down = nn.Linear(hidden, bottleneck_dim, device=device, dtype=dtype)
        self.activation = activation()
        # Up-project: bottleneck → hidden. Zero-init so the adapter
        # is the identity transform at step 1 (no chicken-and-egg
        # training stall - the up output is zero, so the wrapper
        # matches the base output at step 1).
        self.up = nn.Linear(bottleneck_dim, hidden, device=device, dtype=dtype)
        nn.init.zeros_(self.up.weight)
        nn.init.zeros_(self.up.bias)

        # Freeze the base layer so only the adapter trains.
        self.base_layer.weight.requires_grad = False
        if self.base_layer.bias is not None:
            self.base_layer.bias.requires_grad = False

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass: frozen base output + residual adapter output."""
        base_output = self.base_layer(x)
        adapter_output = self.up(self.activation(self.down(base_output)))
        return base_output + adapter_output

    def merge_weights(self) -> None:
        """No-op for adapters.

        Unlike LoRA / IA³ the adapter has no math to fold into the
        base weight - the up-projection being zero means the adapter
        contributes zero to the output. ``merge_weights`` is kept
        for API parity with the other PEFT helpers; it does nothing.
        """
        # Intentionally empty - see docstring.

    def unmerge_weights(self) -> None:
        """No-op for adapters (mirror of :meth:`merge_weights`).

        Kept for API parity; nothing to undo.
        """
        # Intentionally empty - see docstring.

    @property
    def trainable_parameters(self) -> int:
        """Number of trainable adapter parameters (down + up weights + biases)."""
        return self.down.weight.numel() + self.down.bias.numel() + self.up.weight.numel() + self.up.bias.numel()

    def extra_repr(self) -> str:
        return (
            f"in_features={self.base_layer.in_features}, "
            f"out_features={self.base_layer.out_features}, "
            f"bottleneck={self.bottleneck_dim}, "
            f"trainable={self.trainable_parameters}"
        )

trainable_parameters property

trainable_parameters

Number of trainable adapter parameters (down + up weights + biases).

forward

forward(x)

Forward pass: frozen base output + residual adapter output.

源代码位于： src/llm/core/adapter.py

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Forward pass: frozen base output + residual adapter output."""
    base_output = self.base_layer(x)
    adapter_output = self.up(self.activation(self.down(base_output)))
    return base_output + adapter_output

merge_weights

merge_weights()

No-op for adapters.

Unlike LoRA / IA³ the adapter has no math to fold into the base weight - the up-projection being zero means the adapter contributes zero to the output. merge_weights is kept for API parity with the other PEFT helpers; it does nothing.

源代码位于： src/llm/core/adapter.py

def merge_weights(self) -> None:
    """No-op for adapters.

    Unlike LoRA / IA³ the adapter has no math to fold into the
    base weight - the up-projection being zero means the adapter
    contributes zero to the output. ``merge_weights`` is kept
    for API parity with the other PEFT helpers; it does nothing.
    """

unmerge_weights

unmerge_weights()

No-op for adapters (mirror of :meth:merge_weights).

Kept for API parity; nothing to undo.

源代码位于： src/llm/core/adapter.py

def unmerge_weights(self) -> None:
    """No-op for adapters (mirror of :meth:`merge_weights`).

    Kept for API parity; nothing to undo.
    """

apply_adapter

apply_adapter(model, bottleneck_dim=64, target_modules=None)

Apply adapter bottleneck modules to specified linear layers.

参数：

名称	类型	描述	默认
model	Module	The model to adapt (modified in-place).	必需
bottleneck_dim	int	Width of the adapter's hidden dim (passed through to :class:AdapterLinear).	64
target_modules	list[str] | None	List of module-name substring patterns. If None (default), every nn.Linear is wrapped. Pass e.g. ["q_proj", "v_proj"] to wrap only attention projections.	None

返回：

类型	描述
Module	The model with adapters applied (modified in-place).

源代码位于： src/llm/core/adapter.py

def apply_adapter(
    model: nn.Module,
    bottleneck_dim: int = 64,
    target_modules: list[str] | None = None,
) -> nn.Module:
    """Apply adapter bottleneck modules to specified linear layers.

    Args:
        model: The model to adapt (modified in-place).
        bottleneck_dim: Width of the adapter's hidden dim (passed
            through to :class:`AdapterLinear`).
        target_modules: List of module-name substring patterns. If
            ``None`` (default), every ``nn.Linear`` is wrapped. Pass
            e.g. ``["q_proj", "v_proj"]`` to wrap only attention
            projections.

    Returns:
        The model with adapters applied (modified in-place).
    """
    if target_modules is None:
        target_modules = []

    def should_apply(name: str) -> bool:
        if not target_modules:
            return True
        return any(pattern in name for pattern in target_modules)

    replacements = []
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear) and should_apply(name):
            replacements.append((name, module))

    for name, module in replacements:
        adapter = AdapterLinear(module, bottleneck_dim=bottleneck_dim)
        parts = name.split(".")
        parent = model
        for part in parts[:-1]:
            parent = getattr(parent, part)
        setattr(parent, parts[-1], adapter)

    return model

merge_adapter

merge_adapter(model)

No-op for adapters - kept for API parity with LoRA / IA³.

Unlike LoRA / IA³, the adapter has no math to fold into the base weight. The up-projection is zero-initialized, so the adapter contributes zero to the output unless the user trained it.

This function is provided so that apply_adapter / merge_adapter / unmerge_adapter follow the same call pattern as the other PEFT helpers.

源代码位于： src/llm/core/adapter.py

def merge_adapter(model: nn.Module) -> nn.Module:
    """No-op for adapters - kept for API parity with LoRA / IA³.

    Unlike LoRA / IA³, the adapter has no math to fold into the base
    weight. The up-projection is zero-initialized, so the adapter
    contributes zero to the output unless the user trained it.

    This function is provided so that ``apply_adapter`` /
    ``merge_adapter`` / ``unmerge_adapter`` follow the same call
    pattern as the other PEFT helpers.
    """
    for module in model.modules():
        if isinstance(module, AdapterLinear):
            module.merge_weights()
    return model

unmerge_adapter

unmerge_adapter(model)

No-op for adapters - mirror of :func:merge_adapter.

源代码位于： src/llm/core/adapter.py

def unmerge_adapter(model: nn.Module) -> nn.Module:
    """No-op for adapters - mirror of :func:`merge_adapter`."""
    for module in model.modules():
        if isinstance(module, AdapterLinear):
            module.unmerge_weights()
    return model

get_adapter_parameters

get_adapter_parameters(model)

Yield every trainable adapter parameter - down + up weights + biases per wrapper, nothing from the base Linear.

源代码位于： src/llm/core/adapter.py

def get_adapter_parameters(model: nn.Module) -> Iterator[torch.Tensor]:
    """Yield every trainable adapter parameter - down + up weights + biases
    per wrapper, nothing from the base Linear.
    """
    for module in model.modules():
        if isinstance(module, AdapterLinear):
            yield module.down.weight
            assert module.down.bias is not None  # noqa: S101
            yield module.down.bias
            yield module.up.weight
            assert module.up.bias is not None  # noqa: S101
            yield module.up.bias

count_adapter_parameters

count_adapter_parameters(model)

Count trainable vs. total parameters in an adapter-adapted model.

源代码位于： src/llm/core/adapter.py

def count_adapter_parameters(model: nn.Module) -> tuple[int, int]:
    """Count trainable vs. total parameters in an adapter-adapted model."""
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    return trainable, total

disable_adapter

disable_adapter(model)

Disable adapters by zeroing the up-projection (so the adapter contributes zero to the output).

Saves the up-projection weight / bias under _original_up_weight / _original_up_bias so :func:enable_adapter can restore them.

源代码位于： src/llm/core/adapter.py

def disable_adapter(model: nn.Module) -> None:
    """Disable adapters by zeroing the up-projection (so the adapter
    contributes zero to the output).

    Saves the up-projection weight / bias under ``_original_up_weight``
    / ``_original_up_bias`` so :func:`enable_adapter` can restore them.
    """
    for module in model.modules():
        if isinstance(module, AdapterLinear):
            module._original_up_weight = module.up.weight.detach().clone()
            module._original_up_bias = module.up.bias.detach().clone()
            with torch.no_grad():
                module.up.weight.zero_()
                module.up.bias.zero_()

enable_adapter

enable_adapter(model)

Re-enable adapters after :func:disable_adapter - restores the saved up-projection snapshot. No-op if disable_adapter was never called.

源代码位于： src/llm/core/adapter.py

def enable_adapter(model: nn.Module) -> None:
    """Re-enable adapters after :func:`disable_adapter` - restores the
    saved up-projection snapshot. No-op if ``disable_adapter`` was
    never called.
    """
    for module in model.modules():
        orig_weight = getattr(module, "_original_up_weight", None)
        orig_bias = getattr(module, "_original_up_bias", None)
        if isinstance(module, AdapterLinear) and orig_weight is not None and orig_bias is not None:
            with torch.no_grad():
                module.up.weight.copy_(orig_weight)
                module.up.bias.copy_(orig_bias)
            del module._original_up_weight
            del module._original_up_bias

pfeiffer_adapter

Pfeiffer Adapter (Pfeiffer et al. 2020).

Parameter-efficient fine-tuning via bottleneck modules inserted only after FFN/MLP layers, not after attention projections. The variant was introduced in Pfeiffer et al., 2020 — AdapterHub: A Framework for Adapting Transformers, arXiv:2007.07779 — and is the production default in AdapterHub / HuggingFace PEFT, roughly half the parameters of Houlsby 2019 at comparable accuracy on most tasks.

The wrapper class is the same :class:llm.core.adapter.AdapterLinear used by Houlsby — there is no new tensor type. The only difference between Houlsby and Pfeiffer is which linears get wrapped:

• Houlsby: every nn.Linear (attention + MLP)
• Pfeiffer: only the FFN / MLP linears (default fc1 + fc2, matching :class:llm.core.mlp.MLP)

Per the original paper, the adapter is the same bottleneck residual:

``y = base_linear(x) + up(activation(down(base_linear(x))))``

with up zero-initialized so the adapter is the identity at step 1 (no chicken-and-egg training stall).

Per-layer trainable cost: 2 * out_features * bottleneck_dim + out_features + bottleneck_dim — identical to Houlsby's per-layer cost. The parameter savings come from wrapping fewer layers, not from a smaller per-layer footprint. A 2-layer transformer block with hidden=4096 and bottleneck=64 has:

• Houlsby: 4 attention linears + 2 MLP linears = 6 wrappers ≈ 6 * (2 * 4096 * 64 + 4096 + 64) ≈ 3.2M trainable per block
• Pfeiffer: 2 MLP linears = 2 wrappers ≈ 2 * (2 * 4096 * 64 + 4096 + 64) ≈ 1.05M trainable per block

The helper API mirrors the Houlsby one (merge_* / unmerge_* / get_*_parameters / count_*_parameters / disable_*` /enable_*) so swappingadapter→pfeiffer_adapter` in user code is a one-import change. Internally the helpers delegate to the Houlsby implementations — Pfeiffer wrappers **are** :class:llm.core.adapter.AdapterLinear` instances, so there's nothing to distinguish at runtime.

Reference: Pfeiffer et al., 2020 — AdapterHub: A Framework for Adapting Transformers, arXiv:2007.07779. The Compacter (Kronecker decomposition) and MAD-X (cross-lingual modular) variants are deliberate follow-ups; this slice ships the simple FFN-only variant.

apply_pfeiffer_adapter

apply_pfeiffer_adapter(model, bottleneck_dim=64, target_modules=None)

Apply Pfeiffer Adapter — bottleneck residual only after FFN/MLP.

参数：

名称	类型	描述	默认
model	Module	The model to adapt (modified in-place).	必需
bottleneck_dim	int	Width of the adapter's hidden dim. Forwarded to :class:llm.core.adapter.AdapterLinear. Defaults to 64 (the Houlsby 2019 paper convention, also used by the Pfeiffer 2020 reproductions).	64
target_modules	list[str] | None	List of module-name substring patterns used to pick which nn.Linear modules get wrapped. If None (default), the standard FFN/MLP filter ["fc1", "fc2"] is used — matching the layer names in :class:llm.core.mlp.MLP. Pass a custom list to wrap a different subset (e.g. ["q_proj", "v_proj"] is invalid here — Pfeiffer is FFN-only — but you can point at non-standard MLP layer names in custom architectures).	None

返回：

类型	描述
Module	The same model (modified in-place; chainable).

Note

Internally this is a thin delegate to :func:llm.core.adapter.apply_adapter with the FFN-only target filter. The wrapper class is :class:llm.core.adapter.AdapterLinear — Pfeiffer IS Houlsby-on-MLP-only, so no new wrapper code is needed.

Unlike LoRA / IA³ / Prefix Tuning, bottleneck_dim is the knob (rather than rank) and the up projection is zero-initialized so the wrapper is the identity transform at step 1.

源代码位于： src/llm/core/pfeiffer_adapter.py

def apply_pfeiffer_adapter(
    model: nn.Module,
    bottleneck_dim: int = 64,
    target_modules: list[str] | None = None,
) -> nn.Module:
    """Apply Pfeiffer Adapter — bottleneck residual only after FFN/MLP.

    Args:
        model: The model to adapt (modified in-place).
        bottleneck_dim: Width of the adapter's hidden dim. Forwarded
            to :class:`llm.core.adapter.AdapterLinear`. Defaults to
            64 (the Houlsby 2019 paper convention, also used by the
            Pfeiffer 2020 reproductions).
        target_modules: List of module-name substring patterns used to
            pick which ``nn.Linear`` modules get wrapped. If ``None``
            (default), the standard FFN/MLP filter ``["fc1", "fc2"]``
            is used — matching the layer names in
            :class:`llm.core.mlp.MLP`. Pass a custom list to wrap a
            different subset (e.g. ``["q_proj", "v_proj"]`` is
            invalid here — Pfeiffer is FFN-only — but you can point
            at non-standard MLP layer names in custom architectures).

    Returns:
        The same ``model`` (modified in-place; chainable).

    Note:
        Internally this is a thin delegate to
        :func:`llm.core.adapter.apply_adapter` with the FFN-only
        target filter. The wrapper class is
        :class:`llm.core.adapter.AdapterLinear` — Pfeiffer IS
        Houlsby-on-MLP-only, so no new wrapper code is needed.

        Unlike LoRA / IA³ / Prefix Tuning, ``bottleneck_dim`` is the
        knob (rather than ``rank``) and the ``up`` projection is
        zero-initialized so the wrapper is the identity transform at
        step 1.
    """
    if target_modules is None:
        target_modules = list(DEFAULT_PFEIFFER_TARGETS)
    return apply_adapter(model, bottleneck_dim=bottleneck_dim, target_modules=target_modules)

merge_pfeiffer_adapter

merge_pfeiffer_adapter(model)

No-op for Pfeiffer — kept for API parity with LoRA / IA³.

Delegates to :func:llm.core.adapter.merge_adapter. The up projection being zero means the adapter contributes nothing to the output unless the user trained it, so there's nothing to fold into the base weight.

源代码位于： src/llm/core/pfeiffer_adapter.py

def merge_pfeiffer_adapter(model: nn.Module) -> nn.Module:
    """No-op for Pfeiffer — kept for API parity with LoRA / IA³.

    Delegates to :func:`llm.core.adapter.merge_adapter`. The ``up``
    projection being zero means the adapter contributes nothing to
    the output unless the user trained it, so there's nothing to
    fold into the base weight.
    """
    return merge_adapter(model)

unmerge_pfeiffer_adapter

unmerge_pfeiffer_adapter(model)

No-op for Pfeiffer — mirror of :func:merge_pfeiffer_adapter.

源代码位于： src/llm/core/pfeiffer_adapter.py

def unmerge_pfeiffer_adapter(model: nn.Module) -> nn.Module:
    """No-op for Pfeiffer — mirror of :func:`merge_pfeiffer_adapter`."""
    return unmerge_adapter(model)

get_pfeiffer_parameters

get_pfeiffer_parameters(model)

Yield every trainable Pfeiffer parameter.

Delegates to :func:llm.core.adapter.get_adapter_parameters. Since both Houlsby and Pfeiffer produce :class:AdapterLinear wrappers, this helper yields parameters from every adapter wrapper in the model (Pfeiffer alone, Houlsby alone, or both coexisting). For Pfeiffer-only the result is identical to get_adapter_parameters.

源代码位于： src/llm/core/pfeiffer_adapter.py

def get_pfeiffer_parameters(model: nn.Module) -> Iterator[torch.Tensor]:
    """Yield every trainable Pfeiffer parameter.

    Delegates to :func:`llm.core.adapter.get_adapter_parameters`.
    Since both Houlsby and Pfeiffer produce :class:`AdapterLinear`
    wrappers, this helper yields parameters from **every** adapter
    wrapper in the model (Pfeiffer alone, Houlsby alone, or both
    coexisting). For Pfeiffer-only the result is identical to
    ``get_adapter_parameters``.
    """
    return get_adapter_parameters(model)

count_pfeiffer_parameters

count_pfeiffer_parameters(model)

Return (trainable, total) parameter counts.

Delegates to :func:llm.core.adapter.count_adapter_parameters. Same caveat as :func:get_pfeiffer_parameters — counts every adapter wrapper, not just Pfeiffer-targeted ones. Users who mix Pfeiffer and Houlsby on the same model should call the per- method helpers selectively.

源代码位于： src/llm/core/pfeiffer_adapter.py

def count_pfeiffer_parameters(model: nn.Module) -> tuple[int, int]:
    """Return ``(trainable, total)`` parameter counts.

    Delegates to :func:`llm.core.adapter.count_adapter_parameters`.
    Same caveat as :func:`get_pfeiffer_parameters` — counts every
    adapter wrapper, not just Pfeiffer-targeted ones. Users who mix
    Pfeiffer and Houlsby on the same model should call the per-
    method helpers selectively.
    """
    return count_adapter_parameters(model)

disable_pfeiffer_adapter

disable_pfeiffer_adapter(model)

Disable Pfeiffer adapters by zeroing the up projection.

Delegates to :func:llm.core.adapter.disable_adapter. After this call every wrapper's up.weight and up.bias are zero, making the wrapper mathematically the identity on top of the base. The pre-disable up-projection is snapshotted under _original_up_weight / _original_up_bias so :func:enable_pfeiffer_adapter can restore it.

源代码位于： src/llm/core/pfeiffer_adapter.py

def disable_pfeiffer_adapter(model: nn.Module) -> None:
    """Disable Pfeiffer adapters by zeroing the ``up`` projection.

    Delegates to :func:`llm.core.adapter.disable_adapter`. After
    this call every wrapper's ``up.weight`` and ``up.bias`` are
    zero, making the wrapper mathematically the identity on top of
    the base. The pre-disable up-projection is snapshotted under
    ``_original_up_weight`` / ``_original_up_bias`` so
    :func:`enable_pfeiffer_adapter` can restore it.
    """
    disable_adapter(model)

enable_pfeiffer_adapter

enable_pfeiffer_adapter(model)

Re-enable Pfeiffer adapters after :func:disable_pfeiffer_adapter.

Delegates to :func:llm.core.adapter.enable_adapter. No-op if disable_pfeiffer_adapter was never called (the snapshot attribute is the sentinel).

源代码位于： src/llm/core/pfeiffer_adapter.py

def enable_pfeiffer_adapter(model: nn.Module) -> None:
    """Re-enable Pfeiffer adapters after :func:`disable_pfeiffer_adapter`.

    Delegates to :func:`llm.core.adapter.enable_adapter`. No-op if
    ``disable_pfeiffer_adapter`` was never called (the snapshot
    attribute is the sentinel).
    """
    enable_adapter(model)

prefix_tuning

Prefix Tuning module (parameter-efficient fine-tuning, T2 PEFT).

Wraps a frozen attention layer with a trainable prefix that gets prepended to K and V at every forward pass. The prefix lives in a small latent space (prefix_small) and is projected into the K/V dimensions by two reparameterization MLPs - that's the "reparameterized" Prefix Tuning from Li & Liang 2021, which stabilises training versus directly learning the K/V prefix.

Multi-backend: MHA / FlashAttention / MultiLatentAttention all support prefix_kv injection (any base that satisfies the :class:llm.core.attn.base.PrefixCapableAttention Protocol). The wrapper freezes the base attention and only prefix_small + the two reparameterization MLPs receive gradients. After training, :func:fold_reparameterization collapses the MLPs into static prefix_k / prefix_v buffers for inference (one fewer matmul per step, no trainable params at serve time, no risk of training-mode behaviour leaking into deployment).

Reference: Li & Liang, 2021 - Prefix-Tuning: Optimizing Continuous Prompts for Generation, arXiv:2101.00190.

PrefixTuningAttention

Bases: Module

Wrap a frozen attention base with trainable prefix K/V.

The base must satisfy the :class:llm.core.attn.base.PrefixCapableAttention Protocol — i.e. accept a prefix_kv kwarg in its forward. MultiHeadAttention (the original reference impl), FlashAttention, and MultiLatentAttention all qualify.

The trainable parameters are:

• prefix_small: (prefix_len, reparam_hidden) - small latent prefix. Lower-rank than the full K/V dimension so the search space is bounded.
• _reparam_k / _reparam_v: nn.Linear(reparam_hidden, kv_dim) MLPs that project the small prefix into the K and V spaces.

Forward computes prefix K/V via the reparam MLPs (or reads the static buffers if :func:fold_reparameterization has been called), expands to [B, num_kv_heads, prefix_len, head_dim], and dispatches to base_attn(x, prefix_kv=...). The base attention is frozen at construction so only prefix parameters receive gradients.

参数：

名称	类型	描述	默认
base_attn	PrefixCapableAttention	The frozen attention layer to wrap. Must satisfy :class:llm.core.attn.base.PrefixCapableAttention and expose num_kv_heads and head_dim attributes.	必需
prefix_len	int	Number of prefix tokens to prepend to each layer's K and V. Typical values: 10-200.	必需
reparam_hidden	int | None	Width of the reparam MLP's hidden dim. Defaults to kv_dim. Smaller values reduce trainable parameters at the cost of expressivity.	None

源代码位于： src/llm/core/prefix_tuning.py

class PrefixTuningAttention(nn.Module):
    """Wrap a frozen attention base with trainable prefix K/V.

    The base must satisfy the :class:`llm.core.attn.base.PrefixCapableAttention`
    Protocol — i.e. accept a ``prefix_kv`` kwarg in its ``forward``.
    ``MultiHeadAttention`` (the original reference impl),
    ``FlashAttention``, and ``MultiLatentAttention`` all qualify.

    The trainable parameters are:

    - ``prefix_small``: ``(prefix_len, reparam_hidden)`` - small latent
      prefix. Lower-rank than the full K/V dimension so the search space
      is bounded.
    - ``_reparam_k`` / ``_reparam_v``: ``nn.Linear(reparam_hidden, kv_dim)``
      MLPs that project the small prefix into the K and V spaces.

    Forward computes prefix K/V via the reparam MLPs (or reads the
    static buffers if :func:`fold_reparameterization` has been called),
    expands to ``[B, num_kv_heads, prefix_len, head_dim]``, and dispatches
    to ``base_attn(x, prefix_kv=...)``. The base attention is frozen at
    construction so only prefix parameters receive gradients.

    Args:
        base_attn: The frozen attention layer to wrap. Must satisfy
            :class:`llm.core.attn.base.PrefixCapableAttention` and
            expose ``num_kv_heads`` and ``head_dim`` attributes.
        prefix_len: Number of prefix tokens to prepend to each layer's
            K and V. Typical values: 10-200.
        reparam_hidden: Width of the reparam MLP's hidden dim. Defaults
            to ``kv_dim``. Smaller values reduce trainable parameters
            at the cost of expressivity.
    """

    # Buffers set by fold_reparameterization() (deployment path).
    # Absent until folding; ``forward`` checks via hasattr.
    prefix_k: torch.Tensor | None
    prefix_v: torch.Tensor | None

    def __init__(
        self,
        base_attn: PrefixCapableAttention,
        prefix_len: int,
        reparam_hidden: int | None = None,
    ) -> None:
        super().__init__()
        # Protocol gate: any base that satisfies the PrefixCapableAttention
        # protocol (i.e. accepts ``prefix_kv`` in its ``forward``) qualifies.
        # We additionally verify the two attributes the wrapper reads at
        # construction time - ``num_kv_heads`` and ``head_dim``. The
        # ``@runtime_checkable`` decorator only inspects method names, so an
        # unrelated ``nn.Linear`` would otherwise sneak through the isinstance
        # check (it has a ``forward`` method).
        if not isinstance(base_attn, PrefixCapableAttention):
            raise TypeError(
                f"PrefixTuningAttention requires a PrefixCapableAttention base "
                f"(see llm.core.attn.base.PrefixCapableAttention); got "
                f"{type(base_attn).__name__}. The base must (a) expose "
                f"num_kv_heads and head_dim attributes and (b) accept "
                f"prefix_kv= in its forward()."
            )
        if not (hasattr(base_attn, "num_kv_heads") and hasattr(base_attn, "head_dim")):
            raise TypeError(
                f"PrefixTuningAttention base must expose num_kv_heads and head_dim "
                f"attributes; {type(base_attn).__name__} is missing one or both. "
                f"See llm.core.attn.base.PrefixCapableAttention."
            )

        self.base_attn = base_attn
        self.prefix_len = prefix_len
        self.num_kv_heads = base_attn.num_kv_heads
        self.head_dim = base_attn.head_dim
        self.kv_dim = self.num_kv_heads * self.head_dim
        self.reparam_hidden = reparam_hidden if reparam_hidden is not None else self.kv_dim

        # The base attention may live on GPU and/or in fp16/bf16 (serve/export
        # paths, or applying PEFT to an already-placed model). Unlike the
        # linear-based PEFT wrappers (LoRA/IA3/Adapter read
        # ``base_layer.weight.device``), an attention module has no single
        # ``weight`` — derive device/dtype from its first parameter and create
        # the prefix params there, so ``torch.cat(prefix_k, k)`` inside the
        # base attention never hits a device/dtype mismatch (RIL ISS-158).
        first_param = next(self.base_attn.parameters())
        device = first_param.device
        dtype = first_param.dtype

        # Trainable parameters.
        # prefix_small is the "small latent" prefix - lives in
        # ``(prefix_len, reparam_hidden)`` and gets projected into K/V
        # space by the reparam MLPs below.
        self.prefix_small = nn.Parameter(torch.empty(prefix_len, self.reparam_hidden, device=device, dtype=dtype))
        self._reparam_k = nn.Linear(self.reparam_hidden, self.kv_dim, bias=True, device=device, dtype=dtype)
        self._reparam_v = nn.Linear(self.reparam_hidden, self.kv_dim, bias=True, device=device, dtype=dtype)

        # Init: prefix_small and both reparam weight matrices use Kaiming
        # uniform so gradients flow at step 1 (a zero-init for the
        # reparam would make ``d_pk / d_prefix_small = 0`` and stall the
        # prefix path until the reparam learned from somewhere else -
        # a chicken-and-egg problem that delays convergence). Biases are
        # zero. The initial prefix contribution is small but non-zero;
        # the optimizer (typically Adam) drives it toward whatever the
        # task loss asks for.
        nn.init.kaiming_uniform_(self.prefix_small, a=math.sqrt(5))
        nn.init.kaiming_uniform_(self._reparam_k.weight, a=math.sqrt(5))
        nn.init.kaiming_uniform_(self._reparam_v.weight, a=math.sqrt(5))
        nn.init.zeros_(self._reparam_k.bias)
        nn.init.zeros_(self._reparam_v.bias)

        # Freeze base MHA - only prefix params train.
        for p in self.base_attn.parameters():
            p.requires_grad = False

    # --- Internal helpers --------------------------------------------------

    def _project_prefix(self) -> tuple[torch.Tensor, torch.Tensor]:
        """Compute prefix K/V via the reparam MLPs (training path)."""
        pk = self._reparam_k(self.prefix_small)
        pv = self._reparam_v(self.prefix_small)
        return pk, pv

    def _expand_to_attn_shape(
        self, pk: torch.Tensor, pv: torch.Tensor, batch_size: int
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Reshape ``[prefix_len, kv_dim]`` → ``[B, num_kv_heads, prefix_len, head_dim]``.

        The buffer / reparam output lives in 2D ``[prefix_len, kv_dim]``.
        Attention expects 4D ``[B, num_kv_heads, prefix_len, head_dim]``;
        we broadcast across the batch dim (the prefix is shared across
        the batch - every sequence in a batch sees the same prefix).
        """
        # ``pk`` is ``[prefix_len, kv_dim]`` with ``kv_dim == num_kv_heads *
        # head_dim``, laid out row-major as ``[pos][head][dim]`` (position
        # varies slowest across rows, head fastest within a row). Attention
        # expects ``[1, num_kv_heads, prefix_len, head_dim]`` (head varies
        # slowest). A bare ``.view`` would reinterpret the SAME flat memory
        # as ``[head][pos][dim]``, so heads read K/V computed from the wrong
        # prefix positions whenever both ``num_kv_heads > 1`` and
        # ``prefix_len > 1`` (RIL ISS-047). Reshape to ``[pos][head][dim]``
        # then permute the position/head axes to get the correct layout.
        pk = pk.reshape(self.prefix_len, self.num_kv_heads, self.head_dim).permute(1, 0, 2).unsqueeze(0)
        pv = pv.reshape(self.prefix_len, self.num_kv_heads, self.head_dim).permute(1, 0, 2).unsqueeze(0)
        # Broadcast across the batch dim.
        pk = pk.expand(batch_size, -1, -1, -1)
        pv = pv.expand(batch_size, -1, -1, -1)
        return pk, pv

    # --- Forward -----------------------------------------------------------

    def forward(
        self,
        hidden_states: torch.Tensor,
        **kwargs,
    ) -> torch.Tensor:
        """Dispatch to the base MHA with the prefix K/V prepended.

        Forwards every keyword argument (``attn_mask``, ``is_causal``,
        ``kv_cache``, ``use_cache``, ``batch_indices``, ``start_pos``,
        ``paged_kv_cache``, ``layer_idx``) to the base MHA. Any caller-
        supplied ``prefix_kv`` is silently overridden - the wrapper owns
        prefix construction.
        """
        kwargs.pop("prefix_kv", None)

        batch_size = hidden_states.shape[0]
        if hasattr(self, "prefix_k") and hasattr(self, "prefix_v"):
            # Folded: static buffers (no reparam MLPs in play).
            assert self.prefix_k is not None  # noqa: S101
            assert self.prefix_v is not None  # noqa: S101
            pk, pv = self._expand_to_attn_shape(self.prefix_k, self.prefix_v, batch_size)
        else:
            pk, pv = self._project_prefix()
            pk, pv = self._expand_to_attn_shape(pk, pv, batch_size)

        return self.base_attn(
            hidden_states,
            prefix_kv=(pk, pv),
            **kwargs,
        )

forward

forward(hidden_states, **kwargs)

Dispatch to the base MHA with the prefix K/V prepended.

Forwards every keyword argument (attn_mask, is_causal, kv_cache, use_cache, batch_indices, start_pos, paged_kv_cache, layer_idx) to the base MHA. Any caller- supplied prefix_kv is silently overridden - the wrapper owns prefix construction.

源代码位于： src/llm/core/prefix_tuning.py

def forward(
    self,
    hidden_states: torch.Tensor,
    **kwargs,
) -> torch.Tensor:
    """Dispatch to the base MHA with the prefix K/V prepended.

    Forwards every keyword argument (``attn_mask``, ``is_causal``,
    ``kv_cache``, ``use_cache``, ``batch_indices``, ``start_pos``,
    ``paged_kv_cache``, ``layer_idx``) to the base MHA. Any caller-
    supplied ``prefix_kv`` is silently overridden - the wrapper owns
    prefix construction.
    """
    kwargs.pop("prefix_kv", None)

    batch_size = hidden_states.shape[0]
    if hasattr(self, "prefix_k") and hasattr(self, "prefix_v"):
        # Folded: static buffers (no reparam MLPs in play).
        assert self.prefix_k is not None  # noqa: S101
        assert self.prefix_v is not None  # noqa: S101
        pk, pv = self._expand_to_attn_shape(self.prefix_k, self.prefix_v, batch_size)
    else:
        pk, pv = self._project_prefix()
        pk, pv = self._expand_to_attn_shape(pk, pv, batch_size)

    return self.base_attn(
        hidden_states,
        prefix_kv=(pk, pv),
        **kwargs,
    )

apply_prefix_tuning

apply_prefix_tuning(model, prefix_len, reparam_hidden=None, target_modules=None)

Wrap every matching MultiHeadAttention in model with prefix tuning.

Mirrors :func:llm.core.lora.apply_lora and :func:llm.core.adalora.apply_adalora so swapping LoRA → AdaLoRA → Prefix Tuning in user code is a one-import change.

参数：

名称	类型	描述	默认
model	Module	The model to adapt. Modified in-place.	必需
prefix_len	int	Number of prefix tokens per attention layer.	必需
reparam_hidden	int | None	Hidden dim of the reparam MLP. None → defaults to kv_dim (the base attention's full K/V dimension).	None
target_modules	list[str] | None	List of module-name substring patterns. If None, every MultiHeadAttention is wrapped. Otherwise only modules whose qualified name contains any of the patterns are wrapped.	None

返回：

类型	描述
Module	The same model, modified in-place.

源代码位于： src/llm/core/prefix_tuning.py

def apply_prefix_tuning(
    model: nn.Module,
    prefix_len: int,
    reparam_hidden: int | None = None,
    target_modules: list[str] | None = None,
) -> nn.Module:
    """Wrap every matching ``MultiHeadAttention`` in ``model`` with prefix tuning.

    Mirrors :func:`llm.core.lora.apply_lora` and
    :func:`llm.core.adalora.apply_adalora` so swapping LoRA → AdaLoRA →
    Prefix Tuning in user code is a one-import change.

    Args:
        model: The model to adapt. Modified in-place.
        prefix_len: Number of prefix tokens per attention layer.
        reparam_hidden: Hidden dim of the reparam MLP. ``None`` → defaults
            to ``kv_dim`` (the base attention's full K/V dimension).
        target_modules: List of module-name substring patterns. If
            ``None``, every ``MultiHeadAttention`` is wrapped. Otherwise
            only modules whose qualified name contains any of the
            patterns are wrapped.

    Returns:
        The same model, modified in-place.
    """
    if target_modules is None:
        target_modules = []

    def should_apply(name: str) -> bool:
        if not target_modules:
            return True
        return any(pattern in name for pattern in target_modules)

    replacements: list[tuple[str, PrefixCapableAttention]] = []
    for name, module in model.named_modules():
        # Filter on both Protocol satisfaction AND the two attributes the
        # wrapper reads at construction time. The ``@runtime_checkable``
        # decorator on the Protocol only inspects method names, so a
        # generic container (Sequential, custom nn.Module) with a
        # ``forward`` method would otherwise sneak through the Protocol
        # check and only blow up later when we read ``num_kv_heads``.
        if (
            isinstance(module, PrefixCapableAttention)
            and hasattr(module, "num_kv_heads")
            and hasattr(module, "head_dim")
            and should_apply(name)
        ):
            replacements.append((name, module))

    for name, module in replacements:
        wrapper = PrefixTuningAttention(
            base_attn=module,
            prefix_len=prefix_len,
            reparam_hidden=reparam_hidden,
        )
        parts = name.split(".")
        parent = model
        for part in parts[:-1]:
            parent = getattr(parent, part)
        setattr(parent, parts[-1], wrapper)

    return model

get_prefix_parameters

get_prefix_parameters(model)

Yield every trainable prefix parameter (prefix_small + reparam MLPs).

Trainers pass this to the optimizer so only the prefix path is updated - the base MHA weights stay frozen. Yields 5 parameters per wrapped attention:

• prefix_small
• _reparam_k.weight, _reparam_k.bias
• _reparam_v.weight, _reparam_v.bias

A wrapper that has been through :func:fold_reparameterization has HAD these deleted (replaced by the static prefix_k/prefix_v buffers) and yields nothing — the PEFT registry's get_parameters for prefix_tuning feeds :func:llm.core.peft.checkpoint.save_peft, which must not crash with AttributeError on a folded-and-share workflow (RIL ISS-209).

源代码位于： src/llm/core/prefix_tuning.py

def get_prefix_parameters(model: nn.Module) -> Iterator[torch.Tensor]:
    """Yield every trainable prefix parameter (``prefix_small`` + reparam MLPs).

    Trainers pass this to the optimizer so only the prefix path is
    updated - the base MHA weights stay frozen. Yields 5 parameters per
    wrapped attention:

    - ``prefix_small``
    - ``_reparam_k.weight``, ``_reparam_k.bias``
    - ``_reparam_v.weight``, ``_reparam_v.bias``

    A wrapper that has been through :func:`fold_reparameterization` has
    HAD these deleted (replaced by the static ``prefix_k``/``prefix_v``
    buffers) and yields nothing — the PEFT registry's ``get_parameters``
    for ``prefix_tuning`` feeds :func:`llm.core.peft.checkpoint.save_peft`,
    which must not crash with ``AttributeError`` on a folded-and-share
    workflow (RIL ISS-209).
    """
    for module in model.modules():
        if isinstance(module, PrefixTuningAttention):
            if not hasattr(module, "prefix_small"):
                # Folded (prefix_small + reparam MLPs removed); no trainable
                # prefix parameters remain to collect.
                continue
            yield module.prefix_small
            yield module._reparam_k.weight
            assert module._reparam_k.bias is not None  # noqa: S101
            yield module._reparam_k.bias
            yield module._reparam_v.weight
            assert module._reparam_v.bias is not None  # noqa: S101
            yield module._reparam_v.bias

fold_reparameterization

fold_reparameterization(model_or_attn)

Collapse reparam MLPs into static prefix buffers for inference.

After fold:

• prefix_small, _reparam_k, _reparam_v are removed from the wrapper (so the optimizer no longer references them and the model state_dict stops carrying them).
• prefix_k, prefix_v are registered as buffers with the final per-layer K/V values. They are not trainable.
• Forward path skips the reparam MLPs and reads the buffers directly.

Idempotent: calling on an already-folded wrapper is a no-op (the existing buffers are preserved exactly).

Works on either a top-level model (walks every wrapper) or a single :class:PrefixTuningAttention directly. Models with no prefix wrappers are left untouched.

参数：

名称	类型	描述	默认
model_or_attn	Module	A model containing wrapped attention modules, or a single :class:PrefixTuningAttention.	必需

返回：

类型	描述
Module	The same object, modified in-place.

源代码位于： src/llm/core/prefix_tuning.py

def fold_reparameterization(model_or_attn: nn.Module) -> nn.Module:
    """Collapse reparam MLPs into static prefix buffers for inference.

    After fold:

    - ``prefix_small``, ``_reparam_k``, ``_reparam_v`` are removed from
      the wrapper (so the optimizer no longer references them and the
      model state_dict stops carrying them).
    - ``prefix_k``, ``prefix_v`` are registered as buffers with the
      final per-layer K/V values. They are not trainable.
    - Forward path skips the reparam MLPs and reads the buffers directly.

    Idempotent: calling on an already-folded wrapper is a no-op (the
    existing buffers are preserved exactly).

    Works on either a top-level model (walks every wrapper) or a single
    :class:`PrefixTuningAttention` directly. Models with no prefix
    wrappers are left untouched.

    Args:
        model_or_attn: A model containing wrapped attention modules, or
            a single :class:`PrefixTuningAttention`.

    Returns:
        The same object, modified in-place.
    """
    for module in model_or_attn.modules():
        if not isinstance(module, PrefixTuningAttention):
            continue
        # Idempotent: skip already-folded wrappers.
        if hasattr(module, "prefix_k") and hasattr(module, "prefix_v"):
            continue
        # Compute final K/V via reparam MLPs (no_grad so the buffers are
        # constants - they're meant to be the deployment-time values).
        with torch.no_grad():
            pk = module._reparam_k(module.prefix_small).detach().clone()
            pv = module._reparam_v(module.prefix_small).detach().clone()
        # Register as buffers (not Parameters - they're static post-fold).
        module.register_buffer("prefix_k", pk)
        module.register_buffer("prefix_v", pv)
        # Drop the trainable reparam path. Using delattr (rather than
        # setting to None) ensures nn.Module removes the entries from
        # ``_parameters`` / ``_modules`` so the optimizer and state_dict
        # stop referencing them.
        delattr(module, "_reparam_k")
        delattr(module, "_reparam_v")
        delattr(module, "prefix_small")

    return model_or_attn

Component Registry

registry

Component registries backed by runtime.Registry.

set_attention_kv_cache_capability

set_attention_kv_cache_capability(name, supports)

Record whether name supports KV cache.

Called from each attention implementation at import time, alongside @register_attention. Validation in ModelConfig.check_consistency raises if a model declares attn_impl that has no capability record.

源代码位于： src/llm/core/registry.py

def set_attention_kv_cache_capability(name: str, supports: bool) -> None:
    """Record whether ``name`` supports KV cache.

    Called from each attention implementation at import time, alongside
    ``@register_attention``. Validation in ``ModelConfig.check_consistency``
    raises if a model declares ``attn_impl`` that has no capability record.
    """
    ATTENTION_KV_CACHE_CAPABILITY[name] = supports

attention_supports_kv_cache

attention_supports_kv_cache(name)

Return whether the registered attention impl supports KV cache.

Raises KeyError if the impl has not declared its capability — that is a registration bug, not a user error. ModelConfig validates this at config-load time.

源代码位于： src/llm/core/registry.py

def attention_supports_kv_cache(name: str) -> bool:
    """Return whether the registered attention impl supports KV cache.

    Raises ``KeyError`` if the impl has not declared its capability — that is
    a registration bug, not a user error. ``ModelConfig`` validates this at
    config-load time.
    """
    return ATTENTION_KV_CACHE_CAPABILITY[name]

ensure_peft_methods_registered

ensure_peft_methods_registered()

Idempotently register built-in methods and load entry points.

Built-ins are registered before the entry-point load so a plugin that claims a built-in name raises / is silently skipped — the built-in is the source of truth. This matches the convention in :func:llm.export.registry.ensure_exporters_registered and :func:llm.generation.registry.ensure_backends_registered.

源代码位于： src/llm/core/peft/registry.py

def ensure_methods_registered() -> None:
    """Idempotently register built-in methods and load entry points.

    Built-ins are registered **before** the entry-point load so a
    plugin that claims a built-in name raises / is silently skipped —
    the built-in is the source of truth. This matches the convention
    in :func:`llm.export.registry.ensure_exporters_registered` and
    :func:`llm.generation.registry.ensure_backends_registered`.
    """
    global _methods_registered
    if _methods_registered:
        return

    # Double-checked locking: the fast guard above is the hot path; the
    # lock serializes the cold-start race so concurrent callers re-check
    # the flag inside the critical section (RIL ISS-119).
    with _method_registration_lock:
        if _methods_registered:
            return

        for method in iter_builtin_methods():
            # ``Registry.register`` raises on duplicate names. Built-ins
            # use stable names so re-import won't trigger duplicates;
            # third parties load after this point and the entry-point
            # loader defaults to ``overwrite=False`` so plugins claiming
            # built-in names are silently skipped (matching the export
            # convention).
            if method.name not in PEFT_REGISTRY:
                PEFT_REGISTRY.register(method.name, method)

        load_entry_point_registry("llm.peft_methods", PEFT_REGISTRY)
        _methods_registered = True