68 lines
2.3 KiB
Python
68 lines
2.3 KiB
Python
#!/usr/bin/env python
|
|
"""Issue 5 round two: raw-op probe of the cutlass W8A8 int8 path on this stack.
|
|
|
|
Bypasses every layer/loader and tests exactly what CutlassInt8ScaledMMLinearKernel
|
|
calls: ops.scaled_int8_quant (dynamic per-token symmetric) then
|
|
ops.cutlass_scaled_mm, against a float32 reference. A failure here means the
|
|
cutlass int8 kernels themselves are broken on this stack (CUDA 13.4 / sm_86 /
|
|
torch 2.13.0+cu130 / vllm 0.31.0), independent of any checkpoint or loader.
|
|
|
|
Usage: cutlass_ops_probe.py
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
|
|
import torch
|
|
from vllm import _custom_ops as ops
|
|
|
|
|
|
def rel(a: torch.Tensor, b: torch.Tensor) -> float:
|
|
return float((a - b).norm() / b.norm())
|
|
|
|
|
|
def main() -> int:
|
|
torch.manual_seed(2301)
|
|
device = torch.device("cuda")
|
|
results = []
|
|
|
|
for m, k, n in ((64, 256, 128), (512, 4096, 4096), (512, 4096, 12288)):
|
|
w = torch.randn(n, k, device=device)
|
|
w_scale = w.abs().amax(dim=1, keepdim=True).clamp_min(1e-12) / 127.0
|
|
w_q = torch.round(w / w_scale).clamp(-127, 127).to(torch.int8)
|
|
|
|
x = torch.randn(m, k, dtype=torch.bfloat16, device=device)
|
|
x_q, x_s, _ = ops.scaled_int8_quant(x.contiguous(), None, None, symmetric=True)
|
|
|
|
# Reference: dequantized per-token acts @ dequantized per-channel weights.
|
|
ref = (x_q.float() * x_s.float()) @ (w_q.float() * w_scale.float()).T
|
|
|
|
w_q_t = w_q.t()
|
|
out = ops.cutlass_scaled_mm(
|
|
x_q, w_q_t, scale_a=x_s, scale_b=w_scale, out_dtype=torch.bfloat16
|
|
).float()
|
|
|
|
# Also check the quantizer op itself against a float32 reference.
|
|
act_scale_ref = x.float().abs().amax(dim=1, keepdim=True).clamp_min(1e-12) / 127.0
|
|
x_q_ref = torch.round(x.float() / act_scale_ref).clamp(-127, 127).to(torch.int8)
|
|
quant_match = bool(torch.equal(x_q, x_q_ref)) and bool(
|
|
torch.allclose(x_s.float(), act_scale_ref, rtol=1e-3, atol=1e-7)
|
|
)
|
|
|
|
results.append({
|
|
"mk": [m, k, n],
|
|
"quant_op_matches_reference": quant_match,
|
|
"gemm_rel_l2": round(rel(out, ref), 5),
|
|
"out_finite": bool(torch.isfinite(out).all()),
|
|
"scale_a_shape": list(x_s.shape),
|
|
"scale_b_shape": list(w_scale.shape),
|
|
})
|
|
|
|
print(json.dumps(results, indent=1))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|