[claudesquad] update from 'layer-prefill-1' on 08 Jan 26 03:36 CST

tests/test_minference_gpu.py

@@ -0,0 +1,163 @@
+"""
+Needle-in-haystack test with MInference sparse attention.
+
+Tests: MInference sparse prefill on GPU-only path (no CPU offload).
+This validates that MInference's vertical + slash sparse pattern can
+correctly retrieve information from long context.
+"""
+
+import os
+os.environ["NANOVLLM_LOG_LEVEL"] = "INFO"
+
+import argparse
+from nanovllm import LLM, SamplingParams
+from nanovllm.config import SparsePolicyType
+from utils import generate_needle_prompt, check_needle_answer
+
+
+def run_minference_test(
+    model_path: str,
+    max_model_len: int = 16384,
+    input_len: int = 8192,
+    needle_position: float = 0.5,
+    needle_value: str = "7492",
+    adaptive_budget: float = 0.3,
+    max_new_tokens: int = 32,
+    verbose: bool = True,
+) -> bool:
+    """
+    Run needle test with MInference sparse prefill attention.
+
+    Args:
+        model_path: Path to model
+        max_model_len: Maximum model context length
+        input_len: Target input sequence length
+        needle_position: Where to place needle (0.0-1.0)
+        needle_value: The secret value to find
+        adaptive_budget: MInference budget as fraction of seq_len
+        max_new_tokens: Maximum tokens to generate
+        verbose: Print detailed output
+
+    Returns:
+        True if test passed, False otherwise
+    """
+    if verbose:
+        print(f"\n{'='*60}")
+        print(f"MInference Sparse Prefill Test (GPU-only)")
+        print(f"{'='*60}")
+        print(f"Model: {model_path}")
+        print(f"Max model len: {max_model_len}")
+        print(f"Input length: {input_len}")
+        print(f"Needle position: {needle_position:.0%}")
+        print(f"Needle value: {needle_value}")
+        print(f"Adaptive budget: {adaptive_budget}")
+        print(f"{'='*60}\n")
+
+    # Initialize LLM with MInference sparse attention
+    llm = LLM(
+        model_path,
+        enforce_eager=True,
+        max_model_len=max_model_len,
+        max_num_batched_tokens=max_model_len,
+        enable_cpu_offload=False,  # GPU-only
+        sparse_policy=SparsePolicyType.MINFERENCE,
+        minference_adaptive_budget=adaptive_budget,
+    )
+
+    # Generate needle prompt
+    prompt, expected = generate_needle_prompt(
+        tokenizer=llm.tokenizer,
+        target_length=input_len,
+        needle_position=needle_position,
+        needle_value=needle_value,
+    )
+
+    # Generate output
+    sampling_params = SamplingParams(
+        temperature=0.6,
+        max_tokens=max_new_tokens,
+    )
+    outputs = llm.generate([prompt], sampling_params, use_tqdm=True)
+
+    # Check result
+    output_text = outputs[0]["text"]
+    output_token_ids = outputs[0]["token_ids"]
+    passed = check_needle_answer(output_text, expected)
+
+    if verbose:
+        print(f"\n{'='*60}")
+        print(f"Result")
+        print(f"{'='*60}")
+        print(f"Expected: {expected}")
+        print(f"Output tokens ({len(output_token_ids)}): {output_token_ids[:20]}")
+        print(f"Output: {output_text[:200]}...")
+        print(f"Status: {'PASSED' if passed else 'FAILED'}")
+        print(f"{'='*60}\n")
+
+    return passed
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Needle-in-haystack test with MInference sparse prefill"
+    )
+    parser.add_argument(
+        "--model", "-m",
+        type=str,
+        default=os.path.expanduser("~/models/Qwen3-4B-Instruct-2507/"),
+        help="Path to model"
+    )
+    parser.add_argument(
+        "--max-model-len",
+        type=int,
+        default=16 * 1024,
+        help="Maximum model context length"
+    )
+    parser.add_argument(
+        "--input-len",
+        type=int,
+        default=8 * 1024,
+        help="Target input sequence length"
+    )
+    parser.add_argument(
+        "--needle-position",
+        type=float,
+        default=0.5,
+        help="Needle position (0.0=start, 0.5=middle, 1.0=end)"
+    )
+    parser.add_argument(
+        "--needle-value",
+        type=str,
+        default="7492",
+        help="The secret value to hide"
+    )
+    parser.add_argument(
+        "--adaptive-budget",
+        type=float,
+        default=0.3,
+        help="MInference adaptive budget (fraction of seq_len)"
+    )
+    parser.add_argument(
+        "--max-new-tokens",
+        type=int,
+        default=32,
+        help="Maximum tokens to generate"
+    )
+    args = parser.parse_args()
+
+    passed = run_minference_test(
+        model_path=args.model,
+        max_model_len=args.max_model_len,
+        input_len=args.input_len,
+        needle_position=args.needle_position,
+        needle_value=args.needle_value,
+        adaptive_budget=args.adaptive_budget,
+        max_new_tokens=args.max_new_tokens,
+        verbose=True,
+    )
+
+    if passed:
+        print("test_minference_gpu: PASSED")
+    else:
+        print("test_minference_gpu: FAILED")
+        exit(1)

tests/test_needle.py

