qwen3_5 flex: bs=4 parity test for captured graph vs eager batched

tests/qwen3_5_flex_parity_bs4.py

@@ -0,0 +1,124 @@
+"""B=4 parity test so the captured graph bucket actually replays.
+
+Compares flex captured-graph decode vs flex eager-decode (capture
+disabled via ``UNSLOTH_FLEX_QWEN3_5_NO_CAPTURE=1``). Exact token-for-token
+match is the success criterion — both paths use identical math, only
+the kernel-launch wrapper differs.
+
+    CUDA_VISIBLE_DEVICES=2 UNSLOTH_FAST_INFERENCE=1 PYTHONUNBUFFERED=1 \\
+        python -u tests/qwen3_5_flex_parity_bs4.py \\
+        --model Qwen/Qwen3.5-4B --max_new_tokens 16 \\
+        --json_out async_task_outputs/qwen3_5/flex_parity_bs4.json
+"""
+import argparse
+import json
+import os
+import time
+
+
+CHAT_PROMPTS = [
+    "In one sentence, what is Paris?",
+    "What is 23 + 19? Answer in one word.",
+    "Continue: The quick brown fox jumps over",
+    "Name one primary color.",
+]
+
+
+def _run(model, tok, args, capture: bool):
+    class SP:
+        def __init__(self, n): self.max_tokens = n; self.temperature = 0.0
+
+    sp = SP(args.max_new_tokens)
+    rendered = [
+        tok.apply_chat_template(
+            [{"role": "user", "content": p}],
+            tokenize=False, add_generation_prompt=True,
+        )
+        for p in CHAT_PROMPTS
+    ]
+    os.environ["UNSLOTH_FLEX_QWEN3_5_NO_CAPTURE"] = "0" if capture else "1"
+    t0 = time.perf_counter()
+    outs = model.fast_generate(rendered, sampling_params=sp, use_tqdm=False)
+    dt = time.perf_counter() - t0
+    results = []
+    for p, o in zip(CHAT_PROMPTS, outs):
+        results.append({
+            "prompt": p,
+            "text": o.outputs[0].text,
+            "ids": list(o.outputs[0].token_ids),
+        })
+    return dt, results
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model", required=True)
+    ap.add_argument("--max_new_tokens", type=int, default=16)
+    ap.add_argument("--max_seq_length", type=int, default=1024)
+    ap.add_argument("--max_batch_size", type=int, default=4)
+    ap.add_argument("--gpu_memory_utilization", type=float, default=0.6)
+    ap.add_argument("--json_out", required=True)
+    args = ap.parse_args()
+
+    import torch
+    os.environ["UNSLOTH_FAST_INFERENCE"] = "1"
+    import unsloth  # noqa
+    from unsloth import FastLanguageModel
+
+    model, tok_raw = FastLanguageModel.from_pretrained(
+        model_name=args.model,
+        max_seq_length=args.max_seq_length,
+        dtype=torch.bfloat16,
+        load_in_4bit=False,
+        fast_inference=True,
+        max_batch_size=args.max_batch_size,
+        gpu_memory_utilization=args.gpu_memory_utilization,
+    )
+    tok = tok_raw.tokenizer if hasattr(tok_raw, "tokenizer") else tok_raw
+
+    # Eager first (so the graph capture doesn't pollute state).
+    dt_eager, eager = _run(model, tok, args, capture=False)
+    print(f"[eager] gen_s={dt_eager:.2f}")
+    for r in eager:
+        print(f"[eager] {r['prompt'][:40]!r} -> {r['text'][:80]!r}")
+
+    dt_cap, cap = _run(model, tok, args, capture=True)
+    print(f"[cap]   gen_s={dt_cap:.2f}")
+    for r in cap:
+        print(f"[cap]   {r['prompt'][:40]!r} -> {r['text'][:80]!r}")
+
+    # Token-by-token parity.
+    mismatches = []
+    for i, (e, c) in enumerate(zip(eager, cap)):
+        matches = 0
+        for a, b in zip(e["ids"], c["ids"]):
+            if a == b:
+                matches += 1
+            else:
+                break
+        mismatches.append({
+            "prompt_idx": i,
+            "prompt": e["prompt"],
+            "matches": matches,
+            "total": min(len(e["ids"]), len(c["ids"])),
+            "eager_text": e["text"],
+            "cap_text": c["text"],
+        })
+        tag = "OK " if matches == min(len(e["ids"]), len(c["ids"])) else "MM "
+        print(f"[{tag}] P{i} {matches}/{min(len(e['ids']), len(c['ids']))}")
+
+    summary = {
+        "model": args.model,
+        "max_new_tokens": args.max_new_tokens,
+        "eager_gen_s": round(dt_eager, 3),
+        "capture_gen_s": round(dt_cap, 3),
+        "mismatches": mismatches,
+    }
+    os.makedirs(os.path.dirname(args.json_out) or ".", exist_ok=True)
+    with open(args.json_out, "w") as f:
+        json.dump(summary, f, indent=2)
+    print(f"[parity-bs4] wrote {args.json_out}")
+
+
+if __name__ == "__main__":
+    main()