llama.cpp/examples/model-conversion/scripts/causal/run-org-model-multi-token.py

#!/usr/bin/env python3

import argparse
import os
import importlib
from pathlib import Path
import re

from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
import torch
import numpy as np

### If you want to dump RoPE activations, apply this monkey patch to the model
### class from Transformers that you are running (replace apertus.modeling_apertus
### with the proper package and class for your model
### === START ROPE DEBUG ===
# from transformers.models.apertus.modeling_apertus import apply_rotary_pos_emb

# orig_rope = apply_rotary_pos_emb
# torch.set_printoptions(threshold=float('inf'))
# torch.set_printoptions(precision=6, sci_mode=False)

# def debug_rope(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
#     # log inputs
#     summarize(q, "RoPE.q_in")
#     summarize(k, "RoPE.k_in")

#     # call original
#     q_out, k_out = orig_rope(q, k, cos, sin, position_ids, unsqueeze_dim)

#     # log outputs
#     summarize(q_out, "RoPE.q_out")
#     summarize(k_out, "RoPE.k_out")

#     return q_out, k_out

# # Patch it
# import transformers.models.apertus.modeling_apertus as apertus_mod  # noqa: E402
# apertus_mod.apply_rotary_pos_emb = debug_rope
### == END ROPE DEBUG ===

token_counter = {}

def summarize(tensor: torch.Tensor, name: str, max_seq: int = 3, max_vals: int = 3):
    global token, token_counter
    """
    Print a tensor in llama.cpp debug style.

    Supports:
    - 2D tensors (seq, hidden)
    - 3D tensors (batch, seq, hidden)
    - 4D tensors (batch, seq, heads, dim_per_head) via flattening heads × dim_per_head
    - 5D tensors

    Shows first and last max_vals of each vector per sequence position.
    """
    t = tensor.detach().to(torch.float32).cpu()
    ten_shape = t.shape
    while t.ndim > 4:
        t = t.squeeze(0)

    # Determine dimensions
    if t.ndim == 3:
        _, s, _ = t.shape
    elif t.ndim == 2:
        _, s = 1, t.shape[0]
        t = t.unsqueeze(0)
    elif t.ndim == 4:
        _, s, _, _ = t.shape
    
    else:
        print(f"Skipping tensor due to unsupported dimensions: {t.ndim}")
        return

    print(f"ggml_debug: {name} = (f32)  ... = {{{ten_shape}}}")
    print("                                     [")
    print("                                      [")

    # Determine indices for first and last sequences
    first_indices = list(range(min(s, max_seq)))
    last_indices = list(range(max(0, s - max_seq), s))

    # Check if there's an overlap between first and last indices or if we're at the edge case of s = 2 * max_seq
    has_overlap = bool(set(first_indices) & set(last_indices)) or (max_seq * 2 == s)

    # Combine indices
    if has_overlap:
        # If there's overlap, just use the combined unique indices
        indices = sorted(list(set(first_indices + last_indices)))
        separator_index = None
    else:
        # If no overlap, we'll add a separator between first and last sequences
        indices = first_indices + last_indices
        separator_index = len(first_indices)

    for i, si in enumerate(indices):
        # Add separator if needed
        if separator_index is not None and i == separator_index:
            print("                                       ...")

        # Extract appropriate slice
        vec = t[0, si]
        if vec.ndim == 2:  # 4D case: flatten heads × dim_per_head
            flat = vec.flatten().tolist()
        else:  # 2D or 3D case
            flat = vec.tolist()

        # First and last slices
        first = flat[:max_vals]
        last = flat[-max_vals:] if len(flat) >= 2 * max_vals else flat
        first_str = ", ".join(f"{v:12.4f}" for v in first)
        last_str = ", ".join(f"{v:12.4f}" for v in last)

        if len(flat) >= 2 * max_vals:
            print(f"                                       [{first_str}, ..., {last_str}]")
        else:
            print(f"                                       [{last_str}]")

    print("                                      ],")
    print("                                     ]")
    print(f"                                     sum = {t.sum().item():.6f}\n")

    pattern = r"model\.layers\.[0-9]+_out"
    pattern2 = r"recurrent_cache_[0-9]+"
    if re.fullmatch(pattern, name) or re.fullmatch(pattern2, name):
        if name not in token_counter:
            token_counter[name] = 1
        else:
            token_counter[name] = token_counter[name] + 1
        save_tensor(t, f"reference/tensors/org/{name}_{token_counter[name]}.bin")

