From ca9dbce98d2ba61e8f71b8a494ee8728e3ea4837 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 21 Apr 2026 09:54:13 +0000 Subject: [PATCH] benchmarks: add verify_*_numerics scripts to cross-check flex vs vanilla HF One-shot correctness checks for the flex_attention + paged KV path against a vanilla HF `model(input_ids, use_cache=False)` forward on the same prompt. Report max / mean abs diff on last-position logits plus argmax match and top-10 overlap, so numerical drift and semantic equivalence are both visible. Shared approach: load the base model, deep-copy it for the flex path so the attention patching does not mutate the vanilla comparison, run both on the same tokenized prompt, report diffs. `verify_gemma4_numerics.py` mirrors `gemma4_flex_inference.py`'s text-only loader (Gemma4ForConditionalGeneration -> drop vision / audio towers -> language_model into a Gemma4ForCausalLM shell) before deep-copying. The softcap is NOT re-applied on the vanilla logits since `Gemma4ForCausalLM.forward` already applies `final_logit_softcapping`. `verify_qwen3_numerics.py` runs `qwen3_flex_inference.FlexInference` against either Qwen3 or Llama-3.2 via `--model_name`. `fa4_prefill` is disabled because short prompts hit a CuteDSL sm_100 shape mismatch in `handle_block_sparse_empty_tile_correction_sm100`. Results on B200 bf16, 6 to 7 token prompt: | Model | max abs | mean abs | argmax | top-10 | |---------------------------------|---------|----------|--------|--------| | unsloth/Qwen3-4B-Base | 0.313 | 0.088 | yes | 10/10 | | unsloth/Llama-3.2-3B-Instruct | 0.125 | 0.021 | yes | 10/10 | | unsloth/gemma-4-E2B-it | 0.375 | 0.080 | yes | 10/10 | All three land in the same bf16 Triton-flex vs eager-matmul drift band; Gemma-4's extra per-layer-input path and layer_scalar do not widen the gap despite 20 of 35 layers going through the shared-KV link. --- scripts/benchmarks/verify_gemma4_numerics.py | 118 +++++++++++++++++++ scripts/benchmarks/verify_qwen3_numerics.py | 94 +++++++++++++++ 2 files changed, 212 insertions(+) create mode 100644 scripts/benchmarks/verify_gemma4_numerics.py create mode 100644 scripts/benchmarks/verify_qwen3_numerics.py diff --git a/scripts/benchmarks/verify_gemma4_numerics.py b/scripts/benchmarks/verify_gemma4_numerics.py new file mode 100644 index 0000000000..850753cc45 --- /dev/null +++ b/scripts/benchmarks/verify_gemma4_numerics.py @@ -0,0 +1,118 @@ +"""Compare first-token logits between `FlexGemma4Inference._prefill` and +vanilla `Gemma4ForCausalLM.forward` on the same prompt. Intended as a +one-shot correctness check; not part of the benchmark matrix. + +Run: + CUDA_VISIBLE_DEVICES=2 python scripts/benchmarks/verify_gemma4_numerics.py +""" + +from __future__ import annotations + +import copy +import sys +from pathlib import Path + +import torch + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +from gemma4_flex_inference import ( # noqa: E402 + FlexGemma4Inference, + Sequence, + _require_gemma4, +) + + +def main(): + Gemma4ForCausalLM, Gemma4Config, Gemma4TextConfig = _require_gemma4() + from transformers.models.gemma4.modeling_gemma4 import Gemma4ForConditionalGeneration + from transformers import AutoTokenizer + + name = "unsloth/gemma-4-E2B-it" + tok = AutoTokenizer.from_pretrained(name) + if tok.pad_token is None: + tok.pad_token = tok.eos_token + + # --- load once (text shell) and keep a pristine copy for the vanilla pass + full_cfg = Gemma4Config.from_pretrained(name) + text_cfg = full_cfg.text_config + + full = Gemma4ForConditionalGeneration.from_pretrained( + name, dtype = torch.bfloat16, attn_implementation = "eager" + ) + lang = full.model.language_model + full.model.vision_tower = None + full.model.audio_tower = None + full.model.embed_vision = None + full.model.embed_audio = None + + base = Gemma4ForCausalLM(text_cfg) + base.model = lang + base.lm_head.weight = lang.embed_tokens.weight + base = base.to(torch.bfloat16).to("cuda") + base.eval() + del full + + # Deep-copy so Flex's attention patching doesn't mutate the vanilla path. + flex_model = copy.deepcopy(base) + + prompt = "The quick brown fox jumps over" + ids = tok(prompt, return_tensors = "pt")["input_ids"].to("cuda") + print(f"prompt len = {ids.shape[1]}") + + with torch.inference_mode(): + # `Gemma4ForCausalLM.forward` applies `final_logit_softcapping` + # internally, so these logits are already softcapped. + out = base(ids, use_cache = False) + ref_logits = out.logits[0, -1, :].float() + print(f"vanilla last-token logits: mean {ref_logits.mean().item():.4f}, " + f"std {ref_logits.std().item():.4f}, " + f"argmax {int(ref_logits.argmax())} " + f"({tok.decode([int(ref_logits.argmax())])!r})") + + # Flex path. + inf = FlexGemma4Inference( + flex_model, + tok, + max_batch_size = 4, + max_seq_length = 256, + n_pages = 64, + page_size = 64, + max_new_tokens = 1, + decode_kernel_options = {"BLOCK_M": 16, "BLOCK_N": 16}, + prefill_kernel_options = { + "FORCE_USE_FLEX_ATTENTION": True, + "BLOCK_M": 32, + "BLOCK_N": 32, + }, + fa4_prefill = False, + ) + seq = Sequence(text = prompt, max_new_tokens = 1) + inf.tokenize([seq]) + bi = inf.page_table.allocate() + inf.page_table.reserve( + bi, + torch.tensor([bi], device = "cuda", dtype = torch.long), + seq.total_length, + ) + seq.batch_idx = bi + with torch.inference_mode(): + flex_logits = inf._prefill([seq])[0].float() + print(f"flex last-token logits: mean {flex_logits.mean().item():.4f}, " + f"std {flex_logits.std().item():.4f}, " + f"argmax {int(flex_logits.argmax())} " + f"({tok.decode([int(flex_logits.argmax())])!r})") + + diff = (flex_logits - ref_logits).abs() + print(f"max abs diff = {diff.max().item():.4e}") + print(f"mean abs diff = {diff.mean().item():.4e}") + print(f"argmax match = {int(ref_logits.argmax()) == int(flex_logits.argmax())}") + # bf16 ULP is ~1e-2 at magnitude ~5. Report top-10 overlap too. + top_ref = set(ref_logits.topk(10).indices.tolist()) + top_flex = set(flex_logits.topk(10).indices.tolist()) + print(f"top-10 overlap = {len(top_ref & top_flex)} / 10") + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmarks/verify_qwen3_numerics.py b/scripts/benchmarks/verify_qwen3_numerics.py new file mode 100644 index 0000000000..56b91f2398 --- /dev/null +++ b/scripts/benchmarks/verify_qwen3_numerics.py @@ -0,0 +1,94 @@ +"""Compare first-token logits between FlexInference._prefill (qwen3 path) +and vanilla model(input_ids) for Qwen3 and Llama-3.2. + +Run: + CUDA_VISIBLE_DEVICES=2 python scripts/benchmarks/verify_qwen3_numerics.py \ + --model_name unsloth/Qwen3-4B-Base + CUDA_VISIBLE_DEVICES=2 python scripts/benchmarks/verify_qwen3_numerics.py \ + --model_name unsloth/Llama-3.2-3B-Instruct +""" + +from __future__ import annotations + +import argparse +import copy +import sys +from pathlib import Path + +import torch + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +from qwen3_flex_inference import FlexInference, Sequence # noqa: E402 + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--model_name", required = True) + p.add_argument("--prompt", default = "The quick brown fox jumps over") + args = p.parse_args() + + from transformers import AutoModelForCausalLM, AutoTokenizer + + tok = AutoTokenizer.from_pretrained(args.model_name) + if tok.pad_token is None: + tok.pad_token = tok.eos_token + + base = AutoModelForCausalLM.from_pretrained( + args.model_name, dtype = torch.bfloat16, attn_implementation = "eager" + ).to("cuda") + base.eval() + + flex_model = copy.deepcopy(base) + + ids = tok(args.prompt, return_tensors = "pt")["input_ids"].to("cuda") + print(f"prompt len = {ids.shape[1]}") + + with torch.inference_mode(): + out = base(ids, use_cache = False) + ref_logits = out.logits[0, -1, :].float() + print( + f"vanilla last-token logits: mean {ref_logits.mean().item():.4f}, " + f"std {ref_logits.std().item():.4f}, argmax {int(ref_logits.argmax())} " + f"({tok.decode([int(ref_logits.argmax())])!r})" + ) + + inf = FlexInference( + flex_model, + tok, + max_batch_size = 4, + max_seq_length = 256, + n_pages = 64, + page_size = 64, + max_new_tokens = 1, + fa4_prefill = False, + ) + seq = Sequence(text = args.prompt, max_new_tokens = 1) + inf.tokenize([seq]) + bi = inf.page_table.allocate() + inf.page_table.reserve( + bi, + torch.tensor([bi], device = "cuda", dtype = torch.long), + seq.total_length, + ) + seq.batch_idx = bi + with torch.inference_mode(): + flex_logits = inf._prefill([seq])[0].float() + print( + f"flex last-token logits: mean {flex_logits.mean().item():.4f}, " + f"std {flex_logits.std().item():.4f}, argmax {int(flex_logits.argmax())} " + f"({tok.decode([int(flex_logits.argmax())])!r})" + ) + + diff = (flex_logits - ref_logits).abs() + print(f"max abs diff = {diff.max().item():.4e}") + print(f"mean abs diff = {diff.mean().item():.4e}") + print(f"argmax match = {int(ref_logits.argmax()) == int(flex_logits.argmax())}") + top_ref = set(ref_logits.topk(10).indices.tolist()) + top_flex = set(flex_logits.topk(10).indices.tolist()) + print(f"top-10 overlap = {len(top_ref & top_flex)} / 10") + + +if __name__ == "__main__": + main()