From 8bf6eff18f7e6d37e9aa280d755fd84ae0244012 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Wed, 22 Jul 2026 10:58:25 -0700 Subject: [PATCH] [perf]: bank first kernel-level MFU gain via pinned cuBLASLt backward algos --- scripts/train/ltx2_mfu/REPORT.md | 7 +- .../benchmark_fastvideo_train_lt_pinned.py | 295 ++++++++++++++++++ .../train/ltx2_mfu/probes/lt_pinned_ops.cpp | 102 ++++++ .../reports/bf16_kernel_research_plan.md | 11 +- 4 files changed, 411 insertions(+), 4 deletions(-) create mode 100644 scripts/train/ltx2_mfu/harness/benchmark_fastvideo_train_lt_pinned.py create mode 100644 scripts/train/ltx2_mfu/probes/lt_pinned_ops.cpp diff --git a/scripts/train/ltx2_mfu/REPORT.md b/scripts/train/ltx2_mfu/REPORT.md index 175143cf7..318222ad9 100644 --- a/scripts/train/ltx2_mfu/REPORT.md +++ b/scripts/train/ltx2_mfu/REPORT.md @@ -444,7 +444,9 @@ Equivalently, with today's kernels a perfect 75 ms systems endgame caps at **48. ### 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. +`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. + +A first end-to-end trainer gate already confirms the direction. `probes/lt_pinned_ops.cpp` (a pinned-algo `cublasLtMatmul` registry built via `cpp_extension`) plus `harness/benchmark_fastvideo_train_lt_pinned.py` (rank-0 tuning broadcast to all ranks; `ReplicatedLinear` forwards wrapped so only dgrad/wgrad dispatch to pinned algos, forward stays stock nvjet with fused bias) patched all 336 block projections and ran the standard 4x/B2 A/X/B on the 1,965 MHz tray `1632324`: controls 0.978211/0.979068 s (0.088% drift, about 29.52% allocation-local MFU), candidate **0.970947 s / -7.692 ms / -0.786%, +0.234 MFU points**, peak memory +0.13 GiB (the 128 MB Lt workspace). The scratch wrapper pays for an un-fused eager dbias, autograd-Function graph breaks in the compiled blocks, and contiguity copies, so it captures only part of the sweep's margin; the production capture path is an Inductor-level lowering (or quant-method hook) that keeps dbias fused and pins algos inside the compiled backward. Accepted as kernel-direction evidence; the recipe change waits for that integration plus a healthy-allocation gate. 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. @@ -709,6 +711,9 @@ f681ac1b60beb31aa3dd868e9005684810e05ddb26966bd78885141acd7864bb /mnt/pr1630_8x 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 +afb1e48236ae5930ce8c12d730b935b6ad05595da9690821a808c75809f7546f /mnt/pr1630_ltpin_gate2/control_a.log +29785fe9e0eae9aa19bc6bc937a3e2a469c78f68c64402a9e044caf477b6e820 /mnt/pr1630_ltpin_gate2/ltpin.log +aae60722634e8c5a9d1f678ffe17f4c01a45ff131987e4252e7326bfb04fc0df /mnt/pr1630_ltpin_gate2/control_b.log 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/harness/benchmark_fastvideo_train_lt_pinned.py b/scripts/train/ltx2_mfu/harness/benchmark_fastvideo_train_lt_pinned.py new file mode 100644 index 000000000..5a5283620 --- /dev/null +++ b/scripts/train/ltx2_mfu/harness/benchmark_fastvideo_train_lt_pinned.py @@ -0,0 +1,295 @@ +#!/usr/bin/env python3 +"""Scratch gate: pinned cuBLASLt dgrad/wgrad algos for the packed projections. + +Follows research-plan item 2 after the algo sweep found an 8-9% weighted GEMM +band win with bit-exact parity. This driver builds the ``lt_pinned_ops`` C++ +extension, tunes the eight conservative backward cases (dgrad/wgrad where the +sweep beat torch by >2%; forward stays stock nvjet with fused bias), wraps the +matching block-level ``nn.Linear`` modules in an autograd function whose +backward dispatches to the pinned algos, and runs the frozen packed benchmark +harness. Bias gradients fall back to an eager ``sum(0)`` (integration would +fuse BGRADB later). +""" + +from __future__ import annotations + +import ctypes +import os +import statistics + +import torch +import torch.nn.functional as F +from torch.utils.cpp_extension import load + +import benchmark_fastvideo_train_pack_d016 as benchmark +import fastvideo.train.trainer as trainer_module + +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 +WORKSPACE_BYTES = 128 * 1024 * 1024 +MAX_ALGOS = 48 +HIDDEN, FFN = 4096, 16384 +VIDEO_TOKENS, TEXT_TOKENS = 11 * 15 * 26, 1024 +LOCAL_BATCH = int(os.environ.get("FASTVIDEO_BENCH_LOCAL_BATCH_SIZE", "1")) +MIN_WIN_PCT = 2.0 + +_ext = load( + name="lt_pinned_ops", + sources=[os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "probes", "lt_pinned_ops.cpp")], + extra_ldflags=["-lcublasLt"], + with_cuda=True, + verbose=False, +) + + +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 _tune_case(lt, handle, workspace, name, m, n, k, transa, transb, a_shape, b_shape, d_shape, torch_fn): + device = torch.device("cuda") + a = torch.randn(a_shape, device=device, dtype=torch.bfloat16) + b = torch.randn(b_shape, device=device, dtype=torch.bfloat16) + d = torch.empty(d_shape, device=device, dtype=torch.bfloat16) + reference = torch_fn(a, b) + + def check(status, what): + if status != 0: + raise RuntimeError(f"{what} failed: {status}") + + 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, transa), (DESC_TRANSB, transb)): + v = ctypes.c_int32(value) + check(lt.cublasLtMatmulDescSetAttribute(op_desc, attr, ctypes.byref(v), 4), "descSet") + lda, ldb = a_shape[1], b_shape[1] + + def layout(rows, cols, ld): + h = ctypes.c_void_p() + check(lt.cublasLtMatrixLayoutCreate(ctypes.byref(h), CUDA_R_16BF, rows, cols, ctypes.c_int64(ld)), "layout") + return h + + layout_a = layout(lda if transa == OP_T else m, m if transa == OP_T else k, lda) + layout_b = layout(ldb if transb == OP_T else k, k if transb == OP_T else n, 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) + check(lt.cublasLtMatmulAlgoGetHeuristic(handle, op_desc, layout_a, layout_b, layout_d, layout_d, + pref, MAX_ALGOS, results, ctypes.byref(found)), "heuristic") + alpha, beta = ctypes.c_float(1.0), ctypes.c_float(0.0) + stream = ctypes.c_void_p(torch.cuda.current_stream().cuda_stream) + + def run_lt(index): + check(lt.cublasLtMatmul(handle, op_desc, ctypes.byref(alpha), + ctypes.c_void_p(a.data_ptr()), layout_a, + ctypes.c_void_p(b.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, + ctypes.byref(results[index], 0), + ctypes.c_void_p(workspace.data_ptr()), ctypes.c_size_t(WORKSPACE_BYTES), + stream), "ltMatmul") + + def time_fn(fn): + for _ in range(6): + fn() + torch.cuda.synchronize() + events = [(torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)) for _ in range(15)] + 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) + + torch_ms = time_fn(lambda: torch_fn(a, b)) + best = None + for index in range(found.value): + if results[index].state != 0: + continue + try: + run_lt(index) + except RuntimeError: + continue + torch.cuda.synchronize() + if (d.float() - reference.float()).abs().max().item() > 0.0: + continue + lt_ms = time_fn(lambda: run_lt(index)) + if best is None or lt_ms < best[0]: + best = (lt_ms, index) + win_pct = 100.0 * (torch_ms - best[0]) / torch_ms if best else -1.0 + spec = None + if best and win_pct >= MIN_WIN_PCT: + spec = { + "m": m, "n": n, "k": k, "transa": transa == OP_T, "transb": transb == OP_T, + "lda": lda, "ldb": ldb, "d_rows": d_shape[0], "d_cols": d_shape[1], + "algo": bytes(results[best[1]].algo), + } + print(f"LT_PIN_TUNE {name} torch_ms={torch_ms:.4f} best_lt_ms={best[0] if best else -1:.4f} " + f"win_pct={win_pct:.2f} pinned={spec is not None}", flush=True) + return spec + + +def _tune_all_blobs() -> dict[tuple[int, int], dict]: + """Rank-0 tuning: returns per-shape dgrad/wgrad win metadata + algo blobs.""" + lt = None + for so in ("libcublasLt.so.13", "libcublasLt.so.12", "libcublasLt.so"): + try: + lt = ctypes.CDLL(so) + break + except OSError: + continue + if lt is None: + raise OSError("libcublasLt not found") + handle = ctypes.c_void_p() + if lt.cublasLtCreate(ctypes.byref(handle)) != 0: + raise RuntimeError("cublasLtCreate failed") + workspace = torch.empty(WORKSPACE_BYTES, dtype=torch.uint8, device="cuda") + mapping: dict[tuple[int, int], dict] = {} + shapes = [ + ("self_qkv", VIDEO_TOKENS * LOCAL_BATCH, HIDDEN, 3 * HIDDEN), + ("video_dd", VIDEO_TOKENS * LOCAL_BATCH, HIDDEN, HIDDEN), + ("text_kv", TEXT_TOKENS * LOCAL_BATCH, HIDDEN, 2 * HIDDEN), + ("ffn_up", VIDEO_TOKENS * LOCAL_BATCH, HIDDEN, FFN), + ("ffn_down", VIDEO_TOKENS * LOCAL_BATCH, FFN, HIDDEN), + ] + for name, rows, in_features, out_features in shapes: + dgrad_spec = _tune_case( + lt, handle, workspace, f"{name}:dgrad", + in_features, rows, out_features, OP_N, OP_N, + (out_features, in_features), (rows, out_features), (rows, in_features), + lambda a, b: b @ a) + wgrad_spec = _tune_case( + lt, handle, workspace, f"{name}:wgrad", + in_features, out_features, rows, OP_N, OP_T, + (rows, in_features), (rows, out_features), (out_features, in_features), + lambda a, b: b.t() @ a) + mapping[(in_features, out_features)] = {"rows": rows, "dgrad": dgrad_spec, "wgrad": wgrad_spec} + return mapping + + +def _register_from_spec(spec: dict | None) -> int: + if spec is None: + return -1 + blob = torch.frombuffer(bytearray(spec["algo"]), dtype=torch.uint8).clone() + return _ext.register_case(spec["m"], spec["n"], spec["k"], spec["transa"], spec["transb"], + spec["lda"], spec["ldb"], spec["d_rows"], spec["d_cols"], blob) + + +def _tune_and_broadcast() -> dict[tuple[int, int], tuple[int, int, int]]: + import torch.distributed as dist + payload = [None] + if not dist.is_initialized() or dist.get_rank() == 0: + payload = [_tune_all_blobs()] + if dist.is_initialized(): + dist.broadcast_object_list(payload, src=0) + mapping: dict[tuple[int, int], tuple[int, int, int]] = {} + for key, entry in payload[0].items(): + mapping[key] = (entry["rows"], _register_from_spec(entry["dgrad"]), _register_from_spec(entry["wgrad"])) + return mapping + + +@torch.library.custom_op("fastvideo::lt_pinned_dgrad", mutates_args=()) +def lt_pinned_dgrad(grad_out: torch.Tensor, weight: torch.Tensor, case_id: int) -> torch.Tensor: + return _ext.lt_mm(weight.contiguous(), grad_out.contiguous(), case_id) + + +@lt_pinned_dgrad.register_fake +def _(grad_out, weight, case_id): + return grad_out.new_empty((grad_out.shape[0], weight.shape[1])) + + +@torch.library.custom_op("fastvideo::lt_pinned_wgrad", mutates_args=()) +def lt_pinned_wgrad(grad_out: torch.Tensor, saved_input: torch.Tensor, case_id: int) -> torch.Tensor: + return _ext.lt_mm(saved_input.contiguous(), grad_out.contiguous(), case_id) + + +@lt_pinned_wgrad.register_fake +def _(grad_out, saved_input, case_id): + return grad_out.new_empty((grad_out.shape[1], saved_input.shape[1])) + + +class _LtPinnedLinear(torch.autograd.Function): + + @staticmethod + def forward(ctx, x2d, weight, bias, rows, dgrad_id, wgrad_id): + ctx.save_for_backward(x2d, weight) + ctx.meta = (rows, dgrad_id, wgrad_id, bias is not None) + return F.linear(x2d, weight, bias) + + @staticmethod + def backward(ctx, grad_out): + x2d, weight = ctx.saved_tensors + rows, dgrad_id, wgrad_id, has_bias = ctx.meta + grad_out = grad_out.contiguous() + if grad_out.shape[0] == rows and dgrad_id >= 0: + grad_input = lt_pinned_dgrad(grad_out, weight, dgrad_id) + else: + grad_input = grad_out @ weight + if grad_out.shape[0] == rows and wgrad_id >= 0: + grad_weight = lt_pinned_wgrad(grad_out, x2d, wgrad_id) + else: + grad_weight = grad_out.t() @ x2d + grad_bias = grad_out.sum(0) if has_bias else None + return grad_input, grad_weight, grad_bias, None, None, None + + +def _patched_forward(module, rows, dgrad_id, wgrad_id): + + def forward(x): + bias = module.bias if not module.skip_bias_add else None + shape = x.shape + x2d = x.reshape(-1, shape[-1]) + out = _LtPinnedLinear.apply(x2d, module.weight, bias, rows, dgrad_id, wgrad_id) + out = out.reshape(*shape[:-1], out.shape[-1]) + output_bias = module.bias if module.skip_bias_add else None + return out, output_bias + + return forward + + +def _install(method) -> None: + from fastvideo.layers.linear import ReplicatedLinear + + mapping = _tune_and_broadcast() + patched = 0 + for name, module in method.student.transformer.named_modules(): + if not isinstance(module, ReplicatedLinear) or "transformer_blocks" not in name: + continue + key = (module.input_size, module.output_size) + if key not in mapping: + continue + rows, dgrad_id, wgrad_id = mapping[key] + if dgrad_id < 0 and wgrad_id < 0: + continue + module.forward = _patched_forward(module, rows, dgrad_id, wgrad_id) + patched += 1 + print(f"LT_PIN_INSTALLED modules={patched}", flush=True) + + +_original_run = trainer_module.Trainer.run + + +def _run_with_pins(self, method, **kwargs): + _install(method) + return _original_run(self, method, **kwargs) + + +trainer_module.Trainer.run = _run_with_pins + +if __name__ == "__main__": + benchmark.main() diff --git a/scripts/train/ltx2_mfu/probes/lt_pinned_ops.cpp b/scripts/train/ltx2_mfu/probes/lt_pinned_ops.cpp new file mode 100644 index 000000000..bd48b161d --- /dev/null +++ b/scripts/train/ltx2_mfu/probes/lt_pinned_ops.cpp @@ -0,0 +1,102 @@ +// Minimal cuBLASLt pinned-algo GEMM registry for the LTX-2 packed shapes. +// Python owns tuning (probes/bench_ltx2_cublaslt_algo_sweep.py logic) and +// registers one case per (m, n, k, transa, transb) with the winning +// 64-byte cublasLtMatmulAlgo_t blob; lt_mm dispatches by case id with zero +// Python in the hot path. BF16 in/out, FP32 compute, no epilogue. + +#include +#include +#include +#include +#include +#include + +namespace { + +struct PinnedCase { + cublasLtMatmulDesc_t op_desc{}; + cublasLtMatrixLayout_t layout_a{}; + cublasLtMatrixLayout_t layout_b{}; + cublasLtMatrixLayout_t layout_d{}; + cublasLtMatmulAlgo_t algo{}; + int64_t m{}, n{}, k{}; + int64_t d_rows{}, d_cols{}; +}; + +cublasLtHandle_t lt_handle() { + static cublasLtHandle_t handle = [] { + cublasLtHandle_t h; + TORCH_CHECK(cublasLtCreate(&h) == CUBLAS_STATUS_SUCCESS, "cublasLtCreate failed"); + return h; + }(); + return handle; +} + +std::vector& cases() { + static std::vector registry; + return registry; +} + +torch::Tensor& workspace() { + static torch::Tensor ws = torch::empty( + {128 * 1024 * 1024}, + torch::TensorOptions().dtype(torch::kUInt8).device(torch::kCUDA)); + return ws; +} + +void check_lt(cublasStatus_t status, const char* what) { + TORCH_CHECK(status == CUBLAS_STATUS_SUCCESS, what, " failed with status ", static_cast(status)); +} + +int64_t register_case(int64_t m, int64_t n, int64_t k, bool transa, bool transb, + int64_t lda, int64_t ldb, int64_t d_rows, int64_t d_cols, + torch::Tensor algo_bytes) { + TORCH_CHECK(algo_bytes.dtype() == torch::kUInt8 && algo_bytes.numel() == sizeof(cublasLtMatmulAlgo_t), + "algo blob must be 64 uint8 bytes"); + PinnedCase entry; + entry.m = m; entry.n = n; entry.k = k; + entry.d_rows = d_rows; entry.d_cols = d_cols; + check_lt(cublasLtMatmulDescCreate(&entry.op_desc, CUBLAS_COMPUTE_32F, CUDA_R_32F), "descCreate"); + const cublasOperation_t opa = transa ? CUBLAS_OP_T : CUBLAS_OP_N; + const cublasOperation_t opb = transb ? CUBLAS_OP_T : CUBLAS_OP_N; + check_lt(cublasLtMatmulDescSetAttribute(entry.op_desc, CUBLASLT_MATMUL_DESC_TRANSA, &opa, sizeof(opa)), "setA"); + check_lt(cublasLtMatmulDescSetAttribute(entry.op_desc, CUBLASLT_MATMUL_DESC_TRANSB, &opb, sizeof(opb)), "setB"); + const int64_t a_rows = transa ? lda : m; + const int64_t a_cols = transa ? m : k; + const int64_t b_rows = transb ? ldb : k; + const int64_t b_cols = transb ? k : n; + check_lt(cublasLtMatrixLayoutCreate(&entry.layout_a, CUDA_R_16BF, a_rows, a_cols, lda), "layoutA"); + check_lt(cublasLtMatrixLayoutCreate(&entry.layout_b, CUDA_R_16BF, b_rows, b_cols, ldb), "layoutB"); + check_lt(cublasLtMatrixLayoutCreate(&entry.layout_d, CUDA_R_16BF, m, n, m), "layoutD"); + std::memcpy(&entry.algo, algo_bytes.data_ptr(), sizeof(cublasLtMatmulAlgo_t)); + cases().push_back(entry); + return static_cast(cases().size()) - 1; +} + +torch::Tensor lt_mm(torch::Tensor a, torch::Tensor b, int64_t case_id) { + TORCH_CHECK(case_id >= 0 && case_id < static_cast(cases().size()), "unknown lt case"); + TORCH_CHECK(a.is_cuda() && b.is_cuda() && a.dtype() == torch::kBFloat16 && b.dtype() == torch::kBFloat16, + "lt_mm expects CUDA bf16 tensors"); + TORCH_CHECK(a.is_contiguous() && b.is_contiguous(), "lt_mm expects contiguous operands"); + const PinnedCase& entry = cases()[case_id]; + auto d = torch::empty({entry.d_rows, entry.d_cols}, a.options()); + const float alpha = 1.0f; + const float beta = 0.0f; + check_lt(cublasLtMatmul(lt_handle(), entry.op_desc, &alpha, + a.data_ptr(), entry.layout_a, + b.data_ptr(), entry.layout_b, + &beta, + d.data_ptr(), entry.layout_d, + d.data_ptr(), entry.layout_d, + &entry.algo, + workspace().data_ptr(), workspace().numel(), + at::cuda::getCurrentCUDAStream()), "ltMatmul"); + return d; +} + +} // namespace + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { + module.def("register_case", ®ister_case, "register pinned cuBLASLt case"); + module.def("lt_mm", <_mm, "run pinned cuBLASLt matmul"); +} 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 b91afb2e1..46737943b 100644 --- a/scripts/train/ltx2_mfu/reports/bf16_kernel_research_plan.md +++ b/scripts/train/ltx2_mfu/reports/bf16_kernel_research_plan.md @@ -52,9 +52,14 @@ pinning. 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. + bit-exact parity on all winners. A scratch end-to-end gate + (`probes/lt_pinned_ops.cpp` + `harness/benchmark_fastvideo_train_lt_pinned.py`, + 336 projections, dgrad/wgrad only) already banks **-0.786% step / +0.234 MFU + points** at 4x/B2 through the wrapper's own overheads (eager dbias, graph + breaks, contiguity copies). Remaining upside sits in a fused integration: + an Inductor-level mm lowering or ReplicatedLinear quant-method hook that + pins algos inside compiled backward and keeps dbias fused; expected total + +1.9 to +2.2 MFU points at 4x/B2. Then re-gate on a healthy allocation. ## 3. Whole-step wall-minus-union gap (days; expected +0.5-1.0 pp)