from transformers.models.qwen3_next.modeling_qwen3_next import torch_causal_conv1d_update, apply_rotary_pos_emb, l2norm  # noqa: E402
orig_conv1d_update = torch_causal_conv1d_update
orig_rope = apply_rotary_pos_emb
import torch.nn.functional as F  # noqa: E402
import typing  # noqa: E402

def patched_torch_causal_conv1d_update(
    hidden_states,
    conv_state,
    weight,
    bias=None,
    activation=None,
):
    _, hidden_size, seq_len = hidden_states.shape
    state_len = conv_state.shape[-1]
    summarize(hidden_states, "hidden_states_in")
    summarize(conv_state, "conv_state_in")

    hidden_states_new = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype)
    summarize(hidden_states_new, "hidden_states_new")
    summarize(hidden_states_new[:, :, -state_len:], "hidden_states_to_copy")
    summarize(conv_state, "conv_state_pre")
    conv_state.copy_(hidden_states_new[:, :, -state_len:])
    summarize(conv_state, "conv_state_post")
    out = F.conv1d(hidden_states_new, weight.unsqueeze(1), bias, padding=0, groups=hidden_size)
    summarize(out, "out")
    summarize(out[:, :, -seq_len:], "out_proper")
    out = F.silu(out[:, :, -seq_len:])
    summarize(out, "out_silu")
    out = out.to(hidden_states.dtype)
    return out

already_dumped_rope = False

def save_tensor(tensor, filename):
    """Save tensor to binary file with shape information."""
    # Ensure tensors directory exists
    os.makedirs(os.path.dirname(filename), exist_ok=True)
    
    # Convert to numpy and save
    np_array = tensor.detach().cpu().numpy()
    
    # Save shape first (4 int64 values), then data
    with open(filename, 'wb') as f:
        shape = list(np_array.shape)
        while len(shape) < 4:
            shape.insert(0, 0)
        
        # Write shape as int64
        shape_array = np.array(shape, dtype=np.int64)
        f.write(shape_array.tobytes())
        
        # Write data as float32
        np_array_float32 = np_array.astype(np.float32)
        f.write(np_array_float32.tobytes())

