feat: execute FlashKDA on RTX 5090
This commit is contained in:
wuyang
@@ -0,0 +1,454 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Run auditable FlashKDA correctness and latency probes on one CUDA GPU.
|
||||
|
||||
The upstream torch reference is imported from a pinned FlashKDA checkout. The
|
||||
script deliberately uses valid synthetic A_log[H] tensors: it does not resolve
|
||||
the public Kimi K3 checkpoint's A_log[128] versus H=96 inconsistency.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import importlib.metadata
|
||||
import json
|
||||
import math
|
||||
import platform
|
||||
import statistics
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
import flash_kda
|
||||
import flash_kda_C
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--flashkda-source", type=Path, required=True)
|
||||
parser.add_argument("--wheel", type=Path, required=True)
|
||||
parser.add_argument("--output", type=Path, required=True)
|
||||
parser.add_argument("--warmup", type=int, default=20)
|
||||
parser.add_argument("--iters", type=int, default=100)
|
||||
parser.add_argument("--repeats", type=int, default=3)
|
||||
parser.add_argument("--seed", type=int, default=20260729)
|
||||
parser.add_argument("--captured-at", default=None)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def sha256(path: Path) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with path.open("rb") as handle:
|
||||
for block in iter(lambda: handle.read(1024 * 1024), b""):
|
||||
digest.update(block)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def percentile(sorted_values: list[float], fraction: float) -> float:
|
||||
if not sorted_values:
|
||||
return float("nan")
|
||||
index = fraction * (len(sorted_values) - 1)
|
||||
lower = math.floor(index)
|
||||
upper = math.ceil(index)
|
||||
if lower == upper:
|
||||
return sorted_values[lower]
|
||||
weight = index - lower
|
||||
return sorted_values[lower] * (1 - weight) + sorted_values[upper] * weight
|
||||
|
||||
|
||||
def nvidia_smi() -> dict[str, Any]:
|
||||
fields = [
|
||||
"name",
|
||||
"driver_version",
|
||||
"memory.total",
|
||||
"power.limit",
|
||||
"clocks.max.sm",
|
||||
"clocks.max.memory",
|
||||
]
|
||||
output = subprocess.check_output(
|
||||
[
|
||||
"nvidia-smi",
|
||||
f"--query-gpu={','.join(fields)}",
|
||||
"--format=csv,noheader,nounits",
|
||||
],
|
||||
text=True,
|
||||
).strip()
|
||||
values = [value.strip() for value in output.split(",")]
|
||||
return dict(zip(fields, values, strict=True))
|
||||
|
||||
|
||||
def make_inputs(
|
||||
sequence_lengths: list[int],
|
||||
heads: int,
|
||||
seed: int,
|
||||
) -> dict[str, torch.Tensor | float | None]:
|
||||
batch = 1
|
||||
dimension = 128
|
||||
total_tokens = sum(sequence_lengths)
|
||||
sequences = len(sequence_lengths)
|
||||
generator = torch.Generator(device="cuda").manual_seed(seed)
|
||||
|
||||
q = F.normalize(
|
||||
torch.randn(
|
||||
(batch, total_tokens, heads, dimension),
|
||||
dtype=torch.float32,
|
||||
device="cuda",
|
||||
generator=generator,
|
||||
),
|
||||
p=2,
|
||||
dim=-1,
|
||||
).to(torch.bfloat16)
|
||||
k = F.normalize(
|
||||
torch.randn(
|
||||
(batch, total_tokens, heads, dimension),
|
||||
dtype=torch.float32,
|
||||
device="cuda",
|
||||
generator=generator,
|
||||
),
|
||||
p=2,
|
||||
dim=-1,
|
||||
).to(torch.bfloat16)
|
||||
v = torch.randn(
|
||||
(batch, total_tokens, heads, dimension),
|
||||
dtype=torch.bfloat16,
|
||||
device="cuda",
|
||||
generator=generator,
|
||||
)
|
||||
g = torch.randn(
|
||||
(batch, total_tokens, heads, dimension),
|
||||
dtype=torch.bfloat16,
|
||||
device="cuda",
|
||||
generator=generator,
|
||||
)
|
||||
beta = torch.randn(
|
||||
(batch, total_tokens, heads),
|
||||
dtype=torch.bfloat16,
|
||||
device="cuda",
|
||||
generator=generator,
|
||||
)
|
||||
a_log = torch.rand(
|
||||
heads, dtype=torch.float32, device="cuda", generator=generator
|
||||
)
|
||||
dt_bias = torch.rand(
|
||||
(heads, dimension),
|
||||
dtype=torch.float32,
|
||||
device="cuda",
|
||||
generator=generator,
|
||||
)
|
||||
initial_state = torch.randn(
|
||||
(sequences, heads, dimension, dimension),
|
||||
dtype=torch.bfloat16,
|
||||
device="cuda",
|
||||
generator=generator,
|
||||
)
|
||||
cu_seqlens = None
|
||||
if len(sequence_lengths) > 1:
|
||||
offsets = [0]
|
||||
for length in sequence_lengths:
|
||||
offsets.append(offsets[-1] + length)
|
||||
cu_seqlens = torch.tensor(offsets, dtype=torch.long, device="cuda")
|
||||
|
||||
return {
|
||||
"q": q,
|
||||
"k": k,
|
||||
"v": v,
|
||||
"g": g,
|
||||
"beta": beta,
|
||||
"A_log": a_log,
|
||||
"dt_bias": dt_bias,
|
||||
"initial_state": initial_state,
|
||||
"cu_seqlens": cu_seqlens,
|
||||
"scale": 1 / math.sqrt(dimension),
|
||||
}
|
||||
|
||||
|
||||
def run_correctness_case(
|
||||
torch_ref: Callable[..., None],
|
||||
name: str,
|
||||
sequence_lengths: list[int],
|
||||
heads: int,
|
||||
seed: int,
|
||||
) -> dict[str, Any]:
|
||||
inputs = make_inputs(sequence_lengths, heads, seed)
|
||||
q = inputs["q"]
|
||||
assert isinstance(q, torch.Tensor)
|
||||
initial_state = inputs["initial_state"]
|
||||
assert isinstance(initial_state, torch.Tensor)
|
||||
|
||||
out_kernel = torch.zeros_like(q)
|
||||
out_reference = torch.zeros_like(q)
|
||||
state_kernel = torch.zeros_like(initial_state)
|
||||
state_reference = torch.zeros_like(initial_state)
|
||||
common = {
|
||||
"A_log": inputs["A_log"],
|
||||
"dt_bias": inputs["dt_bias"],
|
||||
"lower_bound": -5.0,
|
||||
"initial_state": initial_state.clone(),
|
||||
"cu_seqlens": inputs["cu_seqlens"],
|
||||
}
|
||||
|
||||
start = time.perf_counter()
|
||||
flash_kda.fwd(
|
||||
inputs["q"],
|
||||
inputs["k"],
|
||||
inputs["v"],
|
||||
inputs["g"],
|
||||
inputs["beta"],
|
||||
inputs["scale"],
|
||||
out_kernel,
|
||||
final_state=state_kernel,
|
||||
**common,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
kernel_ms = (time.perf_counter() - start) * 1000
|
||||
|
||||
start = time.perf_counter()
|
||||
torch_ref(
|
||||
inputs["q"],
|
||||
inputs["k"],
|
||||
inputs["v"],
|
||||
inputs["g"],
|
||||
inputs["beta"],
|
||||
inputs["scale"],
|
||||
out_reference,
|
||||
final_state=state_reference,
|
||||
**common,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
reference_ms = (time.perf_counter() - start) * 1000
|
||||
|
||||
output_diff = (out_kernel.float() - out_reference.float()).abs()
|
||||
state_diff = (state_kernel.float() - state_reference.float()).abs()
|
||||
result = {
|
||||
"name": name,
|
||||
"sequence_lengths": sequence_lengths,
|
||||
"heads": heads,
|
||||
"dimension": 128,
|
||||
"kernel_ms_first_measured": kernel_ms,
|
||||
"reference_ms": reference_ms,
|
||||
"output_exact": bool(torch.equal(out_kernel, out_reference)),
|
||||
"state_exact": bool(torch.equal(state_kernel, state_reference)),
|
||||
"output_max_abs_diff": output_diff.max().item(),
|
||||
"output_mean_abs_diff": output_diff.mean().item(),
|
||||
"state_max_abs_diff": state_diff.max().item(),
|
||||
"state_mean_abs_diff": state_diff.mean().item(),
|
||||
}
|
||||
if not result["output_exact"] or not result["state_exact"]:
|
||||
raise AssertionError(f"exact correctness failed: {result}")
|
||||
return result
|
||||
|
||||
|
||||
def bench_events(
|
||||
function: Callable[[], None],
|
||||
warmup: int,
|
||||
iters: int,
|
||||
repeats: int,
|
||||
) -> list[float]:
|
||||
for _ in range(max(warmup, 1)):
|
||||
function()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
elapsed: list[float] = []
|
||||
for _ in range(repeats):
|
||||
starts = [torch.cuda.Event(enable_timing=True) for _ in range(iters)]
|
||||
ends = [torch.cuda.Event(enable_timing=True) for _ in range(iters)]
|
||||
torch.cuda.synchronize()
|
||||
for index in range(iters):
|
||||
starts[index].record()
|
||||
function()
|
||||
ends[index].record()
|
||||
torch.cuda.synchronize()
|
||||
elapsed.extend(
|
||||
start.elapsed_time(end) for start, end in zip(starts, ends, strict=True)
|
||||
)
|
||||
return elapsed
|
||||
|
||||
|
||||
def summarize_latency(
|
||||
values: list[float],
|
||||
total_tokens: int,
|
||||
) -> dict[str, Any]:
|
||||
ordered = sorted(float(value) for value in values)
|
||||
mean = statistics.fmean(ordered)
|
||||
return {
|
||||
"samples": len(ordered),
|
||||
"mean_ms": mean,
|
||||
"min_ms": ordered[0],
|
||||
"p50_ms": percentile(ordered, 0.50),
|
||||
"p95_ms": percentile(ordered, 0.95),
|
||||
"max_ms": ordered[-1],
|
||||
"sequence_tokens_per_second": total_tokens / (mean / 1000),
|
||||
}
|
||||
|
||||
|
||||
def run_benchmark_case(
|
||||
name: str,
|
||||
sequence_lengths: list[int],
|
||||
heads: int,
|
||||
seed: int,
|
||||
warmup: int,
|
||||
iters: int,
|
||||
repeats: int,
|
||||
) -> dict[str, Any]:
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
inputs = make_inputs(sequence_lengths, heads, seed)
|
||||
q = inputs["q"]
|
||||
initial_state = inputs["initial_state"]
|
||||
assert isinstance(q, torch.Tensor)
|
||||
assert isinstance(initial_state, torch.Tensor)
|
||||
out = torch.zeros_like(q)
|
||||
state_bf16 = torch.zeros_like(initial_state)
|
||||
initial_fp32 = initial_state.float()
|
||||
state_fp32 = torch.zeros_like(initial_fp32)
|
||||
common = {
|
||||
"A_log": inputs["A_log"],
|
||||
"dt_bias": inputs["dt_bias"],
|
||||
"lower_bound": -5.0,
|
||||
"cu_seqlens": inputs["cu_seqlens"],
|
||||
}
|
||||
|
||||
def invoke(initial: torch.Tensor | None, final: torch.Tensor | None) -> None:
|
||||
flash_kda.fwd(
|
||||
inputs["q"],
|
||||
inputs["k"],
|
||||
inputs["v"],
|
||||
inputs["g"],
|
||||
inputs["beta"],
|
||||
inputs["scale"],
|
||||
out,
|
||||
initial_state=initial,
|
||||
final_state=final,
|
||||
**common,
|
||||
)
|
||||
|
||||
variants = {
|
||||
"bf16_state": lambda: invoke(initial_state, state_bf16),
|
||||
"no_state": lambda: invoke(None, None),
|
||||
"fp32_state": lambda: invoke(initial_fp32, state_fp32),
|
||||
}
|
||||
timings = {
|
||||
variant: summarize_latency(
|
||||
bench_events(function, warmup, iters, repeats),
|
||||
sum(sequence_lengths),
|
||||
)
|
||||
for variant, function in variants.items()
|
||||
}
|
||||
torch.cuda.synchronize()
|
||||
return {
|
||||
"name": name,
|
||||
"sequence_lengths": sequence_lengths,
|
||||
"total_tokens": sum(sequence_lengths),
|
||||
"heads": heads,
|
||||
"dimension": 128,
|
||||
"warmup": warmup,
|
||||
"iters": iters,
|
||||
"repeats": repeats,
|
||||
"timings": timings,
|
||||
"output_abs_mean_after_last_run": out.float().abs().mean().item(),
|
||||
"peak_allocated_mib": torch.cuda.max_memory_allocated() / 2**20,
|
||||
"peak_reserved_mib": torch.cuda.max_memory_reserved() / 2**20,
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
if not torch.cuda.is_available():
|
||||
raise RuntimeError("CUDA GPU is required")
|
||||
if torch.cuda.get_device_capability(0) < (9, 0):
|
||||
raise RuntimeError("FlashKDA requires SM90 or newer")
|
||||
tests_dir = args.flashkda_source / "tests"
|
||||
sys.path.insert(0, str(tests_dir))
|
||||
from torch_ref import torch_ref
|
||||
|
||||
revision = subprocess.check_output(
|
||||
["git", "-C", str(args.flashkda_source), "rev-parse", "HEAD"],
|
||||
text=True,
|
||||
).strip()
|
||||
torch.manual_seed(args.seed)
|
||||
torch.cuda.manual_seed_all(args.seed)
|
||||
|
||||
correctness_specs = [
|
||||
("one_chunk", [16], 1),
|
||||
("partial_tail", [17], 1),
|
||||
("two_chunks", [32], 1),
|
||||
("multi_chunk_multi_head", [65], 2),
|
||||
("k3_head_count", [17], 96),
|
||||
("varlen_partial_chunks", [17, 31], 2),
|
||||
]
|
||||
correctness = [
|
||||
run_correctness_case(torch_ref, name, lengths, heads, args.seed + index)
|
||||
for index, (name, lengths, heads) in enumerate(correctness_specs)
|
||||
]
|
||||
|
||||
benchmark_specs = [
|
||||
("teaching_scale", [512], 8),
|
||||
("intermediate", [2048], 32),
|
||||
("k3_fixed_shape", [8192], 96),
|
||||
("k3_varlen_shape", [1300, 547, 2048, 963, 271, 3063], 96),
|
||||
]
|
||||
benchmarks = [
|
||||
run_benchmark_case(
|
||||
name,
|
||||
lengths,
|
||||
heads,
|
||||
args.seed + 100 + index,
|
||||
args.warmup,
|
||||
args.iters,
|
||||
args.repeats,
|
||||
)
|
||||
for index, (name, lengths, heads) in enumerate(benchmark_specs)
|
||||
]
|
||||
|
||||
captured_at = args.captured_at or datetime.now(timezone.utc).isoformat()
|
||||
result = {
|
||||
"schema_version": 1,
|
||||
"captured_at": captured_at,
|
||||
"evidence_identity": "X / local execution on deterministic synthetic tensors",
|
||||
"boundary": {
|
||||
"k3_checkpoint_loaded": False,
|
||||
"real_token_hidden_states": False,
|
||||
"a_log_shape_conflict_resolved": False,
|
||||
"benchmark_comparison": "local FlashKDA timings only; author H20/GB200 tables remain separate",
|
||||
},
|
||||
"provenance": {
|
||||
"flashkda_revision": revision,
|
||||
"flashkda_package": importlib.metadata.version("flash-kda"),
|
||||
"wheel_filename": args.wheel.name,
|
||||
"wheel_sha256": sha256(args.wheel),
|
||||
"runner": str(Path(__file__).relative_to(Path.cwd())),
|
||||
},
|
||||
"environment": {
|
||||
"python": platform.python_version(),
|
||||
"platform": platform.platform(),
|
||||
"libc": list(platform.libc_ver()),
|
||||
"torch": torch.__version__,
|
||||
"torch_cuda": torch.version.cuda,
|
||||
"gpu": torch.cuda.get_device_name(0),
|
||||
"capability": list(torch.cuda.get_device_capability(0)),
|
||||
"nvidia_smi": nvidia_smi(),
|
||||
"flash_kda_module": flash_kda.__file__,
|
||||
"flash_kda_extension": flash_kda_C.__file__,
|
||||
},
|
||||
"correctness": {
|
||||
"all_exact": all(
|
||||
case["output_exact"] and case["state_exact"]
|
||||
for case in correctness
|
||||
),
|
||||
"cases": correctness,
|
||||
},
|
||||
"benchmarks": benchmarks,
|
||||
}
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
args.output.write_text(json.dumps(result, indent=2) + "\n")
|
||||
print(json.dumps(result, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user