[perf]: bank first kernel-level MFU gain via pinned cuBLASLt backward algos
This commit is contained in:
SolitaryThinker
@@ -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
|
||||
|
||||
@@ -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()
|
||||
@@ -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 <torch/extension.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <cublasLt.h>
|
||||
#include <vector>
|
||||
#include <cstring>
|
||||
#include <stdexcept>
|
||||
|
||||
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<PinnedCase>& cases() {
|
||||
static std::vector<PinnedCase> 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<int>(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<uint8_t>(), sizeof(cublasLtMatmulAlgo_t));
|
||||
cases().push_back(entry);
|
||||
return static_cast<int64_t>(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<int64_t>(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");
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user