def patched_apply_rope(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
    global already_dumped_rope

    # log inputs
    summarize(q, "RoPE.q_in")
    summarize(k, "RoPE.k_in")
    summarize(cos, "cos")
    summarize(sin, "sin")
    # if q.shape[1] == 2 and k.shape[1] == 1 and k.shape[2] == 1 and not already_dumped_rope:
    #     already_dumped_rope = True
    #     print("Dumping input tensors")
    #     save_tensor(q, "reference/tensors/testrope_q_in.bin")
    #     save_tensor(k, "reference/tensors/testrope_k_in.bin")
    #     save_tensor(cos, "reference/tensors/testrope_cos_in.bin")
    #     save_tensor(sin, "reference/tensors/testrope_sin_in.bin")

    if position_ids:
        summarize(position_ids, "position_ids")
    # print(f"Rotary dim is {cos.unsqueeze(unsqueeze_dim).shape[-1]}")

    # call original
    q_out, k_out = orig_rope(q, k, cos, sin, position_ids, unsqueeze_dim)

    # log outputs
    summarize(q_out, "RoPE.q_out")
    summarize(k_out, "RoPE.k_out")

    return q_out, k_out

def patched_torch_chunk_gated_delta_rule(
    query,
    key,
    value,
    g,
    beta,
    chunk_size=64,
    initial_state=None,
    output_final_state=False,
    use_qk_l2norm_in_kernel=False,
    long=False
):
    torch.set_printoptions(threshold=10_000_000, sci_mode=False, precision=10, linewidth=200)
    initial_dtype = query.dtype
    [ summarize(x, y) for (x, y) in ((query, "q_prenorm"), (key, "k_prenorm")) ]
    if use_qk_l2norm_in_kernel:
        query = l2norm(query, dim=-1, eps=1e-6)
        key = l2norm(key, dim=-1, eps=1e-6)
    [ summarize(x, y) for (x, y) in ((query, "q_orig"), (key, "k_orig"), (value, "v_orig"), (beta, "b_orig"), (g, "g_orig")) ]
    query, key, value, beta, g = [
        x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g)
    ]
    [ summarize(x, y) for (x, y) in ((query, "q_tra"), (key, "k_tra"), (value, "v_tra"), (beta, "b_tra"), (g, "g_tra")) ]
    batch_size, sequence_length, num_heads, k_head_dim = key.shape
    print(f"batch_size = {batch_size}, seq_len = {sequence_length}, num_heads = {num_heads}, k_head_dim = {k_head_dim}")
    v_head_dim = value.shape[-1]
    pad_size = (chunk_size - num_heads % chunk_size) % chunk_size
    print(f"Pad size = {pad_size}, chunk_size = {chunk_size}")
    query = F.pad(query, (0, 0, 0, pad_size))
    key = F.pad(key, (0, 0, 0, pad_size))
    value = F.pad(value, (0, 0, 0, pad_size))
    beta = F.pad(beta, (0, pad_size))
    g = F.pad(g, (0, pad_size))
    [ summarize(x, y) for (x, y) in ((query, "q_pad"), (key, "k_pad"), (value, "v_pad"), (beta, "b_pad"), (g, "g_pad")) ]
    tot_heads = num_heads + pad_size
    scale = 1 / (query.shape[-1] ** 0.5)
    print(f"Scale for delta is {scale} (from {query.shape[-1]})")
    query = query * scale

    summarize(query, "q_scaled")
    summarize(key, "k")
    summarize(beta.unsqueeze(-1), "beta")
    v_beta = value * beta.unsqueeze(-1)
    k_beta = key * beta.unsqueeze(-1)
    summarize(k_beta, "k_beta")
    summarize(v_beta, "v_beta")
    # reshape to chunks
    query, key, value, k_beta, v_beta = [
        x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1]) for x in (query, key, value, k_beta, v_beta)
    ]
    g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size)
    [ summarize(x, y) for (x, y) in ((query, "q_resh"), (k_beta, "k_beta_resh"), (v_beta, "v_beta_resh"), (key, "k_resh"), (value, "v_resh")) ]

    mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=0)

    # chunk decay
    g = g.cumsum(dim=-1)
    summarize(g, "g_cumsum")
    sub = g.unsqueeze(-1) - g.unsqueeze(-2)
    bt1, bt2 = torch.broadcast_tensors(g.unsqueeze(-1), g.unsqueeze(-2))
    summarize(bt1, "bt1")
    summarize(bt2, "bt2")
    summarize(sub, "sub")
    decay_mask = sub.tril()
    summarize(decay_mask, "sub_tril")
    decay_mask = decay_mask.exp()
    summarize(decay_mask, "sub_tril_exp")
    decay_mask = decay_mask.float()
    summarize(decay_mask, "sub_tril_exp_float")
    decay_mask = decay_mask.tril()
    summarize(decay_mask, "decay_mask")
    k_t = key.transpose(-1, -2)
    summarize(k_t, "k_t")
    kmul = k_beta @ k_t
    summarize(kmul, "k_beta @ k_t")
    #if not long:
        #print(f"k_beta @ k_t:\n{kmul[:,:,:,:8,:8]}\n\n")
    kmul_decay = kmul * decay_mask
    summarize(kmul_decay, "(k_beta @ k_t) * decay_mask")
    attn = -(kmul_decay).masked_fill(mask, 0)
    summarize(attn, "attn_in")
    for i in range(1, chunk_size):
        row = attn[..., i, :i].clone()
        sub = attn[..., :i, :i].clone()
        attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2)
        #if i <= num_heads and not long: 
            #print(f"Chunk {i}: row:\n{row}\n\nsub:\n{sub}\nrow_unsq:\n{row.unsqueeze(-1)}\nrow_unsq * sub:\n{row.unsqueeze(-1)*sub}\n")
            #print(f"attn => sum = {attn[..., i, :i].sum()}, tensor: \n{attn[..., i, :i]}\n\n")
    summarize(attn, "attn_chunks")
    attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)
    summarize(attn, "attn_eye")
    
    value = attn @ v_beta
    summarize(value, "value")
        
    k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1))
    summarize(k_cumdecay, "k_cumdecay")
    
    last_recurrent_state = (
        torch.zeros(batch_size, sequence_length, k_head_dim, v_head_dim).to(value)
        if initial_state is None
        else initial_state.to(value)
    )
    core_attn_out = torch.zeros_like(value)
    mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=1)

    # for each chunk
    for i in range(0, tot_heads // chunk_size):
        print(f"\n=== Processing chunk {i} ===")
        q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i]
        summarize(q_i, f"q_i_chunk_{i}")
        summarize(k_i, f"k_i_chunk_{i}")
        summarize(v_i, f"v_i_chunk_{i}")
        
        attn = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0)
        summarize(attn, f"attn_chunk_{i}")
        
        v_prime = (k_cumdecay[:, :, i]) @ last_recurrent_state
        summarize(v_prime, f"v_prime_chunk_{i}")
        
        v_new = v_i - v_prime
        summarize(v_new, f"v_new_chunk_{i}")
        
        attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state
        summarize(attn_inter, f"attn_inter_chunk_{i}")
        
        core_attn_out[:, :, i] = attn_inter + attn @ v_new
        summarize(core_attn_out[:, :, i], f"core_attn_out_chunk_{i}")
        
        g_last = g[:, :, i, -1, None, None].exp()
        summarize(g_last, f"g_last_chunk_{i}")
        
        g_diff_exp = (g[:, :, i, -1, None] - g[:, :, i]).exp()
        last_recurrent_state = (
            last_recurrent_state * g_last
            + (k_i * g_diff_exp[..., None]).transpose(-1, -2) @ v_new
        )
        summarize(last_recurrent_state, f"updated_state_chunk_{i}")

    if not output_final_state:
        last_recurrent_state = None
    core_attn_out = core_attn_out.reshape(core_attn_out.shape[0], core_attn_out.shape[1], -1, core_attn_out.shape[-1])
    core_attn_out = core_attn_out[:, :, :num_heads]
    core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype)
    summarize(core_attn_out, "attn_out")
    if not long:
        print(f"attn_out:\n{core_attn_out}\n\n")
        
    if isinstance(last_recurrent_state, torch.Tensor):
        summarize(last_recurrent_state, "state_out")
        if not long:
            print(f"state_out:\n{last_recurrent_state}\n\n")
    return core_attn_out, last_recurrent_state


