qwen3.5-122B-A10B on DGX Spark: vLLM + DFlash + dense-bandwidth stack, one-shot installer

tools/test_int8_lmhead2.py

@@ -0,0 +1,96 @@
+#!/usr/bin/env python3
+"""v2 harness: proper bf16 GEMV baseline + a single BATCHED int8 w8a16 GEMM kernel
+(one launch for any B, dot-based, pads B to >=16) vs albond's per-row loop, at the
+real lm-head shape and B in {1,5,13} (base + DFlash verify-batch sizes).
+"""
+import time
+
+import torch
+import triton
+import triton.language as tl
+
+V, H = 248320, 3072
+DEV = "cuda"
+torch.manual_seed(0)
+
+W_bf16 = (torch.randn(V, H, device=DEV, dtype=torch.float32) * 0.02).to(torch.bfloat16)
+scales = (W_bf16.float().abs().amax(dim=1) / 127.0).clamp(min=1e-12)
+W_int8 = (W_bf16.float() / scales.unsqueeze(1)).round().clamp(-127, 127).to(torch.int8)
+scales_f16 = scales.to(torch.float16)
+INT8_BYTES = V * H
+BF16_BYTES = V * H * 2
+
+
+def bench(fn, iters=30, warmup=5):
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    t0 = time.perf_counter()
+    for _ in range(iters):
+        fn()
+    torch.cuda.synchronize()
+    return (time.perf_counter() - t0) / iters
+
+
+# ---- batched int8 w8a16 GEMM: out[B,N] = (x[B,K] @ (W_int8[N,K]*s[N]).T) ----
+@triton.jit
+def _k_batched(x_ptr, w_ptr, s_ptr, o_ptr, B, N, K,
+               sxb, sxk, swn, swk, sob, son,
+               BLOCK_B: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
+    pid_n = tl.program_id(0)
+    offs_b = tl.arange(0, BLOCK_B)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    offs_k = tl.arange(0, BLOCK_K)
+    x_ptrs = x_ptr + offs_b[:, None] * sxb + offs_k[None, :] * sxk
+    w_ptrs = w_ptr + offs_n[:, None] * swn + offs_k[None, :] * swk
+    acc = tl.zeros((BLOCK_B, BLOCK_N), dtype=tl.float32)
+    for k in range(0, K, BLOCK_K):
+        km = (offs_k[None, :] + k) < K
+        x = tl.load(x_ptrs, mask=(offs_b[:, None] < B) & km, other=0.0).to(tl.float16)
+        w = tl.load(w_ptrs, mask=(offs_n[:, None] < N) & km, other=0).to(tl.float16)
+        acc += tl.dot(x, w.T)  # [BB,BK] @ [BK,BN] -> [BB,BN]
+        x_ptrs += BLOCK_K * sxk
+        w_ptrs += BLOCK_K * swk
+    s = tl.load(s_ptr + offs_n, mask=offs_n < N, other=0.0).to(tl.float32)
+    acc = acc * s[None, :]
+    o_ptrs = o_ptr + offs_b[:, None] * sob + offs_n[None, :] * son
+    tl.store(o_ptrs, acc.to(tl.float16), mask=(offs_b[:, None] < B) & (offs_n[None, :] < N))
+
+
+def run_batched(x, BLOCK_N=128, BLOCK_K=64, num_warps=4, num_stages=3):
+    B = x.shape[0]
+    BLOCK_B = max(16, triton.next_power_of_2(B))
+    out = torch.empty(B, V, dtype=torch.float16, device=DEV)
+    xf = x.to(torch.float16)
+    grid = ((V + BLOCK_N - 1) // BLOCK_N,)
+    _k_batched[grid](xf, W_int8, scales_f16, out, B, V, H,
+                     xf.stride(0), xf.stride(1), W_int8.stride(0), W_int8.stride(1),
+                     out.stride(0), out.stride(1),
+                     BLOCK_B=BLOCK_B, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
+                     num_warps=num_warps, num_stages=num_stages)
+    return out
+
+
+for B in (1, 5, 13):
+    print(f"\n=== B={B} ===")
+    x = torch.randn(B, H, device=DEV, dtype=torch.bfloat16) * 0.1
+    ref_bf16 = (x.float() @ W_bf16.float().T)
+    ref_deq = (x.float() @ (W_int8.float() * scales.unsqueeze(1)).T)
+    am_floor = (ref_bf16.argmax(-1) == ref_deq.argmax(-1)).float().mean().item()
+    print(f"  quant floor: argmax_match={am_floor*100:.1f}%  maxerr={(ref_deq-ref_bf16).abs().max():.4f}")
+
+    # real bf16 GEMV baseline (stays bf16, reads 1.5GB)
+    dt = bench(lambda: torch.matmul(x, W_bf16.t()))
+    print(f"  {'bf16 GEMV (real baseline)':30s} {'':33s}{dt*1e3:7.3f} ms  {BF16_BYTES/dt/1e9:6.1f} GB/s")
+
+    # batched int8 kernel, a few configs
+    for (bn, bk, nw, ns) in [(128, 64, 4, 3), (256, 64, 8, 3), (128, 128, 4, 3), (64, 128, 4, 3)]:
+        try:
+            out = run_batched(x, bn, bk, nw, ns)
+            err = (out.float() - ref_deq.float()).abs().max().item()
+            am = (out.float().argmax(-1) == ref_bf16.argmax(-1)).float().mean().item()
+            dt = bench(lambda: run_batched(x, bn, bk, nw, ns))
+            print(f"  batched N{bn}/K{bk}/w{nw}/s{ns:<2d} maxerr={err:8.4f} argmax={am*100:5.1f}% "
+                  f"{dt*1e3:7.3f} ms  {INT8_BYTES/dt/1e9:6.1f} GB/s")
+        except Exception as e:
+            print(f"  batched N{bn}/K{bk}/w{nw}/s{ns}: ERR {str(e)[:70]}")