diff --git a/scripts/train/ltx2_mfu/REPORT.md b/scripts/train/ltx2_mfu/REPORT.md index fe2a69112..175143cf7 100644 --- a/scripts/train/ltx2_mfu/REPORT.md +++ b/scripts/train/ltx2_mfu/REPORT.md @@ -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 diff --git a/scripts/train/ltx2_mfu/probes/bench_ltx2_cublaslt_algo_sweep.py b/scripts/train/ltx2_mfu/probes/bench_ltx2_cublaslt_algo_sweep.py new file mode 100644 index 000000000..c93ea14fb --- /dev/null +++ b/scripts/train/ltx2_mfu/probes/bench_ltx2_cublaslt_algo_sweep.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() diff --git a/scripts/train/ltx2_mfu/reports/bf16_kernel_research_plan.md b/scripts/train/ltx2_mfu/reports/bf16_kernel_research_plan.md index 464b59f41..b91afb2e1 100644 --- a/scripts/train/ltx2_mfu/reports/bf16_kernel_research_plan.md +++ b/scripts/train/ltx2_mfu/reports/bf16_kernel_research_plan.md @@ -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)