[perf]: cuBLASLt algo sweep finds 8-9% GEMM band headroom with bit-exact parity
This commit is contained in:
SolitaryThinker
@@ -440,7 +440,11 @@ Today's measurements make the 50% gap quantitative at 4x/B2 (2,450 TFLOP/s conve
|
||||
| realistic best systems (120 ms) | 457.8 ms | 1,585 TF/s (64.7%) | x1.142 |
|
||||
| theoretical-max systems (75 ms) | 502.8 ms | 1,443 TF/s (58.9%) | x1.040 |
|
||||
|
||||
Equivalently, with today's kernels a perfect 75 ms systems endgame caps at **48.34% MFU**: 50% is strictly unreachable in BF16 by systems work alone. The admissible next steps are kernel research — raise the nvjet-bound GEMM band from 62.9% toward 70%+ and/or FA4 from 38.4% toward 50% on SM100 — or the separately-gated NVFP4 projection track. Library-level levers (cuBLAS heuristics/workspace, TorchInductor CUTLASS, TunableOp, attention backends and FA4 split schedules, batch capacity) are all measured and closed above.
|
||||
Equivalently, with today's kernels a perfect 75 ms systems endgame caps at **48.34% MFU**: 50% is strictly unreachable in BF16 by systems work alone. The admissible next steps are kernel research — raise the nvjet-bound GEMM band from 62.9% toward 70%+ and/or FA4 from 38.4% toward 50% on SM100 — or the separately-gated NVFP4 projection track. Torch-dispatch-level levers (workspace, TorchInductor CUTLASS, TunableOp, attention backends and FA4 split schedules, batch capacity) are measured and closed above; native cuBLASLt algo selection is NOT closed — see the sweep below.
|
||||
|
||||
### cuBLASLt heuristic algo pinning: open, quantified opportunity
|
||||
|
||||
`probes/bench_ltx2_cublaslt_algo_sweep.py` (research-plan item 2) enumerates up to 48 `cublasLtMatmulAlgoGetHeuristic` candidates per exact packed shape and training orientation via ctypes, validates each winner bit-exactly against torch in-process, and times bare GEMMs on both sides. On a 1,965 MHz tray (ratios valid; absolute ms are bin-scaled), the occurrence-weighted band improves **-7.98% at both B2 and B3** versus torch's own dispatch, **-9.4%** with per-case best-of, and **-7.7%** restricted to dgrad/wgrad only (the conservative subset that needs no bias-epilogue re-validation; production dgrad is epilogue-free and dbias is a separate compiled reduction). Every winning algo matched torch bit-exactly (`max_abs 0.0`). Case highlights at B2: `video_dd:dgrad` -23.6%, `text_kv:wgrad` -23.2%, `self_qkv:dgrad` -17.9%, `ffn_up:wgrad` -16.5%; the heuristic best loses on `ffn_up:fwd` (+15.2%), so pinning must be per-case best-of. Projected onto the healthy-tray 389.300 ms band this is **-30 to -37 ms/step at 4x/B2, about +1.9 to +2.2 MFU points**, lifting the GEMM band to roughly 68-70% of peak — inside the realistic-systems kernel requirement for 50% above. Next step per the plan: a `torch.library` custom op wrapping `cublasLtMatmul` with pinned per-shape algos (opt-in, fwd cases only after an epilogue-aware re-sweep), then a trainer A/X/B on a healthy allocation.
|
||||
|
||||
A third allocation roll drew rack-059 pair `1632122` (gb-nvl-059-compute05/06), which passed the native health gate with that preamble: all-rank scalar correctness and 458.914 GB/s slowest-rank sustained all-reduce bus bandwidth, matching the healthy 463 GB/s reference, so the MNNVL failure is rack-057-specific. The pair is nevertheless the 1,965 MHz clock bin and its absolute B3 result sits in the degraded-allocation class alongside the `1629519` row: the exact-head A/X/B measured 1.442031/1.443601/1.444706 s for control A / env stack / control B — a 1.443369 s / 30.021699% MFU control midpoint with 0.185% drift, and an env-stack delta of **+0.232 ms / +0.016%, i.e. neutral** (memory unchanged, 161.38 GiB peak allocated in all rows). On a degraded pair the pointwise-tuning win measured at 4x is inside noise, so the environment stack remains accepted for 4x work, is harmless on 8x, and stays **out of `run_current.sh`** until a healthy-class (about 43% baseline) 8x pair re-gates it. This row is baseline-ineligible and does not supersede the stopping-point table.
|
||||
|
||||
@@ -702,6 +706,9 @@ f6bdae8f1e6aef72285ddcb15a0639a28549e67d12e0c2b9db18536255efa8a9 /mnt/pr1630_8x
|
||||
4704d335a688c2eab66ad6eb4e033c12b709938689abac82398e52326105cca6 /mnt/pr1630_8x_envstack_envstack_node1.log
|
||||
10eebf0d0eb9c71e5aca4eee51c66e68454767b0838af61516cce0d4e3a231f4 /mnt/pr1630_8x_envstack_control_b_node0.log
|
||||
f681ac1b60beb31aa3dd868e9005684810e05ddb26966bd78885141acd7864bb /mnt/pr1630_8x_envstack_control_b_node1.log
|
||||
697537479ba7f945983b0fe3e632bab39fdbb9406cf39b584587c4ad16387bea /mnt/pr1630_lt_sweep_b2.log
|
||||
a02fe5338276d8dc3ce31b0af903d80eb71a9980a797e078732893f15807c463 /mnt/pr1630_lt_sweep_b3.log
|
||||
41fc1e8970cbe39431a9cd56773910d3ca6cd0950ef843b5882d6c15a5d66f47 [committed probes/bench_ltx2_cublaslt_algo_sweep.py] /mnt/bench_ltx2_cublaslt_algo_sweep.py
|
||||
7994f41c82ece7a3ab8c324dc3543f99cf4437c4dce3db4d0dbf36bd43b12b74 /mnt/pr1630_video_attn_splits_b2.log
|
||||
bd33740e9159d6ea5a4f02e06e5fc15721fb10d8a3dadd91e20d1608175602f7 /mnt/pr1630_video_attn_splits_b3.log
|
||||
647c2ffb55613bf3ded4f464fc1db1cfbe1ca5bfc50df7ab30c1aeecee3909d6 [committed probes/benchmark_ltx2_video_attention_backends.py with split sweep] /mnt/benchmark_ltx2_video_attention_backends.py
|
||||
|
||||
@@ -0,0 +1,241 @@
|
||||
#!/usr/bin/env python3
|
||||
"""cuBLASLt heuristic-algo sweep for the packed LTX-2 GEMM band (plan item 2).
|
||||
|
||||
Enumerates ``cublasLtMatmulAlgoGetHeuristic`` candidates per exact shape and
|
||||
training orientation (fwd, dgrad, wgrad), times each against torch's own
|
||||
dispatch in the same process, and reports the best-algo margin. Kill
|
||||
criterion from ``reports/bf16_kernel_research_plan.md``: if the best
|
||||
heuristic algo is within 2% of torch/nvjet per weighted band, native algo
|
||||
pinning is closed. Bare GEMMs (no bias epilogue) on both sides; BF16 inputs,
|
||||
FP32 compute/accumulate. Parity-checked against torch per case.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import ctypes
|
||||
import json
|
||||
import statistics
|
||||
|
||||
import torch
|
||||
|
||||
CUDA_R_16BF = 14
|
||||
CUDA_R_32F = 0
|
||||
CUBLAS_COMPUTE_32F = 68
|
||||
OP_N, OP_T = 0, 1
|
||||
DESC_TRANSA, DESC_TRANSB = 3, 4
|
||||
PREF_MAX_WORKSPACE = 1
|
||||
HIDDEN = 4096
|
||||
FFN = 16384
|
||||
VIDEO_TOKENS = 11 * 15 * 26
|
||||
TEXT_TOKENS = 1024
|
||||
WORKSPACE_BYTES = 128 * 1024 * 1024
|
||||
MAX_ALGOS = 48
|
||||
|
||||
|
||||
class HeuristicResult(ctypes.Structure):
|
||||
_fields_ = [
|
||||
("algo", ctypes.c_uint64 * 8),
|
||||
("workspaceSize", ctypes.c_size_t),
|
||||
("state", ctypes.c_int),
|
||||
("wavesCount", ctypes.c_float),
|
||||
("reserved", ctypes.c_int * 4),
|
||||
]
|
||||
|
||||
|
||||
def load_lt() -> ctypes.CDLL:
|
||||
torch.cuda.init()
|
||||
for name in ("libcublasLt.so.13", "libcublasLt.so.12", "libcublasLt.so"):
|
||||
try:
|
||||
return ctypes.CDLL(name)
|
||||
except OSError:
|
||||
continue
|
||||
raise OSError("libcublasLt not found")
|
||||
|
||||
|
||||
def check(status: int, what: str) -> None:
|
||||
if status != 0:
|
||||
raise RuntimeError(f"{what} failed with cublas status {status}")
|
||||
|
||||
|
||||
def build_cases(batch: int) -> list[dict]:
|
||||
m_video = VIDEO_TOKENS * batch
|
||||
m_text = TEXT_TOKENS * batch
|
||||
shapes = [
|
||||
("self_qkv", m_video, HIDDEN, 3 * HIDDEN, 48),
|
||||
("video_dd", m_video, HIDDEN, HIDDEN, 144),
|
||||
("text_kv", m_text, HIDDEN, 2 * HIDDEN, 48),
|
||||
("ffn_up", m_video, HIDDEN, FFN, 48),
|
||||
("ffn_down", m_video, FFN, HIDDEN, 48),
|
||||
]
|
||||
cases = []
|
||||
for name, rows, in_features, out_features, occurrences in shapes:
|
||||
# Row-major training GEMMs expressed as column-major cuBLASLt calls
|
||||
# (documented mapping: fwd y=xW^T, dgrad dx=dyW, wgrad dW=dy^T x).
|
||||
cases.append({
|
||||
"case": f"{name}:fwd", "occurrences": occurrences,
|
||||
"m": out_features, "n": rows, "k": in_features,
|
||||
"transa": OP_T, "transb": OP_N,
|
||||
"a_shape": (out_features, in_features), "b_shape": (rows, in_features),
|
||||
"d_shape": (rows, out_features),
|
||||
"torch_fn": lambda a, b: b @ a.t(),
|
||||
})
|
||||
cases.append({
|
||||
"case": f"{name}:dgrad", "occurrences": occurrences,
|
||||
"m": in_features, "n": rows, "k": out_features,
|
||||
"transa": OP_N, "transb": OP_N,
|
||||
"a_shape": (out_features, in_features), "b_shape": (rows, out_features),
|
||||
"d_shape": (rows, in_features),
|
||||
"torch_fn": lambda a, b: b @ a,
|
||||
})
|
||||
cases.append({
|
||||
"case": f"{name}:wgrad", "occurrences": occurrences,
|
||||
"m": in_features, "n": out_features, "k": rows,
|
||||
"transa": OP_N, "transb": OP_T,
|
||||
"a_shape": (rows, in_features), "b_shape": (rows, out_features),
|
||||
"d_shape": (out_features, in_features),
|
||||
"torch_fn": lambda a, b: b.t() @ a,
|
||||
})
|
||||
return cases
|
||||
|
||||
|
||||
def time_fn(fn, warmup: int, repeats: int) -> float:
|
||||
for _ in range(warmup):
|
||||
fn()
|
||||
torch.cuda.synchronize()
|
||||
events = [(torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)) for _ in range(repeats)]
|
||||
for start, end in events:
|
||||
start.record()
|
||||
fn()
|
||||
end.record()
|
||||
torch.cuda.synchronize()
|
||||
return statistics.median(start.elapsed_time(end) for start, end in events)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--batch", type=int, default=2)
|
||||
parser.add_argument("--warmup", type=int, default=8)
|
||||
parser.add_argument("--repeats", type=int, default=21)
|
||||
args = parser.parse_args()
|
||||
|
||||
torch.manual_seed(20260722)
|
||||
device = torch.device("cuda:0")
|
||||
lt = load_lt()
|
||||
handle = ctypes.c_void_p()
|
||||
check(lt.cublasLtCreate(ctypes.byref(handle)), "ltCreate")
|
||||
workspace = torch.empty(WORKSPACE_BYTES, dtype=torch.uint8, device=device)
|
||||
stream = ctypes.c_void_p(torch.cuda.current_stream().cuda_stream)
|
||||
alpha = ctypes.c_float(1.0)
|
||||
beta = ctypes.c_float(0.0)
|
||||
|
||||
summary = []
|
||||
for case in build_cases(args.batch):
|
||||
a_rm = torch.randn(case["a_shape"], device=device, dtype=torch.bfloat16)
|
||||
b_rm = torch.randn(case["b_shape"], device=device, dtype=torch.bfloat16)
|
||||
d = torch.empty(case["d_shape"], device=device, dtype=torch.bfloat16)
|
||||
reference = case["torch_fn"](a_rm, b_rm)
|
||||
|
||||
# Column-major operand descriptors. Buffers: A is the weight-like
|
||||
# tensor, B the activation-like tensor, both row-major contiguous.
|
||||
m, n, k = case["m"], case["n"], case["k"]
|
||||
op_desc = ctypes.c_void_p()
|
||||
check(lt.cublasLtMatmulDescCreate(ctypes.byref(op_desc), CUBLAS_COMPUTE_32F, CUDA_R_32F), "descCreate")
|
||||
for attr, value in ((DESC_TRANSA, case["transa"]), (DESC_TRANSB, case["transb"])):
|
||||
v = ctypes.c_int32(value)
|
||||
check(lt.cublasLtMatmulDescSetAttribute(op_desc, attr, ctypes.byref(v), 4), "descSet")
|
||||
|
||||
def layout(rows_cm: int, cols_cm: int, ld: int) -> ctypes.c_void_p:
|
||||
handle_ = ctypes.c_void_p()
|
||||
check(lt.cublasLtMatrixLayoutCreate(ctypes.byref(handle_), CUDA_R_16BF, rows_cm, cols_cm,
|
||||
ctypes.c_int64(ld)), "layoutCreate")
|
||||
return handle_
|
||||
|
||||
lda = case["a_shape"][1]
|
||||
ldb = case["b_shape"][1]
|
||||
a_rows_cm = lda if case["transa"] == OP_T else m
|
||||
a_cols_cm = m if case["transa"] == OP_T else k
|
||||
b_rows_cm = ldb if case["transb"] == OP_T else k
|
||||
b_cols_cm = k if case["transb"] == OP_T else n
|
||||
layout_a = layout(a_rows_cm, a_cols_cm, lda)
|
||||
layout_b = layout(b_rows_cm, b_cols_cm, ldb)
|
||||
layout_d = layout(m, n, m)
|
||||
|
||||
pref = ctypes.c_void_p()
|
||||
check(lt.cublasLtMatmulPreferenceCreate(ctypes.byref(pref)), "prefCreate")
|
||||
ws = ctypes.c_size_t(WORKSPACE_BYTES)
|
||||
check(lt.cublasLtMatmulPreferenceSetAttribute(pref, PREF_MAX_WORKSPACE, ctypes.byref(ws), 8), "prefSet")
|
||||
|
||||
results = (HeuristicResult * MAX_ALGOS)()
|
||||
found = ctypes.c_int(0)
|
||||
status = lt.cublasLtMatmulAlgoGetHeuristic(handle, op_desc, layout_a, layout_b, layout_d, layout_d,
|
||||
pref, MAX_ALGOS, results, ctypes.byref(found))
|
||||
check(status, "algoGetHeuristic")
|
||||
|
||||
def run_lt(algo_ref) -> None:
|
||||
check(
|
||||
lt.cublasLtMatmul(handle, op_desc, ctypes.byref(alpha),
|
||||
ctypes.c_void_p(a_rm.data_ptr()), layout_a,
|
||||
ctypes.c_void_p(b_rm.data_ptr()), layout_b,
|
||||
ctypes.byref(beta),
|
||||
ctypes.c_void_p(d.data_ptr()), layout_d,
|
||||
ctypes.c_void_p(d.data_ptr()), layout_d,
|
||||
algo_ref,
|
||||
ctypes.c_void_p(workspace.data_ptr()), ctypes.c_size_t(WORKSPACE_BYTES),
|
||||
stream), "ltMatmul")
|
||||
|
||||
torch_ms = time_fn(lambda: case["torch_fn"](a_rm, b_rm), args.warmup, args.repeats)
|
||||
|
||||
best = None
|
||||
parity_max_abs = None
|
||||
for index in range(found.value):
|
||||
if results[index].state != 0:
|
||||
continue
|
||||
algo_ref = ctypes.byref(results[index], 0)
|
||||
try:
|
||||
run_lt(algo_ref)
|
||||
except RuntimeError:
|
||||
continue
|
||||
torch.cuda.synchronize()
|
||||
diff = (d.float() - reference.float()).abs().max().item()
|
||||
if diff > 0.25:
|
||||
continue
|
||||
lt_ms = time_fn(lambda: run_lt(algo_ref), args.warmup, args.repeats)
|
||||
if best is None or lt_ms < best[0]:
|
||||
best = (lt_ms, index)
|
||||
parity_max_abs = diff
|
||||
|
||||
row = {
|
||||
"case": case["case"],
|
||||
"mnk": [m, n, k],
|
||||
"occurrences": case["occurrences"],
|
||||
"algos_returned": found.value,
|
||||
"torch_ms": round(torch_ms, 4),
|
||||
"best_lt_ms": round(best[0], 4) if best else None,
|
||||
"best_algo_index": best[1] if best else None,
|
||||
"parity_max_abs": parity_max_abs,
|
||||
"lt_vs_torch_pct": round(100.0 * (best[0] - torch_ms) / torch_ms, 2) if best else None,
|
||||
}
|
||||
summary.append(row)
|
||||
print("LT_SWEEP_CASE " + json.dumps(row, sort_keys=True), flush=True)
|
||||
for obj, destroy in ((pref, lt.cublasLtMatmulPreferenceDestroy), (layout_a, lt.cublasLtMatrixLayoutDestroy),
|
||||
(layout_b, lt.cublasLtMatrixLayoutDestroy), (layout_d, lt.cublasLtMatrixLayoutDestroy),
|
||||
(op_desc, lt.cublasLtMatmulDescDestroy)):
|
||||
destroy(obj)
|
||||
|
||||
weighted_torch = sum(r["torch_ms"] * r["occurrences"] for r in summary)
|
||||
weighted_lt = sum((r["best_lt_ms"] or r["torch_ms"]) * r["occurrences"] for r in summary)
|
||||
print(
|
||||
"LT_SWEEP_RESULT " + json.dumps({
|
||||
"batch": args.batch,
|
||||
"weighted_torch_band_ms": round(weighted_torch, 2),
|
||||
"weighted_best_lt_band_ms": round(weighted_lt, 2),
|
||||
"band_delta_pct": round(100.0 * (weighted_lt - weighted_torch) / weighted_torch, 3),
|
||||
"cases": summary,
|
||||
}, sort_keys=True),
|
||||
flush=True,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -50,6 +50,11 @@ pinning.
|
||||
shapes, behind an opt-in config; trainer A/X/B afterward.
|
||||
- Kill criteria: heuristic sweep's best is within 2% of nvjet's current
|
||||
selection (would confirm nvjet is already near-optimal and close this line).
|
||||
- STATUS 2026-07-22: sweep ran (`probes/bench_ltx2_cublaslt_algo_sweep.py`);
|
||||
kill criterion NOT met — weighted band -7.98% at B2 and B3, -9.4% best-of,
|
||||
bit-exact parity on all winners. Integration (pinned-algo custom op + trainer
|
||||
A/X/B) is the highest-priority next action; expected +1.9 to +2.2 MFU points
|
||||
at 4x/B2. See the REPORT campaign section for the full table and caveats.
|
||||
|
||||
## 3. Whole-step wall-minus-union gap (days; expected +0.5-1.0 pp)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user