def patched_torch_recurrent_gated_delta_rule(
    query, key, value, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel=False
):
    initial_dtype = query.dtype
    if use_qk_l2norm_in_kernel:
        query = l2norm(query, dim=-1, eps=1e-6)
        key = l2norm(key, dim=-1, eps=1e-6)
    query, key, value, beta, g = [
        x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g)
    ]
    summarize(query, "q_t")
    summarize(key, "k_t")
    summarize(value, "v_t")
    summarize(beta, "beta_t")
    summarize(g, "g_t")

    batch_size, num_heads, sequence_length, k_head_dim = key.shape
    v_head_dim = value.shape[-1]
    scale = 1 / (query.shape[-1] ** 0.5)
    query = query * scale

    summarize(query, "q_scaled")
    if initial_state is not None:
        summarize(initial_state, "initial_state")

    core_attn_out = torch.zeros(batch_size, num_heads, sequence_length, v_head_dim).to(value)
    last_recurrent_state = (
        torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim).to(value)
        if initial_state is None
        else initial_state.to(value)
    )

    for i in range(sequence_length):
        q_t = query[:, :, i]
        k_t = key[:, :, i]
        v_t = value[:, :, i]
        g_t = g[:, :, i].exp().unsqueeze(-1).unsqueeze(-1)
        summarize(g_t, "g_exp_unsq")
        beta_t = beta[:, :, i].unsqueeze(-1)
        summarize(beta_t, "beta_t_unsq")

        last_recurrent_state = last_recurrent_state * g_t
        summarize(last_recurrent_state, "gated_state")
        k_unsq = k_t.unsqueeze(-1)
        summarize(k_unsq, "k_unsqueeze")
        state_k = last_recurrent_state * k_unsq
        summarize(state_k, "state_k_product")
        kv_mem = state_k.sum(dim=-2)
        summarize(kv_mem, "kv_mem")
        delta = (v_t - kv_mem) * beta_t
        summarize(delta, "delta")
        k_delta = k_t.unsqueeze(-1) * delta.unsqueeze(-2)
        summarize(k_delta, "k_delta")
        last_recurrent_state = last_recurrent_state + k_delta
        summarize(last_recurrent_state, "state_plus_k_delta")
        state_q_prod = last_recurrent_state * q_t.unsqueeze(-1)
        summarize(state_q_prod, "state_q_product")
        core_attn_out[:, :, i] = state_q_prod.sum(dim=-2)
        summarize(core_attn_out, "core_attn_out")

    if not output_final_state:
        last_recurrent_state = None
    core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype)
    return core_attn_out, last_recurrent_state

