初始化项目，由ModelHub XC社区提供模型

Model: hxia7/qwen3-4b-blockdist
Source: Original Platform

eval/scripts/flex_attn_test.py

@@ -0,0 +1,92 @@
+"""Operator-level validation of block attention under FlexAttention.
+
+1) CORRECTNESS: flex block-attention output == reference SDPA-with-additive-block-mask (small T).
+2) STABILITY:   flex is deterministic across repeated runs.
+3) SCALE:       flex runs at long T (O(T) memory) where the eager T x T mask would need terabytes.
+
+This validates the kernel/mask logic your serving stack (vLLM) needs, independent of the HF model.
+"""
+from __future__ import annotations
+import argparse, torch
+import torch.nn.functional as F
+from torch.nn.attention.flex_attention import create_block_mask, flex_attention
+
+NEG = torch.finfo(torch.float32).min
+
+
+def make_block_ids(T, n_blocks, dev):
+    """token 0..s = system(-1); then n_blocks equal blocks; last quarter = query(-2)."""
+    ids = torch.empty(T, dtype=torch.long, device=dev)
+    s = max(1, T // 20)
+    q = T - max(1, T // 4)
+    ids[:s] = -1
+    ids[q:] = -2
+    body = q - s
+    per = max(1, body // n_blocks)
+    for b in range(n_blocks):
+        ids[s + b*per : s + (b+1)*per] = b
+    ids[s + n_blocks*per : q] = n_blocks - 1
+    return ids
+
+
+def block_mask_mod(bids):
+    def mm(b, h, qi, ki):
+        causal = ki <= qi
+        bq = bids[qi]; bk = bids[ki]
+        return causal & ((bk == -1) | (bq == bk) | (bq == -2))
+    return mm
+
+
+def ref_sdpa(q, k, v, bids):
+    T = q.shape[-2]; dev = q.device
+    causal = torch.tril(torch.ones(T, T, dtype=torch.bool, device=dev))
+    bq = bids.view(T, 1); bk = bids.view(1, T)
+    allowed = (((bk == -1) | (bq == bk) | (bq == -2)) & causal).view(1, 1, T, T)  # boolean mask (no NaN)
+    return F.scaled_dot_product_attention(q, k, v, attn_mask=allowed)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--H", type=int, default=8); ap.add_argument("--D", type=int, default=128)
+    ap.add_argument("--T", type=int, default=1024); ap.add_argument("--blocks", type=int, default=8)
+    ap.add_argument("--big-T", type=int, default=16384)
+    args = ap.parse_args()
+    dev = "cuda"; torch.manual_seed(0)
+
+    # (1) correctness at moderate T
+    T = args.T
+    q = torch.randn(1, args.H, T, args.D, device=dev, dtype=torch.bfloat16)
+    k = torch.randn(1, args.H, T, args.D, device=dev, dtype=torch.bfloat16)
+    v = torch.randn(1, args.H, T, args.D, device=dev, dtype=torch.bfloat16)
+    bids = make_block_ids(T, args.blocks, dev)
+    bm = create_block_mask(block_mask_mod(bids), B=None, H=None, Q_LEN=T, KV_LEN=T, device=dev)
+    flex = torch.compile(flex_attention)
+    o_flex = flex(q, k, v, block_mask=bm)
+    o_ref = ref_sdpa(q, k, v, bids)
+    d = (o_flex.float() - o_ref.float()).abs()
+    print(f"(1) CORRECTNESS  T={T} H={args.H} D={args.D} blocks={args.blocks}")
+    print(f"    max|Δ|={d.max():.4f}  mean|Δ|={d.mean():.6f}  (bf16 noise floor ~1e-2) -> "
+          f"{'MATCH' if d.max()<0.05 else 'MISMATCH'}")
+
+    # (2) determinism
+    o2 = flex(q, k, v, block_mask=bm)
+    print(f"(2) STABILITY    flex run twice: max|Δ|={(o_flex.float()-o2.float()).abs().max():.6f} "
+          f"-> {'deterministic' if (o_flex.float()-o2.float()).abs().max()<1e-3 else 'nondeterministic'}")
+
+    # (3) scale — flex at big T (eager T x T mask would be H*T*T*2 bytes)
+    bt = args.big_T
+    qb = torch.randn(1, args.H, bt, args.D, device=dev, dtype=torch.bfloat16)
+    kb = torch.randn(1, args.H, bt, args.D, device=dev, dtype=torch.bfloat16)
+    vb = torch.randn(1, args.H, bt, args.D, device=dev, dtype=torch.bfloat16)
+    bidsb = make_block_ids(bt, args.blocks * 8, dev)
+    bmb = create_block_mask(block_mask_mod(bidsb), B=None, H=None, Q_LEN=bt, KV_LEN=bt, device=dev)
+    torch.cuda.reset_peak_memory_stats()
+    ob = flex(qb, kb, vb, block_mask=bmb); torch.cuda.synchronize()
+    peak = torch.cuda.max_memory_allocated() / 1e9
+    eager_mask_gb = args.H * bt * bt * 2 / 1e9
+    print(f"(3) SCALE        flex at T={bt}: OK, peak mem {peak:.1f} GB, out finite={bool(torch.isfinite(ob).all())}")
+    print(f"    (an eager [1,{args.H},{bt},{bt}] score/mask alone would be ~{eager_mask_gb:.0f} GB -> impossible)")
+
+
+if __name__ == "__main__":
+    main()