feat: optimize ssm with jit and add inference sanity check

Author: Pomilon

aetheris/cli/main.py

@@ -107,6 +107,15 @@ def generate_command(args):
     repetition_penalty = args.repetition_penalty
 
     input_ids = tokenizer.encode(prompt, return_tensors='pt').to(device)
+    
+    # --- INFERENCE SANITY CHECK ---
+    print(f"\n[SANITY CHECK] Inference Tokenizer: {tokenizer.name_or_path}")
+    print(f"[SANITY CHECK] Vocab Size: {tokenizer.vocab_size}")
+    print(f"[SANITY CHECK] Input IDs: {input_ids.tolist()}")
+    decoded_prompt = tokenizer.decode(input_ids[0], skip_special_tokens=False)
+    print(f"[SANITY CHECK] Decoded Prompt: '{decoded_prompt}'\n")
+    # ------------------------------
+    
     generated_ids = input_ids.clone()
     history_ids = set(input_ids[0].tolist())
 

aetheris/modules/ssm.py

@@ -3,6 +3,9 @@
 import torch.nn.functional as F
 from ..config import AetherisConfig
 
+from typing import List
+
+@torch.jit.script
 def selective_scan_native(u: torch.Tensor, delta: torch.Tensor, A: torch.Tensor,
                          B: torch.Tensor, C: torch.Tensor, D: torch.Tensor) -> torch.Tensor:
     """Memory-efficient scan with reduced intermediate tensors."""
@@ -15,7 +18,7 @@ def selective_scan_native(u: torch.Tensor, delta: torch.Tensor, A: torch.Tensor,
     # Use in-place operations where possible
     # FORCE FLOAT32 for state to prevent underflow/overflow in long sequences
     h = torch.zeros(B_size, D_inner, D_state, device=u.device, dtype=torch.float32)
-    ys = []
+    ys: List[torch.Tensor] = []
 
     # Cast inputs to float32 for the scan
     # Note: This increases memory usage slightly but is critical for stability