import transformers.models.qwen3_next.modeling_qwen3_next as qwen_mod  # noqa: E402
qwen_mod.torch_chunk_gated_delta_rule = patched_torch_chunk_gated_delta_rule
qwen_mod.torch_causal_conv1d_update = patched_torch_causal_conv1d_update
qwen_mod.apply_rotary_pos_emb = patched_apply_rope
qwen_mod.torch_recurrent_gated_delta_rule = patched_torch_recurrent_gated_delta_rule

# Store original functions for patching
original_functions = {}

def debug_hook(name):
    def fn(_m, input, output):
        if isinstance(input, torch.Tensor):
            summarize(input, name + "_in")
        elif isinstance(input, (tuple, list)) and len(input) > 0 and isinstance(input[0], torch.Tensor):
            summarize(input[0], name + "_in")
        if isinstance(output, torch.Tensor):
            summarize(output, name + "_out")
        elif isinstance(output, (tuple, list)) and len(output) > 0 and isinstance(output[0], torch.Tensor):
            summarize(output[0], name + "_out")

    return fn

def patch_all_forward_methods(model):
    """Apply monkey patches to all forward methods in the model"""
    for name, module in model.named_modules():
        # Set layer index if applicable
        parts = name.split('.')
        module.layer_idx = -1 # Default invalid value

        if len(parts) > 2 and parts[0] == 'model' and parts[1] == 'layers':
            try:
                module.layer_idx = int(parts[2])  # Convert to integer
            except (ValueError, IndexError):
                module.layer_idx = -1

        # Apply forward hook to log all inputs/outputs
        module.register_forward_hook(debug_hook(name))

        # Additional patches for specific methods in various modules
        if hasattr(module, 'forward'):
            original_forward = module.forward
            def make_patched_forward(orig_forward, mod_name):
                def patched_forward(*args, **kwargs):
                    # Log inputs
                    for i, arg in enumerate(args):
                        if isinstance(arg, torch.Tensor):
                            summarize(arg, f"{mod_name}.forward.arg_{i}_in")

                    # Call original forward
                    result = orig_forward(*args, **kwargs)

                    if mod_name.endswith("linear_attn"):
                        cache = kwargs["cache_params"]
                        nameparts = mod_name.split(".")
                        layer_idx = -1
                        try:
                            layer_idx = int(nameparts[2])
                        except (ValueError, IndexError):
                            print(f"\n\nDEBUG: Failed to calculate layer index for module: {mod_name}\n\n")
                        rec_cache = cache.recurrent_states[layer_idx]
                        if rec_cache is not None:
                            summarize(rec_cache, f"recurrent_cache_{layer_idx}")

                    # Log output
                    if isinstance(result, torch.Tensor):
                        summarize(result, f"{mod_name}.forward.out")
                    elif isinstance(result, (tuple, list)):
                        for i, res in enumerate(result):
                            if isinstance(res, torch.Tensor):
                                summarize(res, f"{mod_name}.forward.out_{i}")

                    return result
                return patched_forward

            module.forward = make_patched_forward(original_forward, name)

