[feat] Added debug tools.

nanovllm/debug/utils.py

@@ -0,0 +1,51 @@
+"""Utility functions for breakpoint alignment debugging."""
+
+import torch
+from nanovllm.utils.context import set_context, reset_context
+
+
+def setup_prefill_context(seq_len: int, device: torch.device):
+    """
+    Set up nanovllm context for prefill alignment testing.
+
+    Args:
+        seq_len: Sequence length
+        device: Target device
+    """
+    cu_seqlens = torch.tensor([0, seq_len], dtype=torch.int32, device=device)
+    slot_mapping = torch.full((seq_len,), -1, dtype=torch.int32, device=device)
+
+    set_context(
+        is_prefill=True,
+        cu_seqlens_q=cu_seqlens,
+        cu_seqlens_k=cu_seqlens,
+        max_seqlen_q=seq_len,
+        max_seqlen_k=seq_len,
+        slot_mapping=slot_mapping,
+        is_chunked_prefill=False,
+    )
+
+
+def setup_decode_context(context_len: int, device: torch.device):
+    """
+    Set up nanovllm context for decode alignment testing.
+
+    Args:
+        context_len: Context length (number of previous tokens)
+        device: Target device
+    """
+    context_lens = torch.tensor([context_len], dtype=torch.int32, device=device)
+    slot_mapping = torch.tensor([-1], dtype=torch.int32, device=device)
+    block_tables = torch.zeros((1, 1), dtype=torch.int32, device=device)
+
+    set_context(
+        is_prefill=False,
+        slot_mapping=slot_mapping,
+        context_lens=context_lens,
+        block_tables=block_tables,
+    )
+
+
+def cleanup_context():
+    """Reset nanovllm context after alignment testing."""
+    reset_context()