update

README.md

@@ -0,0 +1,92 @@
+---
+library_name: transformers
+tags:
+  - custom_generate
+  - sampling
+  - kvcache
+---
+
+# Sampling with KV Cache
+
+## Description
+A clean, hackable implementation of sampling (also called ancestral sampling or multinomial sampling) with full KV cache support. This is a simplified alternative to the complex generation mixin in transformers, designed for readability and ease of modification while maintaining full performance.
+
+The implementation supports both sampling and greedy decoding modes, with optional temperature scaling and top-k/top-p filtering.
+
+## Base model
+- [HuggingFaceTB/SmolLM2-135M-Instruct](https://huggingface.co/HuggingFaceTB/SmolLM2-135M-Instruct)
+
+## Model compatibility
+Most transformer LLM/VLM models trained for causal language modeling.
+
+## Relevant Arguments
+- `temperature` (float): Sampling temperature (default: 1.0, higher = more random)
+- `top_k` (int): Only consider top-k most probable tokens (default: None)
+- `top_p` (float): Only consider tokens with cumulative probability <= top_p (default: None)
+- `do_sample` (bool): Whether to use sampling (True, default) or greedy decoding (False)
+
+### Logits Processing Order
+Logits processors are applied in sequence: `temperature → softmax → top_k → top_p` (same as HuggingFace's `LogitProcessor` system). Temperature scaling occurs before top-p filtering, affecting the probability distribution that top-p operates on.
+
+For example, with `temperature=1.0`, `top_p=0.9` might include tokens A, B, C. With `temperature=0.5`, probability mass is much more concentrated, so `top_p=0.9` might only include token A.
+
+## Outputs
+When `return_dict_in_generate=True`, returns a dictionary with:
+- `sequences`: Generated token IDs
+- `scores`: Log probabilities of sampled tokens (with temperature/sampling modifications)
+- `logprobs`: Original model log probabilities (T=1, no modifications)
+Otherwise, returns a tensor of generated token IDs.
+
+## Example usage
+
+```py
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
+model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct", device_map="auto")
+
+inputs = tokenizer(["The quick brown"], return_tensors="pt").to(model.device)
+
+# Basic sampling
+gen_out = model.generate(**inputs, custom_generate="manueldeprada/sampling_with_kvcache", trust_remote_code=True)
+
+# With temperature
+gen_out = model.generate(**inputs, custom_generate="manueldeprada/sampling_with_kvcache", temperature=0.8, trust_remote_code=True)
+
+# With top-k
+gen_out = model.generate(**inputs, custom_generate="manueldeprada/sampling_with_kvcache", top_k=50, trust_remote_code=True)
+
+# With top-p (nucleus sampling)
+gen_out = model.generate(**inputs, custom_generate="manueldeprada/sampling_with_kvcache", top_p=0.9, trust_remote_code=True)
+
+# Greedy decoding (no sampling)
+gen_out = model.generate(**inputs, custom_generate="manueldeprada/sampling_with_kvcache", do_sample=False, trust_remote_code=True)
+
+# Get detailed output with probabilities
+gen_out = model.generate(
+    **inputs, 
+    custom_generate="manueldeprada/sampling_with_kvcache", 
+    return_dict_in_generate=True,
+    trust_remote_code=True
+)
+print(f"Generated text: {tokenizer.batch_decode(gen_out['sequences'], skip_special_tokens=True)}")
+print(f"Sampling scores: {gen_out['scores']}")
+print(f"Model log probabilities: {gen_out['logprobs']}")
+```
+
+## Algorithm
+1. Initialize KV cache and prepare input sequences
+2. For each generation step:
+   - Get logits from the model for the current sequence
+   - Apply temperature scaling to logits
+   - Optionally apply top-k filtering (keep only top-k tokens)
+   - Optionally apply top-p filtering (nucleus sampling)
+   - Convert to probabilities using softmax
+   - Sample from the probability distribution (or take argmax for greedy)
+   - Append the selected token to the sequence
+   - Update KV cache and track sequence completion
+3. Return generated sequences and probability information
+
+
+
+

config.json

@@ -0,0 +1,29 @@
+{
+  "architectures": [
+    "LlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 0,
+  "eos_token_id": 0,
+  "hidden_act": "silu",
+  "hidden_size": 576,
+  "initializer_range": 0.041666666666666664,
+  "intermediate_size": 1536,
+  "is_llama_config": true,
+  "max_position_embeddings": 8192,
+  "model_type": "llama",
+  "num_attention_heads": 9,
+  "num_hidden_layers": 30,
+  "num_key_value_heads": 3,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-05,
+  "rope_interleaved": false,
+  "rope_scaling": null,
+  "rope_theta": 100000,
+  "tie_word_embeddings": true,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.40.1",
+  "use_cache": true,
+  "vocab_size": 49152
+}

custom_generate/generate.py

@@ -0,0 +1,204 @@
+import torch
+
+
+def next_logits_with_cache_update(model, model_kwargs, input_ids):
+    """
+    Gets the next token logits and updates the KV cache:
+    - Runs the model forward pass
+    - Extracts logits for the last token
+    - Updates the KV cache
+    - Returns updated `model_kwargs` and `logits`
+    
+    Args:
+        model: The language model
+        model_kwargs: Model keyword arguments including KV cache
+        input_ids: Current input token IDs
+    
+    Returns:
+        Updated model_kwargs, logits for the next token
+    """
+    model_inputs = model.prepare_inputs_for_generation(input_ids, **model_kwargs)
+    with torch.no_grad():
+        outputs = model(**model_inputs, return_dict=True)
+    
+    logits = outputs.logits[:, -1].detach()
+    model_kwargs = model._update_model_kwargs_for_generation(
+        outputs, model_kwargs, is_encoder_decoder=model.config.is_encoder_decoder
+    )
+    del outputs
+    return model_kwargs, logits
+
+def init_gen(model_kwargs, model, max_new_tokens, bos_token_id):
+    """
+    Initializes the generation process and prepares the KV cache:
+    - Sets up input sequences and model inputs
+    - Prepares the KV cache for generation
+    - Returns updated `model_kwargs` and `input_ids`
+    
+    Args:
+        model_kwargs: Model keyword arguments
+        model: The language model
+        max_new_tokens: Maximum number of new tokens to generate
+        bos_token_id: Beginning-of-sequence token ID
+    
+    Returns:
+        Model keyword arguments and input token IDs
+    """
+
+    input_ids, model_input_name, model_kwargs = model._prepare_model_inputs(
+        None, bos_token_id, model_kwargs
+    )
+    
+    batch_size = input_ids.shape[0]
+    model._prepare_cache_for_generation(
+        model.generation_config, model_kwargs, None, batch_size, 
+        max_cache_length=max_new_tokens, device=input_ids.device
+    )
+    
+    # Get initial cache position
+    model_kwargs = model._get_initial_cache_position(input_ids.shape[1], input_ids.device, model_kwargs)
+    return model_kwargs, input_ids
+
+def _apply_top_k(ps, model):
+    """Apply top-k filtering to probabilities."""
+    if not hasattr(model, 'generation_config') or not hasattr(model.generation_config, 'top_k'):
+        return ps
+    
+    top_k = model.generation_config.top_k
+    if top_k is None or top_k >= ps.size(-1):
+        return ps
+    
+    indices_to_remove = ps < torch.topk(ps, top_k)[0][..., -1, None]
+    ps[indices_to_remove] = 0.0
+    return ps / ps.sum(dim=-1, keepdim=True)
+
+def _apply_top_p(ps, model):
+    """Apply top-p (nucleus) filtering to probabilities."""
+    if not hasattr(model, 'generation_config') or not hasattr(model.generation_config, 'top_p'):
+        return ps
+    
+    top_p = model.generation_config.top_p
+    if top_p is None or top_p >= 1.0:
+        return ps
+    
+    sorted_probs, sorted_indices = torch.sort(ps, descending=True)
+    cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
+    
+    sorted_indices_to_remove = cumulative_probs > top_p
+    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+    sorted_indices_to_remove[..., 0] = 0
+    
+    indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
+    ps[indices_to_remove] = 0.0
+    return ps / ps.sum(dim=-1, keepdim=True)
+
+def sampling_with_kvcache(model_kwargs, model, eos_token_ids, pad_token_id, bos_token_id, do_sample=True, max_new_tokens=20, temperature=1.0):
+    """
+    Sampling implementation with proper KV caching.
+    
+    Args:
+        prompts: List of input prompts
+        model: The language model
+        max_new_tokens: Maximum number of new tokens to generate
+        eos_token_ids: List of end-of-sequence token IDs
+        pad_token_id: Padding token ID
+        bos_token_id: Beginning-of-sequence token ID
+        max_new_tokens: Maximum number of new tokens to generate
+    
+    Returns:
+        Generated sequences, log probabilities, and metadata
+    """
+    # Initialize the generation process and prepare the KV cache
+    model_kwargs, input_ids = init_gen(model_kwargs, model, max_new_tokens, bos_token_id)
+    batch_size, _ = input_ids.shape
+
+    # Keeps track of which sequences are finished and their lengths
+    active_seqs = input_ids.new_ones((batch_size, 1), dtype=torch.bool)
+    # Modified log probabilities of the sequences
+    scores = torch.zeros((batch_size, max_new_tokens), dtype=model.dtype)
+    # Unfiltered sequence log probabilities (temperature=1, no sampling processors applied)
+    logprobs = torch.zeros((batch_size, max_new_tokens), dtype=model.dtype)
+
+    for i in range(max_new_tokens):
+        # Get the next token probabilities and update the KV cache
+        model_kwargs, logits = next_logits_with_cache_update(model, model_kwargs, input_ids)
+        # Store original model probabilities (temperature=1, no sampling processors applied)
+        model_ps = logits.softmax(-1)
+        
+        # Logit processors (temperature, top-k, top-p). We can chain these!
+        ps = (logits/temperature).softmax(-1)
+        ps = _apply_top_k(ps, model)
+        ps = _apply_top_p(ps, model)
+        
+        # Sample the next token and gather the log probabilities
+        if do_sample: # Sampling
+            next_token_ids = torch.multinomial(ps, 1) * active_seqs + pad_token_id * ~active_seqs
+        else: # Greedy decoding
+            next_token_ids = torch.argmax(ps, dim=-1).unsqueeze(-1) * active_seqs + pad_token_id * ~active_seqs
+        next_token_logprobs = ps.gather(-1, next_token_ids).log()
+        next_token_model_logprobs = model_ps.gather(-1, next_token_ids).log()
+        
+        input_ids = torch.cat([input_ids, next_token_ids], dim=-1)
+        scores[:, i] = (next_token_logprobs * active_seqs).squeeze()
+        logprobs[:, i] = (next_token_model_logprobs * active_seqs).squeeze()
+        
+        active_seqs &= ~torch.isin(next_token_ids, eos_token_ids)
+        if active_seqs.sum() == 0:
+            break
+    return input_ids.detach().cpu(), scores[:,:i+1], logprobs[:,:i+1]
+
+def generate(model, **kwargs):
+    """
+    Sampling strategy - multinomial sampling with temperature and optional top-k/top-p filtering.
+    Simple implementation with proper KV caching support.
+    
+    Args:
+        model: The language model
+        model_kwargs: Model keyword arguments from the tokenizer
+        generation_config: Generation configuration
+        temperature: Sampling temperature (higher = more random)
+        top_k: Only consider top-k most probable tokens
+        top_p: Only consider tokens with cumulative probability <= top_p
+        **kwargs: Additional arguments
+    
+    Returns:
+        Generated token IDs
+    """
+    generation_config = model.generation_config
+    max_new_tokens = kwargs.get('max_new_tokens', generation_config.max_new_tokens)
+    max_new_tokens = 512 if max_new_tokens is None else max_new_tokens
+    do_sample = kwargs.get('do_sample', True)
+    eos_token_ids = kwargs.get('eos_token_ids', generation_config.eos_token_id)
+    if eos_token_ids is None:
+        raise ValueError("Model generation config does not have an EOS token id. You must provide it to generate() with the eos_token_ids argument.")
+    eos_token_ids = torch.as_tensor(eos_token_ids, device=model.device)
+    if eos_token_ids is not None and eos_token_ids.ndim == 0:
+        eos_token_ids = eos_token_ids.unsqueeze(0)
+    
+    pad_token_id = kwargs.get('pad_token_id', generation_config.pad_token_id if generation_config.pad_token_id is not None else eos_token_ids[0])
+    bos_token_id = kwargs.get('bos_token_id', generation_config.bos_token_id)
+    if bos_token_id is None:
+        raise ValueError("Model generation config does not have a BOS token id. You must provide it to generate() with the bos_token_id argument.")
+    temperature = kwargs.get('temperature', 1.0)
+    return_dict = kwargs.get('return_dict_in_generate', False)
+
+    generated_ids, scores, logprobs = sampling_with_kvcache(
+        model_kwargs=kwargs,
+        model=model,
+        eos_token_ids=eos_token_ids,
+        pad_token_id=pad_token_id,
+        bos_token_id=bos_token_id,
+        do_sample=do_sample,
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+    )
+
+    if return_dict:
+        return {
+            "sequences": generated_ids,
+            "scores": scores,
+            "logprobs": logprobs,
+        }
+    else:
+        return generated_ids
+

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 0,
+  "eos_token_id": 0,
+  "transformers_version": "4.40.1"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:80521b40281d6ce74e35c9282c22539e75aa0ac8578892b2a59955ef78d55da1
+size 269060552

special_tokens_map.json

@@ -0,0 +1,42 @@
+{
+  "additional_special_tokens": [
+    "<|endoftext|>",
+    "<|im_start|>",
+    "<|im_end|>",
+    "<repo_name>",
+    "<reponame>",
+    "<file_sep>",
+    "<filename>",
+    "<gh_stars>",
+    "<issue_start>",
+    "<issue_comment>",
+    "<issue_closed>",
+    "<jupyter_start>",
+    "<jupyter_text>",
+    "<jupyter_code>",
+    "<jupyter_output>",
+    "<jupyter_script>",
+    "<empty_output>"
+  ],
+  "bos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,167 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<repo_name>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4": {
+      "content": "<reponame>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "5": {
+      "content": "<file_sep>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "6": {
+      "content": "<filename>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "7": {
+      "content": "<gh_stars>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "8": {
+      "content": "<issue_start>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "9": {
+      "content": "<issue_comment>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "10": {
+      "content": "<issue_closed>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "11": {
+      "content": "<jupyter_start>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "12": {
+      "content": "<jupyter_text>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "13": {
+      "content": "<jupyter_code>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "14": {
+      "content": "<jupyter_output>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "15": {
+      "content": "<jupyter_script>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "16": {
+      "content": "<empty_output>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<|endoftext|>",
+    "<|im_start|>",
+    "<|im_end|>",
+    "<repo_name>",
+    "<reponame>",
+    "<file_sep>",
+    "<filename>",
+    "<gh_stars>",
+    "<issue_start>",
+    "<issue_comment>",
+    "<issue_closed>",
+    "<jupyter_start>",
+    "<jupyter_text>",
+    "<jupyter_code>",
+    "<jupyter_output>",
+    "<jupyter_script>",
+    "<empty_output>"
+  ],
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "model_max_length": 8192,
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>",
+  "vocab_size": 49152
+}