def patch_silu():
    """Patch torch.nn.functional.silu to log inputs and outputs"""
    global original_functions

    if 'silu' not in original_functions:
        original_functions['silu'] = torch.nn.functional.silu

    def patched_silu(input, inplace=False):
        # Log input
        summarize(input, "silu_in")

        # Call original function
        result = original_functions['silu'](input, inplace)

        # Log output
        summarize(result, "silu_out")

        return result

    # Replace the function in the torch.nn.functional module
    torch.nn.functional.silu = patched_silu

def patch_sigmoid():
    """Patch torch.nn.functional.sigmoid to log inputs and outputs"""
    global original_functions

    if 'sigmoid' not in original_functions:
        original_functions['sigmoid'] = torch.nn.functional.sigmoid

    def patched_sigmoid(input):
        # Log input
        summarize(input, "sigmoid_in")

        # Call original function
        result = original_functions['sigmoid'](input)

        # Log output
        summarize(result, "sigmoid_out")

        return result

    # Replace the function in the torch.nn.functional module
    torch.nn.functional.sigmoid = patched_sigmoid


def patch_torch_sigmoid():
    """Patch torch.nn.functional.sigmoid to log inputs and outputs"""
    global original_functions

    if 'torch_sigmoid' not in original_functions:
        original_functions['torch_sigmoid'] = torch.sigmoid

    def patched_torch_sigmoid(input):
        # Log input
        summarize(input, "torch_sigmoid_in")

        # Call original function
        result = original_functions['torch_sigmoid'](input)

        # Log output
        summarize(result, "torch_sigmoid_out")

        return result

    # Replace the function in the torch.nn.functional module
    torch.sigmoid = patched_torch_sigmoid


def patch_pad():
    """Patch torch.nn.functional.pad to log inputs and outputs"""
    global original_functions

    if 'pad' not in original_functions:
        original_functions['pad'] = torch.nn.functional.pad

    def patched_pad(input: torch.Tensor, pad: typing.Sequence[int], mode: str = 'constant', value: float | None = None): # pyright: ignore[reportGeneralTypeIssues]
        # Log input
        summarize(input, "pad_in")
        print(f"Padding shape is {pad}")

        # Call original function
        result = original_functions['pad'](input=input, pad=pad, mode=mode, value=value)

        # Log output
        summarize(result, "pad_out")

        return result

    # Replace the function in the torch.nn.functional module
    torch.nn.functional.pad = patched_pad


