40 lines
1.3 KiB
Python
40 lines
1.3 KiB
Python
"""Check the installed Windows vLLM block-FP8 kernel before loading 27B weights."""
|
|
|
|
import json
|
|
import sys
|
|
import traceback
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
|
|
def main() -> None:
|
|
out = Path(sys.argv[1])
|
|
result = {
|
|
"torch": torch.__version__,
|
|
"device": torch.cuda.get_device_name(0),
|
|
"capability": torch.cuda.get_device_capability(0),
|
|
}
|
|
try:
|
|
import vllm.model_executor.kernels.linear.scaled_mm.triton # noqa: F401
|
|
|
|
a = torch.ones((4, 128), device="cuda", dtype=torch.float8_e4m3fn)
|
|
b = torch.ones((128, 128), device="cuda", dtype=torch.float8_e4m3fn)
|
|
sa = torch.ones((4, 1), device="cuda", dtype=torch.float32)
|
|
sb = torch.ones((1, 1), device="cuda", dtype=torch.float32)
|
|
y = torch.ops.vllm.w8a8_triton_block_scaled_mm_func(
|
|
a, b, sa, sb, [128, 128], torch.bfloat16
|
|
)
|
|
torch.cuda.synchronize()
|
|
result.update(status="pass", shape=list(y.shape), value=float(y[0, 0]))
|
|
except Exception as exc:
|
|
result.update(status="failed", error=repr(exc), traceback=traceback.format_exc())
|
|
out.parent.mkdir(parents=True, exist_ok=True)
|
|
out.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
|
|
if result["status"] != "pass":
|
|
raise SystemExit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|