Add docstrings to some modules and classes (#100)

cacheflow/model_executor/layers/activation.py

@@ -1,3 +1,4 @@
+"""Custom activation functions."""
 import torch
 import torch.nn as nn
 
@@ -5,6 +6,10 @@ from cacheflow import activation_ops
 
 
 class SiluAndMul(nn.Module):
+    """An activation function for SwiGLU.
+
+    The function computes x -> silu(x[:d]) * x[d:] where d = x.shape[1] // 2.
+    """
 
     def __init__(self):
         super().__init__()

cacheflow/model_executor/layers/attention.py

@@ -1,3 +1,4 @@
+"""Multi-head attention."""
 from typing import Optional
 
 import torch
@@ -11,6 +12,32 @@ from cacheflow.model_executor.input_metadata import InputMetadata
 
 
 class GPTCacheFlowAttention(nn.Module):
+    """GPT-style multi-head attention.
+
+    This class takes flattened 1D query, key, and value tensors as input. The
+    input 1D tensors can be split into three parts: the prompt tokens, the
+    generation tokens, and the paddings.
+
+    |<------------------------------------- num_valid_tokens ------------------------------------->|
+    |<--------------- num_prompt_tokens -------------->|<------- num_generation_tokens (M) ------->|
+    |<--prompt_0-->|<--prompt_1-->|...|<--prompt_N-1-->|<--generation_0-->|...|<--generation_M-1-->|<--padding-->|
+
+    The prompts might have different lengths, while the generation tokens always
+    have length 1. The paddings are appended to make the input length a multiple
+    of 8, which is desirable for Tensor Cores.
+
+    The class does the following:
+    1. Perform multi_query_kv_attention for the prompts. This operation does
+        not use the KV cache.
+    2. Wait for the cache operations (e.g., swap, copy) to finish. The cache
+        operations are issued by the cache engine before executing the forward
+        pass of the model, and they are executed asynchronously.
+    3. Reshape and store the input key and value tensors in the KV cache.
+    4. Perform single_query_cached_kv_attention for the generation tokens.
+        This operation reads the previous key and value tensors from the KV
+        cache.
+    5. Output a flattened 1D tensor.
+    """
 
     def __init__(self, scale: float) -> None:
         super().__init__()
@@ -157,7 +184,7 @@ class GPTNeoXCacheFlowAttention(GPTCacheFlowAttention):
         torch_dtype = torch.get_default_dtype()
         cache = cache.to(torch_dtype)
         # Embedding size: [max_position, rotary_dim]
-        self.register_buffer('cos_sin_cache', cache, persistent=False)
+        self.register_buffer("cos_sin_cache", cache, persistent=False)
 
     def forward(
         self,

cacheflow/model_executor/layers/layernorm.py

@@ -1,3 +1,4 @@
+"""Custom normalization layers."""
 import torch
 import torch.nn as nn
 
@@ -5,6 +6,11 @@ from cacheflow import layernorm_ops
 
 
 class RMSNorm(nn.Module):
+    """Root mean square normalization.
+
+    Computes x -> w * x / sqrt(E[x^2] + eps) where w is the learned weight.
+    Refer to https://arxiv.org/abs/1910.07467
+    """
 
     def __init__(
         self,

cacheflow/model_executor/layers/sampler.py

@@ -1,3 +1,4 @@
+"""A layer that samples the next tokens from the model's outputs."""
 from typing import Dict, List, Tuple
 
 import numpy as np
@@ -12,6 +13,19 @@ from cacheflow.sequence import SequenceOutputs
 
 
 class Sampler(nn.Module):
+    """Samples the next tokens from the model's outputs.
+
+    This layer does the following:
+    1. Discard the hidden states that are not used for sampling (i.e., all
+        tokens except the final one in each prompt).
+    2. Compute the logits for the next tokens.
+    3. Apply presence and frequency penalties.
+    4. Apply temperature scaling.
+    5. Apply top-p and top-k truncation.
+    6. Sample the next tokens.
+    Here, each sequence group within the batch can have different sampling
+    parameters (e.g., sampling method, temperature, top-p, top-k, etc.).
+    """
 
     def __init__(self, vocab_size: int) -> None:
         super().__init__()