def save_kv_cache(past_key_values, step_num, data_dir, model_name):
    """Save KV cache tensors for each layer"""
    cache_dir = data_dir / f"kv_cache_step_{step_num}"
    cache_dir.mkdir(exist_ok=True)
    
    # Access past_key_values if available
    if past_key_values is not None:
        for layer_idx, cache_tuple in enumerate(past_key_values):
            if cache_tuple is None:
                print(f"Cache tuple is None for layer {layer_idx} at step {step_num}")
                continue
                
            # Handle different cache formats
            if isinstance(cache_tuple, (tuple, list)) and len(cache_tuple) >= 2:
                key, value = cache_tuple[0], cache_tuple[1]
                
                # Check if key and value are not None
                if key is not None and value is not None:
                    # Save key cache
                    key_filename = cache_dir / f"layer_{layer_idx}_key.bin"
                    key.detach().cpu().numpy().astype(np.float32).tofile(key_filename)
                    
                    # Save value cache
                    value_filename = cache_dir / f"layer_{layer_idx}_value.bin"
                    value.detach().cpu().numpy().astype(np.float32).tofile(value_filename)
                    
                    print(f"Saved KV cache for layer {layer_idx} at step {step_num}: key.shape={key.shape}, value.shape={value.shape}")
                else:
                    print(f"Key or value is None for layer {layer_idx} at step {step_num}")
            else:
                # Handle other cache formats (e.g., recurrent models)
                print(f"Non-standard cache format for layer {layer_idx} at step {step_num}: {type(cache_tuple)}")
                # Save as generic cache if it's a tensor
                if hasattr(cache_tuple, 'detach'):
                    cache_filename = cache_dir / f"layer_{layer_idx}_cache.bin"
                    cache_tuple.detach().cpu().numpy().astype(np.float32).tofile(cache_filename)
                    print(f"Saved generic cache for layer {layer_idx} at step {step_num}: shape={cache_tuple.shape}")
    else:
        print(f"No KV cache available at step {step_num}")


unreleased_model_name = os.getenv("UNRELEASED_MODEL_NAME")

parser = argparse.ArgumentParser(description="Process model with specified path")
parser.add_argument("--model-path", "-m", help="Path to the model")
parser.add_argument("--num-tokens", "-n", type=int, default=5, help="Number of tokens to generate")
parser.add_argument("--prompt", "-p", default="Hello, my name is", help="Input prompt")
parser.add_argument("--save-cache", action="store_true", help="Save KV cache at each step")
args = parser.parse_args()

model_path = os.environ.get("MODEL_PATH", args.model_path)
if model_path is None:
    parser.error(
        "Model path must be specified either via --model-path argument or MODEL_PATH environment variable"
    )

config = AutoConfig.from_pretrained(model_path)

print("Model type:       ", config.model_type)
print("Vocab size:       ", config.vocab_size)
print("Hidden size:      ", config.hidden_size)
print("Number of layers: ", config.num_hidden_layers)
print("BOS token id:     ", config.bos_token_id)
print("EOS token id:     ", config.eos_token_id)

print("Loading model and tokenizer using AutoTokenizer:", model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)
config = AutoConfig.from_pretrained(model_path)

if unreleased_model_name:
    model_name_lower = unreleased_model_name.lower()
    unreleased_module_path = (
        f"transformers.models.{model_name_lower}.modular_{model_name_lower}"
    )
    class_name = f"{unreleased_model_name}ForCausalLM"
    print(f"Importing unreleased model module: {unreleased_module_path}")

    try:
        model_class = getattr(
            importlib.import_module(unreleased_module_path), class_name
        )
        model = model_class.from_pretrained(
            model_path
        )  # Note: from_pretrained, not fromPretrained
    except (ImportError, AttributeError) as e:
        print(f"Failed to import or load model: {e}")
        exit(1)
else:
    model = AutoModelForCausalLM.from_pretrained(
        model_path, device_map="auto", offload_folder="offload"
    )

patch_all_forward_methods(model)
patch_silu()
patch_pad()
patch_sigmoid()
patch_torch_sigmoid()

model_name = os.path.basename(model_path)
# Printing the Model class to allow for easier debugging. This can be useful
# when working with models that have not been publicly released yet and this
# migth require that the concrete class is imported and used directly instead
# of using AutoModelForCausalLM.
print(f"Model class: {model.__class__.__name__}")

device = next(model.parameters()).device
prompt = args.prompt
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)

print(f"Input tokens: {input_ids}")
print(f"Input text: {repr(prompt)}")
print(f"Tokenized: {tokenizer.convert_ids_to_tokens(input_ids[0])}")

data_dir = Path("data")
data_dir.mkdir(exist_ok=True)

# Store all generated tokens and logits
all_generated_tokens = []
all_logits = []