@@ -31,8 +31,13 @@ def run_needle_test(
     max_new_tokens: int = 32,
     enable_cpu_offload: bool = False,
     enable_quest: bool = False,
+    enable_minference: bool = False,
     sparse_topk: int = 8,
     sparse_threshold: int = 4,
+    minference_budget: float = 0.3,
+    minference_vertical: int = 1000,
+    minference_slash: int = 6096,
+    gpu_utilization: float = 0.9,
     verbose: bool = True,
 ) -> bool:
     """
@@ -49,14 +54,25 @@ def run_needle_test(
         max_new_tokens: Maximum tokens to generate
         enable_cpu_offload: Enable CPU offload mode
         enable_quest: Enable Quest sparse attention (decode-only Top-K)
+        enable_minference: Enable MInference sparse prefill (GPU-only)
         sparse_topk: Top-K blocks for Quest
         sparse_threshold: Apply sparse only when blocks > threshold
+        minference_budget: MInference adaptive budget (fraction of seq_len, None=fixed mode)
+        minference_vertical: Fixed vertical_size (only used when budget=None)
+        minference_slash: Fixed slash_size (only used when budget=None)
+        gpu_utilization: GPU memory utilization fraction
         verbose: Print detailed output
 
     Returns:
         True if test passed, False otherwise
     """
-    sparse_policy = SparsePolicyType.QUEST if enable_quest else SparsePolicyType.FULL
+    # Determine sparse policy
+    if enable_minference:
+        sparse_policy = SparsePolicyType.MINFERENCE
+    elif enable_quest:
+        sparse_policy = SparsePolicyType.QUEST
+    else:
+        sparse_policy = SparsePolicyType.FULL
 
     if verbose:
         print(f"\n{'='*60}")
@@ -69,8 +85,14 @@ def run_needle_test(
         print(f"Needle position: {needle_position:.0%}")
         print(f"Needle value: {needle_value}")
         print(f"CPU offload: {enable_cpu_offload}")
-        if enable_cpu_offload:
-            print(f"Sparse policy: {sparse_policy.name} (topk={sparse_topk}, threshold={sparse_threshold})")
+        print(f"Sparse policy: {sparse_policy.name}")
+        if enable_cpu_offload and enable_quest:
+            print(f"  Quest: topk={sparse_topk}, threshold={sparse_threshold}")
+        if enable_minference:
+            if minference_budget is not None:
+                print(f"  MInference: adaptive (budget={minference_budget})")
+            else:
+                print(f"  MInference: fixed (vertical={minference_vertical}, slash={minference_slash})")
         print(f"{'='*60}\n")
 
     # 1. Initialize LLM
@@ -80,12 +102,19 @@ def run_needle_test(
         "max_num_batched_tokens": max_model_len,
         "enable_cpu_offload": enable_cpu_offload,
         "kvcache_block_size": block_size,
+        "gpu_memory_utilization": gpu_utilization,
     }
     if enable_cpu_offload:
         llm_kwargs["num_gpu_blocks"] = num_gpu_blocks
         llm_kwargs["sparse_policy"] = sparse_policy
         llm_kwargs["sparse_topk_blocks"] = sparse_topk
         llm_kwargs["sparse_threshold_blocks"] = sparse_threshold
+    elif enable_minference:
+        # MInference is GPU-only sparse prefill
+        llm_kwargs["sparse_policy"] = sparse_policy
+        llm_kwargs["minference_adaptive_budget"] = minference_budget
+        llm_kwargs["minference_vertical_size"] = minference_vertical
+        llm_kwargs["minference_slash_size"] = minference_slash
 
     llm = LLM(model_path, **llm_kwargs)
 
@@ -186,6 +215,11 @@ if __name__ == "__main__":
         action="store_true",
         help="Enable Quest sparse attention (decode-only Top-K selection)"
     )
+    parser.add_argument(
+        "--enable-minference",
+        action="store_true",
+        help="Enable MInference sparse prefill (GPU-only, vertical+slash pattern)"
+    )
     parser.add_argument(
         "--sparse-topk",
         type=int,
@@ -198,8 +232,35 @@ if __name__ == "__main__":
         default=4,
         help="Apply sparse only when blocks > threshold"
     )
+    parser.add_argument(
+        "--minference-budget",
+        type=float,
+        default=0.3,
+        help="MInference adaptive budget (fraction of seq_len, 0.3=30%% compute, 0=fixed mode)"
+    )
+    parser.add_argument(
+        "--minference-vertical",
+        type=int,
+        default=1000,
+        help="Fixed vertical_size (only used when budget=0)"
+    )
+    parser.add_argument(
+        "--minference-slash",
+        type=int,
+        default=6096,
+        help="Fixed slash_size (only used when budget=0)"
+    )
+    parser.add_argument(
+        "--gpu-utilization",
+        type=float,
+        default=0.9,
+        help="GPU memory utilization (default: 0.9)"
+    )
     args = parser.parse_args()
 
+    # Convert budget=0 to None for fixed mode
+    minference_budget = args.minference_budget if args.minference_budget > 0 else None
+
     passed = run_needle_test(
         model_path=args.model,
         max_model_len=args.max_model_len,
@@ -211,8 +272,13 @@ if __name__ == "__main__":
         max_new_tokens=args.max_new_tokens,
         enable_cpu_offload=args.enable_offload,
         enable_quest=args.enable_quest,
+        enable_minference=args.enable_minference,
         sparse_topk=args.sparse_topk,
         sparse_threshold=args.sparse_threshold,
+        minference_budget=minference_budget,
+        minference_vertical=args.minference_vertical,
+        minference_slash=args.minference_slash,
+        gpu_utilization=args.gpu_utilization,
         verbose=True,
     )