ent
46 lines
1.9 KiB
Python
46 lines
1.9 KiB
Python
#!/usr/bin/env python3
|
|
"""Validate patch_int8_lmhead_v3 landed + the helper runs on GPU. Run AFTER the
|
|
patch in the same container."""
|
|
import torch
|
|
|
|
import vllm.model_executor.layers.logits_processor as lp
|
|
|
|
print("import OK | helpers:",
|
|
hasattr(lp, "_spark_int8_gemm"),
|
|
hasattr(lp, "_spark_int8_lmhead_apply"),
|
|
hasattr(lp, "_spark_k_int8"))
|
|
src = open(lp.__file__).read()
|
|
print("sentinel in _get_logits:", "DGX_SPARK_INT8_LMHEAD_V3: int8 w8a16" in src)
|
|
|
|
V, H = 4096, 512
|
|
torch.manual_seed(0)
|
|
W = torch.randn(V, H, device="cuda") * 0.02
|
|
s = (W.abs().amax(1) / 127).clamp(min=1e-12)
|
|
wi = (W / s.unsqueeze(1)).round().clamp(-127, 127).to(torch.int8).contiguous()
|
|
sf = s.to(torch.float16)
|
|
for B in (1, 5, 13):
|
|
x = torch.randn(B, H, device="cuda", dtype=torch.bfloat16) * 0.1
|
|
out = lp._spark_int8_gemm(x, wi, sf)
|
|
ref = x.float() @ (wi.float() * s.unsqueeze(1)).T
|
|
am = (out.argmax(-1) == ref.argmax(-1)).float().mean().item()
|
|
print(f" B={B}: out{tuple(out.shape)} {out.dtype} argmax={am*100:.0f}% "
|
|
f"maxerr={(out - ref).abs().max().item():.4f}")
|
|
|
|
# Exercise the FULL apply path (print + quantize-once + gemm) with a mock lm_head
|
|
# (vocab > 100k to trigger the int8 path). self is unused -> None.
|
|
print("--- full _spark_int8_lmhead_apply path (mock lm_head, V=131072) ---")
|
|
Vbig = 131072
|
|
|
|
class _MockLMHead:
|
|
pass
|
|
|
|
mh = _MockLMHead()
|
|
mh.weight = (torch.randn(Vbig, H, device="cuda", dtype=torch.float32) * 0.02).to(torch.bfloat16)
|
|
hs = torch.randn(3, H, device="cuda", dtype=torch.bfloat16) * 0.1
|
|
o1 = lp._spark_int8_lmhead_apply(None, mh, hs, None) # first call: quantizes + prints
|
|
o2 = lp._spark_int8_lmhead_apply(None, mh, hs, None) # second call: reuses int8
|
|
refb = hs.float() @ mh.weight.float().T
|
|
am = (o2.argmax(-1) == refb.argmax(-1)).float().mean().item()
|
|
print(f" apply: out{tuple(o2.shape)} {o2.dtype} argmax_vs_bf16={am*100:.0f}% "
|
|
f"int8_ready={getattr(mh, '_spark_int8_ready', None)} weight_kept={mh.weight.numel() > 0}")
|