with torch.no_grad():
    # Initial forward pass
    print(f"\n=== Initial Forward Pass ===")
    outputs = model(input_ids, use_cache=True)
    logits = outputs.logits
    
    # Extract logits for the last token (next token prediction)
    last_logits = logits[0, -1, :].cpu().numpy()
    all_logits.append(last_logits)
    
    print(f"Logits shape: {logits.shape}")
    print(f"Last token logits shape: {last_logits.shape}")
    
    # Generate first token
    next_token_id = np.argmax(last_logits).item()
    all_generated_tokens.append(next_token_id)
    
    # Show top 5 predicted tokens for first step
    top_indices = np.argsort(last_logits)[-5:][::-1]
    print("Top 5 predictions for first token:")
    for idx in top_indices:
        token = tokenizer.decode([idx])
        print(f"  Token {idx} ({repr(token)}): {last_logits[idx]:.6f}")
    
    print(f"Generated token {next_token_id} ({repr(tokenizer.decode([next_token_id]))})")
    
    # Save KV cache if requested
    if args.save_cache:
        save_kv_cache(outputs.past_key_values, 0, data_dir, model_name)
    
    # Prepare for next iteration
    past_key_values = outputs.past_key_values
    current_input = torch.tensor([[next_token_id]], device=device)
    
    # Generate remaining tokens
    for step in range(1, args.num_tokens):
        print(f"\n=== Generation Step {step} ===")
        
        # Forward pass with cache
        outputs = model(
            input_ids=current_input, 
            past_key_values=past_key_values,
            use_cache=True
        )
        
        logits = outputs.logits
        last_logits = logits[0, -1, :].cpu().numpy()
        all_logits.append(last_logits)
        
        # Generate next token
        next_token_id = np.argmax(last_logits).item()
        all_generated_tokens.append(next_token_id)
        
        # Show top 5 predicted tokens for this step
        top_indices = np.argsort(last_logits)[-5:][::-1]
        print(f"Top 5 predictions for step {step}:")
        for idx in top_indices:
            token = tokenizer.decode([idx])
            print(f"  Token {idx} ({repr(token)}): {last_logits[idx]:.6f}")
        
        print(f"Generated token {next_token_id} ({repr(tokenizer.decode([next_token_id]))})")
        
        # Save KV cache if requested
        if args.save_cache:
            save_kv_cache(outputs.past_key_values, step, data_dir, model_name)
        
        # Update for next iteration
        past_key_values = outputs.past_key_values
        current_input = torch.tensor([[next_token_id]], device=device)

# Save results
bin_filename = data_dir / f"pytorch-{model_name}-multi-token.bin"
txt_filename = data_dir / f"pytorch-{model_name}-multi-token.txt"

# Save all logits concatenated
all_logits_array = np.array(all_logits)
all_logits_array.astype(np.float32).tofile(bin_filename)

# Also save as text file for easy inspection
with open(txt_filename, "w") as f:
    f.write(f"Generated tokens: {all_generated_tokens}\n")
    f.write(f"Generated text: {repr(tokenizer.decode(all_generated_tokens))}\n")
    f.write(f"Full sequence: {repr(tokenizer.decode(input_ids[0].tolist() + all_generated_tokens))}\n\n")
    
    for step, logits in enumerate(all_logits):
        f.write(f"=== Step {step} logits ===\n")
        for i, logit in enumerate(logits):
            f.write(f"{i}: {logit:.6f}\n")
        f.write("\n")

print(f"\n=== Generation Complete ===")
print(f"Generated {len(all_generated_tokens)} tokens: {all_generated_tokens}")
print(f"Generated text: {repr(tokenizer.decode(all_generated_tokens))}")
print(f"Full sequence: {repr(tokenizer.decode(input_ids[0].tolist() + all_generated_tokens))}")

print(f"Saved bin logits to: {bin_filename}")
print(f"Saved txt logits to: {txt_filename}")

if args.save_cache:
    print(f"KV cache saved to: {data_dir}/kv_cache_step_*")