Debug a MiniGPT Transformer and Implement KV Cache (PyTorch)

Problem: Debug a MiniGPT Transformer and Implement KV Cache (PyTorch) You are given a PyTorch code skeleton implementing a simplified GPT (miniGPT). The current implementation fails to generate correct text (e.g., garbled output, repetition, or mismatch with expected outputs) and is inefficient at autoregressive inference. Your tasks Debug miniGPT: fix potential bugs in the forward pass, attention, masking, positional embeddings, logits computation, and decoding/sampling so that it generates correct text under the provided test or benchmark. Implement KV cache (follow-up): Cache per-layer attention keys/values during autoregressive decoding. When generating a new token, compute K/V only for the new token and reuse cached history. Ensure outputs match the non-cached implementation under the same seed and decoding settings. I/O (adapt to provided code) Input: prompt tokens (or text), max generation length, decoding params (greedy / temperature / top-k / top-p, etc.). Output: generated token ids / decoded text. Constraints / edge cases Optional: support batching and variable-length prompts. Must handle the causal mask correctly. KV cache must correctly manage shapes, device/dtype, and per-layer organization. Consider cache growth and memory usage.