Add NVFP4 pack geometry profiler
This commit is contained in:
Daniel Maddern
parent
d4c243c9ef
commit
ddc0c5a2b7
4 changed files with 164 additions and 6 deletions
|
|
@ -4,7 +4,7 @@
|
|||
|
||||
torch::Tensor nvfp4_activation_scale_cuda(torch::Tensor input, double divisor);
|
||||
torch::Tensor nvfp4_activation_scale_into_cuda(torch::Tensor input, double divisor, torch::Tensor partials, torch::Tensor output, int64_t blocks, int64_t threads);
|
||||
std::vector<torch::Tensor> quantize_nvfp4_bf16_cuda(torch::Tensor input, torch::Tensor scale, bool pad_16x);
|
||||
std::vector<torch::Tensor> quantize_nvfp4_bf16_cuda(torch::Tensor input, torch::Tensor scale, bool pad_16x, int64_t threads);
|
||||
|
||||
torch::Tensor nvfp4_activation_scale(torch::Tensor input, double divisor) {
|
||||
TORCH_CHECK(input.is_cuda(), "nvfp4_activation_scale expects a CUDA tensor");
|
||||
|
|
@ -29,7 +29,7 @@ torch::Tensor nvfp4_activation_scale_into(torch::Tensor input, double divisor, t
|
|||
return nvfp4_activation_scale_into_cuda(input, divisor, partials, output, blocks, threads);
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> quantize_nvfp4_bf16(torch::Tensor input, torch::Tensor scale, bool pad_16x) {
|
||||
std::vector<torch::Tensor> quantize_nvfp4_bf16(torch::Tensor input, torch::Tensor scale, bool pad_16x, int64_t threads) {
|
||||
TORCH_CHECK(input.is_cuda(), "quantize_nvfp4_bf16 expects a CUDA tensor");
|
||||
TORCH_CHECK(input.is_contiguous(), "quantize_nvfp4_bf16 expects contiguous input");
|
||||
TORCH_CHECK(input.dim() == 2, "quantize_nvfp4_bf16 expects a 2D tensor");
|
||||
|
|
@ -37,7 +37,8 @@ std::vector<torch::Tensor> quantize_nvfp4_bf16(torch::Tensor input, torch::Tenso
|
|||
TORCH_CHECK(scale.is_cuda(), "quantize_nvfp4_bf16 expects a CUDA scale tensor");
|
||||
TORCH_CHECK(scale.device() == input.device(), "quantize_nvfp4_bf16 scale must be on the input device");
|
||||
TORCH_CHECK(scale.numel() == 1, "quantize_nvfp4_bf16 scale must be scalar");
|
||||
return quantize_nvfp4_bf16_cuda(input, scale, pad_16x);
|
||||
TORCH_CHECK(threads == 64 || threads == 128 || threads == 256 || threads == 512, "quantize_nvfp4_bf16 threads must be 64, 128, 256, or 512");
|
||||
return quantize_nvfp4_bf16_cuda(input, scale, pad_16x, threads);
|
||||
}
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
|
|
|
|||
|
|
@ -285,7 +285,7 @@ torch::Tensor nvfp4_activation_scale_into_cuda(torch::Tensor input, double divis
|
|||
return output;
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> quantize_nvfp4_bf16_cuda(torch::Tensor input, torch::Tensor scale, bool pad_16x) {
|
||||
std::vector<torch::Tensor> quantize_nvfp4_bf16_cuda(torch::Tensor input, torch::Tensor scale, bool pad_16x, int64_t threads) {
|
||||
c10::cuda::CUDAGuard device_guard(input.device());
|
||||
const int64_t rows = input.size(0);
|
||||
const int64_t cols = input.size(1);
|
||||
|
|
@ -297,8 +297,7 @@ std::vector<torch::Tensor> quantize_nvfp4_bf16_cuda(torch::Tensor input, torch::
|
|||
auto qdata = torch::empty({q_rows, q_cols}, input.options().dtype(torch::kUInt8));
|
||||
auto block_scale = torch::zeros({scale_rows, scale_cols}, input.options().dtype(torch::kUInt8)).view(torch::kFloat8_e4m3fn);
|
||||
auto stream = at::cuda::getCurrentCUDAStream();
|
||||
constexpr int threads = 256;
|
||||
quantize_nvfp4_bf16_kernel<<<scale_rows, threads, 0, stream>>>(
|
||||
quantize_nvfp4_bf16_kernel<<<scale_rows, static_cast<int>(threads), 0, stream>>>(
|
||||
reinterpret_cast<const uint16_t*>(input.data_ptr()),
|
||||
scale.data_ptr<float>(),
|
||||
qdata.data_ptr<uint8_t>(),
|
||||
|
|
|
|||
|
|
@ -197,6 +197,7 @@ def vortex_native_quantize_nvfp4(
|
|||
tensor,
|
||||
scale,
|
||||
TensorCoreNVFP4Layout.get_padded_shape(orig_shape) != orig_shape,
|
||||
_env_int("H3_NVFP4_PACK_THREADS", 256),
|
||||
),
|
||||
)
|
||||
return _record_timing(
|
||||
|
|
|
|||
157
tools/profile_nvfp4_pack.py
Normal file
157
tools/profile_nvfp4_pack.py
Normal file
|
|
@ -0,0 +1,157 @@
|
|||
"""Profile NVFP4 activation pack kernels without GEMM timing noise."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
import warnings
|
||||
from pathlib import Path
|
||||
|
||||
warnings.filterwarnings("ignore", message="Found GPU0 NVIDIA GB10 which is of cuda capability 12.1.*", category=UserWarning)
|
||||
|
||||
import torch
|
||||
|
||||
import comfy_kitchen as ck
|
||||
from comfy_kitchen.tensor import TensorCoreNVFP4Layout
|
||||
|
||||
from h3_blackwell_runtime.nvfp4_quant import _vortex_scale_extension, nvfp4_activation_scale
|
||||
from tools.profile_nvfp4_linear import module_for_name, representative_inputs
|
||||
|
||||
|
||||
def sync() -> None:
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.synchronize()
|
||||
|
||||
|
||||
def summarize(values: list[float]) -> dict[str, float]:
|
||||
ordered = sorted(values)
|
||||
|
||||
def percentile(percent: float) -> float:
|
||||
if len(ordered) == 1:
|
||||
return ordered[0]
|
||||
rank = (len(ordered) - 1) * percent
|
||||
low = int(rank)
|
||||
high = min(low + 1, len(ordered) - 1)
|
||||
weight = rank - low
|
||||
return ordered[low] * (1.0 - weight) + ordered[high] * weight
|
||||
|
||||
return {
|
||||
"count": len(values),
|
||||
"mean_s": sum(values) / len(values),
|
||||
"p50_s": percentile(0.50),
|
||||
"p90_s": percentile(0.90),
|
||||
"p95_s": percentile(0.95),
|
||||
"p99_s": percentile(0.99),
|
||||
"min_s": ordered[0],
|
||||
"max_s": ordered[-1],
|
||||
}
|
||||
|
||||
|
||||
def timed(fn, iterations: int) -> dict[str, float]:
|
||||
values = []
|
||||
for _ in range(iterations):
|
||||
sync()
|
||||
started = time.perf_counter()
|
||||
fn()
|
||||
sync()
|
||||
values.append(time.perf_counter() - started)
|
||||
return summarize(values)
|
||||
|
||||
|
||||
def bytes_touched(tensor: torch.Tensor, qdata: torch.Tensor, block_scale: torch.Tensor) -> int:
|
||||
return tensor.numel() * tensor.element_size() + qdata.numel() * qdata.element_size() + block_scale.numel()
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--model-path", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
|
||||
parser.add_argument("--output", type=Path, default=Path("/output/h3-blackwell-runtime/benchmarks/nvfp4-pack-profile.json"))
|
||||
parser.add_argument("--width", type=int, default=960)
|
||||
parser.add_argument("--height", type=int, default=544)
|
||||
parser.add_argument("--frames", type=int, default=124)
|
||||
parser.add_argument("--steps", type=int, default=12)
|
||||
parser.add_argument("--sampler-step", type=int, default=1)
|
||||
parser.add_argument("--seed", type=int, default=440407)
|
||||
parser.add_argument("--text-tokens", type=int, default=93)
|
||||
parser.add_argument("--block-index", type=int, default=24)
|
||||
parser.add_argument("--attention", default="sage2")
|
||||
parser.add_argument("--linears", nargs="+", choices=("mlp_fc1", "mlp_fc2", "attn_qkv_proj", "attn_out_proj"), default=("mlp_fc1", "mlp_fc2"))
|
||||
parser.add_argument("--threads", nargs="+", type=int, default=(64, 128, 256, 512))
|
||||
parser.add_argument("--warmup", type=int, default=20)
|
||||
parser.add_argument("--iterations", type=int, default=80)
|
||||
parser.add_argument("--device", default="cuda")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
os.environ.setdefault("H3_NVFP4_SCALE_BACKEND", "vortex")
|
||||
os.environ.setdefault("H3_NVFP4_SCALE_VERSION", "1")
|
||||
|
||||
block, inputs, metadata = representative_inputs(args)
|
||||
extension = _vortex_scale_extension()
|
||||
results = []
|
||||
with torch.inference_mode():
|
||||
for name in args.linears:
|
||||
module = module_for_name(block, name)
|
||||
tensor = inputs[name].reshape(-1, module.in_features).contiguous()
|
||||
scale = nvfp4_activation_scale(tensor).float()
|
||||
pad = TensorCoreNVFP4Layout.get_padded_shape(tuple(tensor.shape)) != tuple(tensor.shape)
|
||||
q_ref, b_ref = ck.quantize_nvfp4(tensor, scale, pad_16x=pad)
|
||||
touched = bytes_touched(tensor, q_ref, b_ref)
|
||||
|
||||
for _ in range(args.warmup):
|
||||
ck.quantize_nvfp4(tensor, scale, pad_16x=pad)
|
||||
ck_timing = timed(lambda: ck.quantize_nvfp4(tensor, scale, pad_16x=pad), args.iterations)
|
||||
results.append(
|
||||
{
|
||||
"name": name,
|
||||
"packer": "ck",
|
||||
"threads": None,
|
||||
"input_shape": list(tensor.shape),
|
||||
"qdata_shape": list(q_ref.shape),
|
||||
"block_scale_shape": list(b_ref.shape),
|
||||
"bytes_touched": touched,
|
||||
"effective_gbps_p50": touched / ck_timing["p50_s"] / 1e9,
|
||||
"effective_gbps_p95": touched / ck_timing["p95_s"] / 1e9,
|
||||
"timing": ck_timing,
|
||||
"qdata_diff": 0,
|
||||
"block_scale_diff": 0,
|
||||
}
|
||||
)
|
||||
|
||||
for threads in args.threads:
|
||||
for _ in range(args.warmup):
|
||||
extension.quantize_nvfp4_bf16(tensor, scale, pad, threads)
|
||||
q_native, b_native = extension.quantize_nvfp4_bf16(tensor, scale, pad, threads)
|
||||
qdiff = int((q_ref != q_native).sum().item())
|
||||
bdiff = int((b_ref.view(torch.uint8) != b_native.view(torch.uint8)).sum().item())
|
||||
timing = timed(lambda: extension.quantize_nvfp4_bf16(tensor, scale, pad, threads), args.iterations)
|
||||
results.append(
|
||||
{
|
||||
"name": name,
|
||||
"packer": "vortex_native",
|
||||
"threads": threads,
|
||||
"input_shape": list(tensor.shape),
|
||||
"qdata_shape": list(q_native.shape),
|
||||
"block_scale_shape": list(b_native.shape),
|
||||
"bytes_touched": touched,
|
||||
"effective_gbps_p50": touched / timing["p50_s"] / 1e9,
|
||||
"effective_gbps_p95": touched / timing["p95_s"] / 1e9,
|
||||
"timing": timing,
|
||||
"qdata_diff": qdiff,
|
||||
"block_scale_diff": bdiff,
|
||||
}
|
||||
)
|
||||
print(name, "threads", threads, "p50_ms", round(timing["p50_s"] * 1000, 3), "p95_ms", round(timing["p95_s"] * 1000, 3), "qdiff", qdiff, "bdiff", bdiff, flush=True)
|
||||
|
||||
output = {"metadata": metadata, "results": results}
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
args.output.write_text(json.dumps(output, indent=2), encoding="utf-8")
|
||||
print(json.dumps(output, indent=2), flush=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Loading…
Add table
Reference in a new issue