diff --git a/CMakeLists.txt b/CMakeLists.txt index 9aa6e393d1..b903f1918f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -782,6 +782,16 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" AND SM70_TURBOMIND_ARCHS) ARCHITECTURES "70" LIBRARIES "${TORCH_INSTALL_PREFIX}/lib/libtorch_python.so" WITH_SOABI) + define_extension_target( + _sm70_sparse_attention_C DESTINATION vllm LANGUAGE CUDA + SOURCES "flash-attention-v100/kernel/h3/forward_sparse.cu" + INCLUDE_DIRECTORIES + "${CUTLASS_INCLUDE_DIR}" + "${cutlass_SOURCE_DIR}/examples/41_fused_multi_head_attention" + COMPILE_FLAGS ${VLLM_GPU_FLAGS} + ARCHITECTURES "70" + LIBRARIES "${TORCH_INSTALL_PREFIX}/lib/libtorch_python.so" + WITH_SOABI) endif() if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") diff --git a/benchmarks/kernels/benchmark_h3_vsa_fp32.py b/benchmarks/kernels/benchmark_h3_vsa_fp32.py new file mode 100644 index 0000000000..92c2e5b076 --- /dev/null +++ b/benchmarks/kernels/benchmark_h3_vsa_fp32.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Acceptance-only sparse FP32 CUDA probe; never registers a runtime backend. + +Run on an externally leased SM70 GPU. An optional captured attention input is +compared with independent gathered FP32 QK/global softmax/PV. Operator timings +are diagnostic and cannot satisfy the complete-denoise performance gate. +""" + +import argparse +import hashlib +import importlib.util +import json +import statistics +from pathlib import Path + +import torch + + +def load_binary(path: Path): + path = path.resolve(strict=True) + spec = importlib.util.spec_from_file_location(path.stem, path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def fp32_reference(q, k, v, block_map, sizes, scale): + """Independent mathematical oracle, with physical padding masked explicitly. + + Match the frozen oracle's grouping (at most eight queries and 256 MiB per + gathered operand). No production sparse kernel or geometry helper is used. + This is an engineering FP32 reference, not the official GPU implementation. + """ + if torch.is_autocast_enabled("cuda"): + raise ValueError("FP32 reference requires CUDA autocast to be disabled") + batch, rows, heads, dim = q.shape + blocks = rows // 64 + flat = [x.permute(0, 2, 1, 3).reshape(batch * heads, rows, dim) for x in (q, k, v)] + maps = block_map.reshape(batch * heads, blocks, blocks) + order = torch.arange(blocks, device=q.device).view(1, 1, blocks).expand_as(maps) + selected = order.masked_fill(~maps, blocks).sort(dim=-1).values + counts = maps.sum(-1).amax(0).tolist() + lanes = torch.arange(64, device=q.device) + bh = torch.arange(batch * heads, device=q.device)[:, None, None] + output = torch.zeros_like(flat[0]) + pos = 0 + while pos < blocks: + keep = int(counts[pos]) + if keep == 0: + raise ValueError("reference requires nonempty selected rows") + chunk = max(1, min(8, 256 * 1024**2 // (batch * heads * keep * 64 * dim * 4))) + end = pos + 1 + while end < min(blocks, pos + chunk) and counts[end] == keep: + end += 1 + index = selected[:, pos:end, :keep] + safe = index.clamp_max(blocks - 1) + tokens = (safe[..., None] * 64 + lanes).flatten(-2) + valid = ( + (index < blocks)[..., None] & (lanes < sizes[safe][..., None]) + ).flatten(-2) + keys, values = [flat[i][bh, tokens].float() for i in (1, 2)] + keys.masked_fill_(~valid[..., None], 0) + values.masked_fill_(~valid[..., None], 0) + queries = flat[0][:, pos * 64 : end * 64].reshape(-1, 64, dim).float() + shape = (batch * heads * (end - pos), keep * 64, dim) + scores = torch.bmm(queries, keys.reshape(shape).transpose(1, 2)) * scale + scores.masked_fill_(~valid.reshape(-1, 1, keep * 64), -float("inf")) + answer = torch.bmm(scores.softmax(-1), values.reshape(shape)) + output[:, pos * 64 : end * 64] = answer.reshape(batch * heads, -1, dim).to( + q.dtype + ) + pos = end + return output.reshape(batch, heads, rows, dim).permute(0, 2, 1, 3).contiguous() + + +def prepare_capture(path): + from vllm.model_executor.models.minimax_h3 import vsa + + data = torch.load(path, weights_only=True, map_location="cpu", mmap=True) + q, k, v = [data[name].cuda() for name in ("q", "k", "v")] + part, sizes, nonpad, _, prefix, video = vsa._get_h3_tile_metadata( + tuple(data["prefix_segments"]), tuple(data["video_shape"]), q.device + ) + tiled = [] + for tensor in (q, k, v): + out = torch.zeros( + q.size(0), len(sizes) * 64, q.size(2), 128, device=q.device, dtype=q.dtype + ) + out[:, nonpad] = tensor[:, part] + tiled.append(out) + q, k, v = tiled + scores = ( + torch.matmul( + vsa._pool_h3_tiles(q, sizes), vsa._pool_h3_tiles(k, sizes).transpose(-2, -1) + ) + * data["scale"] + ) + mask = vsa._build_h3_block_map(scores, prefix, video, data["topk"]) + return (q, k, v, mask, sizes, data["scale"], prefix, data["topk"]), nonpad + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--binary", type=Path) + parser.add_argument("--build-directory", type=Path) + parser.add_argument("--cutlass-root", type=Path) + parser.add_argument("--capture", type=Path) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + if args.binary: + if args.build_directory or args.cutlass_root: + parser.error("choose an existing binary or a build directory/CUTLASS root") + binary = args.binary.resolve(strict=True) + ops = load_binary(binary) + else: + if not args.build_directory or not args.cutlass_root: + parser.error("building requires --build-directory and --cutlass-root") + from torch.utils.cpp_extension import load + + args.build_directory.mkdir(parents=True, exist_ok=True) + ops = load( + name="h3_vsa_cutlass_fp32", + sources=[str(Path(__file__).with_name("h3_vsa_fp32.cu"))], + extra_include_paths=[str(args.cutlass_root / "include")], + extra_cuda_cflags=["-O3", "--ptxas-options=-v"], + build_directory=str(args.build_directory), + verbose=True, + ) + binary = Path(ops.__file__) + record = { + "binary": str(binary), + "binary_sha256": hashlib.sha256(binary.read_bytes()).hexdigest(), + "source_sha256": hashlib.sha256( + Path(__file__).with_name("h3_vsa_fp32.cu").read_bytes() + ).hexdigest(), + "built_from_reported_source": not bool(args.binary), + "torch_version": torch.__version__, + "cuda_version": torch.version.cuda, + "performance_eligible": False, + "scope": "operator diagnostic only", + } + if args.capture: + with torch.inference_mode(): + values, valid = prepare_capture(args.capture) + q, k, v, mask, sizes, scale, _, _ = values + actual = ops.forward(*values) + expected = fp32_reference(q, k, v, mask, sizes, scale) + record["bitwise"] = torch.equal( + actual[:, valid].view(torch.int16), expected[:, valid].view(torch.int16) + ) + if record["bitwise"]: + for _ in range(2): + ops._forward_prevalidated(*values) + times = [] + for _ in range(4): + start, end = [ + torch.cuda.Event(enable_timing=True) for _ in range(2) + ] + start.record() + ops._forward_prevalidated(*values) + end.record() + end.synchronize() + times.append(start.elapsed_time(end)) + record.update(times_ms=times, median_ms=statistics.median(times)) + args.output.write_text(json.dumps(record, indent=2)) + print(json.dumps(record)) + if record.get("bitwise") is False: + raise SystemExit("FP32 comparison failed; candidate is not admitted") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_sm70_exact_row_reduce.py b/benchmarks/kernels/benchmark_sm70_exact_row_reduce.py new file mode 100644 index 0000000000..19736734c4 --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_exact_row_reduce.py @@ -0,0 +1,228 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""TP4 explicit row-plan correctness controls; launch under an owned GPU lease. + +Use torchrun --standalone --nproc_per_node=4 with this script. Measurements +are isolated communication diagnostics, never full-model acceptance. +""" + +import argparse +import hashlib +import importlib.util +import json +import os +import statistics +import sys +import time +from pathlib import Path + +import torch + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--extension", type=Path) + parser.add_argument("--full-shape", action="store_true") + args = parser.parse_args() + from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config + from vllm.distributed import ( + cleanup_dist_env_and_memory, + get_tp_group, + init_distributed_environment, + initialize_model_parallel, + ) + from vllm.model_executor.layers import sm70_collectives as shared + from vllm.video.benchmark import source_provenance + + rank, local = int(os.environ["RANK"]), int(os.environ["LOCAL_RANK"]) + if int(os.environ["WORLD_SIZE"]) != 4: + raise ValueError("This control requires exactly four TP ranks") + torch.cuda.set_device(local) + torch.set_num_threads(4) + torch.manual_seed(5091 + rank) + record = dict( + rank=rank, + state="running", + cases=[], + guards=[], + source=source_provenance(), + scope="isolated operator control; no model acceptance", + ) + if args.extension: + spec = importlib.util.spec_from_file_location( + args.extension.stem, args.extension + ) + extension = importlib.util.module_from_spec(spec) + spec.loader.exec_module(extension) + shared._extension = lambda: extension + sys.modules["onecat_sm70_exact_reduce"] = extension + record["extension"] = { + "path": str(args.extension), + "sha256": hashlib.sha256(args.extension.read_bytes()).hexdigest(), + } + source = ( + Path(shared.__file__).resolve().parents[3] + / "csrc/sm70_turbomind/ops/exact_row_reduce.cu" + ) + record["cuda_source_sha256"] = hashlib.sha256(source.read_bytes()).hexdigest() + record["benchmark_source_sha256"] = hashlib.sha256( + Path(__file__).read_bytes() + ).hexdigest() + with set_current_vllm_config( + VllmConfig(parallel_config=ParallelConfig(tensor_parallel_size=4)) + ): + init_distributed_environment(4, rank, "env://", local, "nccl") + initialize_model_parallel(4) + try: + group = get_tp_group() + for label, shape, budget in ( + ("one-rank-invalid-shape", (3, 3) if rank == 0 else (4, 3), 2**30), + ("one-rank-small-budget", (4, 3), 1 if rank == 0 else 2**30), + ("different-valid-shapes", (4, 3) if rank == 0 else (8, 3), 2**30), + ): + try: + shared.SM70ExactRowReductionPlan( + group, shape, memory_budget_bytes=budget + ) + except ValueError: + record["guards"].append(label) + else: + raise AssertionError(f"Expected collective rejection: {label}") + shapes = [ + (4, 3), + (12, 65), + (68, 257), + (128, 768), + (260, 1024), + (1028, 3072), + ] + if args.full_shape: + shapes.append((34560, 5376)) + for shape in shapes: + started = time.perf_counter() + plan = shared.SM70ExactRowReductionPlan( + group, shape, memory_budget_bytes=4 * 2**30 + ) + calibration_seconds = time.perf_counter() - started + try: + with torch.inference_mode(): + for label, scale in ( + ("ordinary", 1.0), + ("wide", 1e30), + ("subnormal", 1e-40), + ): + storage = ( + torch.randn(shape[0] * shape[1] + 1, device="cuda") + * scale + ) + value = storage[1:].view(shape) + ref = group.all_reduce(value).chunk(4)[rank] + actual = plan.reduce(value) + mismatch = int( + torch.count_nonzero( + ref.view(torch.int32) != actual.view(torch.int32) + ) + ) + record["cases"].append( + dict( + shape=shape, + input=label, + storage_offset=1, + mismatch=mismatch, + calibration_seconds=calibration_seconds, + raw_ipc_bytes=plan.raw_ipc_bytes, + ) + ) + value = torch.full( + shape, + float("inf") if rank < 2 else -float("inf"), + device="cuda", + ) + ref = group.all_reduce(value).chunk(4)[rank] + actual = plan.reduce(value) + record["cases"].append( + dict( + shape=shape, + input="opposing-inf", + mismatch=int( + torch.count_nonzero( + ref.view(torch.int32) + != actual.view(torch.int32) + ) + ), + ) + ) + with torch.cuda.stream(torch.cuda.Stream()): + try: + plan.reduce(value) + except RuntimeError: + record["guards"].append("different-stream") + else: + raise AssertionError("Different stream was accepted") + vote = torch.tensor( + int(all(x["mismatch"] == 0 for x in record["cases"])), + device="cuda", + ) + torch.distributed.all_reduce( + vote, op=torch.distributed.ReduceOp.MIN + ) + if not vote.item(): + raise AssertionError("Native FP32 bits changed") + if shape == (34560, 5376): + value.normal_() + for _ in range(3): + group.all_reduce(value) + plan.reduce(value) + torch.cuda.synchronize() + times = {"native": [], "peer_rows": []} + for repeat in range(7): + for name in ( + ("native", "peer_rows") + if repeat % 2 == 0 + else ("peer_rows", "native") + ): + torch.distributed.barrier(group=group.cpu_group) + torch.cuda.synchronize() + start, end = ( + torch.cuda.Event(enable_timing=True), + torch.cuda.Event(enable_timing=True), + ) + start.record() + output = ( + group.all_reduce(value) + if name == "native" + else plan.reduce(value) + ) + end.record() + end.synchronize() + times[name].append(start.elapsed_time(end)) + del output + record["times_ms"] = times + record["median_ms"] = { + key: statistics.median(values) + for key, values in times.items() + } + finally: + plan.close() + try: + plan.reduce(value) + except RuntimeError: + record["guards"].append("closed-plan") + else: + raise AssertionError("Closed plan was accepted") + print( + json.dumps(dict(rank=rank, shape=shape, state="passed")), flush=True + ) + record["state"] = "passed_operator_control" + except BaseException as error: + record.update(state="failed", error=repr(error)) + raise + finally: + args.output.mkdir(parents=True, exist_ok=True) + (args.output / f"rank-{rank}.json").write_text(json.dumps(record, indent=2)) + cleanup_dist_env_and_memory() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/h3_vsa_fp32.cu b/benchmarks/kernels/h3_vsa_fp32.cu new file mode 100644 index 0000000000..3fa2d1e700 --- /dev/null +++ b/benchmarks/kernels/h3_vsa_fp32.cu @@ -0,0 +1,314 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// Acceptance-only H3 sparse FP32 arithmetic; no runtime backend registration. +// Requires CUTLASS 4.4.2 and SM70. Preserve sequential FMA and global softmax. +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// Scores and probabilities each have this bound. This is an opt-in diagnostic +// budget, not a runtime default. Directly indexed K/V require no per-query +// copy. +constexpr int kMaxQueryBatch = 32; +constexpr int64_t kScoreBudgetBytes = 2LL * 1024 * 1024 * 1024; + +struct Geometry { + const __half *q, *k, *v; + __half* out; + const int *indices, *sizes; + int rows, heads, blocks, prefix, video_keep, group_indices; + int first, count, keep; + float scale; +}; +__device__ int64_t map_offset(const Geometry& g, int group, int query) { + return int64_t(group) * g.group_indices + + (query < g.prefix + ? query * g.blocks + : g.prefix * g.blocks + (query - g.prefix) * g.video_keep); +} +__device__ int64_t input_offset(const Geometry& g, int group, int row, + int channel) { + return (int64_t(group / g.heads) * g.rows + row) * g.heads * 128 + + (group % g.heads) * 128 + channel; +} +__global__ void compact(const bool* mask, int* indices, int blocks, int prefix, + int video_keep, int group_indices) { + int row = blockIdx.x, lane = threadIdx.x, query = row % blocks, + group = row / blocks; + int capacity = query < prefix ? blocks : video_keep; + int64_t offset = + int64_t(group) * group_indices + + (query < prefix ? query * blocks + : prefix * blocks + (query - prefix) * video_keep); + int count = 0; + for (int start = 0; start < blocks; start += 32) { + int key = start + lane; + bool selected = key < blocks && mask[int64_t(row) * blocks + key]; + unsigned ballot = __ballot_sync(0xffffffff, selected); + int pos = count + __popc(ballot & ((1u << lane) - 1)); + if (selected && pos < capacity) indices[offset + pos] = key; + count += __popc(ballot); + } +} + +template +using BaseKernel = typename cutlass::gemm::kernel::DefaultGemmUniversal< + float, cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 1, + float, + typename std::conditional::type, + cutlass::ComplexTransform::kNone, 1, float, cutlass::layout::RowMajor, + float, cutlass::arch::OpClassSimt, cutlass::arch::Sm70, + cutlass::gemm::GemmShape<64, 128, 8>, cutlass::gemm::GemmShape<32, 64, 8>, + cutlass::gemm::GemmShape<1, 1, 1>, + cutlass::epilogue::thread::LinearCombination, + cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, 2, + cutlass::arch::OpMultiplyAdd, cutlass::gemm::SharedMemoryClearOption::kNone, + false, true, false>::GemmKernel; + +template +struct H3IndexedCutlassKernel : BaseKernel { + using Parent = BaseKernel; + using Params = typename Parent::Params; + using SharedStorage = typename Parent::SharedStorage; + CUTLASS_DEVICE void operator()(Params const& input, SharedStorage& storage) { + Params p = input; + p.ptr_gather_B_indices += + int64_t(blockIdx.z) * (PV ? p.problem_size.k() : p.problem_size.n()); + Parent::operator()(p, storage); + } + CUTLASS_DEVICE static void invoke(Params const& p, SharedStorage& storage) { + H3IndexedCutlassKernel op; + op(p, storage); + } +}; +using QK = + cutlass::gemm::device::GemmUniversalBase>; +using PV = + cutlass::gemm::device::GemmUniversalBase>; + +__device__ void convert_vector(uint4 value, float4* dest) { + const __half2* h = reinterpret_cast(&value); + float2 a = __half22float2(h[0]), b = __half22float2(h[1]); + float2 c = __half22float2(h[2]), d = __half22float2(h[3]); + dest[0] = make_float4(a.x, a.y, b.x, b.y); + dest[1] = make_float4(c.x, c.y, d.x, d.y); +} +__global__ void convert_qkv(Geometry g, float* converted, int64_t elements) { + int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= elements / 8) return; + int row = (i * 8 / (g.heads * 128)) % g.rows; + bool valid = row % 64 < g.sizes[row / 64]; + uint4 zero = make_uint4(0, 0, 0, 0); + convert_vector(valid ? reinterpret_cast(g.q)[i] : zero, + reinterpret_cast(converted + i * 8)); + convert_vector(valid ? reinterpret_cast(g.k)[i] : zero, + reinterpret_cast(converted + elements + i * 8)); + convert_vector(valid ? reinterpret_cast(g.v)[i] : zero, + reinterpret_cast(converted + 2 * elements + i * 8)); +} +__global__ void prepare_batch(Geometry g, const float* data, int64_t elements, + float* scores, float* probability, float* result, + int64_t* pointers, int* tokens) { + int batch = blockIdx.x, group = batch / g.count, + query = g.first + batch % g.count; + int batches = gridDim.x, n = g.keep * 64; + int64_t offset = map_offset(g, group, query); + for (int i = threadIdx.x; i < n; i += blockDim.x) + tokens[int64_t(batch) * n + i] = g.indices[offset + i / 64] * 64 + i % 64; + if (threadIdx.x == 0) { + pointers[batch] = + reinterpret_cast(data + input_offset(g, group, query * 64, 0)); + pointers[batches + batch] = reinterpret_cast( + data + elements + input_offset(g, group, 0, 0)); + pointers[2 * batches + batch] = + reinterpret_cast(scores + int64_t(batch) * 64 * n); + pointers[3 * batches + batch] = + reinterpret_cast(probability + int64_t(batch) * 64 * n); + pointers[4 * batches + batch] = reinterpret_cast( + data + 2 * elements + input_offset(g, group, 0, 0)); + pointers[5 * batches + batch] = + reinterpret_cast(result + int64_t(batch) * 64 * 128); + } +} +__global__ void mask_scores(Geometry g, const int* tokens, float* scores) { + int batch = blockIdx.y, key = blockIdx.x, col = threadIdx.x; + int n = g.keep * 64, selected = tokens[int64_t(batch) * n + key * 64] / 64; + if (col < g.sizes[selected]) return; + for (int q = 0; q < 64; ++q) + scores[(int64_t(batch) * 64 + q) * n + key * 64 + col] = -CUDART_INF_F; +} +__global__ void scatter_result(Geometry g, const float* result, int batches) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= batches * 64 * 16) return; + int batch = i / (64 * 16), row = i / 16 % 64, d = i % 16 * 8; + int group = batch / g.count, query = g.first + batch % g.count; + if (row >= g.sizes[query]) return; + float4 a = reinterpret_cast(result + i * 8)[0]; + float4 b = reinterpret_cast(result + i * 8)[1]; + uint4 value; + __half2* h = reinterpret_cast<__half2*>(&value); + h[0] = __floats2half2_rn(a.x, a.y); + h[1] = __floats2half2_rn(a.z, a.w); + h[2] = __floats2half2_rn(b.x, b.y); + h[3] = __floats2half2_rn(b.z, b.w); + *reinterpret_cast( + g.out + input_offset(g, group, query * 64 + row, d)) = value; +} + +torch::Tensor forward_impl(torch::Tensor q, torch::Tensor k, torch::Tensor v, + torch::Tensor mask, torch::Tensor sizes, + double scale, int64_t prefix, int64_t topk, + bool checked) { + TORCH_CHECK(q.is_cuda() && q.dim() == 4 && q.scalar_type() == at::kHalf && + q.is_contiguous() && q.size(0) > 0 && q.size(1) > 0 && + q.size(1) % 64 == 0 && q.size(2) > 0 && q.size(3) == 128, + "indexed FP32 H3 requires contiguous FP16 BSHD128 tiles"); + for (auto& x : {k, v}) + TORCH_CHECK(x.device() == q.device() && x.sizes() == q.sizes() && + x.scalar_type() == at::kHalf && x.is_contiguous() && + !x.requires_grad(), + "matching inference QKV required"); + TORCH_CHECK(!q.requires_grad(), "inference only"); + int64_t blocks = q.size(1) / 64, groups = q.size(0) * q.size(2); + TORCH_CHECK(q.size(1) <= INT_MAX && groups * blocks <= INT_MAX && + groups * kMaxQueryBatch <= 65535, + "indexed H3 exceeds index limits"); + TORCH_CHECK(prefix >= 0 && prefix < blocks && topk > 0, + "invalid H3 prefix/topk"); + TORCH_CHECK(std::isfinite(scale) && scale > 0, + "positive finite scale required"); + TORCH_CHECK(mask.device() == q.device() && mask.dim() == 4 && + mask.scalar_type() == at::kBool && mask.is_contiguous() && + mask.size(0) == q.size(0) && mask.size(1) == q.size(2) && + mask.size(2) == blocks && mask.size(3) == blocks, + "H3 map shape mismatch"); + TORCH_CHECK(sizes.device() == q.device() && sizes.dim() == 1 && + sizes.numel() == blocks && sizes.is_contiguous() && + sizes.scalar_type() == at::kInt, + "H3 sizes shape mismatch"); + c10::cuda::CUDAGuard guard(q.device()); + auto* props = at::cuda::getCurrentDeviceProperties(); + TORCH_CHECK(props->major == 7 && props->minor == 0, "SM70 required"); + int keep = prefix + std::min(topk, blocks - prefix); + if (checked) { + TORCH_CHECK(sizes.min().item() > 0 && sizes.max().item() <= 64, + "sizes must lie in [1,64]"); + auto counts = mask.sum(-1); + TORCH_CHECK((prefix == 0 || + counts.slice(-1, 0, prefix).eq(blocks).all().item()) && + counts.slice(-1, prefix).eq(keep).all().item(), + "map must follow H3 prefix and topk counts"); + TORCH_CHECK(prefix == 0 || mask.slice(-1, 0, prefix).all().item(), + "all prefix keys must be selected"); + } + int64_t per_group = prefix * blocks + (blocks - prefix) * keep; + TORCH_CHECK(per_group <= INT_MAX, "compact map exceeds index limits"); + + TORCH_CHECK(((uintptr_t(q.data_ptr()) | uintptr_t(k.data_ptr()) | + uintptr_t(v.data_ptr())) & + 15) == 0, + "16-byte aligned sources required"); + auto indices = torch::zeros({groups, per_group}, sizes.options()); + auto output = torch::zeros_like(q); + auto stream = at::cuda::getCurrentCUDAStream(); + compact<<>>(mask.data_ptr(), + indices.data_ptr(), blocks, + prefix, keep, per_group); + Geometry g{reinterpret_cast(q.data_ptr()), + reinterpret_cast(k.data_ptr()), + reinterpret_cast(v.data_ptr()), + reinterpret_cast<__half*>(output.data_ptr()), + indices.data_ptr(), + sizes.data_ptr(), + int(q.size(1)), + int(q.size(2)), + int(blocks), + int(prefix), + keep, + int(per_group), + 0, + 0, + 0, + float(scale)}; + auto data = torch::empty({3, q.numel()}, q.options().dtype(at::kFloat)); + convert_qkv<<<(q.numel() / 8 + 255) / 256, 256, 0, stream>>>( + g, data.data_ptr(), q.numel()); + QK qk; + PV pv; + for (int section = 0; section < 2; ++section) { + int begin = section == 0 ? 0 : prefix, end = section == 0 ? prefix : blocks; + g.keep = section == 0 ? blocks : keep; + int64_t score_bytes_per_query = groups * g.keep * 64 * 64 * 4; + TORCH_CHECK(begin == end || score_bytes_per_query <= kScoreBudgetBytes, + "a single query exceeds the diagnostic score budget"); + // Populate the GPU with independent PV queries; do not union their maps. + // Every query still visits its selected keys in the same ascending order. + int chunk = std::max( + 1, std::min(kMaxQueryBatch, + kScoreBudgetBytes / score_bytes_per_query)); + for (int first = begin; first < end; first += chunk) { + g.first = first; + g.count = std::min(chunk, end - first); + int batches = groups * g.count, n = g.keep * 64; + auto scores = + torch::empty({batches, 64, n}, q.options().dtype(at::kFloat)); + auto probability = torch::empty_like(scores); + auto result = + torch::empty({batches, 64, 128}, q.options().dtype(at::kFloat)); + auto pointers = + torch::empty({6, batches}, sizes.options().dtype(at::kLong)); + auto tokens = torch::empty({batches, n}, sizes.options()); + auto* ptr = pointers.data_ptr(); + prepare_batch<<>>( + g, data.data_ptr(), q.numel(), scores.data_ptr(), + probability.data_ptr(), result.data_ptr(), ptr, + tokens.data_ptr()); + QK::Arguments qa(cutlass::gemm::GemmUniversalMode::kArray, {64, n, 128}, + batches, {float(scale), 0.0f}, ptr, ptr + batches, + ptr + 2 * batches, ptr + 2 * batches, 0, 0, 0, 0, + int(q.size(2)) * 128, int(q.size(2)) * 128, n, n, + nullptr, tokens.data_ptr(), nullptr); + TORCH_CHECK(qk(qa, nullptr, stream) == cutlass::Status::kSuccess, + "CUTLASS QK failed"); + mask_scores<<>>( + g, tokens.data_ptr(), scores.data_ptr()); + torch::softmax_out(probability, scores, -1); + PV::Arguments pa(cutlass::gemm::GemmUniversalMode::kArray, {64, 128, n}, + batches, {1.0f, 0.0f}, ptr + 3 * batches, + ptr + 4 * batches, ptr + 5 * batches, ptr + 5 * batches, + 0, 0, 0, 0, n, int(q.size(2)) * 128, 128, 128, nullptr, + tokens.data_ptr(), nullptr); + TORCH_CHECK(pv(pa, nullptr, stream) == cutlass::Status::kSuccess, + "CUTLASS PV failed"); + scatter_result<<<(batches * 64 * 16 + 255) / 256, 256, 0, stream>>>( + g, result.data_ptr(), batches); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + } + return output; +} + +torch::Tensor forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, + torch::Tensor mask, torch::Tensor sizes, double scale, + int64_t prefix, int64_t topk) { + return forward_impl(q, k, v, mask, sizes, scale, prefix, topk, true); +} +torch::Tensor prevalidated(torch::Tensor q, torch::Tensor k, torch::Tensor v, + torch::Tensor mask, torch::Tensor sizes, + double scale, int64_t prefix, int64_t topk) { + return forward_impl(q, k, v, mask, sizes, scale, prefix, topk, false); +} +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("forward", &forward); + m.def("_forward_prevalidated", &prevalidated); +} diff --git a/csrc/sm70_turbomind/ops/diffusion_epilogue.h b/csrc/sm70_turbomind/ops/diffusion_epilogue.h new file mode 100644 index 0000000000..9e53dda799 --- /dev/null +++ b/csrc/sm70_turbomind/ops/diffusion_epilogue.h @@ -0,0 +1,99 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +#pragma once + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace sm70_diffusion { +template +__global__ void scaled_add_rows(Output* output, const float* delta, + const float* scales, int64_t count, + int64_t width, int64_t output_width, + int64_t offset, float alpha) { + for (int64_t index = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; + index < count; index += int64_t(gridDim.x) * blockDim.x) { + const int64_t row = index / width; + const int64_t col = index - row * width; + const int64_t destination = row * output_width + offset + col; + // Retain the explicit FP32 scale-restoration boundary before addition. + const float restored = + scales ? __fmul_rn(delta[index], scales[row]) : delta[index]; + float base; + if constexpr (std::is_same_v) { + base = __half2float(output[destination]); + } else { + base = output[destination]; + } + const float result = __fmaf_rn(alpha, restored, base); + if constexpr (std::is_same_v) { + output[destination] = __float2half_rn(result); + } else { + output[destination] = result; + } + } +} + +inline torch::Tensor scaled_add(torch::Tensor output, torch::Tensor delta, + std::optional scales, + double alpha, int64_t offset) { + TORCH_CHECK(output.is_cuda() && output.dim() == 2 && output.is_contiguous(), + "SM70 scaled addition requires contiguous CUDA [M,N] output"); + TORCH_CHECK(output.scalar_type() == torch::kFloat16 || + output.scalar_type() == torch::kFloat32, + "SM70 scaled addition output must be FP16 or FP32"); + TORCH_CHECK(delta.device() == output.device() && delta.dim() == 2 && + delta.is_contiguous() && + delta.scalar_type() == torch::kFloat32 && + delta.size(0) == output.size(0), + "SM70 scaled addition requires matching FP32 [M,K] delta"); + TORCH_CHECK(offset >= 0 && offset <= output.size(1) && + delta.size(1) <= output.size(1) - offset, + "SM70 scaled addition slice is outside output"); + TORCH_CHECK(std::isfinite(alpha) && + std::abs(alpha) <= std::numeric_limits::max(), + "SM70 scaled addition alpha must be finite FP32"); + TORCH_CHECK(!output.requires_grad() && !delta.requires_grad(), + "SM70 scaled addition is inference-only"); + at::assert_no_overlap(output, delta); + const c10::cuda::CUDAGuard guard(output.device()); + const auto* properties = at::cuda::getCurrentDeviceProperties(); + TORCH_CHECK(properties->major == 7 && properties->minor == 0, + "SM70 scaled addition requires SM70"); + const float* scale_data = nullptr; + if (scales.has_value()) { + TORCH_CHECK( + scales->device() == output.device() && scales->is_contiguous() && + scales->scalar_type() == torch::kFloat32 && + scales->numel() == output.size(0) && !scales->requires_grad(), + "SM70 scaled addition needs one FP32 scale per row"); + at::assert_no_overlap(output, *scales); + scale_data = scales->data_ptr(); + } + if (!delta.numel()) return output; + const int blocks = std::min((delta.numel() + 255) / 256, 65535); + const auto stream = at::cuda::getCurrentCUDAStream(); + if (output.scalar_type() == torch::kFloat16) { + scaled_add_rows<<>>( + reinterpret_cast(output.data_ptr()), + delta.data_ptr(), scale_data, delta.numel(), delta.size(1), + output.size(1), offset, float(alpha)); + } else { + scaled_add_rows<<>>( + output.data_ptr(), delta.data_ptr(), scale_data, + delta.numel(), delta.size(1), output.size(1), offset, float(alpha)); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return output; +} +} // namespace sm70_diffusion diff --git a/csrc/sm70_turbomind/ops/exact_row_reduce.cu b/csrc/sm70_turbomind/ops/exact_row_reduce.cu new file mode 100644 index 0000000000..57e7472379 --- /dev/null +++ b/csrc/sm70_turbomind/ops/exact_row_reduce.cu @@ -0,0 +1,165 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +#include +#include +#include +#include +#include +#include +#include + +struct Peers { + float* data[4]; + unsigned* flags[4]; +}; +constexpr int GRID = 80; +__device__ __forceinline__ void publish(unsigned* p, unsigned value) { + asm volatile("st.release.sys.global.u32 [%0], %1;" ::"l"(p), "r"(value) + : "memory"); +} +__device__ __forceinline__ unsigned acquire(unsigned* p) { + unsigned value; + asm volatile("ld.acquire.sys.global.u32 %0, [%1];" + : "=r"(value) + : "l"(p) + : "memory"); + return value; +} +__device__ __forceinline__ void barrier(Peers peers, int rank, unsigned epoch, + int phase) { + __syncthreads(); + if (threadIdx.x < 4) { + int peer = threadIdx.x; + int location = phase * GRID * 4 + blockIdx.x * 4; + publish(peers.flags[peer] + location + rank, epoch); + while (acquire(peers.flags[rank] + location + peer) != epoch) { + } + } + __syncthreads(); +} +__device__ __forceinline__ float tree(unsigned code, float x0, float x1, + float x2, float x3) { + switch (code) { + case 0: + return __fadd_rn(x0, __fadd_rn(x1, __fadd_rn(x2, x3))); + case 1: + return __fadd_rn(x0, __fadd_rn(__fadd_rn(x1, x2), x3)); + case 2: + return __fadd_rn(x0, __fadd_rn(__fadd_rn(x1, x3), x2)); + case 3: + return __fadd_rn(__fadd_rn(x0, x1), __fadd_rn(x2, x3)); + case 4: + return __fadd_rn(__fadd_rn(x0, x2), __fadd_rn(x1, x3)); + case 5: + return __fadd_rn(__fadd_rn(x0, x3), __fadd_rn(x1, x2)); + case 6: + return __fadd_rn(__fadd_rn(x0, __fadd_rn(x1, x2)), x3); + case 7: + return __fadd_rn(__fadd_rn(__fadd_rn(x0, x1), x2), x3); + case 8: + return __fadd_rn(__fadd_rn(__fadd_rn(x0, x2), x1), x3); + case 9: + return __fadd_rn(__fadd_rn(x0, __fadd_rn(x1, x3)), x2); + case 10: + return __fadd_rn(__fadd_rn(__fadd_rn(x0, x1), x3), x2); + case 11: + return __fadd_rn(__fadd_rn(__fadd_rn(x0, x3), x1), x2); + case 12: + return __fadd_rn(__fadd_rn(x0, __fadd_rn(x2, x3)), x1); + case 13: + return __fadd_rn(__fadd_rn(__fadd_rn(x0, x2), x3), x1); + case 14: + return __fadd_rn(__fadd_rn(__fadd_rn(x0, x3), x2), x1); + default: + return __int_as_float(0x7fffffff); + } +} +__global__ void reduce_rows(Peers peers, const uint8_t* codes, float* output, + int64_t count, int64_t offset, int rank, + unsigned epoch) { + barrier(peers, rank, epoch, 0); + for (int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; i < count; + i += int64_t(gridDim.x) * blockDim.x) { + int64_t at = offset + i; + float x0 = __ldcg(peers.data[0] + at); + float x1 = __ldcg(peers.data[1] + at); + float x2 = __ldcg(peers.data[2] + at); + float x3 = __ldcg(peers.data[3] + at); + output[i] = tree(codes[i], x0, x1, x2, x3); + } + barrier(peers, rank, epoch, 1); +} +std::tuple allocate(int64_t bytes) { + TORCH_CHECK(bytes > 0); + void* ptr = nullptr; + C10_CUDA_CHECK(cudaMalloc(&ptr, bytes)); + cudaIpcMemHandle_t handle; + try { + C10_CUDA_CHECK(cudaIpcGetMemHandle(&handle, ptr)); + C10_CUDA_CHECK(cudaMemset(ptr, 0, bytes)); + } catch (...) { + cudaFree(ptr); + throw; + } + return {reinterpret_cast(ptr), + pybind11::bytes(reinterpret_cast(&handle), sizeof(handle))}; +} +int64_t open_handle(pybind11::bytes bytes) { + std::string value = bytes; + TORCH_CHECK(value.size() == sizeof(cudaIpcMemHandle_t)); + cudaIpcMemHandle_t handle; + std::memcpy(&handle, value.data(), sizeof(handle)); + void* ptr = nullptr; + C10_CUDA_CHECK( + cudaIpcOpenMemHandle(&ptr, handle, cudaIpcMemLazyEnablePeerAccess)); + return reinterpret_cast(ptr); +} +void release(int64_t ptr, bool owner) { + if (owner) + C10_CUDA_CHECK(cudaFree(reinterpret_cast(ptr))); + else + C10_CUDA_CHECK(cudaIpcCloseMemHandle(reinterpret_cast(ptr))); +} +void run(torch::Tensor input, torch::Tensor codes, torch::Tensor output, + std::vector pointers, std::vector flags, int rank, + unsigned epoch) { + TORCH_CHECK(input.is_cuda() && input.scalar_type() == torch::kFloat32 && + input.is_contiguous()); + TORCH_CHECK(codes.is_cuda() && codes.scalar_type() == torch::kUInt8 && + codes.is_contiguous()); + TORCH_CHECK(output.is_cuda() && output.scalar_type() == torch::kFloat32 && + output.is_contiguous()); + TORCH_CHECK(input.device() == codes.device() && + input.device() == output.device()); + TORCH_CHECK(rank >= 0 && rank < 4 && pointers.size() == 4 && + flags.size() == 4 && epoch > 0); + TORCH_CHECK(input.numel() == output.numel() * 4 && + codes.numel() == output.numel()); + c10::cuda::CUDAGuard guard(input.device()); + auto stream = at::cuda::getCurrentCUDAStream(); + const auto* properties = at::cuda::getDeviceProperties(input.get_device()); + TORCH_CHECK( + properties->major == 7 && properties->minor == 0 && + properties->multiProcessorCount >= GRID, + "Exact peer reduction requires an SM70 device with at least 80 SMs"); + Peers peers; + for (int i = 0; i < 4; ++i) { + TORCH_CHECK(pointers[i] && flags[i]); + peers.data[i] = reinterpret_cast(pointers[i]); + peers.flags[i] = reinterpret_cast(flags[i]); + } + C10_CUDA_CHECK(cudaMemcpyAsync(peers.data[rank], input.data_ptr(), + input.nbytes(), cudaMemcpyDeviceToDevice, + stream)); + reduce_rows<<>>( + peers, codes.data_ptr(), output.data_ptr(), + output.numel(), output.numel() * rank, rank, epoch); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("allocate", &allocate); + m.def("open_handle", &open_handle); + m.def("release", &release); + m.def("run", &run); +} diff --git a/csrc/sm70_turbomind/ops/h3_w8a16.cu b/csrc/sm70_turbomind/ops/h3_w8a16.cu index 27a4aeec5a..72383c9387 100644 --- a/csrc/sm70_turbomind/ops/h3_w8a16.cu +++ b/csrc/sm70_turbomind/ops/h3_w8a16.cu @@ -8,6 +8,7 @@ #include #include "h3_column_major_gemm.h" +#include "diffusion_epilogue.h" namespace { __global__ void prepare_fp16_rows(const float* input, half* output, @@ -256,6 +257,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("prepare_fp16", &h3_prepare_fp16); m.def("dequantize", &h3_dequantize); m.def("rotate", &h3_rotate); + m.def("scaled_add_", &sm70_diffusion::scaled_add, pybind11::arg("output"), + pybind11::arg("delta"), pybind11::arg("scales"), pybind11::arg("alpha"), + pybind11::arg("offset") = 0); m.def("gemm", &h3_fp16_gemm, pybind11::arg("input"), pybind11::arg("weight"), pybind11::arg("output_fp32") = false); } diff --git a/docs/design/minimax_h3/CAMPAIGN_RESULTS.md b/docs/design/minimax_h3/CAMPAIGN_RESULTS.md new file mode 100644 index 0000000000..64c8f4de7d --- /dev/null +++ b/docs/design/minimax_h3/CAMPAIGN_RESULTS.md @@ -0,0 +1,42 @@ +# H3 SM70 campaign results + +No configuration has completed the >80 useful TFLOP/s/card, independent official quality and human-review gates. + +Development concentrates on FlashAttention-V100. FI retains its validated implementation; further FI optimization and exhaustive acceptance are paused at the user request. Quality and workflow gates remain unchanged. + +All rows use TP4. Rows marked formal use a complete request warmup plus three unprofiled requests. Other timings are captured cold diagnostics. Original floating-weight controls with legacy FLOP accounting omit throughput. + +| Workflow | Weights | Backend | Denoise seconds | Minimum card TFLOP/s | Native numerical control | Formal performance | +| --- | --- | --- | ---: | ---: | --- | --- | +| FL2V Light4 v1.2_768p | W8A16 | FLASH_ATTN_V100 | 59.749 | 51.939 | passed | failed >80 | +| light4-v10 | W8A16 | FLASH_ATTN_V100 | 61.263 | 50.655 | passed | not measured | +| light4-v11 | W8A16 | FLASH_ATTN_V100 | 61.019 | 50.858 | passed | not measured | +| light4-v01 | W8A16 | FLASH_ATTN_V100 | 61.000 | 50.873 | passed | not measured | +| light8-v10-non768 | W8A16 | FLASH_ATTN_V100 | 120.478 | 51.516 | passed | not measured | +| FL2V Light8 v1.0_768p | W8A16 | FLASH_ATTN_V100 | 121.542 | 51.065 | passed | not measured | +| Ref2V Light8 v1.0_768p, image/video/audio | W8A16 | FLASH_ATTN_V100 | 366.800 | 52.919 | passed | not measured | +| Ref2V Light4 v0.1, image/video/audio | W8A16 | FLASH_ATTN_V100 | 184.564 | 52.586 | passed | not measured | +| FL2V Light4 v1.2_768p, register FI | W8A16 | FLASHINFER_SM70 | 62.321 | 49.795 | passed | failed >80 | +| FlashGen four-step | original floating | FLASH_ATTN_V100 | 59.488 | 51.621 | passed | not measured | +| FastH3 Dense data-free | original floating | FLASH_ATTN_V100 | 56.224 | 54.125 | passed | not measured | +| FastH3 VSA data-free | original floating | FASTVIDEO_VSA | 37.387 | 45.268 | failed | not measured | +| FL2V Light4 v1.2_768p, original native peer rows | original floating | FLASH_ATTN_V100 | 59.324 | 52.312 | passed | not measured | +| Ref2V Light4 v0.1, mixed native peer rows | W8A16 | FLASH_ATTN_V100 | 181.261 | 53.544 | passed | not measured | +| Base H3, no LoRA, 49 updates, native peer rows | original floating | FLASH_ATTN_V100 | 649.973 | 57.353 | passed | not measured | +| FL2V Light4 v1.2_768p, original floating | original floating | FLASH_ATTN_V100 | 66.366 | not measured | passed | not measured | +| FL2V Light4 v1.2_768p, first | W8A16 | FLASH_ATTN_V100 | 66.419 | 50.697 | passed | not measured | +| FL2V Light4 v1.2_768p, last | W8A16 | FLASH_ATTN_V100 | 64.852 | 51.923 | passed | not measured | +| FL2V Light4 v1.2_768p, first-last | W8A16 | FLASH_ATTN_V100 | 69.615 | 52.305 | passed | not measured | +| FL2V Light4 v1.2_768p, native peer rows | W8A16 | FLASH_ATTN_V100 | 58.293 | 53.236 | passed | failed >80 | +| FL2V Light4 v1.2_768p, register FI native peer rows | W8A16 | FLASHINFER_SM70 | 60.970 | 50.899 | passed | failed >80 | + +The VSA failure is against an explicitly labeled FP32 selected-key diagnostic, not the unmodified official GPU kernel. Generation alone does not establish numerical quality. + +- Only explicitly marked formal rows are full-request warmup-plus-three result; other timings are captured cold diagnostics. +- Useful FLOPs exclude padding, duplicate work and skipped sparse/cache work; denominator is the slowest complete-denoise rank. +- Native parity does not establish independent official-model or human audiovisual quality. +- The 720p-family request uses an internal 1280x736/124-frame canvas; Ref2VA has longer conditioning sequences. +- First/last/both keyframes pass native controls for W8A16 Light4 v1.2; remaining adapter/weight keyframes, legal reference combinations, and primary-shape/TP matrix remain incomplete. +- TeaCache and Cache-DiT/SCM currently have separate small-shape lifecycle evidence, not primary >80 or official quality acceptance. + +Exact source/run paths and the evidence index are retained in `campaign-results.json` and `campaign-results.csv`. diff --git a/docs/design/minimax_h3/CONTROL.md b/docs/design/minimax_h3/CONTROL.md index b875f18e60..1be64ee286 100644 --- a/docs/design/minimax_h3/CONTROL.md +++ b/docs/design/minimax_h3/CONTROL.md @@ -1,5 +1,35 @@ # Native MiniMax H3 migration control +## Current campaign direction + +The user authorized integrating the implemented H3 SM70 and VSA source stack +through PR #583 on 2026-09-10. The retained dense workflows use FA; new FI +optimization and workflow expansion remain paused. See +[CURRENT_STATUS.md](CURRENT_STATUS.md) for the dense report and +[VSA_QUALITY_SPEED.md](VSA_QUALITY_SPEED.md) for the subsequent VSA stage and +merge validation. No VSA configuration passes both quality and speed, and no +experimental replacement enters default/AUTO selection. Historical experiments +and the unfinished full campaign remain below for reference. + +New Attention development and exhaustive workflow acceptance concentrate on +FlashAttention-V100 at the user's request. The already validated FlashInfer +implementation remains available, but new FI optimization and exhaustive FI +qualification are paused. The historical parallel-backend decisions below +are retained as evidence, not current work allocation. + +The full-request TP4 four-step native peer-row FA measurement records a +minimum-card median of 53.235745 useful TFLOP/s, 58.293218 seconds denoise, +and CV 0.076959%. It does not meet the >80 gate. Original floating weights, +W8A16, legal adapters/reference inputs, explicit VSA/cache algorithms and +independent official/human quality remain part of the campaign scope. +Only configurations satisfying the unchanged quality and full performance +gates may enter automatic selection. + +See [FA_DEVELOPMENT.md](FA_DEVELOPMENT.md) for current bottleneck evidence and +rejected kernel candidates, and [CAMPAIGN_RESULTS.md](CAMPAIGN_RESULTS.md) +for the per-workflow qualification table. Source integration is authorized; +local operator results are not end-to-end model acceptance. + Latest FlashInfer change: [FLASHINFER_LOCAL_ROTATION.md](FLASHINFER_LOCAL_ROTATION.md). Rotate FP16 rows on their owner before all-gather, avoiding duplicated QKV/MLP diff --git a/docs/design/minimax_h3/CURRENT_STATUS.md b/docs/design/minimax_h3/CURRENT_STATUS.md new file mode 100644 index 0000000000..1328d0162e --- /dev/null +++ b/docs/design/minimax_h3/CURRENT_STATUS.md @@ -0,0 +1,97 @@ +# H3 retained FA delivery status + +Current delivery concentrates on already working H3 dense workflows and the +retained `FLASH_ATTN_V100` implementation. This is already the native H3 +default backend. Slower experimental replacements are excluded. Further +Attention prototype work and new workflow expansion are paused at the user's +request; no GPU experiments remain queued. Native CLI/HTTP APIs remain the +frontend entry point. + +## Formal four-step result + +TP4, V100 SXM2 32GB, single request, 1280x736 internal canvas, 120 requested +frames aligned to 124, 24 FPS, LightX2V v1.2 four-step W8A16. Both backend +runs use the same pageable host/shared VAE policy and native peer reduction. +Each has one complete request warmup and three unprofiled measurements. + +| Metric | Retained FA | Retained FI | +| --- | ---: | ---: | +| Complete denoise median | 58.293218 s | 60.969633 s | +| Denoise divided by four updates | 14.573304 s | 15.242408 s | +| Minimum-card median useful TFLOP/s | 53.235745 | 50.898828 | +| Denoise CV | 0.076959% | 0.069457% | +| Complete request median | 91.071940 s | 89.525699 s | + +FA reduces denoise by 4.389750%. The complete request includes other stages +and is slower in this measurement; an overall FA request speedup is not +established. The tracked live-allocation upper bound including IPC is +20,475,227,136 bytes/card. All >80 gates remain incomplete. + +## Retained coverage + +These are complete native output-preservation controls. Except for the +explicit formal result above, the times below are captured cold diagnostics. +They are not repeated performance acceptance or independent official-model +quality acceptance. + +| Workflow | Weights | Denoise | Numerical control | +| --- | --- | ---: | --- | +| Base H3, no adapter, 49 updates | Original floating | 649.973 s | Passed | +| LightX2V four-step v1.2, peer reduction | Original floating | 59.324 s | Passed | +| Six official FL2V four/eight-step adapters | W8A16 | Four: 58.293-61.263 s; eight: 120.478-121.542 s | Passed | +| Ref2VA four-step, image/video/audio, peer reduction | W8A16 | 181.261 s | Passed | +| Ref2VA eight-step, image/video/audio | W8A16 | 366.800 s | Passed | +| Light4 v1.2 first / last / both keyframes | W8A16 | 66.419 / 64.852 / 69.615 s | Passed | +| FlashGen four-step | Original floating | 59.488 s | Passed | +| FastH3 Dense data-free | Original floating | 56.224 s | Passed | + +The no-adapter original-weight control decreases denoise from 725.819 to +649.973 seconds (10.449687%) and complete request from 835.356 to 700.121 +seconds. Final video/audio latents, all pre-encoding RGB frames and PCM are +bitwise equal. Combined host, projection and residual changes contribute; +this is not an isolated Attention comparison. + +Original floating and W8A16 bases share FP16 projection, scale restoration, +LoRA addition and Attention interfaces. Adapter increments retain unrotated +inputs and enter row reduction before the collective. Residual sharding and +native peer reduction are explicit options, with memory-budget fallback. +No quantization-only or adapter-free restriction selects shared Attention. + +## Reproducing the measured configuration + +Use the native `H3Config` with the model and matching adapter identifiers +from the recorded request contract. The measured TP4 primary configuration +sets `attention_backend="FLASH_ATTN_V100"`, `attention_query_tile=128`, +`fp16_weight_layout="column"`, `residual_sequence_parallel=True`, +`residual_reduction="peer"`, `residual_reduction_memory_gib=4`, +`host_weight_pin_memory=False`, and `share_host_vae_weights=True`. +The mixed Ref4 control uses an explicit 8 GiB communication budget. +Keep each adapter's official sigma, flow-shift and task contract. + +This records explicit measured options, not a universal default or AUTO +promotion. Peer reduction has complete controls for the primary four-step, +original four-step/no-adapter and mixed Ref4 cases. Other rows retain their +measured configurations; the latest peer option is not claimed validated for +every shape/adapter combination. TP1/TP2 capacity checks use layer offload +and ordinary reduction, with no primary-shape >80 claim. + +FastH3 VSA remains outside qualified delivery because its full numerical +diagnostic fails. TeaCache and Cache-DiT/SCM have request-lifecycle and +small-shape GPU evidence; primary-shape official quality and performance are +incomplete. The 243-frame and 15-second runs establish generation/memory +compatibility only. Human audiovisual review and independent official-model +controls remain pending. The user authorized source integration through +PR #583 on 2026-09-10; this does not close those acceptance gates. See +[VSA_QUALITY_SPEED.md](VSA_QUALITY_SPEED.md) for the subsequent VSA stage: +the native path reaches 30.990756 seconds but fails independent quality; +the acceptance-only exact FP32 path passes primary numerical checks at +60.353224 seconds and fails the speed target. Neither enters automatic selection. + +Evidence: [CAMPAIGN_RESULTS.md](CAMPAIGN_RESULTS.md), +[FA_DEVELOPMENT.md](FA_DEVELOPMENT.md), and artifact root +`/data/minimax-h3/sm70-general-20260909/` with +`peer-api-720p-three-runs/performance.json`, +`peer-api-fi-720p-three-runs/performance.json`, +`peer-api-breadth-summary.json`, `original-base-summary.json`, and full +request/source/binary contracts. No new timing is inferred from the scope +change or documentation update. diff --git a/docs/design/minimax_h3/EXACT_ROW_REDUCTION.md b/docs/design/minimax_h3/EXACT_ROW_REDUCTION.md new file mode 100644 index 0000000000..2456f3c2db --- /dev/null +++ b/docs/design/minimax_h3/EXACT_ROW_REDUCTION.md @@ -0,0 +1,189 @@ +# Explicit SM70 local-row reduction + +The shared `SM70ExactRowReductionPlan` interface is experimental and has no +automatic dispatch. H3 exposes an explicit runtime selection; its native +four-step integration passes media preservation and remains below the +formal >80 throughput gate. The ordinary residual path +continues to use FP32 all-reduce followed by a local-row slice. No configuration +has passed the campaign's >80 useful TFLOP/s/card and complete quality gates. + +## Arithmetic and ownership + +A conventional reduce-scatter changes FP32 addition order relative to the +existing all-reduce, and the earlier full H3 control failed numerical gates. +This interface instead classifies the native communicator's addition order +for the actual two-dimensional FP32 shape. Ten fixed finite probes distinguish +all 15 four-input binary addition trees. Unmatched or ambiguous elements reject +setup collectively. Calibration never reads model weights or model activations. + +The CUDA implementation copies each rank's partial input into owned IPC +storage, reads only its destination rows from peers and evaluates the calibrated +tree with rounded FP32 additions. System release/acquire flags establish input +visibility and completion. The grid has 80 blocks; the interface requires +SM70 devices with at least 80 SMs, four distinct peer-accessible devices on +one host, and consistent CUDA visibility. All row-parallel adapter contributions +must already be included in the input. + +The caller supplies an explicit budget covering persistent IPC buffers, +local output, one-byte arithmetic codes and calibration scratch. GPU buffers +belong to the plan, not a global shape cache. Returned tensors alias the plan's +output; consume them before the next call. Call `close()` collectively before +tearing down the TP group. Rank-dependent setup errors are exchanged over the +CPU group, and peer handles close before owners free their allocations. + +The plan rejects another stream/device, autograd inputs, incompatible layouts, +CUDA Graph execution and epoch exhaustion. It does not silently change a +backend or precision. Callers retain their ordinary collective when a plan +is unsuitable. H3 owns one plan per pipeline, reuses it for identical shapes and closes it +collectively on shape changes or worker shutdown. Other models must likewise +own the plan lifecycle explicitly. + +## Validation and measured limits + +Environment: Torch 2.10.0+cu128, CUDA toolkit 12.8.93, NCCL 2.27.5, four leased +V100 SXM2 32GB cards. Evidence root: +`/data/minimax-h3/sm70-general-20260909/exact-peer-reduction/`. + +- The first arithmetic classification covers every element of the real + 34560x5376 projection. Three independent wide-range inputs and signed-zero + controls match the native all-reduce bitwise. +- The source implementation's TP4 control covers seven shapes from 4x3 through + 34560x5376, including non-H3 DiT widths, tails and non-aligned storage offsets. + All four ranks preserve bits for ordinary, wide, subnormal and opposing + infinity inputs: 112 numerical cases. Mismatched shapes, insufficient + per-rank budgets, different streams and closed plans are rejected. +- The prototype's independent four-rank memcheck fixture reports zero errors + on every rank. An earlier NCCL-bearing fixture reported only initialization + `cudaFuncGetAttributes` probes for unsupported kernels; NCCL explicitly + skips that return code in its [corresponding source](https://github.com/NVIDIA/nccl/blob/v2.27.5-1/src/enqueue.cc#L37-L38). + The isolated fixture uses Gloo for coordination and an independent FP32 + arithmetic reference; no CUDA API error suppression was applied. +- Source operator medians, including the full input copy and device barriers: + native 14.561–14.641 ms, peer rows 10.939–10.991 ms. These seven alternating + measurements are communication diagnostics only. +- A separate artifact override at source `3280edbfcc` completes a full denoise + warmup per implementation, then one measurement each: 59.717282 seconds + native versus 58.228244 seconds peer rows, a 2.49348% reduction. Candidate + useful throughput is 53.295148–53.295167 TFLOP/s/card. Every final video/audio + latent bit matches across all four passes, and the baseline also matches + the previously frozen FA query-128 control. Each candidate pass uses 400 + peer reductions. This is not the full-request warmup-plus-three protocol. +- The prototype's full native media control also preserves both final latents, + all 124 RGB frames and PCM bitwise. SSIM and RMS ratio are 1; spectral cosine + exceeds 0.99999999999998. Its captured cold request takes 94.269191 seconds, + including 61.688596 seconds denoise. This is native preservation, not an + independent official-model or human audiovisual review. +- That full request peaks at 19,732,554,240 PyTorch-allocated bytes/card plus + 743,180,800 persistent raw IPC bytes/card. Their sum is 20,475,735,040 bytes; + driver/library overhead is additional. Do not report only the PyTorch number. + +The committed shared interface at `6d2a44b8d0` now also passes a complete +native H3 control using an explicit forward override. Dynamic calibration uses +the current native group and the actual shape, with no saved arithmetic-code +map. Final video/audio latents, all 124 RGB frames and PCM match the frozen +FA query-128 control bitwise; SSIM and RMS ratio are 1. The captured request +records 58.310740 seconds denoise and 99.562953 seconds total. Its contract, +source/binary manifests, `review-native-quality.json` and +`review-native-summary.json` are retained separately from prototype evidence. + +This validates the final shared operator in one H3 request, but does not establish full-request warmup-plus-three performance. +Final native integration validation, formal repeated requests, TP/shape breadth and official/human +quality gates remain incomplete. No AUTO promotion is made. + +Reproduce the operator control with an owned native GPU lease and +`torchrun --standalone --nproc_per_node=4 +benchmarks/kernels/benchmark_sm70_exact_row_reduce.py --output +--full-shape`. The optional `--extension` pins an already-built library; the +report records its SHA256 plus benchmark, CUDA and shared Python source hashes. + +## Native API selection and memory records + +Start `vllm video serve` or `vllm video generate` with +`--residual-sequence-parallel --residual-reduction peer +--residual-reduction-memory-gib 4` to select the explicit candidate. The default +remains `native`; HTTP clients continue to use the existing video API. Selection +does not depend on floating versus W8A16 weights, adapters or task labels. + +TP1 retains its ordinary path. TP2 uses the original all-reduce and local slice. +TP4 creates a shared plan only when its complete calibration/storage requirement +fits the explicit budget; larger shapes use the ordinary collective. A plan +remains valid only for its original communicator and shape. Setup happens on +the first actual projection; its time is included in complete denoise and is +also reported separately in `residual_communication.setup_seconds`. Repeated +requests of the same shape reuse calibration without reading model state. + +Each result records peer/native call counts and the fallback reason. The +`torch_peak_allocated_bytes` and `raw_ipc_peak_bytes` fields remain separate; +`peak_allocated_bytes` is their conservative sum and is marked as an upper +bound when raw IPC storage is present. The performance validator therefore +includes external communication allocations in its memory gate. CUDA driver +and library overhead still require the retained NVML measurements. + +CPU request ownership, budget fallback, config, API and existing residual +regressions pass. The final native selection now has the separate controls and measurements +below. Wider workflow validation remains incomplete. + +## Final native four-step measurements + +At source `ca82c279bc`, the ordinary engine selects peer rows through H3Config, +with no forward replacement. The separate captured native request preserves +both final latents, all 124 RGB frames and PCM bitwise. All four ranks report +400 peer calls, zero native fallbacks and about 0.374 seconds initial plan +setup. `peer-api-native-quality.json` and `peer-api-native-summary.json` retain +this control and its precise configuration. + +`peer-api-720p-three-runs/performance.json` records one complete request warmup +and three unprofiled, uncaptured requests of that same configuration: + +| Measurement | Value | +| --- | ---: | +| Warmup denoise | 59.698783 s | +| Measured denoise | 58.344130 / 58.293218 / 58.234342 s | +| Median useful TFLOP/s/card | 53.235745–53.235764 | +| Denoise coefficient of variation | 0.076959% | +| Complete request | 91.071940 / 90.560870 / 91.905029 s | +| Live allocation upper bound, including IPC | 20,475,227,136 bytes/card | +| Memory gate | Pass | +| >80 throughput gate | Fail | + +The companion `peer-api-formal-telemetry.json` retains 1,995 NVML samples at a +0.25-second interval. Across startup, warmup and measurement, the maximum +sampled device usage is 24,387,256,320 bytes, including allocator caches and +runtime overhead. High-utilization samples have median power of about +278–280 W/card. Telemetry is independent of the CUDA work counters. + +This formal request uses pageable host masters and shared VAE weights; the +older FA query-128 formal control used pinned masters. The latter's shorter +complete-request time must not be presented as a matched comparison of the +communication kernels. The isolated 2.49348% denoise comparison above used +matching host and compute settings. No overall request speedup is claimed +across the different host-memory policies. + +The interface remains explicit. The >80 target, official/human quality, wider +workflow and shape/TP matrix are still incomplete. Initial setup, skipped work, +raw IPC storage and slower end-to-end outcomes are retained in the records. + +## Native backend and workload breadth + +The same native API path with register-probability FI also passes a complete +latent/RGB/PCM bitwise control against its frozen FI baseline. Full-request +warmup plus three unprofiled measurements record denoise +60.986616 / 60.889535 / 60.969633 seconds, median +50.898828-50.898847 useful TFLOP/s/card and CV 0.069457%. Complete requests take +89.859034 / 89.086770 / 89.525699 seconds. The allocation upper bound including +raw IPC is 20,475,227,136 bytes/card. FA and FI peer runs share the pageable +host/shared VAE policy. Both fail the >80 gate. Evidence: +`peer-api-fi-720p-three-runs/performance.json`, +`peer-api-fi-native-quality.json` and `peer-api-fi-formal-telemetry.json`. + +Additional complete native controls preserve video/audio latents, all 124 RGB +frames and PCM bitwise with original floating Light4 weights and W8A16 Ref4 +mixed image/video/audio conditioning. Original Light4 records 59.323955 seconds +denoise and 22,022,771,200 bytes/card allocation upper bound. Ref4 records +181.260984 seconds and 21,572,765,184 bytes/card; its longer reference sequence +uses an explicit 8 GiB reduction budget and 1,491,864,064 raw IPC bytes/card. +Both record 400 peer calls with zero native fallbacks. These are captured cold +quality controls, not formal repeated performance measurements. Evidence: +`peer-api-breadth-summary.json` and both corresponding `*-quality.json` files. +Independent official quality, human review and the full task/weight/shape +matrix remain incomplete. diff --git a/docs/design/minimax_h3/FA_DEVELOPMENT.md b/docs/design/minimax_h3/FA_DEVELOPMENT.md new file mode 100644 index 0000000000..e128af6798 --- /dev/null +++ b/docs/design/minimax_h3/FA_DEVELOPMENT.md @@ -0,0 +1,104 @@ +# FlashAttention development focus + +Current delivery uses the retained FlashAttention-V100 implementation on +already working H3 dense workflows. At the user's latest request, further +investigation of slower Attention prototypes and new workflow expansion are +paused. Validated FlashInfer remains available, with further FI optimization +and exhaustive FI acceptance paused. The unfinished matrix and unchanged +quality/performance gates remain recorded as incomplete. + +See [CURRENT_STATUS.md](CURRENT_STATUS.md) for the retained configuration, +measured workflow coverage and delivery limits. + +With native TP4 peer rows, shared pageable VAE weights and a complete request +warmup plus three unprofiled requests, FA records a minimum-card median +53.235745 useful TFLOP/s and 58.293218 seconds denoise. The corresponding FI +measurements are 50.898828 TFLOP/s and 60.969633 seconds. Both pass native +latent/RGB/PCM preservation and fail the >80 throughput gate. Independent +official and human quality gates remain pending. + +## Evidence guiding the next operator change + +A development-only FA timer samples thread zero in head zero and every 32nd +query CTA. Nine boundary lengths and the actual [1,34551,14,128] captured +input remain bitwise equal to the retained FA implementation. Instrumented +operator median is 150.045700 ms versus 149.676025 ms without instrumentation +(0.247% overhead). Main-shape compilation uses 250 registers and no spills. +The query-64/key-128 diagnostic specialization spills and is not measured. + +| Sampled phase | Fraction of sampled warp spans | +| --- | ---: | +| QK including operand loads and barrier | 38.85% | +| V prefetch and softmax | 18.76% | +| Probability stores and publication | 7.94% | +| PV including operand loads and barrier | 33.20% | +| Iteration join | 1.25% | + +These phase spans include scheduling and waits; they are not whole-kernel +critical-path percentages or formal model performance. The measurement guides +operand and register-lifetime work, without qualifying a faster configuration. + +The primary 34,560-row GEMM diagnostic also compares all 18 returned eligible +zero-workspace, no-split, FP32-accumulation Lt choices across QKV, FC1, output +projection and FC2. Timed alternatives preserve outputs bitwise, and the +existing algorithm 21/tile 24 remains the fastest choice in each shape. +No GEMM plan change is justified by this measurement. + +## Rejected candidates + +- Normal-range hardware exp2 retains the library outside [-126,0]. Boundary + and actual-input outputs are bitwise, but paired median regresses from + 149.407745 to 158.803970 ms. No full-model run or source promotion follows. +- A new explicit WMMA Q16/K64-owner prototype retains FA K128 softmax panels, + FP32 accumulation and output scaling. It lowers register use from 248 to + 128 without spilling and doubles a Q128 CTA to 512 threads. Nine boundary + lengths and actual input are bitwise, but operator median regresses from + 147.507202 to 323.808258 ms. Register count alone is insufficient evidence + of a speedup. This is distinct from the old invalid CUTLASS warp16 shape. +- Adding vector Q/K loads, V lane-exchange transpose and swizzled probability + storage to that prototype spills 160 bytes per thread at the 128-register + limit. The resource gate rejects it before GPU timing. A separate scoped + staging experiment reduces the spill to 116 bytes with a 120-byte stack, + which still fails the resource gate. Neither staging variant is GPU timed. + +Exact code, binary hashes, clocks, numerical results and paired measurements +are retained under `/data/minimax-h3/sm70-general-20260909/` in +`attention-fa-phase-clock`, `gemm-primary-heuristics`, +`attention-fa-normal-exp2`, `attention-fa-warp16-native` and +`attention-fa-warp16-staging`. None is installed as a production replacement. + +The vector-staging Q96 follow-up raises the per-thread register budget to 168. +Its only 4-byte local spill is stored before the key loop and reloaded after +its final back edge, as verified in SASS. This admits a bounded numerical and +performance probe without claiming zero spills. Nine boundaries and the +actual 34,551-token input are bitwise, but median regresses from 148.462585 to +281.825287 ms. This closes the split-key warp16 staging route. + +A separate one-owner FA K128 prototype keeps rounded probability fragments +in registers, preserves the FA K64-half sum order and retains Q across key +panels. It uses 199 registers, no spills and 96 KiB shared memory. Nine +boundaries and the actual input are bitwise, but median is 195.892136 ms +versus 147.519104 ms. Its V-prefetch follow-up compiles to 240 registers with +no spills, but is not GPU timed: the parent exceeds the predeclared 10% +slowdown limit. No full-model run or production promotion follows. + +Evidence: `attention-fa-warp16-staging-q96`, `attention-fa-register-k128`, +`attention-fa-register-k128-vprefetch`, and `fa-vprefetch-after-register.json`. + +The artifact-only D128 tiled-GEMM port also closes without promotion. Padding +only PV's reduction width to 32 fixes its nonaligned-key vector read; the +original failing boundary then passes Compute Sanitizer with zero errors. +Eleven boundary/stress cases and the actual input pass the independent FP32 +operator gate. Actual-input relative L2 is 0.000231 against that oracle and +0.000335 against FA; output is not bitwise equal to FA. Its paired median is +212.961273 ms versus retained FA's 149.299194 ms, so no complete sampling +quality run is justified. + +A bounded Nsight Systems attribution records 209.361943 ms operator wall +time: QK plus tile softmax 142.181160 ms, PV plus probability rescaling +61.613049 ms, and GPU idle/host gaps 0.582682 ms. API durations overlap GPU +execution and cannot be added to those kernel durations. Packing or CPU +scheduling does not explain this regression. A row-maximum epilogue follow-up +completed its CPU build but was stopped at the user's scope change before +any GPU numerical check or timing. Neither artifact replaces production FA. +Evidence: `attention-fa-tiled-gemm-d128` and `attention-fa-tiled-gemm-rowmax`. diff --git a/docs/design/minimax_h3/FLASHINFER_REGISTER_PROBABILITY.md b/docs/design/minimax_h3/FLASHINFER_REGISTER_PROBABILITY.md new file mode 100644 index 0000000000..7c32cb8526 --- /dev/null +++ b/docs/design/minimax_h3/FLASHINFER_REGISTER_PROBABILITY.md @@ -0,0 +1,118 @@ +# FlashInfer SM70 probability fragments in registers + +This change improves the explicitly selected FlashInfer backend. It does not +qualify an AUTO configuration or meet the campaign's >80 useful TFLOP/s/card +gate. The shared `sm70_attention.noncausal_attention` interface dispatches on +SM70 hardware, FP16 tensors and BSHD D128 layout, independently of model, +quantization and adapter identity. + +## Arithmetic and execution + +One warp owns 16 query rows and both logical 32-key halves of the existing +64-key tile. QK/PV accumulators, online maxima and denominators remain FP32. +Each original partial sum, XOR-2/XOR-8 reduction, left/right addition and +per-output MMA K order is retained. Probabilities cross the same FP16 rounding +boundary as before. + +Six 32-bit lane exchanges convert the rounded probability pairs into Volta A +fragments. PV reuses each fragment across eight output fragments. Probabilities +no longer traverse shared memory, and the warp owns its softmax state without +cross-warp barriers. CTA barriers still protect K/V staging and reuse. + +The CTA covers 192 queries with 384 threads. Cooperative prefetch rounds up +to three vectors per thread. Extra complete warps in the final prefetch group +skip rows outside the 64-key tile; the guard is warp-uniform before shuffles. +Valid-length padding, batches, heads and unaligned global-storage loads retain +their original handling. Explicit lane selects avoid addressable local arrays. + +CUDA 12.8 emits 168 registers/thread, no stack or spills, and 84,992 dynamic +shared bytes. These resource counts describe the implementation; they are not +throughput measurements. + +The FP16-accumulator layout shortcut in [FastAttention Appendix B](https://arxiv.org/html/2410.16663v1#A2) +is outside this campaign's precision contract. This implementation exchanges +already rounded probabilities while retaining FP32 accumulation. + +## Development evidence + +Environment: Python 3.12.13, Torch 2.10.0+cu128, CUDA toolkit 12.8.93, +V100 SXM2 32GB. Artifacts are under +`/data/minimax-h3/sm70-general-20260909/`. + +- `attention-fi-register-probability-q192/probe.json`: 17 boundary lengths + and the actual 34,551-token, 14-head H3 capture match the frozen FI binary + bitwise. Sampled FP32-reference relative L2 is 0.000374 on the actual input. +- Seven alternating operator timings on the same GPU give medians + 183.124985 -> 160.701447 ms, 12.245% lower latency. Observed clocks remain + 1425-1432 MHz. This is an operator measurement only. +- Compute Sanitizer 12.8 memcheck, racecheck and synccheck each report zero + errors on 12 boundary cases with two batches, three heads and storage offsets. +- The formatted native build passes 33 GPU numerical, unaligned-storage, + independent-query, graph and non-H3 shared-interface checks. The affected + CPU contracts and provenance/acceptance suite pass 46 checks, with seven + GPU checks skipped in the explicitly device-masked CPU invocation. +- `fi-general-control-quality.json`: the current common prepared/residual + path preserves the previously frozen FI final video/audio latents bitwise. +- `fi-register-denoise-summary.json`: one complete denoise warmup per + implementation, followed by one unprofiled measurement per implementation + in reverse order. Both use the same TP4 W8A16 LightX2V four-step v1.2 weights, + five sigma points, captured conditioning/noise, column layout, exact residual + sharding and zero persistent FP16 weight cache. All four final-latent pairs + match bitwise on every rank. + +| Complete-denoise control | Slowest-rank seconds | Useful TFLOP/s/card | +| --- | ---: | ---: | +| Previous FI kernel, common execution path | 67.303216 | 46.108983-46.109000 | +| Register-probability FI kernel, same path | 62.471266 | 49.675363-49.675382 | + +The isolated kernel change reduces this complete denoise by 7.179%. Candidate +steps take about 15.59-15.64 seconds. The control excludes encoder, VAE and +packaging, and includes only one measurement per implementation. It is not +the required full-request warmup-plus-three acceptance. Existing FA query-128 +remains faster in its separately recorded full-request measurements. + +## Complete native media control + +The formatted native build at `f8b85c681a402ad26aac13adc701687a342881d9` +completes a separate pair of full TP4 W8A16 LightX2V four-step v1.2 requests. +Both implementations use the same current Python source, shared pageable VAE +host storage, zero persistent FP16 cache and the five-second request's +1280x736/124-frame internal canvas. Only the immutable FI binary differs. + +`fi-register-native-quality.json` passes every numerical gate. Final video +and audio latents, all 124 unencoded RGB frames and PCM match bitwise. Video +PSNR is infinite and SSIM is 1; audio spectral cosine is +0.9999999999999756 and RMS ratio is 1. Both requests pass native media and +strict actual-work validation. Their peak allocation is unchanged at +19,501,498,880 bytes/card. Commands, binary hashes and clean source provenance +are in `fi-register-native-pair.json` and `fi-register-native-summary.json`. + +The single captured cold requests take 71.085578 / 63.649644 seconds denoise +and 107.358301 / 100.664110 seconds request (control/candidate). These captures +include different first-use setup costs and are not a formal speed comparison. +Use the matched warmed denoise control above for the isolated 7.179% result. + +## Formal repeated requests + +Source `c69cfc7024460e314e79a0bba37a3b736340bc6e` completes one full native +warmup and three requests without profiler or captures, with the same +media-checked configuration and immutable binary. The warmup takes 63.412712 +seconds denoise and 98.849412 seconds request. + +| Measurement | Denoise seconds | Complete request seconds | +| --- | ---: | ---: | +| 1 | 62.321408 | 91.585107 | +| 2 | 62.339754 | 91.385871 | +| 3 | 62.257260 | 95.401157 | + +Minimum-to-maximum rank median throughput is **49.794813-49.794831 useful +TFLOP/s/card**, using the slowest rank's complete denoise time. Denoise CV is +0.056762%; peak allocation remains 19,501,498,880 bytes/card. The full contract, +per-rank stages/steps, loaded binary hashes, source hashes and NVML samples are +retained in `fi-register-720p-three-runs/`. + +**The >80 performance gate fails.** The earlier FA query-128 configuration's +51.939 TFLOP/s/card remains the campaign's best formal result. The complete +native FI numerical control passes, but independent official references and +human review remain pending. No AUTO selection, precision relaxation or +campaign completion is claimed. diff --git a/docs/design/minimax_h3/GENERAL_SM70.md b/docs/design/minimax_h3/GENERAL_SM70.md new file mode 100644 index 0000000000..7bffd21bdf --- /dev/null +++ b/docs/design/minimax_h3/GENERAL_SM70.md @@ -0,0 +1,261 @@ +# General SM70 H3 acceleration campaign + +Integration: `onecat/main`, base `4f19ef7a20db60bb0685e599bd3f4dd156202eed`. +Owned branch: `codex/v100-h3-general-sm70-20260908-165458`. +Raw evidence: `/data/minimax-h3/sm70-general-20260909/`. + +## Accepted objective + +Accelerate every supported H3 generation task across original floating-point +weights, W8A16, all eight LightX2V four/eight-step adapters, FlashGen, FastH3 +Dense/VSA, TeaCache and Cache-DiT. Keep official task, adapter and sigma +contracts. Extract reusable SM70 DiT operators; a second complete model is +outside this campaign. Frontends continue to call the native vLLM API. + +Use hardware, tensor layout/dtype, precision, adapter and collective capabilities +for dispatch. Preserve FP32 accumulation, residuals and wide-range scaling. +ConvRot applies only to matching rotated weights. LoRA receives the unrotated +activation and joins TP partial sums before reduction. Extend residual sequence +sharding to original weights, Ref2VA, adapters and TP2/TP4; TP1 remains ordinary +execution. AUTO may select only qualified routes. + +Primary performance cases are the original 720p five-second four-step sample, +1344x768/243-frame official workload, and the 15-second boundary on TP4. +Count actual useful rank-local model work over the slowest rank's complete +denoise wall time, excluding padding, redundant work and skipped operations. +Sparse attention counts selected pairs and its actual compression projections; +cache savings are reported separately from arithmetic throughput. Record +actual denoiser calls, executed blocks and per-step/stage timing. + +Acceptance requires one full warmup then three unprofiled measurements, +each rank's median above 80 TFLOP/s and denoise CV <= 5%, plus measured request +latency and memory. Smaller shapes and TP1/TP2 require compatibility checks. +The old 43.914752 TFLOP/s/card sample had 71.315842 s denoise and no warmup; +its equal-work 80 TFLOP/s budget is 39.147719 s. The old 39-frame/20-step +FA/FI diagnostics are not matching speed baselines. + +After auditing duplicated column-LoRA A work, the current four-step numerator +is 3,103,284,010,387,456 useful FLOPs on rank 0. Its corrected 80-TFLOP/s budget +is 38.791050 seconds. Historical numbers above retain their original accounting +and must not be mixed with the corrected formal results below. + +Quality compares each variant against the same weights, initial noise and +official algorithm, including full sampling. Default numerical gates are +video/audio final-latent relative L2 <= 0.01, pre-encoding video PSNR >= 40 dB +and SSIM >= 0.99, audio spectral cosine >= 0.99 and RMS ratio in [0.99, 1.01]. +Complete audiovisual and reference-consistency review remains a separate gate. +Unaccepted drift is not a new oracle; thresholds must not be relaxed to pass. + +## Progress and required evidence + +| Requirement | Implementation | Validation / evidence | +| --- | --- | --- | +| Common FP16 input/GEMM, dense column-major path | Shared operator and explicit dense layout | GPU operator checks pass; original/W8A16 four-step complete controls match bitwise | +| Prepared LoRA input and collective ordering | Implemented, including explicit original basis | GPU prepared/normal results bitwise equal; TP2/TP4 block comparisons pass | +| General residual sequence sharding | TP2/TP4, original/INT8, matching adapters; TP1 no-op | TP2 original Light4 small full control and TP4 Ref8 mixed-reference control match bitwise; broader matrix pending | +| FA and FI kernel optimization | Explicit query-128 and shared epilogues in #581 | Audited FA 47.092, FI 43.990, candidate FA 51.939 TFLOP/s/card; >80 fails | +| All dense task/weight/adapter combinations | Partial mainline support | Full matrix pending | +| FastH3 VSA on SM70 | True block-sparse native API in #583 | Full generation completes; final FP32-math diagnostic fails latent/video gates | +| TeaCache and Cache-DiT | Request-scoped official policies in #584 | TP1/2/4 small forwards and full native cached/lossless/cached lifecycle pass; official quality pending | +| AUTO and native variant APIs | Variant APIs in #583/#584; AUTO pending | No configuration qualified for automatic selection | +| Workflow-specific performance accounting | Actual intervals, blocks, sparse pairs and cache hits in #578/#583/#584 | Strict validators pass; skipped/padded/duplicate work excluded | +| Non-H3 DiT operator reuse | Shared GEMM/input preparation and explicit dense attention | GEMM and non-H3 BSHD attention shapes pass; no second complete model added | +| TP1/TP2 original-weight capacity | Explicit DiT/encoder layer staging | Full small original Light4 generations pass; pinned/pageable TP1 and ordinary/sharded TP2 final media match bitwise | +| >80 TFLOP/s/card, full quality, memory | Not achieved | No qualifying results | +| Draft PRs, matrix report and playable samples | Draft PRs #571/#578/#581/#583/#584 | Original/W8A16 Light4, W8A16 Light8, FlashGen, FastH3 and cache samples retained; full matrix pending | + +## Development record + +- Baseline and active PRs inspected; no overlapping H3 PR is open. Both prior + attention directions and NVENC are merged at the declared base. +- All eight V100s were idle at preflight. Every GPU launch must acquire the + native GPU lease and recheck actual processes; idle observations do not + reserve a device. +- The attachment and complete accepted plan were read. No previous execution + turn existed: the preceding turn produced a plan, not implementation evidence. +- Next: rebuild owned baseline extensions, prove the current numerical route, + then implement shared prepared linear execution and residual sharding. + +### First implementation checkpoint + +- Fresh SM70 extensions built from the integration base; immutable paths and + SHA256 recorded in `baseline-binaries.json`. The baseline runtime is a clean + `git archive` of that SHA. Generated extension aliases are confined to the + artifact bootstrap, preserving package ABI names without copying stale H3 + binaries from another task. +- `prepared-linear-v2.log`: 28 checks passed in the leased GPU batch (24 GPU + checks and four CPU layout checks: shared GEMM, scaling, LoRA, activation, + column-major plans). Prepared LoRA equals normal execution + bitwise for original FP16 and W8A16, scales 0/0.75/-0.5, including explicit + unrotated inputs beside rotated base operands. +- `candidate-tp2.log` and `candidate-tp4.log`: the complete distributed block + case passes on all respective ranks. Cases include both attention backends, + original and INT8 adapted blocks, consecutive blocks, padding and residuals + exceeding FP16 range. This is not complete-model quality evidence. +- `core-regressions-v2.log`: 133 CPU adapter/config/API/conversion checks passed, + one skipped and eight GPU cases deselected. Earlier export/import failures + were missing video-extra dependencies in the new isolated environment. + The environment now matches the retained native model-component versions; + Torch remains 2.10.0+cu128. Both dynamic VAE classes import successfully. +- Failed setup paths retained: combining two distributed fixture lifetimes in + one torchrun produced a Gloo rendezvous error after the first case passed; + separate torchrun invocations resolve it. The first prepared run caught a + removed compatibility export; the alias is restored and all 28 tests pass. +- The first full baseline stopped during VAE construction because the new + environment lacked `diffusers`; no generated result or speed is claimed. + Align the video dependencies with the retained native environment before retry. +- Active LoRA retains the normal unrotated gather in residual sharding. Its + prepared row projections are enabled, and explicit dual-basis column input + is supported, but reducing adapter input/gather overhead remains work. +- Original checkpoint loading rejects non-finite FP16 conversion, while keeping + wide-range FP32 parameters intact. Tests cover BF16 overflow, infinity and NaN. +- Current full baseline: `baseline-720p-v2`, frozen integration Python source, + fresh owned kernels, FA backend, no residual sharding, four-step v1.2 adapter, + no FP16 weight cache, full warmup then one captured quality/timing request. + This is a baseline acquisition, not the formal three-run acceptance. + +### Complete four-step quality localization + +The baseline completed on GPUs 0-3: one full warmup, then 66.206537 s denoise, +47.303750 TFLOP/s/card, four calls and 3,131,817,518,663,680 useful FLOPs/rank. +The first request took 83.139160 s denoise including fresh-cache startup; it is +excluded. Complete-request memory peaked at 18.162186 GiB allocated/card. +The baseline captured all latents and unencoded RGB/PCM; capture overhead is +outside denoise but included in VAE/end-to-end wall time. + +`candidate-720p-residual` took 62.624739 s / 50.009271 TFLOP/s/card after one +warmup. **Rejected:** video/audio latent relative L2 is 0.2957225/0.0565235, +video PSNR 28.312 dB and SSIM 0.878025. Finite media and a 5.41% time reduction +do not satisfy the accepted quality gate. This is not a qualified speedup. + +`candidate-720p-prepared-only` disables residual sharding while retaining every +new prepared/dense execution change. Complete video/audio latents match the +baseline bitwise; pre-encoding RGB matches exactly (SSIM 1.0), and all audio +gates pass. Its single cold request is a quality control, not a speed baseline. +This isolates the full-sampling regression to the residual route rather than +the newly generalized prepared matrix/LoRA path. + +The residual implementation now uses the same full FP32 all-reduce as the +replicated path, then selects local residual rows. This intentionally gives up +the unqualified reduce-scatter communication saving. The distributed oracle +is tightened from a tolerance to bitwise equality; complete-model revalidation +is still required. Tensor-parallel GEMM, local normalization/residual ownership +and adapter support remain active. No >80 result or human acceptance exists. + +`candidate-720p-exact-reduction` validates code `9623a9adb2`: complete video +and audio latents are **bitwise equal** to the frozen mainline baseline, all +124 pre-encoding frames match (SSIM 1.0), and all audio numerical gates pass. +See `exact-reduction-quality.json`. This was a single no-warmup quality request +(66.954238 s denoise); it is not a qualified performance comparison. + +The tightened bitwise block oracle also passes on both TP2 ranks +(`exact-reduction-tp2-v3.log`) and all TP4 ranks (`exact-reduction-tp4.log`). +These results establish the sampled four-step route's numerical preservation, +not all adapters/partitions or human audiovisual acceptance. Draft PR #571 +contains the implementation and remains Draft. + +### Original floating-weight control and host memory + +The first original checkpoint run was interrupted by a host restart and has no +result. The retry (`baseline-720p-original-v2`) was stopped during startup after +host usage reached 187 GiB with less than 1 GiB available and heavy swapping. +No denoise performance was recorded. Original-weight CPU masters per rank are +17,252,698,560 bytes DiT, 13,770,235,360 bytes encoder, 10,415,484,160 bytes video +VAE and 605,306,340 bytes audio VAE, before loader/allocator overhead. + +Two complete controls therefore used identical **pageable** CPU masters for +all four components, with otherwise unchanged sampling, original BF16 checkpoint +converted through the existing FP16/FP32 runtime, LightX2V four-step v1.2 and FA: + +| Quality control | Denoise seconds | Useful TFLOP/s/card | +| --- | ---: | ---: | +| Frozen mainline, ordinary residuals and row layout | 69.044838 | 45.359844 | +| General prepared path, column layout, exact residual sharding | 66.426983 | 47.147453 | + +These are single no-warmup quality requests; the timing difference is **not an +accepted speedup**. Full video/audio latents match bitwise, all 124 decoded RGB +frames match (SSIM 1.0), spectral cosine is effectively 1 and RMS ratio is 1. +See `baseline-720p-original-pageable`, `candidate-720p-original-column` and +`original-column-quality.json`. These prove preservation of frozen mainline for +this original-weight workflow, not independent official or human acceptance. + +The corresponding native deployment option is `host_weight_pin_memory=False` +or `--disable-host-weight-pinning`, applied to the DiT, encoder and both VAEs. +It preserves values, aliases and layouts and changes host residency/transfers +only. Automatic host/GPU memory budgeting remains further work. The recorded +complete controls used an artifact-local stager option before the native flag +was added; they must not be presented as full native-flag generation evidence. +`host-memory-cpu-v1.log`: 43 config/VAE/workflow/API checks pass. +`host-memory-gpu-v1.log`: two pinned/pageable alias-preserving repeated transfer +checks pass. The full original-weight native-flag run is still pending. + +The separate workflow-metrics branch replaces fixed 49-call validation with +actual sigma intervals and step/block counts. Its results and formal FA/FI +comparisons will be recorded independently. No >80 configuration is qualified. + +### Current shared storage and campaign checkpoints + +The later native-flag run at `324f2463c78fb0f69d1546c817e67f3802e52342` +also enables explicit shared VAE host storage. Full original-weight Light4 +latents, RGB and PCM match the original column-weight control bitwise. +TP4 VAE mappings total 11.02 GB physical PSS instead of four physical replicas, +and the engine cleans up its owned files. See [shared host weights](SHARED_HOST_WEIGHTS.md). +This supersedes the pending native-flag status above; warmed request speed +and automatic memory budgeting remain unqualified. + +The accounting branch corrects duplicate column-LoRA A projections. For the +four-step W8A16 sample, useful rank-zero FLOPs are now +3,103,284,010,387,456; the equal-work 80 TFLOP/s budget is approximately +38.79 seconds. Earlier TFLOP/s values in this historical record use the old +numerator and must not be mixed with corrected results. + +The kernel branch's full warmup plus three unprofiled query-128 requests take +59.743807 / 59.748563 / 59.748666 seconds, median 51.939 TFLOP/s/card and +CV 0.003793%. Light4 and Light8 complete native-control comparisons pass +bitwise for video/audio latents, RGB and PCM. Neither establishes independent +official-model or human quality acceptance. The >80 gate still fails. + +The variants branch completes native FlashGen, FastH3 Dense and true VSA +generation. VSA's full FP32 selected-key attention control produces final +video/audio latent L2 errors 0.390670/0.088636, PSNR 24.277 dB and SSIM +0.769774: it is unqualified despite small local operator errors. Unmodified +official FastVideo Triton cannot compile FP16 inputs on this V100 setup; +an independent compatible official runtime remains required. + +Both request caches complete original-weight TP4 cached/lossless/cached +generation at 256x448, 107 internal frames and 49 intervals. TeaCache records +5/0/5 hits and Cache-DiT 34/0/34, with repeated latents/RGB bitwise and PCM +within the declared gates. These are lifecycle checks, not full official +quality or the primary performance matrix. Subsequent branch documents hold +the detailed records; their APIs are not all present in this common-base PR. + +### Original projection conversion audit + +`float-matrix-conversion-audit.json` scans every original attention/MLP matrix +that the native model executes in FP16, including the token refiners: 208 +matrices and 20,038,287,360 values per partition. Protected FP32 normalization, +AdaLN and embedding parameters are outside this conversion. The audit retains +each source tensor's SHA256, shape/dtype, conversion error and underflow count. +Checkpoint revision is `42ed227ee7df40d41602854ae760620d6eb651fe`. + +| Partition | Aggregate relative L2 | Maximum matrix relative L2 | Nonzero values underflowed to zero | Overflow / non-finite input | +| --- | ---: | ---: | ---: | ---: | +| FL2VA | 9.000060e-10 | 4.256286e-9 | 4,055 | 0 / 0 | +| Ref2VA | 9.004251e-10 | 4.208590e-9 | 4,058 | 0 / 0 | + +This explicitly records FP16 subnormal-range loss; conversion is not claimed +to preserve every source bit. No value clipping is performed. The loading +guard still rejects FP16 overflow. These weight-only measurements do not +replace the final latent, video and audio quality gates. + +The official eight LightX2V files are now present and hash-verified against +repository revision `2f015e66b37c585cea9dc4ae6f1850ea8788e742`. Native header +inspection accepts all eight with the recipe's task families, four/eight +intervals, five/nine sigma points, flow shifts and alpha values; see +`official-variants/all-lightx2v-inventory.json`. Download/header validation is +not full GPU generation or acceptance for the remaining adapter versions. + +The kernel branch's Ref2VA eight-step mixed image/video/audio control also +passes frozen-native latent/RGB/PCM comparison bitwise, with 69,325 valid DiT +tokens. Its single captured denoise is 366.799932 seconds; no formal three-run +or independent official quality qualification is claimed. diff --git a/docs/design/minimax_h3/LAYER_WEIGHT_OFFLOAD.md b/docs/design/minimax_h3/LAYER_WEIGHT_OFFLOAD.md new file mode 100644 index 0000000000..c7d7a7fd4e --- /dev/null +++ b/docs/design/minimax_h3/LAYER_WEIGHT_OFFLOAD.md @@ -0,0 +1,84 @@ +# Layer weight offload for capacity-limited H3 execution + +The original model's attention/MLP FP16 matrices alone occupy 40,076,574,720 +bytes globally. A 32-GB V100 cannot hold that complete DiT, even before FP32 +protected weights and activations. The explicit `--weight-offload layer` +deployment option stages individual DiT blocks and Qwen text/vision layers. +`H3Config(weight_offload="layer")` is the Python equivalent. The default +remains whole-component staging for the existing TP4 execution path. + +The plan partitions the existing immutable host snapshot by actual storage +ownership. Private block storage moves to the GPU before its forward and is +released afterwards. Storage shared with other blocks or outer consumers stays +resident for the component context, preserving offsets, strides and aliases. +The first DiT block's normalization and AdaLN projection stay resident because +cache decision probes may call them outside the block's forward. Adapter +buffers follow their owning block. No weight precision or reduction changes. + +Both normal and failed forwards release block storage. The context removes its +hooks on exit, rejects overlapping use and prevents a whole-component load +during layer staging. GPU copy streams/events are reused sequentially from the +original snapshot; block boundaries synchronize. Allocator retention stays +bounded. This mode does not support a fixed persistent FP16 weight-cache list. + +All transfers inside sampling remain in the full denoise denominator. +`dit_layer_weight_staging` / `dit_layer_weight_offload` and the corresponding +encoder fields are host boundary measurements, not isolated GPU transfer +durations. Loading includes host submission and any blocking copies; offload +waits for pending H2D and compute before releasing storage, with no D2H copy. +DiT loading also includes the initial resident setup outside denoise. These +fields must not be summed again into request latency. This is a capacity option, +not an automatic or >80-TFLOP/s configuration. + +## Development validation + +Python 3.12.13, Torch 2.10.0+cu128, CUDA 12.8.93, V100 SXM2 32GB. +Evidence: `/data/minimax-h3/sm70-general-20260909/`. + +- `layer-staging-cpu-v1.log`: 22 ownership, lifetime, failure cleanup, + configuration and service checks pass, three GPU cases deselected. +- `layer-staging-gpu-v1.log`: FP16 and FP32 models with column-major weights, + adapter buffers, cross-component aliases and 65-row tails reproduce + whole-resident results bitwise across three load/forward/offload cycles. + +`layer-staging-cpu-v2.log` additionally passes 35 checks, including explicit +resident consumers and both CLI modes, with three GPU cases deselected. + +Source `aea5a0fc35` completes a TP1 original-weight LightX2V four-step request +on one V100, using column-major matrices and pageable host masters. The first +capacity check uses the smallest legal temporal extent (22 frames) on a +256x448 canvas; it is not a primary performance workload. The full request +passes basic media validation and peaks at 15,473,571,328 allocated GPU bytes. +Denoise takes 126.274887 seconds and the complete request 152.031041 seconds. +Recorded DiT weight loading takes 123.212966 seconds including the initial +3.769427-second resident setup outside denoise. Pageable staging accounts for +most host boundary time in this run. + +The full contract, source/binary hashes, raw latents/RGB/PCM and stages are in +`/home/ymzx/h3-sm70-artifacts-20260909/runs/layer-offload-tp1-original-minimal/`. +The matching pinned-host run completes denoise in 53.840319 seconds and the +request in 67.370839 seconds, with the same GPU peak. Final video/audio latents, +all 22 pre-encoding frames and decoded PCM match the pageable run bitwise; +SSIM is 1 and audio RMS ratio is 1 (`layer-tp1-pinned-quality.json`). Recorded +DiT load/offload boundaries take 35.957164/15.261954 seconds; the latter includes +waiting for asynchronous copies and compute, not device-to-host weight traffic. +Both are captured cold capacity checks, not formal warmed speed measurements. + +The full TP2 original-weight DiT also exceeds 32-GB/card capacity in component +mode: loading fails before denoise after the allocator's bounded retry. No +timing or output is claimed for that failed run. The matching layer policy +completes a full request with exact residual sharding, peaking at +15,449,646,080 allocated bytes/card. Denoise is 90.431030 seconds and request +latency 110.419366 seconds. A separate layer-mode control disables residual +sharding while preserving TP2, original weights, adapter, seed and sampling. +Final video/audio latents, all 22 RGB frames and decoded PCM match bitwise +(`layer-tp2-quality.json`). Both controls use shared pageable VAE masters. +The ordinary-residual cold request takes 166.482792 seconds denoise; variable +host paging and captures preclude a formal performance comparison. + +Artifacts are `layer-offload-tp2-original-component/`, +`layer-offload-tp2-original-layer/` and +`layer-offload-tp2-original-layer-ordinary/` under the same `runs/` root. +Larger TP1/TP2 shapes remain pending. Neither these residency/sharding +comparisons nor the operator tests establish independent official full-model +quality or performance acceptance. diff --git a/docs/design/minimax_h3/NATIVE_VARIANTS.md b/docs/design/minimax_h3/NATIVE_VARIANTS.md new file mode 100644 index 0000000000..a88f9cb143 --- /dev/null +++ b/docs/design/minimax_h3/NATIVE_VARIANTS.md @@ -0,0 +1,120 @@ +# Native H3 variant integration + +Current VSA development and its explicitly agreed 31.3-second stage are +tracked in [VSA_QUALITY_SPEED.md](VSA_QUALITY_SPEED.md). Independent official +GPU-kernel validation and >80 throughput remain future objectives; the +existing full FP32-control quality failure has not been reclassified. + +This branch extends the prepared execution and workflow accounting stack +(#571 / #578) with the pinned official FastH3 VSA algorithm. Performance above +80 useful TFLOP/s/card, independent official full-sampling quality and human +audiovisual review remain incomplete. + +## Sparse execution contract + +Select `attention_backend="FASTVIDEO_VSA"` with an explicit official FastH3 +VSA adapter and `vsa_topk=64` in `H3Config`. The equivalent native CLI options +are `--attention-backend FASTVIDEO_VSA --fastvideo-vsa-topk 64`. Existing video +HTTP APIs consume the configured engine; no additional frontend is required. + +The loader validates the artifact's declared inventory and injects all 50 main +DiT compression projections before normal TP loading and CPU staging. The +token refiner remains dense, as in the release. Sparse artifacts cannot run on +a dense backend, dense artifacts cannot silently switch to VSA, and FastH3 +still requires original floating weights, T2VA, four intervals and its own +per-modality shifted DMD2 schedule. Fused adapters remain immutable per engine. + +The geometry and learned compression follow `fastvideo_vsa.py` from Omni +`7be014bce6374f06c95b703763bdbac4c6198f31`. Prefix segments form independent +64-token chunks; target video forms 4x4x4 tiles. Prefix queries select every +valid key block, while video queries select all prefix blocks plus top-k video +blocks. Edge tiles retain their actual lengths. The learned compressed output +is added using the same pooled FP32 scores and explicit gate; no sigmoid or +replacement with a dense attention call is introduced. + +The shared `sm70_sparse_attention.block_sparse_attention` interface accepts +pre-tiled FP16 operands, a boolean block map and int32 valid block lengths. +A warp compacts each map row in ascending block order. The SM70 kernel executes +only those blocks and masks all nonterminal padding holes. Its CUDA extension +is included in the SM70 wheel targets and also supports the existing CUTLASS +source-development workflow. No external FastVideo kernel package is required. + +## Accounting and isolation + +Per-layer and per-step records retain actual selected block/token-pair counts, +pooled compression FLOPs and avoided attention FLOPs. The numerator includes +selected valid pairs, actual compression and the gate projections. Padding +and unselected pairs do not inflate throughput. The acceptance evaluator checks +the declared geometry against every layer and step, keeps algorithm savings +separate, and still requires a full warmup and three unprofiled requests. + +Geometry caches contain immutable indices only, keyed by shape, prefix and +device. Scores, pooled activations and gate values are created for each call. +This implementation does not enable TeaCache or Cache-DiT; those request-state +policies and AUTO selection remain separate campaign work. + +## Validation record + +Python 3.12.13, Torch 2.10.0+cu128, CUDA toolkit 12.8.93, V100 SXM2 32GB. +Evidence root: `/data/minimax-h3/sm70-general-20260909/`. + +- `sparse-gpu-v1.log`: 8 operator checks pass against an independent masked + FP32 reference, with NaN-poisoned padding holes, different batch/head maps, + dense-equivalent arithmetic and invalid-input rejection. +- `vsa-geometry-cpu-v1.log`: 13 geometry, top-k and learned-gate checks pass. +- `vsa-integration-cpu-v1.log`: 169 existing API/workflow/adapter checks pass, + one GPU case skipped. +- `vsa-fusion-cpu-v1.log`: 33 checks pass, one GPU case skipped; this includes + all 50 gate assignments, TP1/TP2/TP4 native loaders and host snapshots. +- `vsa-dense-and-accounting-gpu.log`: 37 checks pass, including the rebuilt + dense specialization and a real sparse DiT block with independent FLOP + arithmetic, bitwise hook transparency and hook cleanup. +- `vsa-accounting-cpu-v1.log`: 71 checks pass, two GPU cases deselected. The + evaluator rejects invented padding, block counts, compression and savings. + +Full native FlashGen and FastH3 Dense five-second GPU generation completed on +the preceding combined source `5195f31b8d`, with original weights and native +pageable host masters. Basic media checks pass; those cold captured requests +are not speed or independent quality acceptance. + +Full native VSA generation on source `61c57e36ae` also completed with the +official data-free adapter (SHA256 +`42dc502a2078f166c396a1fa75f29728d1844363652d345d5ef3e2b444ed6470`). +This TP4 cold captured request used the same 1280x736, five-second, seed-42 +prompt. Complete denoising took 37.387443 seconds; the slowest rank per step +took 12.211990, 8.383499, 8.374612 and 8.412264 seconds. Peak allocated memory +was 20,987,568,640 bytes/card. All four ranks passed the native sparse-work +validator, and basic video/audio checks passed. Results and raw captures are +in `/home/ymzx/h3-sm70-artifacts-20260909/runs/fasth3-vsa-720p-native/`; +`fasth3-vsa-native-summary.json` records the compact audit. + +Actual useful throughput was 45.2679–45.2773 TFLOP/s/card. Approximately +1.484e15 skipped attention FLOPs/card are reported separately, not credited +to throughput. This single cold request is not the warmup-plus-three-run +performance gate. Independent official full-sampling comparison and human +review remain pending; a sampled frame contains several ducks despite the +prompt specifying one, so basic media checks do not establish prompt fidelity. +No complete workflow is yet qualified for AUTO or the >80 target. + +## Full sparse math control + +`vsa-fp32-control` runs two requests in one native engine: the sparse kernel, +then the pinned official Omni frontend with independent selected-key FP32 +QK/softmax/PV. Initial video/audio rows and text inputs are bitwise equal; the +native request's final latents also reproduce the first bringup bitwise. The +official frontend alone matches the port bitwise in operator checks. + +Small local kernel differences amplify during full sampling. Against this +FP32 math control, final video/audio latent relative L2 is 0.390670 / 0.088636, +video PSNR is 24.2768 dB and SSIM is 0.769774. Audio spectral cosine is 0.992129 +and RMS ratio is 1.001159. The required latent and video gates fail; do not +qualify this configuration or relax thresholds. Full results are retained in +`vsa-fp32-control-quality.json` at the evidence root. + +This control substitutes the sparse math operation and keeps the native model; +it is not an independent official full-model or official GPU-kernel reference. +The unchanged FastVideo Triton source at `a943220c115228ade5d57b3bab9a6a87fd600a10` +fails to compile FP16 inputs because probabilities are cast to BF16 while V is +FP16. `official-triton-sm70-probe.json` retains that failure. Matching the +official kernel's intended precision and locating full-sampling amplification +remain required quality work. diff --git a/docs/design/minimax_h3/README.md b/docs/design/minimax_h3/README.md index 62442e55b1..ffcd0d02f8 100644 --- a/docs/design/minimax_h3/README.md +++ b/docs/design/minimax_h3/README.md @@ -1,12 +1,11 @@ # Native MiniMax H3 (development) -The latest optional residual-sharded FlashInfer development run completes -39 frames and 20 denoise updates in 61.538397 seconds, with unchanged -video/audio latents and fresh MP4. See -[FLASHINFER_LOCAL_ROTATION.md](FLASHINFER_LOCAL_ROTATION.md) for local ConvRot, -the current bottleneck trace, validation and remaining gates. -The numerical-drift overlap experiment is rejected; see -[FLASHINFER_OVERLAP.md](FLASHINFER_OVERLAP.md). +The retained dense workflows use FlashAttention-V100; see +[CURRENT_STATUS.md](CURRENT_STATUS.md) for measured coverage. The subsequent +[VSA quality/speed stage](VSA_QUALITY_SPEED.md) reaches 30.990756 seconds on the +native path with failing independent quality, or 60.353224 seconds with passing +primary numerical checks in the acceptance-only FP32 diagnostic. No VSA +configuration passes the combined stage or enters default/AUTO selection. This is an in-progress native integration. Full checkpoint video quality and four-card 80 useful TFLOPS acceptance are not yet established. The control log @@ -83,6 +82,15 @@ its downloaded model directory. Frozen revisions and port licenses are recorded in the model package's `UPSTREAM.md`. FFmpeg and FFprobe must be on `PATH` for reference-video/audio processing. +For hosts that cannot hold every component's pinned CPU copy, use +`--disable-host-weight-pinning` (Python: `host_weight_pin_memory=False`). This +keeps pageable immutable masters for the DiT, text encoder and both VAEs. +It preserves weights, layouts and computation, while allowing the OS to page +inactive components; transfers and request startup can be slower. The original +TP4 deployment holds about 157 GiB of CPU masters before loader temporaries and +allocator caches, so all-pinned startup exhausted the tested 188 GiB host. +See [GENERAL_SM70.md](GENERAL_SM70.md) for the complete floating-weight control. + Use `--image first.png --keyframe-indices 0`, `--image last.png --keyframe-indices -1`, or two `--image` arguments with `--keyframe-indices 0 -1` for FL2VA. A Ref2VA instance uses `--partition ref2va` and the matching transformer @@ -123,13 +131,13 @@ The measured 39-frame/20-update cache list and native cuBLASLt configuration are documented in [FLASHINFER_TO50.md](FLASHINFER_TO50.md), including exact commands, output comparison and the still-incomplete 50-second target. -For the experimental FlashInfer TP4 INT8 FL2VA route, add -`--residual-sequence-parallel` to shard FP32 residual rows and reduce TP -communication. It defaults off and changes floating-point reduction order. -The unchanged 39-frame/20-update denoise measures 66.863312 seconds; automatic -checks pass, while five-axis human quality review and the 50-second target -remain open. See [FLASHINFER_RESIDUAL.md](FLASHINFER_RESIDUAL.md) for output -differences, GPU tests, current hardware counters and rollback. +The earlier experimental FlashInfer residual route changed reduction order; +its historical results and unaccepted differences are recorded in +[FLASHINFER_RESIDUAL.md](FLASHINFER_RESIDUAL.md). Current residual sharding uses +the ordinary full FP32 all-reduce before selecting local residual rows. It +preserves the complete four-step control bitwise, and does not claim the old +reduce-scatter communication saving. It remains explicitly enabled with +`--residual-sequence-parallel` and requires wider quality/performance validation. The subsequent probability-tile layout change reduces this same optional route to 65.804661 seconds and preserves its video/audio outputs bitwise. @@ -148,16 +156,27 @@ For INT8 MLPs, the native path combines FP32 SiLU/product evaluation and power-of-two FP16 input preparation in one kernel. The following projection restores the scale in FP32 before the normal TP reduction. This removes the large FP32 activation intermediate without lowering arithmetic precision or -adding a persistent cache. CPU, unquantized and FP32-input paths keep their -existing implementation. +adding a persistent cache. Prepared inputs now also support original floating +weights and active LoRA. LoRA restores the first A projection's row scale before +preparing B, retaining both rounding boundaries and pre-reduction delta addition. +The shared FP16 GEMM and preparation operators live outside the H3 model package. +Original floating projections can opt into `--fp16-weight-layout column`; the +default remains `row`. Layout preparation preserves logical weight coordinates +and storage size. Full-workflow qualification is tracked in [GENERAL_SM70.md](GENERAL_SM70.md). For the measured TP4 FL2VA INT8 development workload, optionally add `--residual-sequence-parallel` to `video generate` or `video serve`. This keeps FP32 residual rows sharded across ranks and gathers normalized FP16 inputs at the existing precision boundary. Both native SM70 attention backends support -this option; it rejects BF16, Ref2VA, adapters and non-TP4 configurations. It defaults to -false. Normalized FP16 rows are now rotated locally before all-gather, avoiding +this option. It now accepts original weights, Ref2VA and matching adapters on +TP2/TP4; TP1 uses ordinary execution. It defaults to false. Runtime guards still +require one request, FP32 residuals, aligned metadata and no Ulysses hooks. +For projections that can consume only a rotated input, normalized FP16 rows +are rotated locally before all-gather, avoiding four copies of the same ConvRot work while preserving gathered bits and GEMMs. +An active adapter also needs the unrotated input, so the ordinary gather is +retained there until a validated combined preparation removes that dependency. +The following measurements predate the generalized route and do not qualify it: The native FlashAttention 39-frame/20-update run measures 58.190385 s and 6.195001125 GiB peak denoise allocation/card, versus 62.804019 s and 6.566735744 GiB without the flag, with zero persistent cache in both cases. @@ -167,6 +186,12 @@ the sharded route changes FP32 reduction order from the default and its human quality review remains pending. Omit the flag to restore the default. These are short development measurements, not the 243-frame acceptance run. +The generalized four-step/full-canvas check exposed substantial latent drift +from reduce-scatter's FP32 addition order. Residual sharding now performs the +same full all-reduce as the replicated path before selecting each rank's rows. +This preserves the reference sum order and does not claim a communication-volume +reduction. The old reduce-scatter speed figures above are not qualified outputs. + Outputs include `video.mp4`, original decoded `audio.wav`, `run.json`, sampled `nvml.jsonl`, `quality.json` and frame screenshots. Automatic checks do not replace the five-axis human quality review. Useful TFLOPS use actual local diff --git a/docs/design/minimax_h3/SHARED_ATTENTION.md b/docs/design/minimax_h3/SHARED_ATTENTION.md new file mode 100644 index 0000000000..281fd34c71 --- /dev/null +++ b/docs/design/minimax_h3/SHARED_ATTENTION.md @@ -0,0 +1,47 @@ +# Shared SM70 attention interface + +`vllm.model_executor.layers.sm70_attention.noncausal_attention` exposes both +native SM70 dense attention implementations without importing the H3 model +package. H3's dense facade uses the same function. Historical extension ABI +names and H3 loader exports remain compatible with existing wheels; a default +call still invokes the original four-argument entrypoint. + +```python +from vllm.model_executor.layers.sm70_attention import noncausal_attention + +output = noncausal_attention( + q, k, v, scale=128**-0.5, backend="FLASH_ATTN_V100", query_tile=128 +) +``` + +Inputs are CUDA FP16 BSHD MHA tensors with head dimension 128; accumulation and +softmax remain FP32. No model name, quantization label or adapter identity is +required. FlashAttention supports different Q/K lengths, strided storage and +explicit query/key tiles. FlashInfer requires matching lengths and receives +contiguous inputs. Invalid scales, including values overflowing FP32, fail +before native loading. Unsupported dtype, hardware and shapes fail at the +CUDA entrypoint. There is no silent BF16/FP32 conversion or backend substitution. + +Callers own masks, suffix padding and sparse geometry. The H3 facade still +slices valid tokens before dispatch and restores zero suffix padding; VSA +keeps its separate prefix, selected-block and learned-gate implementation. +AUTO qualification is a separate unfinished campaign requirement. + +## Validation + +Environment: Python 3.12.13, Torch 2.10.0+cu128, CUDA 12.8.93, V100 SXM2 32GB. +No CUDA source or binary changes accompany this interface extraction. + +- `shared-attention-gpu.log`: 12 GPU checks pass. Both native backends and both + FA query geometries preserve direct-entrypoint results bitwise for BSHD + shapes `(1, 1537, 24, 128)` and `(2, 65, 8, 128)`, including strided inputs + and a later increase in the online maximum. Sampled FP32 relative L2 is + below 0.001. Additional checks cover unequal Q/K lengths, dtype rejection, + H3 poisoned suffix padding and CUDA Graph replay with new input values. +- `shared-attention-cpu.log`: 18 scale, deployment, service and residency + checks pass; nine GPU cases are deselected. + +These tests exercise non-H3 DiT operator shapes, not a second complete model. +They establish interface preservation, not a new end-to-end speed result or +independent official quality acceptance. Raw logs and binary hashes are under +`/data/minimax-h3/sm70-general-20260909/`. diff --git a/docs/design/minimax_h3/SHARED_HOST_WEIGHTS.md b/docs/design/minimax_h3/SHARED_HOST_WEIGHTS.md new file mode 100644 index 0000000000..a1c0e6ca72 --- /dev/null +++ b/docs/design/minimax_h3/SHARED_HOST_WEIGHTS.md @@ -0,0 +1,72 @@ +# Shared immutable host VAE weights + +`--share-host-vae-weights --disable-host-weight-pinning` enables an explicit +TP2/TP4 memory policy. The Python equivalent is +`H3Config(share_host_vae_weights=True, host_weight_pin_memory=False)` through +`H3Engine`. TP1 uses ordinary storage. The option is off by default. + +The native video and audio VAEs keep replicated FP32 host masters even though +their computation is distributed. In the measured original-weight TP4 setup, +each rank retains 10,415,484,160 video-VAE bytes and 605,306,340 audio-VAE bytes. +Sharing one physical replica instead of four can remove 33,062,371,500 bytes +of duplicated host storage. This does not change GPU weights, arithmetic, +reference encoding or the VAE compute topology. + +Each engine creates an owned temporary directory under `/dev/shm`. Rank zero +writes aligned storage groups and checksums; every worker validates its own +weights/layouts against the snapshot and maps the bytes privately. Private +mappings share clean physical pages and isolate accidental CPU writes. +Shapes, strides, storage offsets, mixed-dtype aliases and empty tensors are +preserved. GPU load/offload still copies the exact original storage bytes. +The parent removes its directory only after its workers stop, including failed +startup. Unrelated shared-memory files are not modified. + +The policy requires pageable masters: ordinary PyTorch pinning would copy the +mapping into a separate pinned allocation per rank, defeating the memory +saving. CUDA host registration and automatic policy selection are not added. +The measured VAE snapshot needs approximately 11.02 GB of tmpfs capacity plus +small alignment/metadata overhead. Space is reserved before mapping writes. + +## Evidence and current limits + +The original-weight Cache-DiT lifecycle trace separates allocation and copies. +First DiT device allocation takes 0.18–0.19 s, while copies take 27.7–43.2 s +with 142,577–268,365 major faults/rank. Later copies take 6.5–8.2 s. The memory +policy targets duplicated host storage and cold paging; it is not evidence of +higher denoising TFLOP/s. + +`shared-host-cpu-v1.log`: 19 host-storage and native service checks pass, +including nonzero offsets, transposed and strided views, mixed-dtype aliases, +private writes, mismatched replicas, corrupt snapshots and owned cleanup. +`shared-host-gpu.log`: exact storage roundtrips pass across three GPU +load/offload cycles. + +Source `324f2463c78fb0f69d1546c817e67f3802e52342` additionally completes a full +TP4 original-weight LightX2V four-step request: column-major FP16 weights, +FP32 residual sharding, frozen FA binaries, 1280x736/124 internal frames, +seed 42, five sigma points and no persistent FP16 weight cache. All final +video/audio latents, 124 RGB frames and decoded PCM match the prior original +column-weight control bitwise (`shared-host-quality.json`). PSNR is infinity, +SSIM and audio RMS ratio are 1; all numerical preservation gates pass. + +The four workers' VAE mappings total 44,083,544,064 RSS bytes but only +11,020,886,016 PSS bytes, with zero private mapped pages at startup. Each rank +maps the same two files and accounts for one quarter of their physical pages. +This verifies sharing of the approximately 11.02 GB replica and eliminates +the three redundant replicas; mapping alignment adds small overhead to the +raw tensor sizes above. The engine removes its owned directory after shutdown. + +The single captured cold request takes 66.365689 s in denoise and 126.846697 s +overall, with 21,979,466,240 peak allocated GPU bytes/card. DiT staging still +takes 8.44–26.53 s across ranks after startup paging. This is not a warmed +speed comparison or proof that all host paging has been eliminated. Full +independent official quality, human review and performance qualification +remain pending; AUTO is not enabled. + +`shared-host-summary.json` records the physical mapping totals and stages. +The complete contract, source/binary hashes, startup smaps snapshot and raw +media are retained under +`/home/ymzx/h3-sm70-artifacts-20260909/runs/shared-host-original-light4/`. + +Evidence root: `/data/minimax-h3/sm70-general-20260909/`. Runtime: +Python 3.12.13, Torch 2.10.0+cu128, CUDA 12.8.93, V100 SXM2 32GB. diff --git a/docs/design/minimax_h3/SM70_EPILOGUES.md b/docs/design/minimax_h3/SM70_EPILOGUES.md new file mode 100644 index 0000000000..59c4f05b2b --- /dev/null +++ b/docs/design/minimax_h3/SM70_EPILOGUES.md @@ -0,0 +1,372 @@ +# Shared SM70 projection epilogues + +This development branch is stacked on the common prepared execution (#571) +and workflow accounting (#578) branches. No configuration has passed the +campaign's >80 useful TFLOP/s/card and complete official quality gates. + +The separate [FI register-probability update](FLASHINFER_REGISTER_PROBABILITY.md) +now also has full native media preservation and a formal warmup-plus-three +result of 49.795 useful TFLOP/s/card. It remains below the FA query-128 result +and the campaign target. Both use the same shared projection interface. + +## Measured problem and implementation + +The matching four-step FA denoise profile spends 6.890 seconds in miscellaneous +elementwise kernels, including FP32 row-scale restoration, adapter additions +and output casts. Attention, GEMM and communication separately consume +31.439, 16.041 and 8.453 seconds. These are profiled service diagnostics; +the unprofiled audited baseline is 47.091839–47.091855 useful TFLOP/s/card. + +`sm70_diffusion.fp16_linear_add` prepares the projection input at the existing +FP16 boundary and keeps GEMM accumulation in FP32. A shared CUDA epilogue +restores row scales with a rounded FP32 multiplication, adds the scaled delta +with the same FP32 fused multiply-add as PyTorch, and stores FP16 or FP32. +It writes only the requested output slice. Delta/output and scale/output +storage overlap is rejected, including differently typed views of one buffer. + +The H3 adapter uses this interface when output hardware, precision and layout +support it and adapter slices are disjoint. Overlapping contributions retain +the original FP32 buffer until the final cast. Wheels without the new extension +ABI retain the ordinary path. No quantization label or adapter name controls +dispatch. The first LoRA projection still uses unrotated inputs; row-parallel +increments still join the partial result before the collective. + +## Development evidence + +Environment: Python 3.12.13, Torch 2.10.0+cu128, CUDA toolkit 12.8.93, +V100 SXM2 32GB. All GPU tests use an owned native lease. Evidence root: +`/data/minimax-h3/sm70-general-20260909/`. + +- `epilogue-integration-v1.log`: 43 checks pass, including explicit comparison + with the old unfused adapter path for original/W8A16 bases, FP16/FP32 output, + prepared/ordinary inputs, negative scales, wide intermediates, untouched + slices, overlap fallback, unaligned storage and CUDA Graph replay. +- `epilogue-cpu-v1.log`: 94 adapter, workflow and strict acceptance checks pass. +- `epilogue-micro.json`: paired postprocessing-only medians for 34,560 rows, + output width 5,376: three FP16 QKV slices 7.524352 -> 2.015232 ms; one FP32 + row projection 4.562944 -> 3.094528 ms. Results match bitwise. These timings + exclude GEMM and do not establish complete-request speedup. +- `epilogue-binaries.json` retains the first tested binary. The strengthened + alias check is in `epilogue-binaries-v3.json`; validation is recorded separately. + +The final alias guard passed on GPU0 (`epilogue-alias-v3-gpu0.log`). An earlier +GPU4 attempt was refused by an existing lease before starting the test. + +## Complete four-step control and measurements + +Source `7ea8908d83827dd8d82c34ba6a60b2beaa8057d6`, TP4, LightX2V four-step v1.2, +W8A16, FA, exact residual sharding, no persistent FP16 weight cache, and the +original 1280x736/124-frame internal canvas for the five-second sample: + +- `epilogue-quality.json`: final video/audio latents and all pre-encoding RGB + frames match frozen mainline bitwise; video SSIM 1, spectral cosine 1 and RMS + ratio 1. This is numerical preservation, not independent official acceptance. +- `epilogue-720p-three-runs/performance.json`: one full warmup (64.373997 seconds + denoise) followed by three complete requests without profiler or captures. + Denoise times are 62.245159 / 62.183194 / 62.162648 seconds; CV 0.056387%. + Every rank reports median **49.905491–49.905510 useful TFLOP/s**. The declared + >80 gate fails and remains incomplete. +- Complete request times are 84.363265 / 91.992465 / 88.620830 seconds. Peak + allocation remains 19,501,498,880 bytes per card. Exact source/kernel hashes, + per-step records and NVML samples are retained beside each run. +- The 62.183194-second median is 5.64% below the audited original FA baseline + (65.898529 seconds). This measures the **combined** prepared/residual/epilogue + changes, not an isolated attribution to this CUDA epilogue. A separate matched + profile is necessary for attribution. + +Independent official reference, full audiovisual review, other adapters and +the complete shape/TP matrix remain required. No AUTO selection is qualified. + +## Explicit attention query geometry + +`attention_query_tile=128` / `--attention-query-tile 128` opts into a 128-query +FlashAttention-V100 CTA. The default remains 64 and retains the previous call +ABI. Both sizes use the same 32x64 warp arithmetic and key-tile selection. +The option applies independently of model weights and adapters; other attention +backends reject this explicit tiling option rather than ignoring it. + +The separate prototype retains exact outputs at nine boundary lengths and +the actual 34,551-token Q/K/V capture. Full four-step video/audio latents also +match frozen mainline bitwise (`q128-profile-quality.json`). A matching pair of +full-denoise profiles is retained in `epilogue-profile-breakdown/` and +`q128-epilogue-profile-breakdown/`; profiler timings are not acceptance results. +The public kernel/CLI implementation additionally passes 69 GPU tail, storage, +cross-attention-length and graph checks; 53 strengthened comparisons also +require exact equality between query geometries and reject invalid query tiles. + +The native option at source `6b39f1c23ca6834e9beead89b4cdf57548101e97` passes +the complete latent/RGB/PCM comparison (`query-tile-quality.json`): both final +latents and all 124 frames match frozen mainline bitwise; SSIM 1 and all audio +gates pass. The same explicit configuration completed one full warmup plus +three unprofiled requests (`query-tile-720p-three-runs/performance.json`): + +| Configuration | Median denoise seconds | Useful TFLOP/s/card | Denoise CV | +| --- | ---: | ---: | ---: | +| Audited original FA baseline | 65.898529 | 47.091839–47.091855 | 0.055668% | +| Prepared/residual/epilogue, query tile 64 | 62.183194 | 49.905491–49.905510 | 0.056387% | +| Same path, explicit query tile 128 | 59.748563 | 51.939038–51.939057 | 0.003793% | + +The last three denoise times are 59.743807 / 59.748563 / 59.748666 seconds. +Complete request times are 81.805001 / 82.780855 / 84.272880 seconds and peak +allocation remains 19,501,498,880 bytes/card. The 128-query setting reduces +median denoise by 3.92% relative to the same prepared 64-query configuration, +and the combined change reduces it by 9.33% relative to the original baseline. +These measurements cover one five-second workflow only. **The >80 gate still +fails; official reference, human review and the wider matrix remain pending.** + +## Eight-step numerical preservation + +The same explicit query-128/prepared/residual/epilogue configuration at +`bd5e1898265eb1783fcc413de321125230fbe594` also completed a matched TP4 +LightX2V eight-step FL2V v1.0_768p comparison with frozen mainline `4f19ef7`. +Both runs use W8A16, seed 42, the same five-second request, nine sigma points, +flow shift 6 and audio flow shift 3. The official adapter SHA256 is +`9b0efe3613b43a84e30febaa43af27432ea9d0711eac7bba904b2556b175f6d4`. + +`light8-720p-quality.json` passes every declared numerical gate: both final +latents, all 124 RGB frames and decoded PCM match bitwise. This extends +preservation evidence beyond the four-step adapter, using the same shared +operators without an adapter-specific dispatch exception. It remains a frozen +native control, not independent official-model acceptance. + +The captured cold requests took 141.160054 and 121.541891 seconds in denoise; +complete request times were 174.840862 and 203.651622 seconds respectively. +These single captured runs have different staging conditions and no full +warmup, so they do not establish a formal performance result or an overall +request speedup. Source hashes, binary manifests and commands are recorded +in `light8-pair.json`. The eight-step >80 gate remains unmeasured. + +An additional 16-row warp experiment (`attention-warp16/hypothesis.json`) +was rejected at compilation: CUTLASS Volta MMA requires a multiple of its +interleaved tile shape. No GPU run or production change followed. Supporting +that geometry requires new MMA and accumulator iterators, not another +configuration-only benchmark of the rejected shape. + +## Mixed-reference eight-step control + +Source `2063b09f2d75c4a63af3f90a3b7803744ffc6e02` completes Ref2VA with the +official eight-step v1.0_768p adapter, W8A16 Ref2VA base, seed 42 and one image, +one 2.5-second reference video plus one standalone audio reference. The video +start time is zero. The output remains 1280x736/124 internal frames for the +five-second request. This control uses 69,325 valid DiT tokens and a 10,273-token +Qwen presentation, exercising mixed reference indices and padded residual rows. + +`ref8-mixed-quality.json` passes all gates against frozen native mainline: +video/audio latents, all RGB frames and PCM are bitwise equal; PSNR infinity, +SSIM 1, RMS ratio 1. Candidate settings include prepared execution, exact +residual sharding, query tile 128 and explicit shared pageable VAE host weights. +The frozen control uses its ordinary query-64 path and pinned host masters. +Host residency changes byte ownership/transfers, not GPU arithmetic. + +Single captured cold denoise times are 421.173489 seconds for the frozen +control and 366.799932 seconds for the candidate; request times are +586.794155 and 454.461091 seconds. Peak GPU allocation is 19,623,684,096 and +19,615,279,104 bytes/card respectively. The candidate reports corrected useful +throughput 52.919225–52.919231 TFLOP/s/card. These runs lack full warmup and +three measurements and have different host staging policies, so they do not +qualify either performance or attribution to an individual optimization. +Independent official quality, audiovisual/reference review and other reference +combinations remain pending. Raw contracts and media are retained in +`/home/ymzx/h3-sm70-artifacts-20260909/runs/ref8-mixed-720p-{baseline,candidate}/`. + +Two additional CTA-barrier coalescing candidates preserve bitwise edge and +34,551-token results. The second passes 20 synccheck and racecheck geometries +with zero errors/hazards. Their paired operator gains are only 0.3–0.6%, so +neither is retained or promoted to a full-request performance claim. Evidence: +`attention-barrier-coalesce/`, `attention-barrier-coalesce-v2/` and the +`barrier-v2-*.log` files under the campaign root. + +Further isolated operand probes are also rejected: + +- `attention-fi-p-reuse/` interchanges PV loops in the retained FlashInfer + Q128/K64 kernel to share a loaded P fragment across four output fragments. + Nine boundary shapes and the real 34,551-token input remain bitwise equal, + but paired latency regresses from 183.285767 to 185.328644 ms. +- `attention-rescale-identity/` uses a warp-uniform check to skip accumulator + multiplication when every applicable online-softmax scale equals one. + The same numerical cases remain bitwise equal; 150.328323 to 149.099518 ms + is less than 1% and includes clock variation, so it is not retained. + +Neither probe changes production kernels or establishes a full-denoise gain. +Their source, build logs, hashes and actual-input timings remain in the +campaign artifact root to prevent repeating unchanged experiments. + +## Larger canvas and duration compatibility + +Source `9d2489fc4e2f32ea500ec13478bd68ef9000c1cc` completes two additional +TP4 W8A16 LightX2V four-step requests with prepared execution, exact residual +sharding, query tile 128 and shared pageable VAE masters. Both use seed 42, +the original paper-boat prompt, five sigma points and flow shifts 6/3. + +| Requested shape | Actual frames | Denoise seconds | Request seconds | Peak GPU allocation bytes/card | Useful TFLOP/s/card | +| --- | ---: | ---: | ---: | ---: | ---: | +| 1344x768, 243 frames | 243 | 203.431920 | 254.138019 | 24,341,115,904 | 52.603224–52.603229 | +| 1344x768, 15 seconds | 362 | 401.026288 | 469.032456 | 28,777,495,040 | 53.561244–53.561247 | + +The 15-second request resolves to the model's 362-frame aligned output; it is +not claimed to be an exactly 15.000-second encoded clip. Both complete native +media validation and strict actual-work checks, and remove their owned shared +weight directories after shutdown. Full captures, source/binary manifests, +per-rank stages and NVML samples are retained under +`/home/ymzx/h3-sm70-artifacts-20260909/runs/official-243-frame/` and +`boundary-15-second/`; `large-canvas-summary.json` summarizes the evidence. + +These first captured requests establish shape and memory compatibility only. +They have no matched quality reference, full warmup or three post-warmup +measurements. Both are below 80 and remain unqualified. The independent +official reference and human review gates are also pending. + +## Remaining FL2V Turbo versions + +Source `be89a26d1c` completes four more matched frozen-mainline comparisons on +TP4 at the same 1280x736/124-frame internal canvas for a five-second request. +All use W8A16, seed 42, pageable host masters and no fixed FP16 weight cache. +The candidate uses prepared execution, exact residual sharding, shared LoRA +epilogues and query tile 128. The frozen `4f19ef7` control uses its ordinary +query-64 path. Each artifact retains its official alpha and flow shift. + +| Official artifact | Intervals / sigma points | Video shift / alpha | Candidate denoise seconds | Request seconds | +| --- | ---: | ---: | ---: | ---: | +| FL2V four-step v1.0_768p | 4 / 5 | 6 / 128 | 61.263482 | 96.285706 | +| FL2V four-step v1.1_768p | 4 / 5 | 6 / 128 | 61.018722 | 96.491105 | +| FL2V four-step v0.1 | 4 / 5 | 12 / 8 | 61.000133 | 94.938178 | +| FL2V eight-step v1.0 (non-768p) | 8 / 9 | 12 / 8 | 120.478234 | 156.924670 | + +For every pair, final video/audio latents, all 124 pre-encoding RGB frames and +decoded PCM match bitwise. SSIM and RMS ratio are 1. Strict per-rank workload +validators pass, including actual intervals, block counts and duplicate-work +exclusion. Peak allocation is 19,501,498,880 bytes/card for the four-step cases +and 19,502,023,168 bytes/card for eight-step. Single-request useful throughput +is approximately 50.65–51.52 TFLOP/s/card, below 80. + +These are captured cold quality controls, not warmed three-run performance. +Together with the existing v1.2 four-step and v1.0_768p eight-step controls, +all six official FL2V Turbo artifacts now have a complete W8A16 T2VA numerical +preservation result. Original-weight and keyframe combinations remain separate +pending coverage. The subsequent Ref2V four-step control is recorded below. +No independent official quality or human acceptance is inferred. + +Evidence: `remaining-turbo-pairs.json`, `remaining-turbo-summary.json`, +`light4-v{10,11,01}-720p-quality.json`, `light8-v10-non768-720p-quality.json` +and the corresponding captured runs under the campaign's artifact root. + +## Four-step mixed references and complete adapter inventory + +Source `c69cfc7024460e314e79a0bba37a3b736340bc6e` completes the official +Ref2V four-step v0.1 adapter with a W8A16 Ref2VA base, one image, one +2.5-second video and one standalone audio reference. The video starts at zero; +seed 42, five sigma points, video/audio shifts 12/3 and alpha 8 are retained. +The candidate uses the same general FA query-128 path as the other adapters. + +`ref4-mixed-quality.json` passes all declared numerical gates against frozen +native `4f19ef7`: final video/audio latents, all 124 RGB frames and PCM match +bitwise, PSNR is infinite, SSIM is 1 and RMS ratio is 1. Spectral cosine is +0.9999999999999695. Both native generations complete. The candidate's strict +actual-work checks pass; its complete denoise is 184.563686 seconds, request +269.210367 seconds, and peak allocation 19,613,711,360 bytes/card. Corrected +useful throughput is 52.585556-52.585562 TFLOP/s/card. + +The frozen control takes 214.108869 seconds denoise and 321.250580 seconds +request with the same peak allocation. These are captured cold requests with +different host VAE sharing policies; they are not formal speed acceptance. +The old control script did not embed Git metadata. The separate +`baseline-source-audit.json` verifies all 2,483 tracked `vllm` files in the +frozen archive against `4f19ef7`, without changing the historical contract. + +All eight official LightX2V artifacts now have complete native numerical +preservation evidence: six FL2V adapters on T2VA and both Ref2V adapters with +mixed references. This is not the full task/weight/reference cross-product, +independent official quality, human review or >80 acceptance. Exact paths, +source identities and timing scope are retained in `ref4-mixed-pair.json`, +`ref4-mixed-summary.json` and the campaign result index. + +## First, last and both-frame controls + +Source `4ed70419e7f42c6e9f4625f7fa92cf8dea2126d8` also completes all three +FL2VA keyframe modes using the official Light4 v1.2_768p adapter and W8A16, +with one immutable engine per implementation. The candidate uses the same +general FA query-128/prepared/exact-residual/epilogue path. First and last +images are the retained frames 0 and 123 of the campaign sample, selected +with indices `[0]`, `[-1]` and `[0,-1]` respectively. + +| Constraint | Frozen-control denoise seconds | Candidate denoise seconds | Candidate request seconds | Candidate useful TFLOP/s/card, minimum | +| --- | ---: | ---: | ---: | ---: | +| First frame | 77.275887 | 66.419303 | 106.123582 | 50.697486 | +| Last frame | 71.537050 | 64.852263 | 102.142091 | 51.922501 | +| First and last frames | 77.299085 | 69.615145 | 108.431893 | 52.304580 | + +Every pair passes complete numerical preservation: final video/audio latents, +all 124 decoded frames and PCM match bitwise; SSIM and RMS ratio are 1. +Peak allocation is unchanged within each pair: 19,512,957,952 bytes/card +for a single image and 19,522,796,544 bytes/card for both images. + +These are captured requests with different first-use state, reference lengths +and host VAE sharing policies. They establish full execution and native +preservation, not a formal performance comparison or independent verification +of reference fidelity. Original-weight/other-adapter keyframes and the full +legal Ref2VA combination matrix remain pending. + +Evidence: `keyframe-pairs.json`, `keyframe-summary.json`, and +`keyframe-{first,last,first-last}-quality.json`. The candidate contract retains +shared-operator, model and video-source hashes plus clean Git provenance. + +## Original floating FlashGen and FastH3 Dense controls + +Both four-interval T2VA variants now have complete native comparisons with +frozen mainline `4f19ef7a20db60bb0685e599bd3f4dd156202eed`. Each pair uses +original floating FL2VA weights, its matching official adapter, seed 42, +TP4, flow shifts 12/3 and the same 1280x736/124-frame internal canvas. +Candidate source and individual file hashes are retained in each contract; +the frozen source audit covers all 2,483 tracked package files with no mismatch. + +Both final video/audio latents, all 124 pre-encoding RGB frames and PCM match +bitwise for both variants. SSIM and RMS ratio are 1; spectral cosine exceeds +0.99999999999996. This establishes native numerical preservation for these +configurations, not independent official-model or human quality acceptance. + +| Variant | Baseline denoise / request seconds | Candidate denoise / request seconds | Candidate useful TFLOP/s/card | Candidate peak bytes/card | +| --- | ---: | ---: | ---: | ---: | +| FlashGen four-step | 73.290174 / 162.679376 | 59.488178 / 121.697365 | 51.620900 | 20,781,940,224 | +| FastH3 Dense data-free | 61.574164 / 121.787763 | 56.223567 / 114.632281 | 54.124949 | 20,023,082,496 | + +These are captured cold controls. Both arms use pageable host weights; the +candidate also shares host VAE weights and enables prepared column weights, +exact residual sharding, shared epilogues and explicit FA query tile 128. +The measurements combine these changes and do not isolate a kernel effect. +Neither variant has completed a full warmup plus three unprofiled requests. + +Evidence: `original-variant-pairs.json`, `original-variant-summary.json`, +`flashgen-original-quality.json`, `fasth3-dense-original-quality.json` and +`baseline-source-audit.json` in the campaign artifact root. Complete media and +contracts reside in the corresponding `*-original-{baseline,candidate}` runs. +The [campaign table](CAMPAIGN_RESULTS.md) separates these diagnostic timings +from formal acceptance measurements. Every configuration remains unqualified. + +## Original floating weights without an adapter + +The complete default-sampling T2VA native control uses original floating +weights, no LoRA, 50 sigma points and 49 actual updates. Both requests retain +seed 42, the same prompt and 1280x736/124-frame internal canvas. The frozen +`4f19ef7` control uses ordinary residuals and row-major floating weights; +the candidate at runtime source `570be8d407` uses prepared column-major +projections, FA query128 and explicit native peer rows. Shared host VAE +masters and pageable staging are recorded separately in the run contracts. + +All declared numerical gates pass: video/audio latents, all 124 pre-encoding +RGB frames and PCM are bitwise equal; SSIM and RMS ratio are 1. Denoise falls +from 725.819030 to 649.973214 seconds (10.449687%). Complete captured requests +are 835.355602 and 700.121402 seconds. Candidate actual-work validation passes +on all ranks and records 57.353041 useful TFLOP/s/card. Tracked allocation upper bounds +are 21,633,302,016 bytes/card for the baseline and 20,998,705,664 for the +candidate, including the candidate's raw IPC memory. + +These are single captured cold requests, not warmup-plus-three performance +acceptance. The combined result covers shared operators, residual layout and +host residency; it does not isolate an Attention-only or host-policy-only +speedup. Independent official reference, continuous human audiovisual review +and >80 acceptance remain incomplete. Four sampled baseline frames were +inspected, which does not replace those quality gates. + +Evidence: `original-base-pair.json`, `original-base-summary.json`, +`original-base-no-lora-original-quality.json` and the corresponding original +base run directories under the campaign artifact roots. diff --git a/docs/design/minimax_h3/VSA_QUALITY_SPEED.md b/docs/design/minimax_h3/VSA_QUALITY_SPEED.md new file mode 100644 index 0000000000..4e17771e0d --- /dev/null +++ b/docs/design/minimax_h3/VSA_QUALITY_SPEED.md @@ -0,0 +1,387 @@ +# VSA quality and 31.3-second stage + +FastH3 VSA Data-Free on original floating weights and V100 TP4 has no +configuration that passes the joint 31.3-second quality/speed stage. The native +FP16 sparse path reaches a formal 30.990756-second median but fails independent +FP32 quality. The acceptance-only exact FP32 path passes the primary numerical +gates at a formal 60.353224-second median; human audiovisual review and the +extended coverage matrix remain incomplete. Neither result authorizes AUTO +promotion. Official GPU-kernel validation and >80 useful TFLOP/s/card remain +separate unfinished objectives. No quality threshold or official sampling/ +selection rule is relaxed. + +## Main integration + +On 2026-09-10 the user authorized integrating #583 and its dependencies +(#571, #578 and #581) into `main`. The synchronization base is +`24220ca0eb4a02b2376cf48feb430bb4d5c5c3d7`. Source integration does not change the +numerical/performance or human-review status above and does not register the +FP32 diagnostic in default/AUTO selection. Historical benchmark records retain +their original source and binary identities. + +The merge retains main's mmap host weights, media progress and physical-device +mask handling alongside prepared FP16 execution, layer offload, shared VAE host +storage, work counters and VSA layout/diagnostic changes. Tests combine mmap +masters with layerwise adapter/alias roundtrips on the GPU. No numerical CUDA +kernel is changed by this synchronization. + +The merged candidate passes the complete CPU video suite: 359 passed and +251 GPU/opt-in checks skipped. The focused leased-V100 integration suite passes +73 checks, including mmap/layer staging, aliases, VSA geometry/layout, strict +sparse validation and work counters; all 24 QK/RoPE GPU checks also pass. +Ten existing GPU-only cases now skip explicitly when CUDA is unavailable. +Source-integration logs and GPU lease records are retained under +`/home/ymzx/h3-sm70-artifacts-20260909/vsa-merge-main-20260910/`. +These integration checks do not replace complete model quality or performance +acceptance. All applicable pre-commit checks pass before source publication. + +## Frozen baseline and diagnosis + +The source baseline is `970c5fb3f86d59440a3431e53029e98f8d690778`, which merges +the existing shared-kernel dependency into the owned VSA worktree. Commands, +source/binary hashes, input captures and diagnostic artifacts are retained in +`/home/ymzx/h3-sm70-artifacts-20260909/vsa-quality-speed-20260910/`. + +The native host policy now matches the retained Dense configuration: pageable +shared VAE masters, prepared floating columns and explicit native TP4 peer +rows with a 4 GiB communication budget. A complete primary capture preserves +the previous VSA final video/audio latents, all pre-encoding RGB frames and +PCM bitwise. Its 34.895545-second denoise is a cold diagnostic including input +capture, not formal performance or a quality repair. + +A separate profiled request records a 32.039245-second denoise NVTX span on +the slowest rank. Sparse Attention accounts for 7.343837 seconds and GEMM for +15.466638 seconds. The initial parser classified peer `reduce_rows` in other +kernels; its 3.654961 seconds must be included in communication, alongside +1.727143 seconds of other collectives. Profiler start/stop overhead is outside +this NVTX span but inside native stage timing; neither time is an unprofiled +acceptance result. No NCU occupancy/utilization claim is made. + +The fixed-input diagnostic shows probability rounding contributes to error, +but a compensated-PV prototype only reduces actual operator relative L2 from +0.000203199 to 0.000135391. It fails its numerical admission criterion and is +not installed or run as a full-sampling candidate. Further localization keeps +QKV, selected blocks and gates fixed; FP32 output and QK/PV diagnostics are +separate from production precision. + +## Deferred work accounting and internal validation + +Dynamic selected-pair and selected-block counts stay as int64 device scalars +while layers execute. The request-owned counter retains their producing +tensors and step/layer association, then copies all counts to CPU together +after complete denoise. Public result fields remain Python integers and are +reconciled across steps, layers and totals. Counter completion remains inside +the complete-denoise timer and before the final TP barrier. Closing or failing +a request removes hooks and releases pending tensors. + +H3-owned geometry and a nonempty mask produced by the official top-k/prefix +construction use a private prevalidated CUDA entrypoint. Device, layout, +dtype, shape, alignment and index-limit checks remain. The general sparse +operator still validates block-size values and rejects empty query rows. +Older wheels use its checked entrypoint until rebuilt; no user flag bypasses +validation. Kernel arithmetic is unchanged. + +Validation so far: 73 CPU checks pass, 2 GPU checks skip in the CPU run; +23 leased-GPU checks pass, including real dense/sparse DiT work accounting, +tail/padding controls, strict public input rejection, exact public/private +output equality and absence of host scalar reads in the private entrypoint. +The rebuilt sparse kernel retains 215 registers and zero spills. Nine further +checks pass under both Compute Sanitizer memcheck and synccheck with zero +reported errors. The +complete primary capture preserves final video/audio latents, pre-encoding +RGB and PCM bitwise. Selected blocks, valid pairs and useful FLOPs also match +at every rank, layer and step. Its 34.458830-second cold captured denoise is +diagnostic only. This preserves the old native output, whose independent +FP32 quality still fails; it does not establish engineering quality. + +## Full-sampling amplification diagnosis + +A separate artifact-only prototype resets the FP32 PV accumulator for each +64-key block and adds a scaled low probability component. Fixed-input relative +L2 improves from 0.000203199 to 0.000046140, but this still misses the prototype's +fivefold numerical improvement criterion. It is not installed. A full run was +used specifically to localize amplification, not as acceptance or performance +evidence. The centered-exponent variant adds no useful numerical improvement. + +The independent FP32 reference was run again and reproduces the original final +video and audio latents bitwise. With identical initial tensors, the block-local +candidate's video latent errors after the four steps are 0.003839, 0.016306, +0.056423 and 0.369510; final audio error is 0.050502. Both final latent gates fail. +Selection first differs in the second layer on all ranks. By the last layer of +the first step, 73.2--78.4% of query blocks have at least one changed selected +block. This is the fraction of affected queries, not the fraction of replaced +keys. Fixed-reference-map controls separate continuous arithmetic error from +dynamic-selection amplification. They cannot be shipped. Fixed reference maps +still give final video/audio errors of 0.189996/0.016187 with the native kernel +and 0.175657/0.014001 with compensated block-local PV. Exact FP32 prefix queries +with dynamic video selection give 0.401763/0.086767. All three fail; neither +fixed routing nor prefix precision alone solves the problem. + +The native frontend using the same FP32 sparse operator matches the independent +official frontend bitwise on real inputs, including the official transport-only +partner block. The block-local diagnostic before final FP16 rounding has relative +L2 0.000005753 against FP32 (prefix 0.000015489, video 0.000005443). A further +artifact isolates exact FP32 QK/softmax and compensated 64-key PV partials; no +production sparse-math change is admitted from these local numbers. + +## Exact layout fusion + +One kernel gathers Q/K/V directly into their final padded tiles. Another +combines learned compression with the sparse output and restores the original +row order, preserving separate FP16 multiply/add rounding, including overflow. +The original pooling, top-k, sparse traversal and work reductions are retained. +An older extension or strided/unaligned inputs use the existing Python layout. + +Only geometry indices are cached across requests. A denoise context owns QKV +scratch storage separately for each device/stream and releases it on success or +failure. Nested contexts restore their parent. Every returned output is freshly +allocated, so a later layer cannot overwrite an earlier result. + +The artifact prototype preserves real primary outputs and useful counts bitwise. +The corrected comparison includes work counting on both sides: median 41.969666 +ms baseline versus 37.358593 ms fused over seven paired operator measurements. +This is an operator result, not complete-denoise acceptance. The installed source +passes 30 GPU tests covering layout, sparse math and geometry; its eight layout +checks also pass memcheck and synccheck with zero errors. The complete primary +capture preserves final video/audio latents, all 124 pre-encoding RGB frames and +PCM bitwise; every rank/layer/step work count is identical. Its captured denoise +is 32.989041 seconds, which is diagnostic only. Independent FP32 engineering +quality remains incomplete; VSA is not promoted into default or AUTO selection. + +## Stage timing evaluation + +After separate full native benchmarks for VSA and Dense, run: + +```bash +python -m vllm.video.vsa_acceptance --vsa VSA_BENCHMARK_DIR \ + --dense DENSE_BENCHMARK_DIR --output stage-performance.json +``` + +The tool reuses strict same-session warmup, three-request and rank/step/layer +accounting validation. It additionally requires the primary dimensions, TP4, +top-k 64, official FastH3 four-step sigma positions and matched request, GPU, +host-weight, VAE, communication and export settings. Only backend, query tile, +top-k and adapter path may differ between the two algorithm controls. It checks +the slowest rank's median denoise <=31.3 seconds, CV <=5%, and a strictly lower +complete-request median than Dense. The future 80-TF criterion is explicitly +reported separately. Passing this timing tool does not establish weight +identity, numerical/human quality, other-shape coverage or official hardware +validation. + +## Matched formal timing: speed gate passes, joint stage incomplete + +Source `450f9b9bc6458e31ff83150f33b79df31a89732d` with the retained +`native-layout-binaries.json` completed one full warmup and three consecutive +unprofiled, uncaptured single requests per backend. Both controls use the same +original H3 weights, request, TP4 GPUs, pageable shared VAE masters, prepared +FP16 columns, 4 GiB native peer-row communication budget and libx264 export. +Each uses its corresponding official FastH3 Data-Free adapter. VSA uses top-k +64 and query tile 64; Dense FA uses query tile 128. + +| Measurement | VSA, existing FP16 sparse math | Dense FA | +| --- | ---: | ---: | +| Denoise run 1 (s), slowest rank | 30.948464 | 52.886819 | +| Denoise run 2 (s), slowest rank | 30.999521 | 52.873835 | +| Denoise run 3 (s), slowest rank | 30.990756 | 52.858062 | +| Median denoise (s) | 30.990756 | 52.873835 | +| Denoise CV | 0.071956% | 0.022239% | +| Median complete request (s) | 68.157817 | 87.426512 | +| Peak allocation bound incl. raw IPC (GiB/card) | 20.843828 | 19.554280 | +| Effective model TFLOP/s/card, lowest rank median | 54.611483 | 57.553944 | + +Per-step medians of the maximum rank GPU time are 7.730810, 7.748937, +7.747062 and 7.753540 seconds for VSA, versus 13.183612, 13.219172, +13.240814 and 13.227989 seconds for Dense. These event timings do not replace +the complete-denoise wall-clock criterion, which includes counter completion. +Avoided sparse work is not included in useful throughput. + +`native-layout-stage-performance.json` passes all three timing checks. The +underlying sparse arithmetic still fails the independent FP32 reference: +video/audio latent relative L2 is 0.390670/0.088636. This timing pass cannot be +combined with the quality pass of a different, slower kernel. No configuration +has yet passed the joint stage. `vsa-layout-formal-telemetry.json` and its Dense +counterpart retain per-device clocks, power, memory and utilization samples; +these span the complete campaign and are not denoise-only utilization figures. + +## Exact FP32 sparse CUDA candidate + +The compensated-PV experiment with exact FP32 QK/global softmax still fails +full sampling (video/audio latent errors 0.367458/0.074789), despite a fivefold +local improvement. Admission now requires real-input FP32 parity before another +full sample; a small local relative error is not sufficient evidence. + +Sequential FP32 FFMA QK and PV match cuBLAS FP32 on seven real query subsets. +A directly indexed CUDA implementation avoids gathered K/V copies, compacts +selected blocks in ascending order, preserves prefix-dense/video-sparse queries +and uses a global FP32 softmax. Its 64-by-64 output-tile variant matches all +609 blocks of the actual primary operator bitwise, including five boundary, +extreme-score and poisoned-padding cases. Six small cases pass memcheck and +synccheck with zero errors. + +The complete dynamic-selection seed-42 sample also matches the frozen reference +video/audio latents and all decoded RGB frames bitwise. PSNR is infinite, +SSIM is 1, audio spectral cosine is 0.999999999999945 and RMS ratio is +0.999999998643185. This is a numerical pass only: human review remains pending. +Captured denoise is approximately 109.6 seconds, so this kernel fails speed. + +The 64-by-64 operator takes approximately 407 ms. Nsight Systems attributes +158.81 ms to QK, 215.65 ms to PV and 32.37 ms to global softmax. Shared-memory +padding (411 ms) and a 4-by-4 register tile (437 ms) preserve bitwise output but +are slower and rejected. A CUTLASS SIMT candidate preserves the five boundary +cases and full primary operator bitwise at approximately 233 ms. Six small +cases pass memcheck and synccheck with zero errors. Nsight Systems attributes +75.76 ms to QK, 116.46 ms to PV and 32.46 ms to global softmax. Its complete +seed-42 sample preserves initial video/audio noise, text tensors, every step +and final latents bitwise. All 124 RGB frames are identical; audio spectral +cosine is 0.999999999999945 and RMS ratio is 0.999999998643185. Captured denoise +is 74.545147 seconds and complete request is 225.097778 seconds; the latter +includes cold staging and capture and is not a matched formal measurement. +The primary numerical gate passes, but speed and human review do not. No +precision candidate changes the runtime backend or sampling algorithm. + +Remaining acceptance work: one configuration satisfying both quality and speed, +seeds 43/44, 243-frame and 15-second boundaries, TP1/TP2 compatibility, representative +Dense/LightX2V/Ref2VA output regressions, and human audiovisual review. Official +original-hardware comparison remains explicitly deferred. + +## Reproducing the precision diagnostic + +The acceptance-only source is `benchmarks/kernels/h3_vsa_fp32.cu`. It has no +runtime backend registration. Build it separately with CUTLASS 4.4.2, CUDA +12.8.93, Torch 2.10.0+cu128, Python 3.12.13 and SM70: + +```bash +CUDA_HOME=/path/to/cuda-12.8 TORCH_CUDA_ARCH_LIST=7.0 MAX_JOBS=2 \ + .venv/bin/python benchmarks/kernels/benchmark_h3_vsa_fp32.py \ + --build-directory /path/to/diagnostic-build \ + --cutlass-root /path/to/cutlass-4.4.2 --output /path/to/build.json +``` + +After acquiring a GPU lease, run the regression checks against that exact +binary. For sanitizer checks, put `compute-sanitizer --tool memcheck +--error-exitcode 86` or `--tool synccheck --error-exitcode 86` before Python: + +```bash +H3_VSA_FP32_DIAGNOSTIC=/path/to/diagnostic-build/h3_vsa_cutlass_fp32.so \ + .venv/bin/python -m pytest -q tests/video/test_h3_vsa_fp32_diagnostic.py +``` + +The same benchmark accepts `--binary`, `--capture attention-input-rank-0.pt` +and `--output`. It compares the complete captured operator with independent +FP32 math before measuring it. Reports retain binary/source hashes and mark +operator measurements ineligible for complete-denoise acceptance. Supplying +an existing binary does not assert that it was built from the current source. +The diagnostic binary must never be substituted into a formal runtime timing +record without recording the actual operator override. + +The initial versioned diagnostic passes 20 leased-GPU regression checks, +including six mathematical/boundary cases and fourteen public input rejection +cases. The complete captured primary operator also matches FP32 bitwise. All +20 checks pass memcheck and synccheck with zero errors. Its build source hash is +`f21ae1df63032be02853c2968dacdc001511bb5729942c8b006e316a0771d8fc`. +The full quality capture used the arithmetic-identical artifact build +`cafdc331dfbc2218e6e982a6ce8b4d8c4cd562ad45e4106ca5968de4fb911349`; +the separately built versioned extension is +`ef4a9b8ebee7091b71951692d46313a1b59c322f1970d1925759a285c1d8cb96`. +The source body is unchanged apart from inlining includes, using the base +CUTLASS header directly, formatting and comments. These binary identities +must remain distinct in retained measurement records. + +## Batch independent FP32 queries + +The initial prefix PV launch has only 14 threadblocks on an 80-SM V100. Its +query-chunk limit was inherited from the gathered-K/V oracle, although the +CUDA implementation indexes shared converted Q/K/V directly. The diagnostic +now submits at most 32 independent query blocks together, with a 2 GiB limit +on each score/probability allocation and a matching CUDA grid-size guard. +Requests with even one query exceeding that allocation limit are rejected. +Each query retains its own compact selected indices and ascending key order; +there is no selection union, changed probability precision or split-K reduction. + +Source `cb33adbdd821a77a8f7f5575f0a6acb68fa5bb606d2b41dd02657d458e448e56` +builds binary +`8a0192009af23b25d559b91e9155385266f9574c6e92adba3570d67b539f5cd6`. +The exact versioned binary passes all 20 GPU, memcheck and synccheck checks, +including a 38-block case crossing the new query-batch boundary, with zero +sanitizer errors. The complete captured primary operator remains bitwise FP32, +with median 172.439 ms versus 232.740 ms before batching. + +The complete primary sample preserves the bytes of initial inputs, all four +step outputs and final video/audio latents against the frozen reference. All +124 RGB frames match; PSNR is infinite, SSIM is 1, audio spectral cosine is +0.999999999999945 and RMS ratio is 0.999999998643185. Human review is pending. +Captured denoise is 61.826728 seconds, complete request 164.385431 seconds, +and peak allocation bound including raw IPC is 24,884,304,896 bytes/card. +These captured times are diagnostic; the original and batched captures are +not a formal warm-session comparison. + +The separate one-warmup/three-request benchmark completed. Its provenance +explicitly names `CUTLASS_SIMT_FP32_BATCHED_DIAGNOSTIC` and the binary/source +hashes: the native sampler, work counters and timers are unchanged, but the +acceptance tool replaces the sparse operator. This remains an opt-in diagnostic +with no runtime/default/AUTO registration. The joint 31.3-second quality/speed +stage remains incomplete. + +## Exact FP32 formal result and remaining bottlenecks + +Source `3ae9087b28` and binary `8a019200...539f5cd6` completed a full warmup +followed by three consecutive unprofiled, uncaptured requests in one engine. +Every run records the actual diagnostic sparse operator override. The warmup +(90.458786-second denoise) is excluded from the following measurements. + +| Measurement | Run 1 | Run 2 | Run 3 | +| --- | ---: | ---: | ---: | +| Slowest-rank denoise (s) | 60.317440 | 60.355101 | 60.353224 | +| Complete request (s) | 163.296343 | 140.666471 | 96.415224 | +| Rank-0 text/media encoding (s) | 31.438246 | 25.516768 | 7.559961 | +| Rank-0 DiT staging/weight preparation (s) | 51.678684 | 25.174796 | 9.212433 | +| Rank-0 VAE (s) | 14.851771 | 14.801774 | 14.235778 | +| Rank-0 packaging (s) | 1.838711 | 1.963504 | 1.896965 | + +The denoise median is **60.353224 seconds**, CV **0.028716%**, and the complete +request median is **140.666471 seconds**. The stage evaluator passes CV but +fails both <=31.3 seconds and beating the retained Dense complete-request +median of 87.426512 seconds. The host preparation times are visibly unsettled +in the first two measured requests; the full-request difference must not be +attributed entirely to the sparse kernel. All three prescribed requests remain +in the result, and the fastest request is not substituted for their median. +The retained Dense and native VSA controls have matching configuration but were +measured earlier; background host state was not controlled across campaigns. + +The four per-step maximum-rank GPU-time medians are 15.071570, 15.072747, +15.091930 and 15.108280 seconds. The lowest-rank effective model throughput +median is **28.042311 TFLOP/s/card**. Peak allocation including raw IPC is +24,882,522,624 bytes/card (23.173655 GiB). Avoided work is reported separately, +not added to useful throughput. + +NVML samples are retained for the complete campaign. Restricting to samples +with GPU utilization >=90%, per-card median SM clocks are 1387, 1470, 1485 and +1470 MHz; median powers are 256.597, 261.452, 246.914 and 257.504 W. This selection +includes other GPU-active phases and must not be labeled denoise-only power or +utilization. The complete series and ranges are in the telemetry artifact. + +A separate Nsight Systems trace of the fixed captured operator records 20 QK +and 20 PV launches after batching, versus 84 each before. QK takes 73.09 ms, PV +60.24 ms and global softmax 32.08 ms. Conversion, score masking, pointer/index +preparation and output scattering account for about 7.76 ms. FP32 QK/PV and +normalization still dominate this operator; this is not a complete-denoise +category breakdown. Further layout-only changes cannot account for the large +remaining speed gap. No new Dense FA prototype or FI route was introduced. + +The final operator, quality and timing records are respectively +`cutlass-batched-versioned-capture.json`, +`cutlass-batched-full-quality-metrics.json`, +`vsa-fp32-batched-three-runs/`, and `fp32-batched-stage-performance.json` under +the retained artifact root. `review-media.json` records playable candidate and +reference MP4/WAV paths and hashes; both MP4 files pass complete FFmpeg decode. +They are H.264/AAC, 1280x736 at 24 fps, 5.175 seconds after the frozen request's +frame alignment. Assistant inspection of frames 0/61/123 notes several ducks +in the background, also present in the identical FP32 reference. Object-count +consistency and complete audiovisual quality require human review; the user +has been asked to review the playable sample. No human pass is recorded. + +Seeds 43/44, 243 frames, the 15-second boundary, TP1/TP2 complete compatibility, +and representative Dense/LightX2V/Ref2VA output regression remain **not +completed**. The implementation, numerical primary evidence and formal speed +failure are delivered in PR #583; the requested combined acceptance remains +**not completed**, with no default/AUTO promotion. diff --git a/docs/design/minimax_h3/WORKFLOW_METRICS.md b/docs/design/minimax_h3/WORKFLOW_METRICS.md new file mode 100644 index 0000000000..1d15aa299a --- /dev/null +++ b/docs/design/minimax_h3/WORKFLOW_METRICS.md @@ -0,0 +1,121 @@ +# H3 workflow performance measurements + +The native engine records the actual video/audio sigma sequences, selected +attention backends, useful unpadded sequence length and completed block calls. +The evaluator derives denoiser calls from the runtime intervals. LightX2V's +5/9 API points and FlashGen/FastH3's four-interval API therefore share the same +accounting without a hardcoded 49-call assumption. Normal request validation +still enforces each adapter's legal tasks and sampling settings. + +Each rank reports `denoise_steps` with useful FLOPs, completed DiT calls, +executed blocks, CUDA event time and CPU enqueue time. CUDA event spans include +dependent communication, stream waits and host feeding gaps; they are not sums +of individual kernel execution times. No per-step synchronization is added. +GEMM/attention/communication service attribution still requires a separate +profiler run. The complete synchronized denoise wall time remains the throughput +denominator, using the slowest rank for every rank's numerator. + +For a complete measurement, reuse a native run's `config` and `request` JSON: + +```bash +.venv/bin/python -m vllm.video.benchmark \ + --contract h3-contract.json --output h3-measurements-new +``` + +Run this command without a profiler. It acquires the native GPU lease, uses one +persistent engine, generates one complete warmup video/audio request and three +consecutive measured requests, and saves all four results plus +`performance.json`. The output directory must be new. NVML records memory, +power, clocks, temperature, utilization and running processes alongside each +request. Source revision/file hashes, Torch/CUDA versions and actual loaded H3 +kernel paths/SHA256 hashes accompany the results. Deployment and request JSON +retain model revision, checkpoint/adapter identifiers, sampling and TP settings. + +The evaluator requires a complete warmup from the same engine session, exact +configuration/shape consistency, all TP ranks and all scheduled steps. Missing +blocks, inconsistent per-step/per-layer/total work, profiled runs and quality +captures are rejected. For existing result files: + +```python +import json +from pathlib import Path +from vllm.video.metrics import evaluate_performance + +root = Path("h3-measurements-new") +warmup = json.loads((root / "warmup/run.json").read_text()) +runs = [json.loads((root / f"run-{i}/run.json").read_text()) for i in (1, 2, 3)] +report = evaluate_performance(runs, warmup=warmup) +``` + +Passing performance means every rank's three-run median is strictly above +80 useful TFLOP/s and complete-denoise CV is at most 5%. Peak memory (30 GiB +allocated/card budget), end-to-end latency and quality status are reported +separately. The selected shape and TP size remain in the report; passing a +shorter workload or TP1 does not complete the 243-frame TP4 or 15-second cases. +Quality must separately pass the accepted numerical and audiovisual review. + +This checkpoint counts **dense uncached execution**. It records zero sparse +blocks/cache hits for those actual routes and rejects sparse/cache descriptors +until the corresponding execution counter exists. This is not VSA, TeaCache or +Cache-DiT support. Padding, ConvRot, dequantization and repeated output rows remain +excluded from useful FLOPs. Algorithmic work savings must be accounted separately +when those variants are implemented. + +Column-parallel LoRA A projections have the same weights and input on every +rank. `dense_tp_lora_v2` attributes those input rows once across the TP group, +including uneven tails. Identical replicas are reported in +`redundant_denoise_flops`, `redundant_flops_by_layer` and each step's +`redundant_flops`; they never increase useful throughput. Row-parallel A consumes +distinct input shards, so its resulting partial B products remain useful work. +Legacy records without an accounting version are rejected until explicitly +audited; the measurement source/time must remain unchanged in any such audit. + +## Development evidence + +Integration base: `4f19ef7a20db60bb0685e599bd3f4dd156202eed` (`onecat/main`). +Owned branch: `codex/v100-h3-workflow-metrics-20260909-031302`. +Artifacts: `/data/minimax-h3/sm70-general-20260909/`. + +- `metrics-regressions-v2.log`: 81 CPU acceptance/service/workflow/API checks + pass. These include four/eight/base/DMD2 schedule semantics, TP1/2/4 and + rejection of incomplete warmup, missing steps and inconsistent work counts. +- `metrics-block-gpu-v2.log`: real SM70 DiT block test passes, with an independent + FLOP formula that excludes suffix padding, two measured calls, + bitwise output preservation and hook cleanup before another request. The first + GPU5 attempt was rejected by an existing lease and ran no GPU test; GPU0 was + acquired after the original-weight control released it. +- Full model instrumentation validation and performance measurements are pending. + No configuration has met the campaign's >80 TFLOP/s and full quality gates. + +### Complete four-step controls and first three-run baseline + +`metrics-720p-quality` validates runtime `82362a4312`, W8A16 + LightX2V4 v1.2, +TP4 GPUs0-3, internal 1280x736/124 frames for the five-second 720p sample. +All four ranks recorded four complete calls and 52 blocks/call. Full video/audio +latents and decoded RGB/PCM match frozen mainline bitwise; SSIM is 1.0 and the +audio numerical gates pass (`metrics-quality.json`). This is instrumentation +regression evidence, not independent official or human quality acceptance. + +`fa-720p-three-runs` completes one full native warmup and three unprofiled +requests with the same deployment/sampling on `82362a4312`. Denoise times are +65.880184, 65.898529 and 65.965556 seconds; CV is 0.055668%. End-to-end times +are 93.173206, 95.494415 and 88.343595 seconds. Peak allocation is +19,501,498,880 bytes/card. Warmup alone also retained fresh encoder/denoise input +tensors for later independent diagnostics; measured requests contain no captures. + +The original counter included identical column-A replicas. Its reported +47.524847 TFLOP/s/card is superseded by the explicit header-shape audit in +`fa-720p-three-runs/audited-counts/`. Original files, times and source fingerprints +are retained. Rank0 useful work is 3,103,284,010,387,456 FLOPs after excluding +28,533,508,276,224 replicated A FLOPs (0.9111% of the old numerator). +Audited median throughput is **47.091839–47.091855 TFLOP/s/card**, depending on +the uneven row tail. Performance remains below 80. The audit script checks all +retained per-layer counts against immutable LightX2V A/B header shapes; it is +not a new run of the revised runtime counter. + +`metrics-lora-count-cpu.log`: 49 CPU tests pass, including unique logical column +adapter FLOPs across TP1/2/4 and one-row/97-row/34551-row tails. Two additional +legacy/inconsistent-redundancy rejection cases pass in +`metrics-legacy-rejection.log`. The revised counter will be exercised in the +matching FlashInfer full measurements. Sparse/cache and complete official +workflow/quality/performance coverage remain open. diff --git a/flash-attention-v100/kernel/h3/fmha.h b/flash-attention-v100/kernel/h3/fmha.h index db35f0ece1..4890f24b88 100644 --- a/flash-attention-v100/kernel/h3/fmha.h +++ b/flash-attention-v100/kernel/h3/fmha.h @@ -449,12 +449,17 @@ struct H3FMHAKernel { int keys; int heads; float scale; + const int* block_indices = nullptr; + const int* block_counts = nullptr; + const int* block_sizes = nullptr; + int blocks = 0; static constexpr bool causal = false; }; /// Executes one GEMM - CUTLASS_DEVICE - void operator()(DirectParams const& params, SharedStorage& shared_storage) { + template + CUTLASS_DEVICE void operator()(DirectParams const& params, + SharedStorage& shared_storage) { auto& m_prime = shared_storage.m_prime; auto& s_prime = shared_storage.s_prime; [[maybe_unused]] auto& si = shared_storage.after_mm0.si; @@ -488,7 +493,8 @@ struct H3FMHAKernel { static_assert(kKeepOutputInRF); ElementOAccum* ptr_O_accum = nullptr; const int num_queries = - TileParams::num_queries(threadblock_idx, problem_size0); + Sparse ? params.block_sizes[threadblock_idx] + : TileParams::num_queries(threadblock_idx, problem_size0); auto createOutputIter = [&](int col) -> typename MM1::OutputTileIterator { using OutputTileIterator = typename MM1::OutputTileIterator; @@ -516,12 +522,22 @@ struct H3FMHAKernel { const int num_keys = TileParams::num_keys(threadblock_idx, problem_size0, params.causal); - for (int32_t iter_key_start = 0; iter_key_start < num_keys; - iter_key_start += kKeysPerBlock) { + const int block_row = group * params.blocks + threadblock_idx; + const int key_tiles = + Sparse ? params.block_counts[block_row] + : (num_keys + kKeysPerBlock - 1) / kKeysPerBlock; + for (int key_tile = 0; key_tile < key_tiles; ++key_tile) { + const int selected_block = + Sparse ? params.block_indices[int64_t(block_row) * params.blocks + + key_tile] + : key_tile; + const int iter_key_start = selected_block * kKeysPerBlock; + const int keys_remaining = Sparse ? params.block_sizes[selected_block] + : num_keys - iter_key_start; int32_t problem_size_0_m = cutlass::fast_min((int32_t)kQueriesPerBlock, num_queries); - int32_t problem_size_0_n = cutlass::fast_min((int32_t)kKeysPerBlock, - num_keys - iter_key_start); + int32_t problem_size_0_n = + cutlass::fast_min((int32_t)kKeysPerBlock, keys_remaining); int32_t const& problem_size_0_k = problem_size0.k(); int32_t const& problem_size_1_n = problem_size1.n(); int32_t const& problem_size_1_k = problem_size_0_n; @@ -591,7 +607,7 @@ struct H3FMHAKernel { (warp_id() / MM0::Mma::WarpCount::kM)}; // Mask out last if causal - if (params.causal && num_keys - iter_key_start <= kKeysPerBlock) { + if (params.causal && keys_remaining <= kKeysPerBlock) { auto lane_offset = MM0::AccumLambdaIterator::get_lane_offset( lane_id(), warp_id(), iteratorC_tile_offset); int32_t last_col; @@ -609,9 +625,9 @@ struct H3FMHAKernel { }, [&](int accum_m) {}); } - // DISPATCH_BOOL(iter_key_start == 0, kIsFirst, ([&] { + // DISPATCH_BOOL(key_tile == 0, kIsFirst, ([&] { // DISPATCH_BOOL( - // num_keys - iter_key_start >= kKeysPerBlock, + // keys_remaining >= kKeysPerBlock, // kFullColumns, // ([&] { // // Update `mi` from accum stored in registers @@ -628,7 +644,7 @@ struct H3FMHAKernel { // lane_id(), // thread_id(), // warp_id(), - // num_keys - iter_key_start, + // keys_remaining, // iteratorC_tile_offset, // kSupportsBias ? 1.0f : params.scale); // })); @@ -646,18 +662,18 @@ struct H3FMHAKernel { } // Update `mi` from accum stored in registers // Also does accum[i] <- exp(accum[i] - mi) - if (num_keys - iter_key_start >= kKeysPerBlock) { + if (keys_remaining >= kKeysPerBlock) { iterative_softmax( accum_o, accum, mi, m_prime, s_prime, out_rescale, shared_storage.addition_storage, lane_id(), thread_id(), - warp_id(), num_keys - iter_key_start, iter_key_start == 0, - iteratorC_tile_offset, kSupportsBias ? 1.0f : params.scale); + warp_id(), keys_remaining, key_tile == 0, iteratorC_tile_offset, + kSupportsBias ? 1.0f : params.scale); } else { iterative_softmax( accum_o, accum, mi, m_prime, s_prime, out_rescale, shared_storage.addition_storage, lane_id(), thread_id(), - warp_id(), num_keys - iter_key_start, iter_key_start == 0, - iteratorC_tile_offset, kSupportsBias ? 1.0f : params.scale); + warp_id(), keys_remaining, key_tile == 0, iteratorC_tile_offset, + kSupportsBias ? 1.0f : params.scale); } // Output results to shared-memory @@ -725,10 +741,9 @@ struct H3FMHAKernel { if (!kKeepOutputInRF) { MM1::Mma::drain_cp_asyncs(); DISPATCH_BOOL( - iter_key_start == 0, kIsFirst, ([&] { + key_tile == 0, kIsFirst, ([&] { DISPATCH_BOOL( - (iter_key_start + kKeysPerBlock) >= num_keys, kIsLast, - ([&] { + key_tile + 1 >= key_tiles, kIsLast, ([&] { using DefaultEpilogue = typename MM1::DefaultEpilogue; using DefaultOp = typename MM1::DefaultConfig::EpilogueOutputOp; diff --git a/flash-attention-v100/kernel/h3/forward.cu b/flash-attention-v100/kernel/h3/forward.cu index acd5f96212..74218bd055 100644 --- a/flash-attention-v100/kernel/h3/forward.cu +++ b/flash-attention-v100/kernel/h3/forward.cu @@ -17,29 +17,30 @@ namespace { using Half = cutlass::half_t; constexpr int kHeadDim = 128; -constexpr int kQueries = 64; -template -using KernelFor = typename cutlass::gemm::kernel::H3FMHA< - Half, cutlass::arch::Sm70, true, kQueries, Keys, kHeadDim>::FMHAKernel; +template +using KernelFor = + typename cutlass::gemm::kernel::H3FMHA::FMHAKernel; -template -__global__ __launch_bounds__(128, 1) void h3_flash_v100_d128( - typename KernelFor::DirectParams params) { +template +__global__ __launch_bounds__(Queries * 2, 1) void h3_flash_v100_d128( + typename KernelFor::DirectParams params) { extern __shared__ __align__(16) unsigned char storage[]; if constexpr (Fixed) { params.heads = 14; params.queries = 12323; params.keys = 12323; } - KernelFor kernel; + KernelFor kernel; kernel(params, - *reinterpret_cast::SharedStorage*>(storage)); + *reinterpret_cast::SharedStorage*>( + storage)); } -template +template void launch_attention(at::Tensor const& q, at::Tensor const& k, at::Tensor const& v, at::Tensor& output, float scale) { - using Kernel = KernelFor; + using Kernel = KernelFor; typename Kernel::DirectParams params{ reinterpret_cast(q.data_ptr()), reinterpret_cast(k.data_ptr()), @@ -50,18 +51,38 @@ void launch_attention(at::Tensor const& q, at::Tensor const& k, int(q.size(2)), scale}; if constexpr (Keys == 128) { - // 34,304 bytes/block: allow two resident blocks without extra global - // storage. + // The 64-query tile uses 34,304 bytes/block. Prefer full shared-memory + // capacity for both query geometries without extra global storage. C10_CUDA_CHECK(cudaFuncSetAttribute( - h3_flash_v100_d128, + h3_flash_v100_d128, cudaFuncAttributePreferredSharedMemoryCarveout, 100)); } - h3_flash_v100_d128 - <<, + cudaFuncAttributeMaxDynamicSharedMemorySize, + sizeof(typename Kernel::SharedStorage))); + } + h3_flash_v100_d128 + <<>>(params); } +template +void dispatch_attention(at::Tensor const& q, at::Tensor const& k, + at::Tensor const& v, at::Tensor& output, float scale, + int selected) { + if (selected == 128) { + if (q.size(1) == 12323 && k.size(1) == 12323 && q.size(2) == 14) + launch_attention(q, k, v, output, scale); + else + launch_attention(q, k, v, output, scale); + } else { + launch_attention(q, k, v, output, scale); + } +} + at::Tensor aligned_contiguous(const at::Tensor& tensor) { auto result = tensor.contiguous(); // contiguous() may preserve a contiguous view with an unaligned offset. @@ -72,7 +93,8 @@ at::Tensor aligned_contiguous(const at::Tensor& tensor) { } // namespace at::Tensor h3_flash_attention_forward(at::Tensor q, at::Tensor k, at::Tensor v, - double scale, int key_tile) { + double scale, int key_tile, + int query_tile) { TORCH_CHECK(q.is_cuda() && q.dim() == 4 && q.scalar_type() == at::kHalf, "H3 FlashAttention-V100 requires CUDA FP16 BSND tensors"); TORCH_CHECK(k.device() == q.device() && v.device() == q.device() && @@ -91,12 +113,14 @@ at::Tensor h3_flash_attention_forward(at::Tensor q, at::Tensor k, at::Tensor v, "H3 FlashAttention-V100 is an inference-only operator"); TORCH_CHECK(key_tile == 0 || key_tile == 64 || key_tile == 128, "H3 attention key tile must be 0, 64 or 128"); + TORCH_CHECK(query_tile == 64 || query_tile == 128, + "H3 attention query tile must be 64 or 128"); const c10::cuda::CUDAGuard guard(q.device()); auto* properties = at::cuda::getCurrentDeviceProperties(); TORCH_CHECK(properties->major == 7 && properties->minor == 0, "H3 FlashAttention-V100 requires SM70"); int64_t groups64 = q.size(0) * q.size(2); - int64_t blocks64 = ((q.size(1) + kQueries - 1) / kQueries) * groups64; + int64_t blocks64 = ((q.size(1) + query_tile - 1) / query_tile) * groups64; TORCH_CHECK(q.size(1) <= INT_MAX && k.size(1) <= INT_MAX && groups64 <= 65535 && blocks64 <= INT_MAX, "H3 attention shape exceeds kernel index limits"); @@ -108,13 +132,10 @@ at::Tensor h3_flash_attention_forward(at::Tensor q, at::Tensor k, at::Tensor v, // self-attention reuses each Q fragment across twice as many keys. int selected = key_tile ? key_tile : (q.size(1) >= 1024 && k.size(1) >= 1024 ? 128 : 64); - if (selected == 128) { - if (q.size(1) == 12323 && k.size(1) == 12323 && q.size(2) == 14) - launch_attention<128, true>(q, k, v, output, float(scale)); - else - launch_attention<128>(q, k, v, output, float(scale)); - } else - launch_attention<64>(q, k, v, output, float(scale)); + if (query_tile == 128) + dispatch_attention<128>(q, k, v, output, float(scale), selected); + else + dispatch_attention<64>(q, k, v, output, float(scale), selected); C10_CUDA_KERNEL_LAUNCH_CHECK(); return output; } @@ -122,5 +143,5 @@ at::Tensor h3_flash_attention_forward(at::Tensor q, at::Tensor k, at::Tensor v, PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("forward", &h3_flash_attention_forward, pybind11::arg("q"), pybind11::arg("k"), pybind11::arg("v"), pybind11::arg("scale"), - pybind11::arg("key_tile") = 0); + pybind11::arg("key_tile") = 0, pybind11::arg("query_tile") = 64); } diff --git a/flash-attention-v100/kernel/h3/forward_sparse.cu b/flash-attention-v100/kernel/h3/forward_sparse.cu new file mode 100644 index 0000000000..5a26daa612 --- /dev/null +++ b/flash-attention-v100/kernel/h3/forward_sparse.cu @@ -0,0 +1,156 @@ +// SPDX-License-Identifier: BSD-3-Clause +// SPDX-FileCopyrightText: Copyright contributors to the 1Cat-vLLM project + +#include +#include +#include +#include + +#include +#include + +#include +#include + +#include "default_fmha.h" +#include "vsa_layout.h" + +namespace { +using Half = cutlass::half_t; +using Kernel = + typename cutlass::gemm::kernel::H3FMHA::FMHAKernel; + +// One warp compacts one mask row in ascending key-block order. No sorting, +// padded key arithmetic, host mask transfer or atomic update is necessary. +__global__ void pack_block_map(const bool* mask, int* indices, int* counts, + int rows, int blocks) { + const int row = blockIdx.x; + const int lane = threadIdx.x; + int count = 0; + for (int start = 0; start < blocks; start += 32) { + const int block = start + lane; + const bool selected = block < blocks && mask[int64_t(row) * blocks + block]; + const unsigned ballot = __ballot_sync(0xffffffff, selected); + const unsigned before = (1u << lane) - 1; + if (selected) + indices[int64_t(row) * blocks + count + __popc(ballot & before)] = block; + count += __popc(ballot); + } + if (lane == 0) counts[row] = count; +} + +__global__ __launch_bounds__(128, 1) void sparse_attention( + Kernel::DirectParams params) { + extern __shared__ __align__(16) unsigned char storage[]; + Kernel kernel; + kernel.template operator()( + params, *reinterpret_cast(storage)); +} +} // namespace + +torch::Tensor sparse_forward_impl(torch::Tensor q, torch::Tensor k, + torch::Tensor v, torch::Tensor block_map, + torch::Tensor block_sizes, double scale, + bool validate_values) { + TORCH_CHECK(q.is_cuda() && q.dim() == 4 && q.scalar_type() == at::kHalf && + q.is_contiguous() && q.size(0) > 0 && q.size(1) > 0 && + q.size(1) % 64 == 0 && q.size(2) > 0 && q.size(3) == 128, + "SM70 sparse attention requires contiguous FP16 [B,64*N,H,128]"); + for (const auto& operand : {k, v}) { + TORCH_CHECK( + operand.device() == q.device() && operand.sizes() == q.sizes() && + operand.scalar_type() == q.scalar_type() && operand.is_contiguous(), + "SM70 sparse Q/K/V must share shape, device, dtype and layout"); + } + TORCH_CHECK(!q.requires_grad() && !k.requires_grad() && !v.requires_grad(), + "SM70 sparse attention is inference-only"); + TORCH_CHECK(std::isfinite(scale) && scale > 0 && + scale <= std::numeric_limits::max(), + "SM70 sparse attention scale must be finite and positive"); + const int64_t blocks = q.size(1) / 64; + const int64_t groups = q.size(0) * q.size(2); + TORCH_CHECK( + q.size(1) <= INT_MAX && groups <= 65535 && groups * blocks <= INT_MAX, + "SM70 sparse attention shape exceeds index limits"); + TORCH_CHECK(block_map.device() == q.device() && block_map.dim() == 4 && + block_map.scalar_type() == at::kBool && + block_map.is_contiguous() && block_map.size(0) == q.size(0) && + block_map.size(1) == q.size(2) && + block_map.size(2) == blocks && block_map.size(3) == blocks, + "SM70 sparse attention block map must be bool [B,H,N,N]"); + TORCH_CHECK(block_sizes.device() == q.device() && block_sizes.dim() == 1 && + block_sizes.size(0) == blocks && + block_sizes.is_contiguous() && + block_sizes.scalar_type() == at::kInt, + "SM70 sparse attention block sizes must be int32 [N]"); + const c10::cuda::CUDAGuard guard(q.device()); + auto* properties = at::cuda::getCurrentDeviceProperties(); + TORCH_CHECK(properties->major == 7 && properties->minor == 0, + "SM70 sparse attention requires SM70"); + if (validate_values) { + TORCH_CHECK(block_sizes.min().item() > 0 && + block_sizes.max().item() <= 64, + "SM70 sparse attention block sizes must be in [1,64]"); + TORCH_CHECK( + block_map.any(-1).all().item(), + "SM70 sparse attention requires a selected key for every query block"); + } + // Aligned strides alone do not guarantee an aligned contiguous storage view. + if (reinterpret_cast(q.data_ptr()) % 16) q = q.clone(); + if (reinterpret_cast(k.data_ptr()) % 16) k = k.clone(); + if (reinterpret_cast(v.data_ptr()) % 16) v = v.clone(); + auto indices = torch::empty({groups * blocks, blocks}, block_sizes.options()); + auto counts = torch::empty({groups * blocks}, block_sizes.options()); + auto output = torch::zeros_like(q); + auto stream = at::cuda::getCurrentCUDAStream(); + pack_block_map<<>>( + block_map.data_ptr(), indices.data_ptr(), + counts.data_ptr(), int(groups * blocks), int(blocks)); + Kernel::DirectParams params{reinterpret_cast(q.data_ptr()), + reinterpret_cast(k.data_ptr()), + reinterpret_cast(v.data_ptr()), + reinterpret_cast(output.data_ptr()), + int(q.size(1)), + int(k.size(1)), + int(q.size(2)), + float(scale), + indices.data_ptr(), + counts.data_ptr(), + block_sizes.data_ptr(), + int(blocks)}; + sparse_attention<<>>(params); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return output; +} + +torch::Tensor sparse_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, + torch::Tensor block_map, torch::Tensor block_sizes, + double scale) { + return sparse_forward_impl(q, k, v, block_map, block_sizes, scale, true); +} + +// Private H3 route: the owner constructs immutable sizes from validated host +// geometry and creates a nonempty mask by construction. Keep all device, +// shape, dtype, alignment and indexing checks; only value reductions are +// omitted. +torch::Tensor sparse_forward_prevalidated(torch::Tensor q, torch::Tensor k, + torch::Tensor v, + torch::Tensor block_map, + torch::Tensor block_sizes, + double scale) { + return sparse_forward_impl(q, k, v, block_map, block_sizes, scale, false); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("_h3_tile_qkv_prevalidated", &h3_vsa_layout::tile); + m.def("_h3_gate_untile_prevalidated", &h3_vsa_layout::finish); + m.def("forward", &sparse_forward, pybind11::arg("q"), pybind11::arg("k"), + pybind11::arg("v"), pybind11::arg("block_map"), + pybind11::arg("block_sizes"), pybind11::arg("scale")); + m.def("_forward_prevalidated", &sparse_forward_prevalidated, + pybind11::arg("q"), pybind11::arg("k"), pybind11::arg("v"), + pybind11::arg("block_map"), pybind11::arg("block_sizes"), + pybind11::arg("scale")); +} diff --git a/flash-attention-v100/kernel/h3/vsa_layout.h b/flash-attention-v100/kernel/h3/vsa_layout.h new file mode 100644 index 0000000000..faf8d4bf2a --- /dev/null +++ b/flash-attention-v100/kernel/h3/vsa_layout.h @@ -0,0 +1,140 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +#pragma once +// Private primitives for H3-owned, validated geometry. No activation is cached +// here. +#include +#include +#include +#include +#include +#include +#include + +namespace h3_vsa_layout { + +__global__ void tile_three(const uint4* q, const uint4* k, const uint4* v, + const int* source_rows, uint4* out, int64_t vectors, + int64_t source_tokens, int64_t tiled_tokens, + int64_t row_vectors, int64_t batches) { + int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= vectors) return; + int64_t lane = i % row_vectors; + int64_t row = (i / row_vectors) % tiled_tokens; + int64_t batch = (i / (row_vectors * tiled_tokens)) % batches; + int tensor = i / (row_vectors * tiled_tokens * batches); + int original = source_rows[row]; + const uint4* input = tensor == 0 ? q : tensor == 1 ? k : v; + uint4 value = make_uint4(0, 0, 0, 0); + if (original >= 0 && original < source_tokens) + value = input[(batch * source_tokens + original) * row_vectors + lane]; + out[i] = value; +} + +__global__ void gate_untiling(const uint4* sparse, const uint4* compressed, + const uint4* gate, const int* tiled_rows, + uint4* output, int64_t vectors, + int64_t source_tokens, int64_t tiled_tokens, + int64_t row_vectors) { + int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= vectors) return; + int64_t lane = i % row_vectors; + int64_t row = (i / row_vectors) % source_tokens; + int64_t batch = i / (row_vectors * source_tokens); + int tiled = tiled_rows[row]; + if (tiled < 0 || tiled >= tiled_tokens) { + output[i] = make_uint4(0, 0, 0, 0); + return; + } + uint4 sv = sparse[(batch * tiled_tokens + tiled) * row_vectors + lane]; + uint4 cv = + compressed[(batch * (tiled_tokens / 64) + tiled / 64) * row_vectors + + lane]; + uint4 gv = gate[i]; + uint4 answer; + auto* a = reinterpret_cast<__half2*>(&answer); + const auto* s = reinterpret_cast(&sv); + const auto* c = reinterpret_cast(&cv); + const auto* g = reinterpret_cast(&gv); +#pragma unroll + for (int j = 0; j < 4; ++j) a[j] = __hadd2_rn(s[j], __hmul2_rn(c[j], g[j])); + output[i] = answer; +} + +void check_operand(const torch::Tensor& x, const torch::Tensor& q) { + TORCH_CHECK(x.device() == q.device() && x.scalar_type() == at::kHalf && + x.is_contiguous() && uintptr_t(x.data_ptr()) % 16 == 0, + "H3 layout input must be aligned contiguous FP16 on one device"); +} + +torch::Tensor tile(torch::Tensor q, torch::Tensor k, torch::Tensor v, + torch::Tensor rows, torch::Tensor out) { + TORCH_CHECK(q.is_cuda() && q.dim() == 4 && q.size(3) == 128 && + q.size(0) > 0 && q.size(1) > 0 && q.size(2) > 0, + "H3 layout requires nonempty CUDA BSHD with D128"); + check_operand(q, q); + check_operand(k, q); + check_operand(v, q); + check_operand(out, q); + TORCH_CHECK(q.sizes() == k.sizes() && q.sizes() == v.sizes(), + "QKV shapes differ"); + TORCH_CHECK(rows.device() == q.device() && rows.scalar_type() == at::kInt && + rows.dim() == 1 && rows.is_contiguous() && rows.numel() > 0 && + rows.numel() % 64 == 0, + "H3 source map must be contiguous int32 tile rows"); + TORCH_CHECK(out.dim() == 5 && out.size(0) == 3 && out.size(1) == q.size(0) && + out.size(2) == rows.numel() && out.size(3) == q.size(2) && + out.size(4) == 128, + "H3 tiled output shape mismatch"); + at::assert_no_overlap(out, rows); + at::assert_no_overlap(out, q); + at::assert_no_overlap(out, k); + at::assert_no_overlap(out, v); + c10::cuda::CUDAGuard guard(q.device()); + int64_t n = out.numel() / 8; + tile_three<<<(n + 255) / 256, 256, 0, at::cuda::getCurrentCUDAStream()>>>( + reinterpret_cast(q.data_ptr()), + reinterpret_cast(k.data_ptr()), + reinterpret_cast(v.data_ptr()), rows.data_ptr(), + reinterpret_cast(out.data_ptr()), n, q.size(1), rows.numel(), + q.size(2) * 16, q.size(0)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} + +torch::Tensor finish(torch::Tensor sparse, torch::Tensor compressed, + torch::Tensor gate, torch::Tensor rows) { + TORCH_CHECK(sparse.is_cuda() && sparse.dim() == 4 && sparse.size(3) == 128 && + sparse.size(0) > 0 && sparse.size(1) > 0 && + sparse.size(1) % 64 == 0 && sparse.size(2) > 0, + "H3 sparse output requires nonempty CUDA tile geometry"); + check_operand(sparse, sparse); + check_operand(compressed, sparse); + check_operand(gate, sparse); + TORCH_CHECK(gate.dim() == 4 && gate.size(1) > 0 && + gate.size(0) == sparse.size(0) && + gate.size(2) == sparse.size(2) && gate.size(3) == 128, + "H3 gate geometry mismatch"); + TORCH_CHECK(compressed.dim() == 4 && compressed.size(0) == sparse.size(0) && + compressed.size(1) == sparse.size(1) / 64 && + compressed.size(2) == sparse.size(2) && + compressed.size(3) == 128, + "H3 compressed geometry mismatch"); + TORCH_CHECK(rows.device() == sparse.device() && + rows.scalar_type() == at::kInt && rows.dim() == 1 && + rows.is_contiguous() && rows.numel() == gate.size(1), + "H3 untiling map mismatch"); + c10::cuda::CUDAGuard guard(sparse.device()); + auto out = torch::empty_like(gate); + int64_t n = out.numel() / 8; + gate_untiling<<<(n + 255) / 256, 256, 0, at::cuda::getCurrentCUDAStream()>>>( + reinterpret_cast(sparse.data_ptr()), + reinterpret_cast(compressed.data_ptr()), + reinterpret_cast(gate.data_ptr()), rows.data_ptr(), + reinterpret_cast(out.data_ptr()), n, gate.size(1), sparse.size(1), + sparse.size(2) * 16); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} + +} // namespace h3_vsa_layout diff --git a/flashinfer-sm70/csrc/h3_noncausal_sm70.cu b/flashinfer-sm70/csrc/h3_noncausal_sm70.cu index 00fffc314f..d086c3eabd 100644 --- a/flashinfer-sm70/csrc/h3_noncausal_sm70.cu +++ b/flashinfer-sm70/csrc/h3_noncausal_sm70.cu @@ -9,28 +9,20 @@ namespace fi = flashinfer::attention::sm70; namespace { constexpr int D = 128; -constexpr int BQ = 128; +constexpr int BQ = 192; constexpr int BK = 64; -// Two warps share 16 query rows. Each owns two K16 score fragments, so -// expanding BK does not double the CTA's thread count. -constexpr int KEY_WARPS = 2; +// One warp owns Q16 and both logical K32 halves, retaining FP32 arithmetic. +constexpr int KEY_WARPS = 1; constexpr int KEY_FRAGMENTS = BK / (KEY_WARPS * 16); constexpr int OUTPUT_FRAGMENTS = D / (KEY_WARPS * 16); constexpr int VLD = BK + 8; constexpr int THREADS = (BQ / 16) * KEY_WARPS * 32; -constexpr int PREFETCH_VECTORS = BK * D / (THREADS * 8); -static_assert(THREADS * PREFETCH_VECTORS * 8 == BK * D); -static_assert(BQ / 16 < 16); // Barrier 0 joins the CTA; 1..8 join query pairs. +constexpr int PREFETCH_VECTORS = (BK * D + THREADS * 8 - 1) / (THREADS * 8); +static_assert(THREADS * PREFETCH_VECTORS * 8 >= BK * D); +static_assert(KEY_WARPS == 1 && BK == 64); constexpr int shared_bytes() { - constexpr int QLD = D + 8, PLD = BK + 4; - return (BQ * D + BK * QLD + D * VLD + BQ * PLD) * 2 + - (BQ * KEY_WARPS * 2 + BQ * 2) * 4; -} - -// QK maxima and probability rows are consumed only by the matching pair. -// Keep CTA-wide barriers around K/V staging and tile consumption. -__device__ __forceinline__ void sync_query_pair(int query_group) { - asm volatile("bar.sync %0, 64;" ::"r"(query_group + 1) : "memory"); + constexpr int QLD = D + 8; + return (BQ * D + BK * QLD + D * VLD) * 2; } __device__ __forceinline__ int q_swizzle(int row) { @@ -49,19 +41,28 @@ __device__ __forceinline__ void load_q_fragment(fi::AFragment& fragment, values[1] = *reinterpret_cast(base + ((col + 8) ^ mask)); } -// A 68-half P stride makes accumulator pair stores conflict-free. Odd rows -// remain 8-byte aligned, so use 64-bit loads rather than WMMA's 128-bit loads. -__device__ __forceinline__ void load_p_fragment(fi::AFragment& fragment, - const half* source, int row, - int col) { +// Convert FP16-rounded probability accumulator pairs to the existing Volta +// A fragment layout. Exchange row ownership across lane bit 1, then concatenate +// the two eight-column halves across lane bit 3. No arithmetic in the exchange. +__device__ __forceinline__ void load_probability_fragment( + fi::AFragment& fragment, const unsigned* pairs) { const int lane = threadIdx.x & 31; - const int physical_row = - row + (lane & 3) + ((lane & 16) >> 2) + ((lane & 4) << 1); - const half* base = source + physical_row * (BK + 4) + col; - auto* values = reinterpret_cast(fragment.x); + const int row_bit = (lane >> 1) & 1; + const unsigned own0 = row_bit ? pairs[1] : pairs[0]; + const unsigned own1 = row_bit ? pairs[3] : pairs[2]; + const unsigned other0 = + __shfl_xor_sync(0xffffffff, row_bit ? pairs[0] : pairs[1], 2); + const unsigned other1 = + __shfl_xor_sync(0xffffffff, row_bit ? pairs[2] : pairs[3], 2); + const unsigned local[4] = {row_bit ? other0 : own0, row_bit ? own0 : other0, + row_bit ? other1 : own1, row_bit ? own1 : other1}; + auto* output = reinterpret_cast(fragment.x); #pragma unroll - for (int i = 0; i < 4; ++i) - values[i] = *reinterpret_cast(base + i * 4); + for (int part = 0; part < 4; ++part) { + const unsigned opposite = __shfl_xor_sync(0xffffffff, local[part], 8); + output[part] = (lane & 8) ? opposite : local[part]; + output[part + 4] = (lane & 8) ? local[part] : opposite; + } } __global__ __launch_bounds__(THREADS, @@ -69,15 +70,13 @@ __global__ __launch_bounds__(THREADS, const half* v, half* output, int length, int heads, float scale) { - constexpr int QLD = D + 8, PLD = BK + 4; + constexpr int QLD = D + 8; extern __shared__ __align__(32) unsigned char raw[]; half* qs = reinterpret_cast(raw); half* ks = qs + BQ * D; half* vs = ks + BK * QLD; - half* probabilities = vs + D * VLD; - float* scores = reinterpret_cast(probabilities + BQ * PLD); - float* maximum = scores + BQ * KEY_WARPS * 2; - float* denominator = maximum + BQ; + float running_max[2] = {-INFINITY, -INFINITY}; + float running_sum[2] = {0.f, 0.f}; const int tid = threadIdx.x, warp = tid / 32; const int warp_q = warp / KEY_WARPS, warp_k = warp % KEY_WARPS; const int lane = tid % 32; @@ -101,10 +100,6 @@ __global__ __launch_bounds__(THREADS, row < length ? q[base + int64_t(row) * heads * D + i % D] : __float2half(0.f); } - if (tid < BQ) { - maximum[tid] = -INFINITY; - denominator[tid] = 0.f; - } __syncthreads(); for (int start = 0; start < length; start += BK) { if (start == 0) { @@ -135,76 +130,73 @@ __global__ __launch_bounds__(THREADS, fi::mma_sync_m16n16k16_row_col_f16f16f32(qk[n], qa, kb); } } + unsigned probability_pairs[KEY_FRAGMENTS][4]; { - // Volta distributes each accumulator row across lanes differing in - // bits 1 and 3. Reduce its 16 columns in registers, then combine only - // the two warp partials through shared memory. - float row_max[2] = {-INFINITY, -INFINITY}; + // Preserve the original two logical K32 partials and their FP32 sum + // order even though one warp now owns both halves of this K64 tile. + float row_max[2][2] = {{-INFINITY, -INFINITY}, {-INFINITY, -INFINITY}}; #pragma unroll for (int n = 0; n < KEY_FRAGMENTS; ++n) { #pragma unroll for (int i = 0; i < qk[n].num_elements; ++i) { - const int col = warp_k * (BK / KEY_WARPS) + n * 16 + fragment_col + - (i & 1) + ((i >> 2) & 1) * 4; + const int col = n * 16 + fragment_col + (i & 1) + ((i >> 2) & 1) * 4; + const int row = (i >> 1) & 1; qk[n].x[i] = start + col < length ? qk[n].x[i] * scale : -INFINITY; - row_max[(i >> 1) & 1] = fmaxf(row_max[(i >> 1) & 1], qk[n].x[i]); + row_max[n / 2][row] = fmaxf(row_max[n / 2][row], qk[n].x[i]); } } #pragma unroll - for (int r = 0; r < 2; ++r) { - row_max[r] = - fmaxf(row_max[r], __shfl_xor_sync(0xffffffff, row_max[r], 2)); - row_max[r] = - fmaxf(row_max[r], __shfl_xor_sync(0xffffffff, row_max[r], 8)); - const int row = warp_q * 16 + fragment_row + r * 2; - if ((lane & 10) == 0) scores[row * KEY_WARPS + warp_k] = row_max[r]; - } - sync_query_pair(warp_q); - float new_max[2], row_sum[2] = {0.f, 0.f}; + for (int half = 0; half < 2; ++half) { #pragma unroll - for (int r = 0; r < 2; ++r) { - const int row = warp_q * 16 + fragment_row + r * 2; - new_max[r] = maximum[row]; + for (int row = 0; row < 2; ++row) { + auto& maximum = row_max[half][row]; + maximum = fmaxf(maximum, __shfl_xor_sync(0xffffffff, maximum, 2)); + maximum = fmaxf(maximum, __shfl_xor_sync(0xffffffff, maximum, 8)); + } + } + float new_max[2], row_sum[2][2] = {{0.f, 0.f}, {0.f, 0.f}}; #pragma unroll - for (int w = 0; w < KEY_WARPS; ++w) - new_max[r] = fmaxf(new_max[r], scores[row * KEY_WARPS + w]); - register_alpha[r] = __expf(maximum[row] - new_max[r]); + for (int row = 0; row < 2; ++row) { + new_max[row] = + fmaxf(fmaxf(running_max[row], row_max[0][row]), row_max[1][row]); + register_alpha[row] = __expf(running_max[row] - new_max[row]); } #pragma unroll for (int n = 0; n < KEY_FRAGMENTS; ++n) { #pragma unroll for (int i = 0; i < qk[n].num_elements; ++i) { - const int r = (i >> 1) & 1; - const int row = warp_q * 16 + fragment_row + r * 2; - const int col = warp_k * (BK / KEY_WARPS) + n * 16 + fragment_col + - (i & 1) + ((i >> 2) & 1) * 4; - const float p = __expf(qk[n].x[i] - new_max[r]); - probabilities[row * PLD + col] = __float2half_rn(p); - row_sum[r] += p; + const int row = (i >> 1) & 1; + const float p = __expf(qk[n].x[i] - new_max[row]); + qk[n].x[i] = p; + row_sum[n / 2][row] += p; } - } - float* partial_sums = scores + BQ * KEY_WARPS; #pragma unroll - for (int r = 0; r < 2; ++r) { - row_sum[r] += __shfl_xor_sync(0xffffffff, row_sum[r], 2); - row_sum[r] += __shfl_xor_sync(0xffffffff, row_sum[r], 8); - const int row = warp_q * 16 + fragment_row + r * 2; - if ((lane & 10) == 0) - partial_sums[row * KEY_WARPS + warp_k] = row_sum[r]; + for (int i = 0; i < 4; ++i) { + union PackedPair { + half2 value; + unsigned bits; + } pair; + pair.value = __floats2half2_rn(qk[n].x[2 * i], qk[n].x[2 * i + 1]); + probability_pairs[n][i] = pair.bits; + } } - sync_query_pair(warp_q); - if (warp_k == 0 && (lane & 10) == 0) { #pragma unroll - for (int r = 0; r < 2; ++r) { - const int row = warp_q * 16 + fragment_row + r * 2; - float sum = 0.f; + for (int half = 0; half < 2; ++half) { #pragma unroll - for (int w = 0; w < KEY_WARPS; ++w) - sum += partial_sums[row * KEY_WARPS + w]; - denominator[row] = denominator[row] * register_alpha[r] + sum; - maximum[row] = new_max[r]; + for (int row = 0; row < 2; ++row) { + auto& sum = row_sum[half][row]; + sum += __shfl_xor_sync(0xffffffff, sum, 2); + sum += __shfl_xor_sync(0xffffffff, sum, 8); } } +#pragma unroll + for (int row = 0; row < 2; ++row) { + float sum = 0.f; + sum += row_sum[0][row]; + sum += row_sum[1][row]; + running_sum[row] = running_sum[row] * register_alpha[row] + sum; + running_max[row] = new_max[row]; + } } // Issue the next K/V global loads while the current V tile is consumed. union StagedVector { @@ -219,7 +211,7 @@ __global__ __launch_bounds__(THREADS, const int tile_row = (tile / (D / 32)) * 8 + (lane >> 2); const int next_row = start + BK + tile_row; const int next_col = (tile % (D / 32)) * 32 + (lane & 3) * 8; - if (next_row < length) { + if (tile_row < BK && next_row < length) { const int64_t position = base + int64_t(next_row) * heads * D + next_col; if ((reinterpret_cast(k) % 16 == 0) && @@ -247,19 +239,19 @@ __global__ __launch_bounds__(THREADS, } #pragma unroll for (int part = 0; part < OUTPUT_FRAGMENTS; ++part) { - const int col = warp_k * (D / KEY_WARPS) + part * 16; - const int row = warp_q * 16; - auto& pv = accumulators[part]; #pragma unroll - for (int i = 0; i < pv.num_elements; ++i) - pv.x[i] *= register_alpha[(i >> 1) & 1]; + for (int i = 0; i < accumulators[part].num_elements; ++i) + accumulators[part].x[i] *= register_alpha[(i >> 1) & 1]; + } +#pragma unroll + for (int kv = 0; kv < KEY_FRAGMENTS; ++kv) { + fi::AFragment pa; + load_probability_fragment(pa, probability_pairs[kv]); #pragma unroll - for (int kv = 0; kv < BK; kv += 16) { - fi::AFragment pa; + for (int part = 0; part < OUTPUT_FRAGMENTS; ++part) { fi::QKBFragment vb; - load_p_fragment(pa, probabilities, row, kv); - fi::load_qk_b_fragment(vb, vs + col * VLD + kv, VLD); - fi::mma_sync_m16n16k16_row_col_f16f16f32(pv, pa, vb); + fi::load_qk_b_fragment(vb, vs + part * 16 * VLD + kv * 16, VLD); + fi::mma_sync_m16n16k16_row_col_f16f16f32(accumulators[part], pa, vb); } } __syncthreads(); @@ -269,6 +261,7 @@ __global__ __launch_bounds__(THREADS, const int tile = (tid + n * THREADS) / 32; const int tile_row = (tile / (D / 32)) * 8 + (lane >> 2); const int next_col = (tile % (D / 32)) * 32 + (lane & 3) * 8; + if (tile_row >= BK) continue; *reinterpret_cast(ks + tile_row * QLD + next_col) = next_k[n].packed; // Transpose four rows with exact 32-bit lane exchanges. Each lane @@ -292,7 +285,8 @@ __global__ __launch_bounds__(THREADS, __shfl_xor_sync(0xffffffff, transposed[2 * j], 8); const unsigned other1 = __shfl_xor_sync(0xffffffff, transposed[2 * j + 1], 8); - const unsigned local = transposed[2 * j + ((lane & 8) >> 3)]; + const unsigned local = + ((lane & 8) ? transposed[2 * j + 1] : transposed[2 * j]); const unsigned other = (lane & 8) ? other1 : other0; const uint2 vector = (lane & 8) ? make_uint2(other, local) : make_uint2(local, other); @@ -313,7 +307,7 @@ __global__ __launch_bounds__(THREADS, (i & 1) + ((i >> 2) & 1) * 4; if (q_start + row < length) output[base + int64_t(q_start + row) * heads * D + col] = - __float2half_rn(pv.x[i] / denominator[row]); + __float2half_rn(pv.x[i] / running_sum[(i >> 1) & 1]); } } } diff --git a/setup.py b/setup.py index e225d15506..d6d8953c55 100644 --- a/setup.py +++ b/setup.py @@ -1250,6 +1250,7 @@ def _read_requirements(filename: str) -> list[str]: ext_modules.append(CMakeExtension(name="vllm._h3_w8a16_C")) ext_modules.append(CMakeExtension(name="vllm._h3_flashinfer_C")) ext_modules.append(CMakeExtension(name="vllm._h3_flashattn_C")) + ext_modules.append(CMakeExtension(name="vllm._sm70_sparse_attention_C")) build_sm70_fa2 = _cuda_arch_contains(7, 0) and not _cuda_arch_at_least(8, 0) if _cuda_arch_at_least(8, 0) or build_sm70_fa2: ext_modules.append(CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa2_C")) diff --git a/tests/video/test_h3_acceptance.py b/tests/video/test_h3_acceptance.py index 080e7842e9..75de37b554 100644 --- a/tests/video/test_h3_acceptance.py +++ b/tests/video/test_h3_acceptance.py @@ -9,62 +9,274 @@ from vllm.video.metrics import evaluate_performance -def measurements(): - run = { - "config": asdict(H3Config()), +def measurements(calls=49, api_steps=50, tp=4): + workload = { + "work_accounting": "dense_tp_lora_v2", + "partition": "fl2va", + "task": "t2va", + "adapter": None, + "video_sigmas": [1 - i / calls for i in range(calls + 1)], + "audio_sigmas": [1 - i / calls for i in range(calls + 1)], + "blocks_per_call": 52, + "used_length": 34551, + "attention_algorithm": "dense", + "cache_algorithm": None, + "actual_backends": ["FLASH_ATTN_V100"], + } + total = 6_000_000_000_000_000 + quot, rem = divmod(total, calls) + run: dict = { + "config": asdict(H3Config(tensor_parallel_size=tp)), "request": asdict(H3Request()), - "gpus": [0, 1, 2, 3], - "measurement": {"profiled": False, "warmup_runs": 1}, + "gpus": list(range(tp)), + "engine_session_id": "test-session", + "request_index": 0, + "measurement": {"profiled": False, "warmup": False, "capture": False}, + "end_to_end_seconds": 90.0, "ranks": [ { "rank": rank, - "dit_calls": 49, - "useful_denoise_flops": 6_000_000_000_000_000, + "dit_calls": calls, + "useful_denoise_flops": total, + "denoise_flops_by_layer": {"example": total}, + "redundant_denoise_flops": 0, + "redundant_flops_by_layer": {}, + "denoise_workload": deepcopy(workload), + "denoise_executed_blocks": {str(i): calls for i in range(52)}, + "denoise_steps": [ + { + "index": i, + "dit_calls": 1, + "executed_blocks": 52, + "useful_flops": quot + (i < rem), + "redundant_flops": 0, + "gpu_seconds": 1.0, + "cpu_enqueue_seconds": 0.1, + "sparse_blocks": 0, + "cache_hits": 0, + } + for i in range(calls) + ], "peak_allocated_bytes": 29 * 1024**3, - "stage_seconds": {"denoise": 60.0 if rank == 3 else 50.0}, + "stage_seconds": {"denoise": 60.0 if rank == tp - 1 else 50.0}, } - for rank in range(4) + for rank in range(tp) ], } - return [deepcopy(run) for _ in range(3)] + run["request"]["sampling"]["num_inference_steps"] = api_steps + runs = [deepcopy(run) for _ in range(3)] + for index, measurement in enumerate(runs, 1): + measurement["request_index"] = index + warmup = deepcopy(run) + warmup["measurement"]["warmup"] = True + return warmup, runs + + +def sparse_measurements(): + warmup, runs = measurements(calls=4, api_steps=4) + heads, gated, calls = 14, 50, 4 + selected_pairs = heads * (33 + 32 * 17) + selected_blocks = heads * (3 + 2 * 2) + compression = 4 * heads * 3**2 * 128 + dense_pairs = heads * 33**2 + for run in (warmup, *runs): + run["config"].update(attention_backend="FASTVIDEO_VSA", vsa_topk=1) + for rank in run["ranks"]: + rank["denoise_workload"].update( + work_accounting="sparse_tp_v1", + attention_algorithm="vsa", + used_length=33, + actual_backends=["FASTVIDEO_VSA", "FLASH_ATTN_V100"], + sparse_config=dict( + topk=1, + gated_blocks=gated, + heads=heads, + head_size=128, + prefix_segments=[1], + video_shape=[1, 4, 8], + ), + ) + rank["denoise_sparse_work_by_layer"] = {} + for i in range(gated): + name = f"blocks.{i}.attn.attention" + rank["denoise_sparse_work_by_layer"][name] = dict( + head_size=128, + heads=heads, + selected_blocks=calls * selected_blocks, + selected_token_pairs=calls * selected_pairs, + compression_flops=calls * compression, + dense_token_pairs=calls * dense_pairs, + ) + flops = 4 * calls * selected_pairs * 128 + rank["denoise_flops_by_layer"][name] = flops + rank["denoise_flops_by_layer"][name + ".compression"] = ( + calls * compression + ) + rank["denoise_flops_by_layer"]["example"] -= flops + calls * compression + for step in rank["denoise_steps"]: + step.update( + sparse_blocks=gated * selected_blocks, + sparse_token_pairs=gated * selected_pairs, + sparse_compression_flops=gated * compression, + attention_avoided_flops=4 + * 128 + * gated + * (dense_pairs - selected_pairs), + ) + return warmup, runs + + +def test_sparse_acceptance_keeps_algorithm_savings_out_of_numerator(): + warmup, runs = sparse_measurements() + report = evaluate_performance(runs, warmup=warmup) + assert report["rank_median_tflops"] == [100.0] * 4 + assert report["performance_passed"] + assert all(n > 0 for n in report["attention_avoided_flops_by_run_and_rank"][0]) + + +@pytest.mark.parametrize( + "invalid", + [ + "padded_pairs", + "block_count", + "missing_gate", + "compression", + "savings", + "backend", + ], +) +def test_sparse_acceptance_rejects_invented_or_incomplete_work(invalid): + warmup, runs = sparse_measurements() + rank = runs[0]["ranks"][0] + layers = rank["denoise_sparse_work_by_layer"] + name = next(iter(layers)) + if invalid == "padded_pairs": + layers[name]["selected_token_pairs"] = layers[name]["selected_blocks"] * 64**2 + elif invalid == "block_count": + rank["denoise_steps"][0]["sparse_blocks"] += 1 + elif invalid == "missing_gate": + layers.pop(name) + elif invalid == "compression": + layers[name]["compression_flops"] += 1 + elif invalid == "savings": + rank["denoise_steps"][0]["attention_avoided_flops"] += 1 + else: + rank["denoise_workload"]["actual_backends"] = ["FLASH_ATTN_V100"] + with pytest.raises(ValueError): + evaluate_performance(runs, warmup=warmup) def test_all_ranks_use_slowest_rank_wall_time_and_strict_threshold(): - runs = measurements() - report = evaluate_performance(runs) + warmup, runs = measurements() + report = evaluate_performance(runs, warmup=warmup) assert report["rank_median_tflops"] == [100.0] * 4 assert report["performance_passed"] for run in runs: - run["ranks"][2]["useful_denoise_flops"] = 4_800_000_000_000_000 - assert not evaluate_performance(runs)["performance_passed"] + run["ranks"][-1]["stage_seconds"]["denoise"] = 75.0 + assert not evaluate_performance(runs, warmup=warmup)["performance_passed"] + + +@pytest.mark.parametrize("calls,api_steps", [(49, 50), (4, 5), (8, 9), (4, 4)]) +@pytest.mark.parametrize("tp", [1, 2, 4]) +def test_uses_actual_intervals_for_all_dense_schedules(calls, api_steps, tp): + warmup, runs = measurements(calls, api_steps, tp) + for run in (warmup, *runs): + run["request"]["sampling"].update(width=1280, height=736, num_frames=120) + report = evaluate_performance(runs, warmup=warmup) + assert len(report["rank_median_tflops"]) == tp + assert len(report["workflow"]["video_sigmas"]) == calls + 1 + assert report["performance_passed"] + assert report["quality_status"] == "requires_separate_numerical_and_human_review" @pytest.mark.parametrize( - "invalid", ["profiled", "seed", "rank", "calls", "excluded", "residual"] + "invalid", + [ + "profiled", + "seed", + "rank", + "calls", + "excluded", + "residual", + "capture", + "warmup", + "session", + "index", + "missing_step", + "work", + "block", + "layer_work", + "negative_rank_time", + "nan_sigma", + "short_audio", + "sparse", + "cache", + "missing_memory", + "warmup_incomplete", + "missing_warmup", + "nan_step_time", + "legacy_work", + "redundant_work", + ], ) -def test_rejects_incomparable_measurements(invalid): - runs = measurements() - if invalid == "profiled": - runs[1]["measurement"]["profiled"] = True +def test_rejects_incomplete_or_incomparable_measurements(invalid): + warmup, runs = measurements(4, 5) + run, rank = runs[1], runs[1]["ranks"][0] + if invalid in ("profiled", "capture", "warmup"): + run["measurement"][invalid] = True elif invalid == "seed": - runs[1]["request"]["sampling"]["seed"] = 2026 + run["request"]["sampling"]["seed"] = 2026 elif invalid == "rank": - runs[1]["ranks"][3]["rank"] = 2 + run["ranks"][3]["rank"] = 2 elif invalid == "excluded": - runs[1]["timing_valid"] = False + run["timing_valid"] = False elif invalid == "residual": - runs[1]["config"]["residual_sequence_parallel"] = True + run["config"]["residual_sequence_parallel"] = True + elif invalid == "session": + run["engine_session_id"] = "restarted" + elif invalid == "index": + run["request_index"] = 0 + elif invalid == "missing_step": + rank["denoise_steps"].pop() + elif invalid == "work": + rank["useful_denoise_flops"] += 1 + elif invalid == "block": + rank["denoise_executed_blocks"]["0"] -= 1 + elif invalid == "layer_work": + rank["denoise_flops_by_layer"]["example"] += 1 + elif invalid == "negative_rank_time": + rank["stage_seconds"]["denoise"] = -1.0 + elif invalid == "nan_sigma": + rank["denoise_workload"]["video_sigmas"][1] = float("nan") + elif invalid == "short_audio": + rank["denoise_workload"]["audio_sigmas"].pop() + elif invalid == "sparse": + rank["denoise_workload"]["attention_algorithm"] = "vsa" + elif invalid == "cache": + rank["denoise_workload"]["cache_algorithm"] = "teacache" + elif invalid == "missing_memory": + rank["peak_allocated_bytes"] = 0 + elif invalid == "warmup_incomplete": + warmup["ranks"][0]["dit_calls"] = 3 + elif invalid == "missing_warmup": + warmup["measurement"]["warmup"] = False + elif invalid == "nan_step_time": + rank["denoise_steps"][0]["gpu_seconds"] = float("nan") + elif invalid == "legacy_work": + rank["denoise_workload"].pop("work_accounting") + elif invalid == "redundant_work": + rank["redundant_denoise_flops"] = 1 else: - runs[1]["ranks"][0]["dit_calls"] = 48 + rank["dit_calls"] = 48 with pytest.raises(ValueError): - evaluate_performance(runs) + evaluate_performance(runs, warmup=warmup) def test_memory_gate_and_variability_are_reported_separately(): - runs = measurements() + warmup, runs = measurements() runs[0]["ranks"][1]["peak_allocated_bytes"] = 31 * 1024**3 runs[0]["ranks"][3]["stage_seconds"]["denoise"] = 80.0 - report = evaluate_performance(runs) + report = evaluate_performance(runs, warmup=warmup) assert not report["memory_passed"] assert report["denoise_cv"] > 0.05 assert not report["performance_passed"] diff --git a/tests/video/test_h3_column_major.py b/tests/video/test_h3_column_major.py index 54dc445001..9f84832b62 100644 --- a/tests/video/test_h3_column_major.py +++ b/tests/video/test_h3_column_major.py @@ -7,6 +7,7 @@ import pytest import torch +from vllm.model_executor.layers import sm70_diffusion from vllm.model_executor.models.minimax_h3 import cuda_ops from vllm.model_executor.models.minimax_h3.quantization import ( DiffusionInt8ConvRotConfig, @@ -73,7 +74,9 @@ def test_real_tp4_projection_uses_exact_zero_workspace_plan(n, k, fp32): @cuda def test_missing_plan_falls_back_without_changing_math(monkeypatch): monkeypatch.setattr( - cuda_ops, "_column_major_plan", lambda *args: SimpleNamespace(supported=False) + sm70_diffusion, + "_column_major_plan", + lambda *args: SimpleNamespace(supported=False), ) x = torch.randn(17, 256, dtype=torch.float16, device="cuda") w = torch.randn(65, 256, dtype=torch.float16, device="cuda") diff --git a/tests/video/test_h3_fasth3.py b/tests/video/test_h3_fasth3.py index 024959105d..7bf0e35f98 100644 --- a/tests/video/test_h3_fasth3.py +++ b/tests/video/test_h3_fasth3.py @@ -26,7 +26,9 @@ from vllm.model_executor.models.minimax_h3.pipeline import MiniMaxH3Pipeline -def artifact(tmp_path, *, mutate=None, metadata_updates=None, dtype=torch.bfloat16): +def artifact( + tmp_path, *, mutate=None, metadata_updates=None, dtype=torch.bfloat16, vsa=False +): generator = torch.Generator().manual_seed(519) def rand(*shape): @@ -80,11 +82,20 @@ def rand(*shape): name: (weight.float() + deltas[name]).to(dtype) if name in deltas else weight for name, weight in base.items() } + if vsa: + for i in range(50): + gate = rand(8, 8) + tensors[f"transformer_blocks.{i}.attn.to_gate_compress.set_weight"] = gate + expected[f"blocks.{i}.attn.to_gate_compress.weight"] = gate if mutate: mutate(tensors) metadata = { "format": "fastvideo-lora-v2", - "finetuned_model": "FastVideo/FastVideo-FastH3-Dense-4-step-v1", + "finetuned_model": ( + "FastVideo/FastVideo-FastH3-4-step-v1" + if vsa + else "FastVideo/FastVideo-FastH3-Dense-4-step-v1" + ), "base_model": "MiniMaxAI/MiniMax-H3", "rank": "64", "low_rank_tensors": str( @@ -93,7 +104,9 @@ def rand(*shape): "diff_tensors": str( sum(name.endswith((".diff", ".diff_b")) for name in tensors) ), - "set_weight_tensors": "0", + "set_weight_tensors": str( + sum(name.endswith(".set_weight") for name in tensors) + ), **(metadata_updates or {}), } path = tmp_path / "adapter_model.safetensors" @@ -267,7 +280,8 @@ def test_fasth3_api_defaults_and_unavailable_dynamic_lora(tmp_path): @pytest.mark.parametrize("tp", [1, 2, 4]) -def test_fused_weights_enter_native_tp_loaders_and_host_snapshot(tmp_path, tp): +@pytest.mark.parametrize("vsa", [False, True]) +def test_fused_weights_enter_native_tp_loaders_and_host_snapshot(tmp_path, tp, vsa): from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config from vllm.model_executor.layers.linear import ( ColumnParallelLinear, @@ -277,7 +291,7 @@ def test_fused_weights_enter_native_tp_loaders_and_host_snapshot(tmp_path, tp): from vllm.model_executor.models.minimax_h3.residency import PinnedModuleStager from vllm.model_executor.models.minimax_h3.transformer import MiniMaxH3DiTModel - path, base, expected = artifact(tmp_path) + path, base, expected = artifact(tmp_path, vsa=vsa) fused = dict(FastH3Fusion(path, head_dim=2).apply(base.items())) # The native loader converts serialized [head, Q/K/V, channel] to Q/K/V. qkv = expected["blocks.0.attn.qkv_proj.weight"] @@ -347,6 +361,13 @@ def test_fused_weights_enter_native_tp_loaders_and_host_snapshot(tmp_path, tp): model.video_patch_proj = ColumnParallelLinear( 8, 8, bias=True, params_dtype=torch.float32 ) + if vsa: + for block in model.blocks: + if not hasattr(block, "attn"): + block.attn = nn.Module() + block.attn.to_gate_compress = ColumnParallelLinear( + 8, 8, bias=False, params_dtype=torch.float16 + ) loaded = MiniMaxH3DiTModel.load_weights(model, fused.items()) assert loaded == set(fused) - {"untouched.weight"} for name, param in model.named_parameters(): @@ -373,3 +394,27 @@ def test_fused_weights_enter_native_tp_loaders_and_host_snapshot(tmp_path, tp): stager._restore_masters() for name, param in model.named_parameters(): torch.testing.assert_close(param, before[name], atol=0, rtol=0) + + +def test_vsa_gates_require_complete_inventory_and_matching_backend(tmp_path): + path, base, expected = artifact(tmp_path, vsa=True) + spec = inspect_adapter(path, "fl2va") + assert spec.requires_vsa + fused = dict(FastH3Fusion(path, head_dim=2).apply(base.items())) + for name, value in expected.items(): + torch.testing.assert_close(fused[name], value, rtol=0, atol=0) + with pytest.raises(H3InputError, match="FASTVIDEO_VSA"): + sampling_for_deployment(H3Config(lora_path=str(path))) + sampling = sampling_for_deployment( + H3Config(lora_path=str(path), attention_backend="FASTVIDEO_VSA") + ) + assert sampling.num_inference_steps == 4 + path, _, _ = artifact( + tmp_path, + vsa=True, + mutate=lambda tensors: tensors.pop( + "transformer_blocks.49.attn.to_gate_compress.set_weight" + ), + ) + with pytest.raises(H3InputError, match="every main-block compression gate"): + inspect_adapter(path, "fl2va") diff --git a/tests/video/test_h3_flashattn.py b/tests/video/test_h3_flashattn.py index b05ad2f142..6bdeb9ae5f 100644 --- a/tests/video/test_h3_flashattn.py +++ b/tests/video/test_h3_flashattn.py @@ -26,7 +26,8 @@ def test_flashattn_rejects_batch_head_grid_overflow(): "length", [1, 31, 32, 33, 63, 64, 65, 96, 97, 127, 128, 129, 12323] ) @pytest.mark.parametrize("key_tile", [64, 128]) -def test_flashattn_d128_mha_tails_and_online_rescaling(length, key_tile): +@pytest.mark.parametrize("query_tile", [64, 128]) +def test_flashattn_d128_mha_tails_and_online_rescaling(length, key_tile, query_tile): torch.manual_seed(42) q, k, v = [ torch.randn(2, length, 2, 128, device="cuda", dtype=torch.float16) @@ -37,14 +38,18 @@ def test_flashattn_d128_mha_tails_and_online_rescaling(length, key_tile): k[:, length // 2 :] *= 4 rows = torch.linspace(0, length - 1, min(length, 65), device="cuda").long() expected = chunked_attention_reference(q[:, rows], k, v, scale=128**-0.5) - actual = flashattn_extension().forward(q, k, v, 128**-0.5, key_tile) + actual = flashattn_extension().forward(q, k, v, 128**-0.5, key_tile, query_tile) assert torch.isfinite(actual).all() torch.testing.assert_close(actual[:, rows], expected, atol=0.002, rtol=0.03) + if query_tile == 128: + control = flashattn_extension().forward(q, k, v, 128**-0.5, key_tile, 64) + torch.testing.assert_close(actual, control, atol=0, rtol=0) @pytest.mark.parametrize("layout", ["offset", "strided"]) @pytest.mark.parametrize("key_tile", [64, 128]) -def test_flashattn_d128_storage_and_different_q_k_lengths(layout, key_tile): +@pytest.mark.parametrize("query_tile", [64, 128]) +def test_flashattn_d128_storage_and_different_q_k_lengths(layout, key_tile, query_tile): torch.manual_seed(42) tensors = [] for length, offset in [(33, 1), (65, 3), (65, 5)]: @@ -59,12 +64,12 @@ def test_flashattn_d128_storage_and_different_q_k_lengths(layout, key_tile): tensors.append(value) q, k, v = tensors expected = chunked_attention_reference(q, k, v, scale=128**-0.5) - actual = flashattn_extension().forward(q, k, v, 128**-0.5, key_tile) + actual = flashattn_extension().forward(q, k, v, 128**-0.5, key_tile, query_tile) torch.testing.assert_close(actual, expected, atol=0.002, rtol=0.03) def test_flashattn_dispatch_slices_poisoned_padding(monkeypatch): - from vllm.model_executor.models.minimax_h3 import cuda_ops + from vllm.model_executor.layers import sm70_attention as cuda_ops native = flashattn_extension() calls = [] @@ -98,7 +103,8 @@ def forward(self, q, k, v, scale): @pytest.mark.parametrize("key_tile", [64, 128]) -def test_flashattn_graph_replay_uses_new_values(key_tile): +@pytest.mark.parametrize("query_tile", [64, 128]) +def test_flashattn_graph_replay_uses_new_values(key_tile, query_tile): ops = flashattn_extension() q, k, v = [ torch.randn(1, 129, 2, 128, device="cuda", dtype=torch.float16) @@ -108,11 +114,11 @@ def test_flashattn_graph_replay_uses_new_values(key_tile): stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(stream): for _ in range(3): - ops.forward(q, k, v, 128**-0.5, key_tile) + ops.forward(q, k, v, 128**-0.5, key_tile, query_tile) torch.cuda.current_stream().wait_stream(stream) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - actual = ops.forward(q, k, v, 128**-0.5, key_tile) + actual = ops.forward(q, k, v, 128**-0.5, key_tile, query_tile) for _ in range(2): q.normal_() k.normal_() @@ -122,7 +128,8 @@ def test_flashattn_graph_replay_uses_new_values(key_tile): torch.testing.assert_close(actual, expected, atol=0.002, rtol=0.03) -def test_flashattn_fixed_shape_graph_replay(): +@pytest.mark.parametrize("query_tile", [64, 128]) +def test_flashattn_fixed_shape_graph_replay(query_tile): ops = flashattn_extension() q, k, v = [ torch.randn(1, 12323, 14, 128, device="cuda", dtype=torch.float16) @@ -132,11 +139,11 @@ def test_flashattn_fixed_shape_graph_replay(): stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(stream): for _ in range(3): - ops.forward(q, k, v, 128**-0.5) + ops.forward(q, k, v, 128**-0.5, 0, query_tile) torch.cuda.current_stream().wait_stream(stream) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - actual = ops.forward(q, k, v, 128**-0.5) + actual = ops.forward(q, k, v, 128**-0.5, 0, query_tile) rows = torch.tensor([0, 31, 32, 63, 64, 127, 128, 12322], device="cuda") for _ in range(2): v.normal_() @@ -155,3 +162,5 @@ def test_flashattn_rejects_changed_head_dimension_and_scale(): ops.forward(q, q, q, float("nan")) with pytest.raises(RuntimeError, match="key tile"): ops.forward(q, q, q, 128**-0.5, 32) + with pytest.raises(RuntimeError, match="query tile"): + ops.forward(q, q, q, 128**-0.5, 64, 32) diff --git a/tests/video/test_h3_host_memory.py b/tests/video/test_h3_host_memory.py index 38fbd442ad..c6065698c4 100644 --- a/tests/video/test_h3_host_memory.py +++ b/tests/video/test_h3_host_memory.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project - +import argparse import ast import errno import os @@ -21,6 +21,170 @@ ) +@pytest.mark.parametrize("mode", ["generate", "serve"]) +def test_pageable_weight_masters_are_an_explicit_deployment_option(mode): + from vllm.entrypoints.cli.video import VideoSubcommand + + parser = argparse.ArgumentParser() + VideoSubcommand().subparser_init(parser.add_subparsers()) + assert parser.parse_args(["video", mode]).host_weight_pin_memory + assert not parser.parse_args( + ["video", mode, "--disable-host-weight-pinning"] + ).host_weight_pin_memory + assert not H3Config(host_weight_pin_memory=False).host_weight_pin_memory + with pytest.raises(H3InputError, match="boolean"): + H3Config(host_weight_pin_memory="false") + + +@pytest.mark.parametrize("pin_memory", [True, False]) +@pytest.mark.parametrize("kind", ["MiniMaxH3VideoVAE", "MiniMaxH3AudioVAE"]) +def test_both_vae_stagers_honor_the_host_policy(monkeypatch, pin_memory, kind): + from vllm.model_executor.models.minimax_h3 import vae + + remote = nn.Module() + remote.model = nn.Linear(2, 2) + expected = remote.model.weight.clone() + calls = [] + monkeypatch.setattr( + vae, "_load_component_config", lambda path: {"sample_rate": 44100} + ) + monkeypatch.setattr(vae, "_load_remote_component", lambda *args: remote) + monkeypatch.setattr( + vae, + "PinnedModuleStager", + lambda module, device, **kwargs: calls.append((module, device, kwargs)), + ) + wrapper = getattr(vae, kind)( + "test", + device=torch.device("cuda"), + load_device=torch.device("cpu"), + pin_memory=pin_memory, + ) + assert calls == [(remote, torch.device("cuda"), {"pin_memory": pin_memory})] + assert wrapper.model.weight.dtype == torch.float32 + torch.testing.assert_close(wrapper.model.weight, expected, rtol=0, atol=0) + + +def _alias_module(): + raw = torch.arange(257, dtype=torch.float32) + module = nn.Module() + module.weight = nn.Parameter(raw[8:104].reshape(8, 12).t()) + module.register_buffer("strided", raw[31:67:2]) + module.register_buffer("byte_alias", raw.view(torch.uint8)[9:63]) + module.other = nn.Parameter( + torch.arange(77, dtype=torch.float16).reshape(7, 11).t() + ) + module.register_buffer("empty", torch.empty(0)) + return module + + +def _cpu_snapshot(module): + from vllm.model_executor.models.minimax_h3.residency import PinnedModuleStager + + stager = PinnedModuleStager.__new__(PinnedModuleStager) + stager._groups = stager._snapshot_groups((module,), pin_memory=False) + stager.loaded = False + return stager + + +def test_shared_snapshot_preserves_offsets_aliases_and_private_writes(tmp_path): + left, right = _alias_module(), _alias_module() + a, b = _cpu_snapshot(left), _cpu_snapshot(right) + before = {name: value.clone() for name, value in left.state_dict().items()} + directory = tmp_path / "snapshot" + a._write_shared_groups(directory) + a._read_shared_groups(directory) + b._read_shared_groups(directory) + for module in (left, right): + for name, value in module.state_dict().items(): + torch.testing.assert_close(value, before[name], atol=0, rtol=0) + assert module.weight.stride() == (1, 12) + assert module.strided.stride() == (2,) + assert ( + module.weight.untyped_storage().data_ptr() + == module.strided.untyped_storage().data_ptr() + ) + with torch.no_grad(): + right.weight.add_(7) + torch.testing.assert_close(left.weight, before["weight"], atol=0, rtol=0) + # A private mapping must not modify a later reader or the shared snapshot. + third = _alias_module() + _cpu_snapshot(third)._read_shared_groups(directory) + torch.testing.assert_close(third.weight, before["weight"], atol=0, rtol=0) + a._restore_masters() + for name, value in left.state_dict().items(): + torch.testing.assert_close(value, before[name], atol=0, rtol=0) + + +@pytest.mark.parametrize("corrupt", ["layout", "weights", "file"]) +def test_shared_snapshot_rejects_different_replicas(tmp_path, corrupt): + left, right = _alias_module(), _alias_module() + a = _cpu_snapshot(left) + directory = tmp_path / "snapshot" + a._write_shared_groups(directory) + if corrupt == "layout": + right.weight = nn.Parameter(right.weight.t()) + elif corrupt == "weights": + with torch.no_grad(): + right.other.add_(1) + else: + data = directory / "weights.bin" + with data.open("r+b") as f: + f.seek(257) + f.write(b"\xff") + with pytest.raises(ValueError, match="shared component"): + _cpu_snapshot(right)._read_shared_groups(directory) + + +def test_shared_vae_configuration_requires_pageable_storage(): + with pytest.raises(H3InputError, match="pageable"): + H3Config(share_host_vae_weights=True) + assert H3Config( + share_host_vae_weights=True, host_weight_pin_memory=False + ).share_host_vae_weights + assert not H3Config().share_host_vae_weights + + +def test_engine_cleans_only_its_shared_directory(tmp_path): + import tempfile + + from vllm.video.engine import H3Engine + + untouched = tmp_path / "other" + untouched.mkdir() + engine = H3Engine.__new__(H3Engine) + engine._closed = False + engine.workers = [] + engine.connections = [] + engine._gpu_lease = None + engine._shared_weights = tempfile.TemporaryDirectory(dir=tmp_path, prefix="owned-") + directory = engine._shared_weights.name + engine.close() + from pathlib import Path + + assert not Path(directory).exists() + assert untouched.is_dir() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a leased GPU") +def test_gpu_shared_snapshot_roundtrips_all_storage_views(tmp_path): + from vllm.model_executor.models.minimax_h3.residency import PinnedModuleStager + + model = _alias_module() + before = {name: value.clone() for name, value in model.state_dict().items()} + stager = PinnedModuleStager(model, torch.device("cuda"), pin_memory=False) + stager.share_cpu_storage(tmp_path / "shared") + for _ in range(3): + stager.load() + for name, value in model.state_dict().items(): + assert value.device.type == "cuda" + torch.testing.assert_close(value.cpu(), before[name], atol=0, rtol=0) + stager.offload() + for name, value in model.state_dict().items(): + assert value.device.type == "cpu" + torch.testing.assert_close(value, before[name], atol=0, rtol=0) + + def aliased_module(): module = nn.Module() values = torch.arange(128, dtype=torch.float32).reshape(16, 8) diff --git a/tests/video/test_h3_layer_staging.py b/tests/video/test_h3_layer_staging.py new file mode 100644 index 0000000000..7a8024e0fa --- /dev/null +++ b/tests/video/test_h3_layer_staging.py @@ -0,0 +1,181 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Bounded layer ownership, failure cleanup and exact device transitions.""" + +import argparse + +import pytest +import torch +from torch import nn + +from vllm.model_executor.models.minimax_h3.config import H3Config, H3InputError +from vllm.model_executor.models.minimax_h3.residency import ( + BoundedAllocatorCache, + LayerwiseModuleStager, + MMapHostWeights, + PinnedModuleStager, +) + + +class Block(nn.Module): + def __init__(self, size=8): + super().__init__() + self.weight = nn.Parameter(torch.randn(size, size) / size) + self.register_buffer("adapter_a", torch.randn(2, size) / size) + self.register_buffer("adapter_b", torch.randn(size, 2) / size) + self.fail = False + + def forward(self, x): + if self.fail: + raise RuntimeError("block failure") + return nn.functional.linear(x, self.weight) + 0.25 * nn.functional.linear( + nn.functional.linear(x, self.adapter_a), self.adapter_b + ) + + +class Model(nn.Module): + def __init__(self, size=8): + super().__init__() + self.blocks = nn.ModuleList([Block(size), Block(size)]) + # Aliased storage spanning a block and the outer model stays resident. + self.register_buffer("shared", self.blocks[0].adapter_a[0]) + + def forward(self, x): + for block in self.blocks: + x = block(x) + return x + self.shared + + +def cpu_snapshot(model): + snapshot = PinnedModuleStager.__new__(PinnedModuleStager) + snapshot._groups = snapshot._snapshot_groups((model,), pin_memory=False) + snapshot.loaded = False + snapshot._device_storages = [] + snapshot.device = torch.device("cpu") + snapshot.cache_retention = None + return snapshot + + +def test_layer_storage_partition_and_rejected_overlap(): + model = Model() + snapshot = cpu_snapshot(model) + plan = LayerwiseModuleStager(snapshot, model.blocks) + groups = [g for s in (*plan.stagers, plan.resident) for g in s._groups] + assert len(groups) == len({id(g) for g in groups}) == len(snapshot._groups) + shared = [g for g in plan.resident._groups if len(g.bindings) == 2] + assert len(shared) == 1 + assert {id(b.target) for b in shared[0].bindings} == { + id(model.shared), + id(model.blocks[0].adapter_a), + } + with pytest.raises(ValueError, match="unique"): + LayerwiseModuleStager(snapshot, [model.blocks[0]] * 2) + with pytest.raises(ValueError, match="nested"): + LayerwiseModuleStager(snapshot, [model, model.blocks[0]]) + protected = LayerwiseModuleStager( + snapshot, model.blocks, resident_modules=(model.blocks[1],) + ) + assert not protected.stagers[1]._groups + assert {id(b.target) for g in protected.resident._groups for b in g.bindings} >= { + id(t) for t in model.blocks[1].parameters() + } + + +def test_layer_forward_lifetime_failure_and_reentry(monkeypatch): + monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None) + monkeypatch.setattr( + BoundedAllocatorCache, "release_if_needed", lambda *a, **k: False + ) + monkeypatch.setattr( + PinnedModuleStager, "_load_once", lambda self: setattr(self, "loaded", True) + ) + model = Model().eval() + snapshot = cpu_snapshot(model) + plan = LayerwiseModuleStager(snapshot, model.blocks) + seen = [] + for index, block in enumerate(model.blocks): + # Observe after the plan's pre-hook runs, inside the actual forward. + original = block.forward + + def observe(x, index=index, original=original): + assert plan.resident.loaded + assert [s.loaded for s in plan.stagers] == [i == index for i in range(2)] + seen.append(index) + return original(x) + + block.forward = observe + for fail in (False, True, False): + model.blocks[1].fail = fail + with plan.on_device(): + with pytest.raises(RuntimeError, match="idle host"), plan.on_device(): + pass + with pytest.raises(RuntimeError, match="overlaps"): + snapshot.load() + if fail: + with pytest.raises(RuntimeError, match="block failure"): + model(torch.ones(3, 8)) + else: + assert torch.isfinite(model(torch.ones(3, 8))).all() + assert not snapshot._layerwise_active + assert all(not s.loaded for s in (*plan.stagers, plan.resident)) + assert all( + not b._forward_pre_hooks and not b._forward_hooks for b in model.blocks + ) + assert seen == [0, 1] * 3 + assert plan.loaded_bytes > 0 + + +def test_layer_policy_rejects_incompatible_fixed_cache(): + assert H3Config().weight_offload == "component" + assert H3Config(weight_offload="layer").weight_offload == "layer" + with pytest.raises(H3InputError, match="offload"): + H3Config(weight_offload="automatic") + with pytest.raises(H3InputError, match="fixed GPU weight cache"): + H3Config(weight_offload="layer", fp16_cache_layers=("blocks.0",)) + + +@pytest.mark.parametrize("mode", ["generate", "serve"]) +def test_layer_policy_cli(mode): + from vllm.entrypoints.cli.video import VideoSubcommand + + parser = argparse.ArgumentParser() + VideoSubcommand().subparser_init(parser.add_subparsers()) + assert parser.parse_args(["video", mode]).weight_offload == "component" + assert ( + parser.parse_args(["video", mode, "--weight-offload", "layer"]).weight_offload + == "layer" + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("dtype", [torch.float16, torch.float32]) +@pytest.mark.parametrize("mapped", [False, True]) +def test_gpu_layerwise_adapter_and_alias_roundtrips(dtype, mapped, tmp_path): + torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False + torch.manual_seed(96) + model = Model(128).to(device="cuda", dtype=dtype).eval() + # Module.to can split aliases; establish the actual mixed-owner alias on GPU. + model.shared = model.blocks[0].adapter_a[0] + for block in model.blocks: + block.weight.data = block.weight.t().contiguous().t() + inputs = [torch.randn(65, 128, device="cuda", dtype=dtype) for _ in range(3)] + with torch.inference_mode(): + expected = [model(x).clone() for x in inputs] + snapshot = PinnedModuleStager( + model, + torch.device("cuda"), + pin_memory=False, + host_backing=MMapHostWeights(tmp_path) if mapped else None, + ) + plan = LayerwiseModuleStager(snapshot, model.blocks) + for x, reference in zip(inputs, expected): + with plan.on_device(): + actual = model(x) + torch.testing.assert_close(actual, reference, atol=0, rtol=0) + assert all(p.device.type == "cpu" for p in model.parameters()) + if mapped: + assert all(p.untyped_storage().filename for p in model.parameters()) + assert ( + model.shared.untyped_storage().data_ptr() + == model.blocks[0].adapter_a.untyped_storage().data_ptr() + ) diff --git a/tests/video/test_h3_numerics.py b/tests/video/test_h3_numerics.py index 6153fa30df..1d39f8543d 100644 --- a/tests/video/test_h3_numerics.py +++ b/tests/video/test_h3_numerics.py @@ -178,7 +178,7 @@ def test_attention_padding_excludes_poisoned_suffix(used, padded): @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") -@pytest.mark.parametrize("length", [127, 128, 129, 12323]) +@pytest.mark.parametrize("length", [127, 128, 129, 191, 192, 193, 12323]) def test_flashinfer_online_softmax_across_tiles_and_batches(length): from vllm.model_executor.models.minimax_h3.cuda_ops import flashinfer_extension @@ -197,7 +197,9 @@ def test_flashinfer_online_softmax_across_tiles_and_batches(length): @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") -@pytest.mark.parametrize("length", [31, 32, 33, 63, 64, 65, 127, 128, 129]) +@pytest.mark.parametrize( + "length", [31, 32, 33, 63, 64, 65, 127, 128, 129, 191, 192, 193, 385] +) def test_flashinfer_prefetch_tail_and_unaligned_storage(length): from vllm.model_executor.models.minimax_h3.cuda_ops import flashinfer_extension @@ -218,7 +220,7 @@ def test_flashinfer_prefetch_tail_and_unaligned_storage(length): @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") -@pytest.mark.parametrize("length", [129, 257]) +@pytest.mark.parametrize("length", [129, 193, 257, 385]) def test_flashinfer_query_groups_have_independent_softmax_state(length): from vllm.model_executor.models.minimax_h3.cuda_ops import flashinfer_extension @@ -435,7 +437,8 @@ def test_encoder_uses_functional_all_reduce_return(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") -def test_staging_preserves_aliased_weights_across_repeated_transfers(): +@pytest.mark.parametrize("pin_memory", [True, False]) +def test_staging_preserves_aliased_weights_across_repeated_transfers(pin_memory): from torch import nn from vllm.model_executor.models.minimax_h3.residency import PinnedModuleStager @@ -445,7 +448,7 @@ def test_staging_preserves_aliased_weights_across_repeated_transfers(): module.weight = nn.Parameter(backing) module.register_buffer("view", backing[3:7, 1:5]) expected = module.view.clone() - stager = PinnedModuleStager(module, torch.device("cuda")) + stager = PinnedModuleStager(module, torch.device("cuda"), pin_memory=pin_memory) for _ in range(2): stager.load() assert module.weight.is_cuda and module.view.is_cuda @@ -455,7 +458,8 @@ def test_staging_preserves_aliased_weights_across_repeated_transfers(): ) torch.testing.assert_close(module.view.cpu(), expected) stager.offload() - assert module.weight.is_pinned() and module.view.is_pinned() + assert module.weight.is_pinned() is pin_memory + assert module.view.is_pinned() is pin_memory torch.testing.assert_close(module.view, expected) diff --git a/tests/video/test_h3_prepared_linear.py b/tests/video/test_h3_prepared_linear.py new file mode 100644 index 0000000000..6c5a3cd4bc --- /dev/null +++ b/tests/video/test_h3_prepared_linear.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Independent operand, adapter and output-precision checks for shared GEMM.""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.model_executor.layers.sm70_diffusion import ( + fp16_gemm_input, + fp16_linear_prepared, +) +from vllm.model_executor.models.minimax_h3.lora import TurboLinearMethod, lora_scale +from vllm.model_executor.models.minimax_h3.quantization import ( + DiffusionInt8ConvRotConfig, + FP32OutputLinearMethod, + Int8ConvRotLayerConfig, + Int8ConvRotLinearMethod, +) + + +@pytest.mark.parametrize("value", [100000.0, float("inf"), float("nan")]) +def test_original_weight_loader_rejects_fp16_overflow(value): + from vllm.model_executor.models.minimax_h3.transformer import MiniMaxH3DiTModel + + model = torch.nn.Module() + model.register_parameter( + "weight", torch.nn.Parameter(torch.empty(2, 3, dtype=torch.float16), False) + ) + checkpoint = torch.full((2, 3), value, dtype=torch.bfloat16) + with pytest.raises(ValueError, match="finite FP16"): + MiniMaxH3DiTModel.load_weights(model, [("weight", checkpoint)]) + + +def test_original_weight_loader_keeps_sensitive_fp32_parameters(): + from vllm.model_executor.models.minimax_h3.transformer import MiniMaxH3DiTModel + + model = torch.nn.Module() + model.register_parameter("weight", torch.nn.Parameter(torch.empty(2, 3), False)) + checkpoint = torch.full((2, 3), 100000.0, dtype=torch.bfloat16) + MiniMaxH3DiTModel.load_weights(model, [("weight", checkpoint)]) + torch.testing.assert_close(model.weight, checkpoint.float(), atol=0, rtol=0) + + +@pytest.mark.parametrize("layout", ["row", "column"]) +def test_dense_layout_keeps_logical_weights_and_wide_output(layout): + torch.manual_seed(7) + # Deliberately not an H3 projection shape. + weight = torch.randn(96, 192, dtype=torch.float16) + layer = SimpleNamespace(weight=weight.clone(), h3_fp16_weight_layout=layout) + method = FP32OutputLinearMethod() + for _ in range(2): + method.process_weights_after_loading(layer) + torch.testing.assert_close(layer.weight, weight, atol=0, rtol=0) + assert layer.weight.untyped_storage().nbytes() == weight.numel() * 2 + x = torch.randn(33, 192) * 100000 + values, scale = fp16_gemm_input(x) + result = method.apply_prepared(layer, values, scale) + reference = (values.float() @ weight.float().T) * scale + torch.testing.assert_close(result, reference, atol=0, rtol=0) + assert torch.isfinite(result).all() + assert result.abs().max() > torch.finfo(torch.float16).max + with pytest.raises(ValueError, match="unrotated"): + method.apply_prepared(layer, values, scale, input_is_rotated=True) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires SM70") +@pytest.mark.parametrize("quantized", [False, True]) +@pytest.mark.parametrize("adapter_scale", [0.0, 0.75, -0.5]) +def test_gpu_prepared_adapter_preserves_wide_intermediates( + quantized, adapter_scale, monkeypatch +): + import vllm.model_executor.models.minimax_h3.lora as adapter + from vllm.model_executor.models.minimax_h3.cuda_ops import w8a16_extension + + torch.manual_seed(42) + ops = w8a16_extension() + if quantized: + base = Int8ConvRotLinearMethod( + DiffusionInt8ConvRotConfig(), Int8ConvRotLayerConfig(True), prefix="probe" + ) + weight = torch.randint(-7, 8, (384, 256), device="cuda", dtype=torch.int8) + else: + base = FP32OutputLinearMethod() + weight = torch.randn(384, 256, device="cuda", dtype=torch.float16) * 0.03 + layer = SimpleNamespace( + weight=weight, + weight_scale=torch.full((384,), 0.01, device="cuda"), + h3_output_fp32=True, + h3_lora_a_0=torch.randn(8, 256, device="cuda", dtype=torch.float16), + h3_lora_b_0=torch.randn(384, 8, device="cuda", dtype=torch.float16), + ) + method = TurboLinearMethod(base, [(0, 0, 384)], 0.125) + x = torch.randn(33, 256, device="cuda") * 100000 + values, scale = fp16_gemm_input(x) + token = lora_scale.set(adapter_scale) + try: + with monkeypatch.context() as context: + context.setattr(adapter, "supports_fused_scaled_add", lambda _: False) + expected = method.apply(layer, x) + actual = method.apply_prepared(layer, values, scale) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + assert actual.dtype == torch.float32 and torch.isfinite(actual).all() + if quantized: + rotated = ops.rotate(values) + actual = method.apply_prepared( + layer, rotated, scale, input_is_rotated=True, original_input=values + ) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + if adapter_scale: + with pytest.raises(ValueError, match="original unrotated"): + method.apply_prepared(layer, rotated, scale, input_is_rotated=True) + finally: + lora_scale.reset(token) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires SM70") +def test_gpu_shared_linear_handles_non_h3_shapes_and_fp32_outputs(): + from vllm.model_executor.models.minimax_h3.cuda_ops import w8a16_extension + + torch.manual_seed(2026) + for m, n, k in ((33, 96, 192), (257, 1024, 768)): + x = torch.randn(m, k, device="cuda", dtype=torch.float16) + w = torch.randn(n, k, device="cuda", dtype=torch.float16) + expected = w8a16_extension().gemm(x, w, True) + for weight in (w, w.T.contiguous().T): + actual = fp16_linear_prepared(x, weight, output_fp32=True) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) diff --git a/tests/video/test_h3_provenance.py b/tests/video/test_h3_provenance.py new file mode 100644 index 0000000000..63ee157cc1 --- /dev/null +++ b/tests/video/test_h3_provenance.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import hashlib + +import vllm +from vllm.video.benchmark import source_provenance + + +def test_provenance_tracks_shared_operators_outside_model_directory( + tmp_path, monkeypatch +): + package = tmp_path / "vllm" + layers = package / "model_executor/layers" + layers.mkdir(parents=True) + monkeypatch.setattr(vllm, "__file__", str(package / "__init__.py")) + contents = { + "sm70_diffusion.py": b"gemm revision 1", + "sm70_attention.py": b"dense attention revision 1", + "sm70_sparse_attention.py": b"sparse attention revision 1", + } + for name, content in contents.items(): + (layers / name).write_bytes(content) + before = source_provenance()["sources_sha256"] + for name, content in contents.items(): + assert ( + before[f"model_executor/layers/{name}"] + == hashlib.sha256(content).hexdigest() + ) + + changed = "model_executor/layers/sm70_attention.py" + (package / changed).write_bytes(b"dense attention revision 2") + after = source_provenance()["sources_sha256"] + assert after[changed] != before[changed] + assert {k: v for k, v in before.items() if k != changed} == { + k: v for k, v in after.items() if k != changed + } + + +def test_provenance_tracks_loaded_generic_sm70_binary(tmp_path, monkeypatch): + import sys + from types import SimpleNamespace + + from vllm.video.metrics import loaded_kernel_provenance + + binary = tmp_path / "exact_reduce.so" + binary.write_bytes(b"generic collective binary") + monkeypatch.setitem( + sys.modules, "onecat_sm70_exact_reduce", SimpleNamespace(__file__=str(binary)) + ) + result = loaded_kernel_provenance() + assert result[str(binary)] == hashlib.sha256(binary.read_bytes()).hexdigest() diff --git a/tests/video/test_h3_qk_norm_rope.py b/tests/video/test_h3_qk_norm_rope.py index fc9fcdab89..af9602293a 100644 --- a/tests/video/test_h3_qk_norm_rope.py +++ b/tests/video/test_h3_qk_norm_rope.py @@ -84,6 +84,7 @@ def inputs(tokens, heads, amplitude, weight_dtype): return q, k, *weights, rope, 1e-5 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("tokens,heads", [(1, 1), (65, 3), (129, 14), (12323, 14)]) @pytest.mark.parametrize( "amplitude,weight_dtype", [(1, torch.float16), (2000, torch.float32)] @@ -98,6 +99,7 @@ def test_rounding_and_partial_rotation(tokens, heads, amplitude, weight_dtype): torch.testing.assert_close(result, reference, rtol=0, atol=0) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @torch.inference_mode() def test_graph_replay_reads_changed_inputs_and_rope(): args = inputs(65, 3, 1000, torch.float16) @@ -120,6 +122,7 @@ def test_graph_replay_reads_changed_inputs_and_rope(): torch.testing.assert_close(result, reference, rtol=0, atol=0) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @torch.inference_mode() def test_other_rotary_width_keeps_reference_path(): args = list(inputs(65, 3, 1, torch.float16)) diff --git a/tests/video/test_h3_residual_collectives.py b/tests/video/test_h3_residual_collectives.py new file mode 100644 index 0000000000..33f1527839 --- /dev/null +++ b/tests/video/test_h3_residual_collectives.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + +import pytest +import torch + +from vllm.model_executor.models.minimax_h3 import collectives +from vllm.model_executor.models.minimax_h3.config import H3Config, H3InputError + + +@pytest.mark.parametrize("tp", [1, 2, 4]) +@pytest.mark.parametrize("partition", ["fl2va", "ref2va"]) +def test_explicit_peer_option_keeps_tp_and_partition_compatibility(tp, partition): + config = H3Config( + tensor_parallel_size=tp, + partition=partition, + residual_sequence_parallel=True, + residual_reduction="peer", + ) + assert config.residual_reduction == "peer" + assert H3Config().residual_reduction == "native" + + +@pytest.mark.parametrize("budget", [0, -1, True, float("nan"), float("inf"), 1e308]) +def test_reject_invalid_communication_budget(budget): + with pytest.raises(H3InputError): + H3Config(residual_reduction_memory_gib=budget) + + +def test_peer_option_requires_residual_rows(): + with pytest.raises(H3InputError, match="residual sequence"): + H3Config(residual_reduction="peer") + + +def test_reuse_request_accounting_shape_eviction_and_budget_fallback(monkeypatch): + plans = [] + + class Plan: + raw_ipc_bytes = 64 + + @staticmethod + def required_memory_bytes(shape): + return shape[0] * shape[1] * 4 + + def __init__(self, group, shape, *, memory_budget_bytes): + self.shape = shape + self.closed = False + plans.append(self) + + def reduce(self, value): + return (value * 4).chunk(4)[1] + + def close(self): + self.closed = True + + monkeypatch.setattr(collectives, "SM70ExactRowReductionPlan", Plan) + group = SimpleNamespace(world_size=4, rank_in_group=1, all_reduce=lambda x: x * 4) + owner = collectives.H3ResidualReduction(group, memory_budget_bytes=100) + value = torch.arange(12, dtype=torch.float32).view(4, 3) + expected = (value * 4).chunk(4)[1] + assert torch.equal(owner.reduce(value), expected) + owner.begin_request() + assert owner.snapshot()["peer_calls"] == 0 + assert owner.snapshot()["raw_ipc_peak_bytes"] == 64 + assert torch.equal(owner.reduce(value), expected) + assert len(plans) == 1 + assert owner.snapshot()["setup_seconds"] == 0 + large = torch.ones(16, 3) + assert torch.equal(owner.reduce(large), torch.full((4, 3), 4.0)) + assert plans[0].closed + assert owner.snapshot()["native_calls"] == 1 + assert owner.snapshot()["fallback_reason"] + owner.begin_request() + assert owner.snapshot()["raw_ipc_peak_bytes"] == 0 + owner.close() + + +def test_tp2_uses_ordinary_reduction_without_plan(monkeypatch): + def unexpected(*args, **kwargs): + raise AssertionError("TP2 must not construct a TP4 peer plan") + + monkeypatch.setattr(collectives, "SM70ExactRowReductionPlan", unexpected) + group = SimpleNamespace(world_size=2, rank_in_group=1, all_reduce=lambda x: x * 2) + owner = collectives.H3ResidualReduction(group, memory_budget_bytes=100) + assert torch.equal(owner.reduce(torch.ones(6, 3)), torch.full((3, 3), 2.0)) + assert owner.snapshot()["native_calls"] == 1 diff --git a/tests/video/test_h3_residual_parallel.py b/tests/video/test_h3_residual_parallel.py index dd2f4818a1..ffb7be7bfa 100644 --- a/tests/video/test_h3_residual_parallel.py +++ b/tests/video/test_h3_residual_parallel.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""TP4 residual tests: launch the GPU cases with torchrun --nproc-per-node=4.""" +"""Residual tests: launch the GPU cases with torchrun --nproc-per-node=2 or 4.""" import argparse import os @@ -10,7 +10,7 @@ import pytest import torch -from vllm.model_executor.models.minimax_h3.config import H3Config, H3InputError +from vllm.model_executor.models.minimax_h3.config import H3Config from vllm.model_executor.models.minimax_h3.transformer import MiniMaxH3DiTBlock @@ -35,14 +35,13 @@ def test_residual_parallel_accepts_native_sm70_backends(backend): {"lora_path": "adapter.safetensors"}, ], ) -def test_residual_parallel_rejects_unvalidated_deployments(change): +def test_residual_parallel_accepts_compatible_deployments(change): config = H3Config( transformer_path="int8.safetensors", attention_backend="FLASHINFER_SM70", residual_sequence_parallel=True, ) - with pytest.raises(H3InputError, match="residual sequence parallelism"): - replace(config, **change) + assert replace(config, **change).residual_sequence_parallel @pytest.mark.parametrize("mode", ["generate", "serve"]) @@ -89,8 +88,9 @@ def test_residual_parallel_rejects_invalid_rows_before_collectives( @pytest.fixture def tp4_group(): - if os.environ.get("WORLD_SIZE") != "4": - pytest.skip("requires torchrun --nproc-per-node=4 on a leased GPU group") + world_size = int(os.environ.get("WORLD_SIZE", "1")) + if world_size not in (2, 4): + pytest.skip("requires torchrun with TP2/TP4 on a leased GPU group") from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config from vllm.distributed import ( cleanup_dist_env_and_memory, @@ -104,10 +104,10 @@ def tp4_group(): torch.accelerator.set_device_index(local_rank) torch.set_num_threads(2) with set_current_vllm_config( - VllmConfig(parallel_config=ParallelConfig(tensor_parallel_size=4)) + VllmConfig(parallel_config=ParallelConfig(tensor_parallel_size=world_size)) ): - init_distributed_environment(4, rank, "env://", local_rank, "nccl") - initialize_model_parallel(4) + init_distributed_environment(world_size, rank, "env://", local_rank, "nccl") + initialize_model_parallel(world_size) try: yield get_tp_group() finally: @@ -121,10 +121,20 @@ def test_gpu_residual_parallel_matches_replicated_blocks(tp4_group): for valid in (32, 33, 131): for backend in ("FLASH_ATTN_V100", "FLASHINFER_SM70"): _check_replicated_blocks(tp4_group, valid, backend) + for quantized in (False, True): + for backend in ("FLASH_ATTN_V100", "FLASHINFER_SM70"): + _check_replicated_blocks( + tp4_group, 131, backend, quantized=quantized, adapted=True + ) -def _check_replicated_blocks(group, valid, backend): +def _check_replicated_blocks(group, valid, backend, *, quantized=False, adapted=False): from vllm.model_executor.models.minimax_h3.attention import attention_backend + from vllm.model_executor.models.minimax_h3.lora import TurboLinearMethod, lora_scale + from vllm.model_executor.models.minimax_h3.quantization import ( + DiffusionInt8ConvRotConfig, + Int8ConvRotLinearMethod, + ) from vllm.model_executor.models.minimax_h3.transformer import ( MiniMaxH3DiTArchConfig, ) @@ -137,11 +147,19 @@ def _check_replicated_blocks(group, valid, backend): adaln_out_features=18 * 512, ) token = attention_backend.set(backend) + quant = None + if quantized: + quant = DiffusionInt8ConvRotConfig( + layer_configs={ + f"blocks.0.{name}": {"format": "int8_tensorwise", "convrot": True} + for name in ("attn.qkv_proj", "attn.out_proj", "mlp.fc1", "mlp.fc2") + } + ) try: - baseline = MiniMaxH3DiTBlock(arch, None, prefix="blocks.0").cuda().eval() + baseline = MiniMaxH3DiTBlock(arch, quant, prefix="blocks.0").cuda().eval() candidate = ( MiniMaxH3DiTBlock( - arch, None, prefix="blocks.0", residual_sequence_parallel=True + arch, quant, prefix="blocks.0", residual_sequence_parallel=True ) .cuda() .eval() @@ -151,13 +169,41 @@ def _check_replicated_blocks(group, valid, backend): # Rank-dependent projection weights exercise real TP partial sums. torch.manual_seed(1234 + group.rank_in_group) for name, parameter in baseline.named_parameters(): - if "norm" in name: + if parameter.dtype == torch.int8: + parameter.random_(-7, 8) + elif name.endswith("weight_scale"): + parameter.fill_(0.005) + elif "norm" in name: parameter.fill_(1) else: parameter.normal_(0, 0.03) + if adapted: + for block in (baseline, candidate): + for layer in block.modules(): + method = getattr(layer, "quant_method", None) + if not getattr(method, "supports_prepared_fp16", False): + continue + n, k = layer.weight.shape + layer.register_buffer( + "h3_lora_a_0", + torch.randn(8, k, dtype=torch.float16, device="cuda") * 0.01, + ) + layer.register_buffer( + "h3_lora_b_0", + torch.randn(n, 8, dtype=torch.float16, device="cuda") * 0.01, + ) + layer._sm70_f16_forbidden = True + layer.quant_method = TurboLinearMethod(method, [(0, 0, n)], 1.0) candidate.load_state_dict(baseline.state_dict()) + for block in (baseline, candidate): + for layer in block.modules(): + method = getattr(layer, "quant_method", None) + if isinstance(method, TurboLinearMethod): + method = method.base + if isinstance(method, Int8ConvRotLinearMethod): + method.process_weights_after_loading(layer) torch.manual_seed(42) - total = (valid + 3) // 4 * 4 + total = (valid + group.world_size - 1) // group.world_size * group.world_size x = torch.randn(total, 512, device="cuda", dtype=torch.float32) x[::5, 0] = 70000 # Residuals must not pass through an FP16 collective. kwargs = dict( @@ -168,17 +214,21 @@ def _check_replicated_blocks(group, valid, backend): max_seqlen=valid, packed_total=total, ) - rows = total // 4 + rows = total // group.world_size actual = x.narrow(0, group.rank_in_group * rows, rows).clone() expected = x.clone() - for _ in range(2): - expected = baseline(expected, **kwargs) - actual = candidate(actual, **kwargs) + scale_token = lora_scale.set(0.75 if adapted else 0.0) + try: + for _ in range(2): + expected = baseline(expected, **kwargs) + actual = candidate(actual, **kwargs) + finally: + lora_scale.reset(scale_token) actual = group.all_gather(actual, dim=0) assert actual.dtype == torch.float32 assert torch.isfinite(actual).all() assert actual.abs().max() > torch.finfo(torch.float16).max - torch.testing.assert_close(actual, expected, rtol=2e-3, atol=3e-4) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) # Modifying padding must not change any valid attention output. if valid < total: changed = x.clone() diff --git a/tests/video/test_h3_vsa_fp32_diagnostic.py b/tests/video/test_h3_vsa_fp32_diagnostic.py new file mode 100644 index 0000000000..5c725a280a --- /dev/null +++ b/tests/video/test_h3_vsa_fp32_diagnostic.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Opt-in tests for the acceptance-only CUDA kernel; no runtime registration.""" + +import os +from pathlib import Path + +import pytest +import torch + +from benchmarks.kernels.benchmark_h3_vsa_fp32 import fp32_reference, load_binary + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or not os.environ.get("H3_VSA_FP32_DIAGNOSTIC"), + reason="requires a leased SM70 GPU and H3_VSA_FP32_DIAGNOSTIC binary", +) + + +@pytest.fixture(scope="module") +def ops(): + return load_binary(Path(os.environ["H3_VSA_FP32_DIAGNOSTIC"])) + + +def make_case(lengths, prefix, topk, amplitude=1): + torch.manual_seed(821) + blocks = len(lengths) + sizes = torch.tensor(lengths, device="cuda", dtype=torch.int32) + q, k, v = [ + torch.randn(2, blocks * 64, 3, 128, device="cuda", dtype=torch.float16) + * amplitude + for _ in range(3) + ] + for i, length in enumerate(lengths): + for tensor in (q, k, v): + tensor[:, i * 64 + length : (i + 1) * 64] = float("nan") + # Independent mask construction: prefix queries see every block; video + # queries see every prefix block and their own top-k video blocks. + scores = torch.randn(2, 3, blocks - prefix, blocks - prefix, device="cuda") + selected = scores.topk(min(topk, blocks - prefix), dim=-1).indices + mask = torch.zeros(2, 3, blocks, blocks, device="cuda", dtype=torch.bool) + mask[:, :, :prefix] = True + mask[:, :, :, :prefix] = True + mask[:, :, prefix:, prefix:].scatter_(-1, selected, True) + return [q, k, v, mask, sizes, 128**-0.5, prefix, topk] + + +@pytest.mark.parametrize( + "lengths,prefix,topk,amplitude", + [ + ((1,), 0, 1, 1), + ((63, 1), 1, 1, 1), + ((17, 64, 3, 64, 1), 2, 1, 1), + ((17, 64, 3, 64, 1), 2, 99, 1), + ((64, 3, 1), 1, 1, 16), + ((64,) * 33 + (17, 64, 3, 64, 1), 2, 8, 1), + ], +) +@torch.inference_mode() +def test_exact_fp32_sparse_math(ops, lengths, prefix, topk, amplitude): + args = make_case(lengths, prefix, topk, amplitude) + expected = fp32_reference(*args[:6]) + actual = ops.forward(*args) + private = ops._forward_prevalidated(*args) + valid = torch.cat( + [torch.arange(i * 64, i * 64 + n, device="cuda") for i, n in enumerate(lengths)] + ) + assert torch.equal( + actual[:, valid].view(torch.int16), expected[:, valid].view(torch.int16) + ) + assert torch.equal(actual.view(torch.int16), private.view(torch.int16)) + assert torch.isfinite(actual).all() + for i, length in enumerate(lengths): + assert not actual[:, i * 64 + length : (i + 1) * 64].count_nonzero() + + +@pytest.mark.parametrize( + "invalid", + [ + "empty_row", + "missing_prefix_key", + "zero_size", + "large_size", + "size_dtype", + "map_dtype", + "map_shape", + "q_dtype", + "q_stride", + "q_grad", + "prefix", + "topk", + "scale", + "unaligned", + ], +) +def test_public_entry_rejects_invalid_inputs(ops, invalid): + args = make_case((64, 17, 64), 1, 1) + q, _, _, mask, sizes, _, _, _ = args + if invalid == "empty_row": + mask[:, :, 1] = False + elif invalid == "missing_prefix_key": + mask[:, :, 1, 0] = False + mask[:, :, 1, 1:] = True # Correct count, wrong prefix selection. + elif invalid == "zero_size": + sizes[1] = 0 + elif invalid == "large_size": + sizes[1] = 65 + elif invalid == "size_dtype": + args[4] = sizes.long() + elif invalid == "map_dtype": + args[3] = mask.int() + elif invalid == "map_shape": + args[3] = mask[..., :-1].contiguous() + elif invalid == "q_dtype": + args[0] = q.float() + elif invalid == "q_stride": + args[0] = q.transpose(1, 2) + elif invalid == "q_grad": + q.requires_grad_() + elif invalid == "prefix": + args[6] = 3 + elif invalid == "topk": + args[7] = 0 + elif invalid == "scale": + args[5] = float("nan") + elif invalid == "unaligned": + args[0] = torch.empty(q.numel() + 1, device=q.device, dtype=q.dtype)[ + 1: + ].view_as(q) + with pytest.raises(RuntimeError): + ops.forward(*args) + torch.accelerator.synchronize() diff --git a/tests/video/test_h3_vsa_geometry.py b/tests/video/test_h3_vsa_geometry.py new file mode 100644 index 0000000000..6cff00017b --- /dev/null +++ b/tests/video/test_h3_vsa_geometry.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""H3 geometry, selection and learned-gate checks without model weights.""" + +import pytest +import torch + +from vllm.model_executor.models.minimax_h3 import vsa + + +@pytest.mark.parametrize("prefix", [(1, 65, 2), (97, 414), ()]) +@pytest.mark.parametrize("shape", [(1, 1, 1), (5, 6, 7), (4, 4, 8)]) +def test_geometry_round_trip_and_segment_boundaries(prefix, shape): + partition, sizes, non_pad, untile, prefix_blocks, video_blocks = ( + vsa._get_h3_tile_metadata(prefix, shape, torch.device("cpu")) + ) + source = torch.arange(sum(prefix) + shape[0] * shape[1] * shape[2]) + tiled = torch.full((sizes.numel() * 64,), -1) + tiled[non_pad] = source[partition] + torch.testing.assert_close(tiled[untile], source) + assert (sizes > 0).all() and (sizes <= 64).all() + assert sizes.sum() == source.numel() + assert prefix_blocks + video_blocks == sizes.numel() + boundaries = torch.tensor(prefix).cumsum(0).tolist() + for block in range(prefix_blocks): + rows = tiled[block * 64 : block * 64 + sizes[block]].tolist() + assert all(not rows[0] < end <= rows[-1] for end in boundaries) + assert max(rows) < sum(prefix) + for block in range(prefix_blocks, sizes.numel()): + rows = tiled[block * 64 : block * 64 + sizes[block]] - sum(prefix) + coords = torch.stack( + ( + rows // (shape[1] * shape[2]), + rows // shape[2] % shape[1], + rows % shape[2], + ), + dim=1, + ) + assert torch.equal(coords[0] // 4, (coords // 4).amin(0)) + assert torch.equal(coords[0] // 4, (coords // 4).amax(0)) + + +def test_prefix_scores_do_not_consume_video_topk_budget(): + scores = torch.zeros(1, 2, 5, 5) + scores[..., :2] = 100000 + scores[..., 4] = 1 + mask = vsa._build_h3_block_map(scores, 2, 3, 1) + assert mask[:, :, :2].all() + assert mask[..., :2].all() + assert mask[:, :, 2:, 4].all() + assert not mask[:, :, 2:, 2:4].any() + assert vsa._build_h3_block_map(scores, 2, 3, 99).all() + + +@pytest.mark.parametrize("gate_value", [0.0, 2.0]) +def test_learned_compression_and_exact_sparse_work(gate_value, monkeypatch): + def selected_blocks(q, k, v, block_map, block_sizes, *, scale): + # Check the sparse operator receives holes as zeros, not repeated tokens. + for i, size in enumerate(block_sizes): + for operand in (q, k, v): + assert not operand[:, i * 64 + size : (i + 1) * 64].count_nonzero() + return torch.ones_like(q) + + monkeypatch.setattr(vsa, "block_sparse_attention", selected_blocks) + x = torch.zeros(1, 8, 2, 128, dtype=torch.float16) + output, work = vsa.h3_vsa_attention( + x, + x, + torch.full_like(x, 3), + prefix_segments=(1, 2), + video_shape=(1, 1, 5), + gate_compress=torch.full_like(x, gate_value), + topk=1, + scale=128**-0.5, + ) + torch.testing.assert_close(output, torch.full_like(x, 1 + 3 * gate_value)) + assert work["prefix_blocks"] == 2 and work["video_blocks"] == 2 + assert work["compression_flops"] == 4 * 2 * 4**2 * 128 + # Two prefix queries per head see all four blocks; two video queries see + # both prefix blocks and exactly one video block. Token counts use real + # lengths even if tied top-k scores choose either video tile. + assert work["selected_blocks"] == 2 * (2 * 4 + 2 * 3) + assert work["selected_token_pairs"] in (2 * (3 * 8 + 5 * 4), 2 * (3 * 8 + 5 * 7)) + + +def test_missing_gate_is_rejected_before_sparse_launch(monkeypatch): + def forbidden(*args, **kwargs): + pytest.fail("invalid geometry must not launch a sparse kernel") + + monkeypatch.setattr(vsa, "block_sparse_attention", forbidden) + x = torch.zeros(1, 1, 2, 128, dtype=torch.float16) + with pytest.raises(ValueError, match="learned compression gate"): + vsa.h3_vsa_attention( + x, + x, + x, + prefix_segments=(), + video_shape=(1, 1, 1), + gate_compress=None, + topk=1, + scale=128**-0.5, + ) diff --git a/tests/video/test_h3_vsa_layout.py b/tests/video/test_h3_vsa_layout.py new file mode 100644 index 0000000000..2a48a00063 --- /dev/null +++ b/tests/video/test_h3_vsa_layout.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Exact H3 layout fusion, request ownership and fallback checks.""" + +import pytest +import torch + +from vllm.model_executor.models.minimax_h3 import vsa + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires a leased SM70 GPU" +) + + +@pytest.mark.parametrize( + "prefix,grid,batch,heads,topk", + [ + ((3, 5), (1, 3, 7), 2, 3, 1), + ((65,), (7, 5, 11), 1, 2, 2), + ((), (1, 1, 1), 1, 1, 1), + ((1, 65, 2), (5, 6, 7), 1, 2, 100), + ], +) +def test_fused_layout_preserves_output_and_work( + prefix, grid, batch, heads, topk, monkeypatch +): + torch.manual_seed(814) + rows = sum(prefix) + grid[0] * grid[1] * grid[2] + q, k, v, gate = [ + torch.randn(batch, rows, heads, 128, device="cuda", dtype=torch.float16) + for _ in range(4) + ] + assert vsa._layout_ops(q, k, v, gate) is not None + kwargs = dict( + prefix_segments=prefix, + video_shape=grid, + gate_compress=gate, + topk=topk, + scale=128**-0.5, + ) + with monkeypatch.context() as patch: + patch.setattr(vsa, "_layout_ops", lambda *args: None) + expected, expected_work = vsa.h3_vsa_attention(q, k, v, **kwargs) + with vsa.h3_vsa_workspace(): + actual, work = vsa.h3_vsa_attention(q, k, v, **kwargs) + saved = actual.clone() + scratch = vsa._layout_scratch( + q, vsa._get_h3_tile_metadata(prefix, grid, q.device)[1].numel() * 64 + ) + vsa.h3_vsa_attention(q * 0.5, k, v, **kwargs) + assert torch.equal(actual.view(torch.int16), saved.view(torch.int16)) + assert scratch is next(iter(vsa._layout_buffers.get().values())) + assert vsa._layout_buffers.get() is None + assert torch.equal(actual.view(torch.int16), expected.view(torch.int16)) + assert all(work[key] == value for key, value in expected_work.items()) + + +def test_workspace_isolates_streams_nested_requests_and_failure(): + q = torch.empty(1, 8, 2, 128, device="cuda", dtype=torch.float16) + with vsa.h3_vsa_workspace(): + first = vsa._layout_scratch(q, 64) + assert first is vsa._layout_scratch(q, 64) + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + assert first is not vsa._layout_scratch(q, 64) + with ( + pytest.raises(RuntimeError, match="request failed"), + vsa.h3_vsa_workspace(), + ): + assert first is not vsa._layout_scratch(q, 64) + raise RuntimeError("request failed") + assert first is vsa._layout_scratch(q, 64) + assert vsa._layout_buffers.get() is None + with vsa.h3_vsa_workspace(): + assert first is not vsa._layout_scratch(q, 64) + + +def test_layout_falls_back_for_views_and_older_wheels(monkeypatch): + q = torch.empty(1, 8, 2, 128, device="cuda", dtype=torch.float16) + assert vsa._layout_ops(q[:, ::2]) is None + unaligned = torch.empty(q.numel() + 1, device=q.device, dtype=q.dtype)[1:].view_as( + q + ) + assert vsa._layout_ops(unaligned) is None + assert vsa._layout_ops(q.requires_grad_()) is None + monkeypatch.setattr(vsa, "sparse_extension", lambda: object()) + assert vsa._layout_ops(q.detach()) is None + + +def test_private_layout_rejects_invalid_shapes_and_aliasing(): + q = torch.empty(1, 64, 2, 128, device="cuda", dtype=torch.float16) + ops = vsa.sparse_extension() + rows = torch.arange(64, device=q.device, dtype=torch.int32) + out = torch.empty(3, *q.shape, device=q.device, dtype=q.dtype) + with pytest.raises(RuntimeError, match="source map"): + ops._h3_tile_qkv_prevalidated(q, q, q, rows[:-1], out) + with pytest.raises(RuntimeError, match="single memory location|overlap"): + ops._h3_tile_qkv_prevalidated(out[0], q, q, rows, out) + with pytest.raises(RuntimeError, match="compressed geometry"): + ops._h3_gate_untile_prevalidated(q, q, q, rows) + + +def test_gate_fusion_preserves_two_fp16_roundings(): + torch.manual_seed(819) + sparse = torch.randn(1, 128, 2, 128, device="cuda", dtype=torch.float16) + compressed = torch.randn(1, 2, 2, 128, device="cuda", dtype=torch.float16) + rows = torch.tensor([0, 1, 17, 64, 65, 100, 127], device="cuda", dtype=torch.int32) + gate = torch.randn(1, 7, 2, 128, device="cuda", dtype=torch.float16) + # Include large finite products and overflow: the fusion must not clamp. + gate[:, 0] = 65504 + gate[:, 1] = -0.0 + expected = sparse[:, rows.long()] + compressed[:, (rows // 64).long()] * gate + actual = vsa.sparse_extension()._h3_gate_untile_prevalidated( + sparse, compressed, gate, rows + ) + assert torch.equal(actual.view(torch.int16), expected.view(torch.int16)) diff --git a/tests/video/test_h3_vsa_stage.py b/tests/video/test_h3_vsa_stage.py new file mode 100644 index 0000000000..ebb8092291 --- /dev/null +++ b/tests/video/test_h3_vsa_stage.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Stage acceptance cannot substitute shorter, captured or mismatched runs.""" + +import pytest + +from tests.video.test_h3_acceptance import measurements +from vllm.model_executor.models.minimax_h3.fasth3 import FastH3Spec +from vllm.model_executor.models.minimax_h3.sigma_schedule import DMD2SigmaSchedule +from vllm.video.vsa_acceptance import evaluate_vsa_stage + + +def stage_measurements(): + result = {} + schedule = DMD2SigmaSchedule(FastH3Spec().base_schedule) + for algorithm in ("vsa", "dense"): + warmup, runs = measurements(calls=4, api_steps=4) + result[algorithm] = dict(warmup=warmup, runs=runs) + for run in (warmup, *runs): + run["config"].update( + attention_backend="FASTVIDEO_VSA" + if algorithm == "vsa" + else "FLASH_ATTN_V100", + vsa_topk=64 if algorithm == "vsa" else None, + lora_path=algorithm + "-datafree", + ) + run["request"]["sampling"].update( + width=1280, height=736, num_frames=120, fps=24 + ) + run["end_to_end_seconds"] = 50 if algorithm == "vsa" else 70 + for rank in run["ranks"]: + rank["stage_seconds"]["denoise"] = 31.3 if rank["rank"] == 3 else 30 + rank["useful_denoise_flops"] = 1_700_000_000_000_000 + rank["denoise_flops_by_layer"] = dict( + example=rank["useful_denoise_flops"] + ) + for step in rank["denoise_steps"]: + step["useful_flops"] = rank["useful_denoise_flops"] // 4 + rank["denoise_workload"].update( + adapter="FastH3Spec", + video_sigmas=schedule.shifted_sigmas(12), + audio_sigmas=schedule.shifted_sigmas(3), + ) + if algorithm == "dense": + continue + rank["denoise_workload"].update( + work_accounting="sparse_tp_v1", + attention_algorithm="vsa", + actual_backends=["FASTVIDEO_VSA", "FLASH_ATTN_V100"], + sparse_config=dict( + topk=64, + prefix_segments=[97, 414], + video_shape=[37, 23, 40], + gated_blocks=50, + heads=14, + head_size=128, + ), + ) + selected_pairs = 14 * (511 * 34551 + 34040 * (511 + 64 * 64)) + selected_blocks = 14 * (9 * 609 + 600 * 73) + dense_pairs = 14 * 34551**2 + compression = 4 * 14 * 609**2 * 128 + rank["denoise_sparse_work_by_layer"] = {} + for layer in range(50): + name = f"blocks.{layer}.attn.attention" + rank["denoise_sparse_work_by_layer"][name] = dict( + head_size=128, + heads=14, + selected_blocks=4 * selected_blocks, + selected_token_pairs=4 * selected_pairs, + compression_flops=4 * compression, + dense_token_pairs=4 * dense_pairs, + ) + flops = 4 * 4 * selected_pairs * 128 + rank["denoise_flops_by_layer"][name] = flops + rank["denoise_flops_by_layer"][name + ".compression"] = ( + 4 * compression + ) + rank["denoise_flops_by_layer"]["example"] -= flops + 4 * compression + for step in rank["denoise_steps"]: + step.update( + sparse_blocks=50 * selected_blocks, + sparse_token_pairs=50 * selected_pairs, + sparse_compression_flops=50 * compression, + attention_avoided_flops=4 + * 128 + * 50 + * (dense_pairs - selected_pairs), + ) + return result + + +def test_stage_uses_slowest_rank_and_keeps_80_tf_goal_separate(): + data = stage_measurements() + report = evaluate_vsa_stage(data["vsa"], data["dense"]) + assert report["stage_performance_passed"] + assert report["denoise_median_seconds"] == 31.3 + assert not report["measurements"]["vsa"]["future_80_tflops_passed"] + assert report["quality_status"] == "requires_independent_numerical_and_human_review" + + +@pytest.mark.parametrize("failure", ["denoise", "request", "variance"]) +def test_stage_rejects_failed_timing_even_with_three_complete_runs(failure): + data = stage_measurements() + for i, run in enumerate(data["vsa"]["runs"]): + if failure == "denoise": + run["ranks"][3]["stage_seconds"]["denoise"] = 31.300001 + elif failure == "request": + run["end_to_end_seconds"] = 70 + else: + for rank in run["ranks"]: + rank["stage_seconds"]["denoise"] = (25, 30, 35)[i] + assert not evaluate_vsa_stage(data["vsa"], data["dense"])[ + "stage_performance_passed" + ] + + +@pytest.mark.parametrize("failure", ["capture", "host", "short", "schedule", "topk"]) +def test_stage_rejects_noncomparable_or_changed_workloads(failure): + data = stage_measurements() + if failure == "capture": + data["vsa"]["runs"][0]["measurement"]["capture"] = True + for name, item in data.items(): + for run in (item["warmup"], *item["runs"]): + if failure == "host" and name == "dense": + run["config"]["host_weight_pin_memory"] = not run["config"][ + "host_weight_pin_memory" + ] + if failure == "short": + run["request"]["sampling"]["num_frames"] = 60 + if failure == "schedule": + for rank in run["ranks"]: + rank["denoise_workload"]["video_sigmas"][1] = 0.97 + if failure == "topk" and name == "vsa": + run["config"]["vsa_topk"] = 32 + with pytest.raises(ValueError): + evaluate_vsa_stage(data["vsa"], data["dense"]) diff --git a/tests/video/test_h3_work_counter.py b/tests/video/test_h3_work_counter.py new file mode 100644 index 0000000000..f50808818a --- /dev/null +++ b/tests/video/test_h3_work_counter.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm.video.metrics import DenoiseWorkCounter, lora_work + + +@pytest.mark.parametrize("rows", [1, 97, 34551]) +@pytest.mark.parametrize("tp", [1, 2, 4]) +def test_column_lora_a_is_counted_once_across_ranks_including_tails(rows, tp): + useful = redundant = 0 + for rank in range(tp): + layer = SimpleNamespace( + tp_size=tp, + tp_rank=rank, + h3_lora_a_0=torch.empty(3, 7), + h3_lora_b_0=torch.empty(8 // tp, 3), + ) + work, repeated = lora_work(layer, [(0, 0, 8 // tp)], rows, replicated_a=True) + useful += work + redundant += repeated + # The logical adapter is one 7->3->8 projection regardless of TP size. + assert useful == 2 * rows * (7 * 3 + 3 * 8) + assert redundant == 2 * rows * (7 * 3) * (tp - 1) + + +@pytest.mark.parametrize("tp", [1, 2, 4]) +def test_row_lora_partial_products_are_distinct_work(tp): + layer = SimpleNamespace( + tp_size=tp, + tp_rank=0, + h3_lora_a_0=torch.empty(3, 8 // tp), + h3_lora_b_0=torch.empty(11, 3), + ) + work, redundant = lora_work(layer, [(0, 0, 11)], 13, replicated_a=False) + assert work == 2 * 13 * (3 * (8 // tp) + 11 * 3) + assert redundant == 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a leased GPU") +@pytest.mark.parametrize("sparse", [False, True]) +@torch.inference_mode() +def test_real_block_work_excludes_padding_and_preserves_outputs( + dist_init, default_vllm_config, sparse +): + from vllm.model_executor.models.minimax_h3.attention import ( + VideoTokenLayout, + VideoTokenSpan, + attention_backend, + ) + from vllm.model_executor.models.minimax_h3.transformer import ( + MiniMaxH3DiTArchConfig, + MiniMaxH3DiTBlock, + ) + + arch = MiniMaxH3DiTArchConfig( + hidden_size=512, + num_attention_heads=4, + ffn_hidden_size=1024, + adaln_curve_grid=2, + adaln_out_features=18 * 512, + ) + + class Model(nn.Module): + def __init__(self): + super().__init__() + self.block = MiniMaxH3DiTBlock(arch, None, prefix="block") + if sparse: + self.block.attn.enable_vsa_gate(1) + + def forward(self, x, **kwargs): + return self.block(x, **kwargs) + + token = attention_backend.set("FASTVIDEO_VSA" if sparse else "FLASH_ATTN_V100") + try: + model = Model().cuda().eval() + finally: + attention_backend.reset(token) + torch.manual_seed(412) + for name, parameter in model.named_parameters(): + if "norm" in name: + parameter.fill_(1) + else: + parameter.normal_(0, 0.02) + valid, padded = 33, 64 + x = torch.randn(padded, 512, device="cuda") + kwargs = dict( + t_emb=torch.randn(1, 8, device="cuda"), + combined_indices=torch.arange(padded, device="cuda") % 3, + rope_table=torch.randn(padded, 96, device="cuda"), + cu_seqlens=torch.tensor([0, valid], device="cuda", dtype=torch.int32), + max_seqlen=valid, + packed_total=padded, + ) + if sparse: + kwargs.update( + video_layout=VideoTokenLayout( + used_len=valid, + video_spans=(VideoTokenSpan(1, (1, 4, 8), "target"),), + ), + vsa_prefix_segments=(1,), + ) + expected = model(x, **kwargs) + # Independent geometry: QKV + output + gate/up + down + AdaLN + QK/PV. + flops = ( + 2 * valid * (3 * 512 * 512 + 512 * 512 + 2 * 1024 * 512 + 512 * 1024) + + 2 * 8 * (18 * 512) + + 4 * 4 * valid * valid * 128 + ) + if sparse: + # A one-token dense prefix and two 16-token video tiles. Each video + # query selects the prefix and exactly one video tile; four heads. + flops -= 4 * 4 * valid * valid * 128 + flops += 4 * 4 * (33 + 32 * 17) * 128 + flops += 2 * valid * 512 * 512 # learned gate projection + flops += 4 * 4 * 3**2 * 128 # pooled QK and pooled PV + with DenoiseWorkCounter( + model, used_length=valid, video_outputs=valid, audio_outputs=1 + ) as counter: + for index in range(2): + with counter.step(index): + actual = model(x, **kwargs) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + torch.accelerator.synchronize() + steps = counter.finish_steps() + assert counter.calls == 2 + assert counter.blocks == {"block": 2} + assert counter.flops == 2 * flops + assert sum(counter.by_layer.values()) == counter.flops + assert [step["useful_flops"] for step in steps] == [flops, flops] + assert all(step["gpu_seconds"] > 0 for step in steps) + assert [step["sparse_blocks"] for step in steps] == ( + [4 * (3 + 2 * 2)] * 2 if sparse else [0, 0] + ) + if sparse: + record = counter.sparse_by_layer["block.attn.attention"] + assert record["selected_token_pairs"] == 2 * 4 * (33 + 32 * 17) + assert record["dense_token_pairs"] == 2 * 4 * 33**2 + assert record["compression_flops"] == 2 * 4 * 4 * 3**2 * 128 + assert sum(step["attention_avoided_flops"] for step in steps) == 2097152 + model(x, **kwargs) + assert counter.calls == 2 # Hooks must not leak into the following request. diff --git a/tests/video/test_sm70_attention.py b/tests/video/test_sm70_attention.py new file mode 100644 index 0000000000..cedad0e16f --- /dev/null +++ b/tests/video/test_sm70_attention.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared attention contracts and non-H3 DiT shapes, without model imports.""" + +import pytest +import torch + +from vllm.model_executor.layers import sm70_attention as ops + + +@pytest.mark.parametrize("scale", [0, -1, float("inf"), float("nan"), 1e100]) +def test_invalid_scale_does_not_load_native_code(monkeypatch, scale): + def unexpected_load(): + pytest.fail("invalid scale must fail before extension loading") + + monkeypatch.setattr(ops, "flashattn_extension", unexpected_load) + monkeypatch.setattr(ops, "flashinfer_extension", unexpected_load) + for backend in ("FLASH_ATTN_V100", "FLASHINFER_SM70"): + with pytest.raises(ValueError, match="scale"): + ops.noncausal_attention(None, None, None, scale=scale, backend=backend) + + +def test_explicit_backend_and_geometry_contract(): + for backend, tile in (("AUTO", 64), ("FLASHINFER_SM70", 128)): + with pytest.raises(ValueError): + ops.noncausal_attention( + None, None, None, scale=0.1, backend=backend, query_tile=tile + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires SM70 GPU") +@pytest.mark.parametrize( + "backend,tile", + [("FLASH_ATTN_V100", 64), ("FLASH_ATTN_V100", 128), ("FLASHINFER_SM70", 64)], +) +@pytest.mark.parametrize("batch,length,heads", [(1, 1537, 24), (2, 65, 8)]) +def test_gpu_non_h3_shapes_preserve_native_output(backend, tile, batch, length, heads): + torch.manual_seed(206) + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False + # Strided token storage and non-H3 head counts exercise the shared contract. + q, k, v = [ + torch.randn(batch, length * 2, heads, 128, device="cuda", dtype=torch.float16)[ + :, ::2 + ] + for _ in range(3) + ] + k[:, length // 2 :] *= 4 + scale = 128**-0.5 + actual = ops.noncausal_attention( + q, k, v, scale=scale, backend=backend, query_tile=tile + ) + native = ( + ops.flashattn_extension() + if backend == "FLASH_ATTN_V100" + else ops.flashinfer_extension() + ) + args = (scale, 0, tile) if tile != 64 else (scale,) + expected = native.forward(q.contiguous(), k.contiguous(), v.contiguous(), *args) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + rows = torch.linspace(0, length - 1, min(length, 33), device="cuda").long() + qh, kh, vh = (x.transpose(1, 2).float() for x in (q[:, rows], k, v)) + reference = (((qh @ kh.transpose(-1, -2)) * scale).softmax(-1) @ vh).transpose(1, 2) + assert torch.isfinite(actual).all() + relative_l2 = (actual[:, rows].float() - reference).norm() / reference.norm() + assert relative_l2 < 0.001 + with pytest.raises(RuntimeError, match="FP16"): + ops.noncausal_attention( + q.float(), + k.float(), + v.float(), + scale=scale, + backend=backend, + query_tile=tile, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires SM70 GPU") +def test_gpu_cross_attention_contract(): + torch.manual_seed(207) + q = torch.randn(2, 33, 8, 128, device="cuda", dtype=torch.float16) + k, v = [ + torch.randn(2, 129, 8, 128, device="cuda", dtype=torch.float16) + for _ in range(2) + ] + actual = ops.noncausal_attention( + q, k, v, scale=0.1, backend="FLASH_ATTN_V100", query_tile=128 + ) + expected = ops.flashattn_extension().forward(q, k, v, 0.1, 0, 128) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + with pytest.raises(RuntimeError, match="matching"): + ops.noncausal_attention(q, k, v, scale=0.1, backend="FLASHINFER_SM70") diff --git a/tests/video/test_sm70_collectives.py b/tests/video/test_sm70_collectives.py new file mode 100644 index 0000000000..6d63a8bca6 --- /dev/null +++ b/tests/video/test_sm70_collectives.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Host checks for the explicit, non-automatic SM70 collective interface.""" + +import numpy as np +import pytest + +from vllm.model_executor.layers.sm70_collective_calibration import PROBES +from vllm.model_executor.layers.sm70_collectives import _layout + + +@pytest.mark.parametrize( + "shape", [(0, 8), (3, 8), (4, 0), (-4, 8), (4,), (True, 8), (4, 2**63)] +) +def test_reject_invalid_layout(shape): + with pytest.raises((TypeError, ValueError)): + _layout(shape) + + +def test_budget_covers_ipc_output_codes_and_calibration(): + shape, raw, resident, calibration = _layout((34560, 5376)) + assert shape == (34560, 5376) + assert raw == 743_180_800 + assert resident > raw + assert calibration > resident + 2 * 34560 * 5376 * 4 + + +def _trees(indices): + if len(indices) == 1: + return [indices[0]] + result: list[tuple] = [] + # Anchor the first leaf on the left to remove commutative duplicates. + for mask in range(1, (1 << len(indices)) - 1, 2): + left = tuple(x for i, x in enumerate(indices) if mask & (1 << i)) + right = tuple(x for i, x in enumerate(indices) if not mask & (1 << i)) + result.extend((a, b) for a in _trees(left) for b in _trees(right)) + return result + + +def _evaluate(tree, values): + if isinstance(tree, int): + return np.float32(values[tree]) + return np.float32(_evaluate(tree[0], values) + _evaluate(tree[1], values)) + + +def test_fixed_probes_cover_all_fp32_addition_trees(): + trees = _trees((0, 1, 2, 3)) + assert len(trees) == 15 + actual = { + tuple(int(_evaluate(tree, values).view(np.uint32)) for values, _ in PROBES) + for tree in trees + } + stored = {tuple(bits[i] for _, bits in PROBES) for i in range(15)} + assert len(stored) == 15 + assert actual == stored diff --git a/tests/video/test_sm70_scaled_add.py b/tests/video/test_sm70_scaled_add.py new file mode 100644 index 0000000000..d99e93f37e --- /dev/null +++ b/tests/video/test_sm70_scaled_add.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + +import pytest +import torch + +from vllm.model_executor.layers.sm70_diffusion import sm70_extension + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires a leased SM70 GPU" +) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.float32]) +@pytest.mark.parametrize("alpha", [0.0, 0.0625, 0.75, -0.5, 1.0]) +@pytest.mark.parametrize("scaled", [False, True]) +def test_scaled_add_retains_fp32_rounding_and_untouched_slices(dtype, alpha, scaled): + torch.manual_seed(610) + rows, columns, width, offset = 17, 43, 33, 3 + storage = torch.randn(rows * columns + 1, device="cuda", dtype=dtype) + actual = storage[1:].reshape(rows, columns) + base = actual.clone() + delta = torch.randn(rows, width, device="cuda") * 100 + scales = ( + torch.ldexp( + torch.ones(rows, 1, device="cuda"), + (torch.arange(rows, device="cuda") - 8)[:, None], + ) + if scaled + else None + ) + restored = delta if scales is None else delta * scales + expected = base.float().clone() + expected[:, offset : offset + width].add_(restored, alpha=alpha) + result = sm70_extension().scaled_add_(actual, delta, scales, alpha, offset) + assert result.data_ptr() == actual.data_ptr() + torch.testing.assert_close(actual, expected.to(dtype), rtol=0, atol=0) + torch.testing.assert_close(actual[:, :offset], base[:, :offset], rtol=0, atol=0) + torch.testing.assert_close( + actual[:, offset + width :], base[:, offset + width :], rtol=0, atol=0 + ) + + +def test_scaled_add_rejects_unsafe_memory_aliases_and_bad_bounds(): + data = torch.zeros(65, device="cuda") + output = data[:-1].reshape(8, 8) + overlapping = data[1:].reshape(8, 8) + op = sm70_extension().scaled_add_ + with pytest.raises(RuntimeError): + op(output, overlapping, None, 1.0, 0) + with pytest.raises(RuntimeError): + op(output.view(torch.float16), output, None, 1.0, 0) + with pytest.raises(RuntimeError): + op(output, torch.ones_like(output), data[:8], 1.0, 0) + with pytest.raises(RuntimeError, match="outside output"): + op(output, torch.ones_like(output), None, 1.0, 1) + with pytest.raises(RuntimeError, match="finite FP32"): + op(output, torch.ones_like(output), None, float("nan"), 0) + + +def test_scaled_add_graph_reads_current_operands(): + output = torch.zeros(16, 192, device="cuda", dtype=torch.float16) + delta = torch.randn(16, 192, device="cuda") + scales = torch.ones(16, 1, device="cuda") * 8 + op = sm70_extension().scaled_add_ + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + op(output, delta, scales, 0.75, 0) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(output, delta, scales, 0.75, 0) + for factor in (1.0, -2.0): + output.fill_(factor) + delta.mul_(factor) + expected = output.float().add(delta * scales, alpha=0.75).half() + graph.replay() + torch.testing.assert_close(output, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("quantized", [False, True]) +@pytest.mark.parametrize("output_fp32", [False, True]) +@pytest.mark.parametrize("overlapping", [False, True]) +def test_adapter_fusion_matches_unfused_rounding( + quantized, output_fp32, overlapping, monkeypatch +): + import vllm.model_executor.models.minimax_h3.lora as adapter + from vllm.model_executor.models.minimax_h3.quantization import ( + DiffusionInt8ConvRotConfig, + FP16LinearMethod, + FP32OutputLinearMethod, + Int8ConvRotLayerConfig, + Int8ConvRotLinearMethod, + ) + + torch.manual_seed(913) + if quantized: + base = Int8ConvRotLinearMethod( + DiffusionInt8ConvRotConfig(), Int8ConvRotLayerConfig(True), prefix="probe" + ) + weight = torch.randint(-7, 8, (384, 256), device="cuda", dtype=torch.int8) + else: + base = FP32OutputLinearMethod() if output_fp32 else FP16LinearMethod() + weight = torch.randn(384, 256, device="cuda", dtype=torch.float16) * 0.03 + layer = SimpleNamespace( + weight=weight, + weight_scale=torch.full((384,), 0.01, device="cuda"), + h3_output_fp32=output_fp32, + ) + parts = [] + for i in range(3): + setattr(layer, f"h3_lora_a_{i}", torch.randn(8, 256, device="cuda").half()) + setattr(layer, f"h3_lora_b_{i}", torch.randn(96, 8, device="cuda").half()) + parts.append((i, 32 if overlapping else 16 + i * 112, 96)) + method = adapter.TurboLinearMethod(base, parts, 0.125) + inputs = torch.randn(33, 256, device="cuda", dtype=torch.float16) + fused_add = adapter.fp16_linear_add + calls = [] + + def record(*args, **kwargs): + calls.append(kwargs["offset"]) + return fused_add(*args, **kwargs) + + monkeypatch.setattr(adapter, "fp16_linear_add", record) + token = adapter.lora_scale.set(-0.75) + try: + with monkeypatch.context() as context: + context.setattr(adapter, "supports_fused_scaled_add", lambda _: False) + expected = method.apply(layer, inputs) + actual = method.apply(layer, inputs) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert actual.dtype == (torch.float32 if output_fp32 else torch.float16) + assert calls == ([] if overlapping else [16, 128, 240]) + calls.clear() + actual = method.apply_prepared(layer, inputs, None) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert calls == ([] if overlapping else [16, 128, 240]) + finally: + adapter.lora_scale.reset(token) diff --git a/tests/video/test_sm70_sparse_attention.py b/tests/video/test_sm70_sparse_attention.py new file mode 100644 index 0000000000..4cf1e82b7b --- /dev/null +++ b/tests/video/test_sm70_sparse_attention.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Independent masked FP32 reference for real sparse SM70 execution.""" + +import pytest +import torch + +from vllm.model_executor.layers.sm70_sparse_attention import block_sparse_attention + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires a leased SM70 GPU" +) + + +def _reference(q, k, v, block_map, sizes): + result = torch.zeros_like(q) + for batch in range(q.shape[0]): + for head in range(q.shape[2]): + for query_block, query_size in enumerate(sizes): + indices = [ + torch.arange(i * 64, i * 64 + size, device=q.device) + for i, size in enumerate(sizes) + if block_map[batch, head, query_block, i] + ] + selected = torch.cat(indices) + start = query_block * 64 + queries = q[batch, start : start + query_size, head].float() + keys = k[batch, selected, head].float() + values = v[batch, selected, head].float() + scores = (queries @ keys.T) * 128**-0.5 + result[batch, start : start + query_size, head] = ( + scores.softmax(-1) @ values + ).half() + return result + + +@pytest.mark.parametrize("sizes", [(1,), (63, 1), (17, 64, 3, 64, 1)]) +@pytest.mark.parametrize("dense", [False, True]) +def test_sparse_masks_and_nonterminal_edges(sizes, dense): + torch.manual_seed(711) + blocks = len(sizes) + q, k, v = [ + torch.randn(2, blocks * 64, 3, 128, device="cuda", dtype=torch.float16) + for _ in range(3) + ] + # Poison every padding region, including holes before later valid blocks. + for i, size in enumerate(sizes): + for operand in (q, k, v): + operand[:, i * 64 + size : (i + 1) * 64] = float("nan") + block_map = torch.ones(2, 3, blocks, blocks, device="cuda", dtype=torch.bool) + if not dense and blocks > 1: + block_map[:, :, 1:, ::2] = False + block_map[:, :, 1:, -1] = True + block_map[1, 2, 0, :-1] = False + expected = _reference(q, k, v, block_map, sizes) + actual = block_sparse_attention( + q, + k, + v, + block_map, + torch.tensor(sizes, device="cuda", dtype=torch.int32), + scale=128**-0.5, + ) + assert torch.isfinite(actual).all() + assert (actual.float() - expected.float()).norm() / expected.float().norm() < 0.001 + torch.testing.assert_close(actual, expected, rtol=0.003, atol=0.003) + for i, size in enumerate(sizes): + assert not torch.count_nonzero(actual[:, i * 64 + size : (i + 1) * 64]) + + +def test_sparse_rejects_empty_rows_and_invalid_sizes_before_launch(): + x = torch.zeros(1, 128, 2, 128, device="cuda", dtype=torch.float16) + mask = torch.ones(1, 2, 2, 2, device="cuda", dtype=torch.bool) + sizes = torch.tensor([64, 1], device="cuda", dtype=torch.int32) + mask[:, :, 1] = False + with pytest.raises(RuntimeError, match="selected key"): + block_sparse_attention(x, x, x, mask, sizes, scale=128**-0.5) + mask.fill_(True) + sizes[1] = 65 + with pytest.raises(RuntimeError, match=r"\[1,64\]"): + block_sparse_attention(x, x, x, mask, sizes, scale=128**-0.5) + + +def test_sparse_all_blocks_matches_existing_64_key_arithmetic(): + from vllm.model_executor.models.minimax_h3.cuda_ops import flashattn_extension + + torch.manual_seed(509) + q, k, v = [ + torch.randn(1, 256, 2, 128, device="cuda", dtype=torch.float16) + for _ in range(3) + ] + expected = flashattn_extension().forward(q, k, v, 128**-0.5, 64) + actual = block_sparse_attention( + q, + k, + v, + torch.ones(1, 2, 4, 4, device="cuda", dtype=torch.bool), + torch.full((4,), 64, device="cuda", dtype=torch.int32), + scale=128**-0.5, + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_h3_owned_sparse_route_preserves_output_without_host_scalar_reads(): + from vllm.model_executor.layers.sm70_sparse_attention import ( + _h3_block_sparse_attention, + sparse_extension, + ) + + assert hasattr(sparse_extension(), "_forward_prevalidated") + torch.manual_seed(814) + q, k, v = [ + torch.randn(1, 192, 2, 128, device="cuda", dtype=torch.float16) + for _ in range(3) + ] + sizes = torch.tensor([17, 64, 3], device="cuda", dtype=torch.int32) + mask = torch.ones(1, 2, 3, 3, device="cuda", dtype=torch.bool) + mask[:, :, 1:, 1] = False + expected = block_sparse_attention(q, k, v, mask, sizes, scale=128**-0.5) + _h3_block_sparse_attention(q, k, v, mask, sizes, scale=128**-0.5) + torch.accelerator.synchronize() + with torch.profiler.profile( + activities=[torch.profiler.ProfilerActivity.CPU] + ) as prof: + actual = _h3_block_sparse_attention(q, k, v, mask, sizes, scale=128**-0.5) + assert not any( + event.key == "aten::_local_scalar_dense" for event in prof.key_averages() + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + with pytest.raises(RuntimeError, match="block map"): + _h3_block_sparse_attention(q, k, v, mask[:, :, :2], sizes, scale=128**-0.5) diff --git a/vllm/entrypoints/cli/video.py b/vllm/entrypoints/cli/video.py index 01e1477d20..8dcd9cba12 100644 --- a/vllm/entrypoints/cli/video.py +++ b/vllm/entrypoints/cli/video.py @@ -36,18 +36,59 @@ def subparser_init(self, subparsers): mode.add_argument("--tensor-parallel-size", "-tp", type=int, default=4) mode.add_argument( "--attention-backend", - choices=("FLASH_ATTN_V100", "FLASHINFER_SM70", "TORCH_SDPA"), + choices=( + "FLASH_ATTN_V100", + "FLASHINFER_SM70", + "TORCH_SDPA", + "FASTVIDEO_VSA", + ), default="FLASH_ATTN_V100", ) + mode.add_argument("--fastvideo-vsa-topk", type=int, default=64) mode.add_argument("--fp16-weight-cache-gib", type=float, default=0) + mode.add_argument( + "--attention-query-tile", type=int, choices=(64, 128), default=64 + ) + mode.add_argument( + "--weight-offload", + choices=("component", "layer"), + default="component", + help="Stage complete components or individual DiT/encoder layers", + ) + mode.add_argument( + "--share-host-vae-weights", + action="store_true", + help="Share immutable pageable VAE masters across TP workers", + ) + mode.add_argument( + "--disable-host-weight-pinning", + dest="host_weight_pin_memory", + action="store_false", + help="Keep weight masters pageable when pinned copies exceed host RAM", + ) mode.add_argument("--fp16-cache-layer", action="append", default=[]) mode.add_argument( "--int8-weight-layout", choices=["row", "column"], default="column" ) + mode.add_argument( + "--fp16-weight-layout", choices=["row", "column"], default="row" + ) mode.add_argument( "--residual-sequence-parallel", action="store_true", - help=("Experimental FP32 residual sharding for TP4 FL2VA INT8"), + help="Experimental FP32 residual sharding for TP2/TP4; TP1 is a no-op", + ) + mode.add_argument( + "--residual-reduction", + choices=("native", "peer"), + default="native", + help="Explicit TP4 SM70 row reduction; requires residual sharding", + ) + mode.add_argument( + "--residual-reduction-memory-gib", + type=float, + default=4.0, + help="Communication setup and buffer budget; larger shapes use native", ) mode.add_argument("--output-dir", type=Path, default=Path("h3-output")) mode.add_argument( @@ -112,12 +153,20 @@ def cmd(args): transformer_path=args.transformer_path, tensor_parallel_size=args.tensor_parallel_size, attention_backend=args.attention_backend, + vsa_topk=args.fastvideo_vsa_topk, + attention_query_tile=args.attention_query_tile, fp16_weight_cache_gib=args.fp16_weight_cache_gib, fp16_cache_layers=tuple(args.fp16_cache_layer), lora_path=args.lora_path, int8_weight_layout=args.int8_weight_layout, + fp16_weight_layout=args.fp16_weight_layout, residual_sequence_parallel=args.residual_sequence_parallel, + residual_reduction=args.residual_reduction, + residual_reduction_memory_gib=args.residual_reduction_memory_gib, video_encoder=args.video_encoder, + host_weight_pin_memory=args.host_weight_pin_memory, + share_host_vae_weights=args.share_host_vae_weights, + weight_offload=args.weight_offload, host_memory_mode=args.host_memory_mode, host_memory_directory=args.host_memory_directory, ) diff --git a/vllm/model_executor/layers/sm70_attention.py b/vllm/model_executor/layers/sm70_attention.py new file mode 100644 index 0000000000..ecf654d838 --- /dev/null +++ b/vllm/model_executor/layers/sm70_attention.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Explicit model-independent SM70 FP16 non-causal D128 attention. + +The extension ABI retains historical H3 names. Inputs are BSHD; no model, +quantization, adapter or sampler identity participates in dispatch. +""" + +import math +import os +from functools import lru_cache +from importlib import import_module +from pathlib import Path + + +@lru_cache(maxsize=1) +def flashinfer_extension(): + try: + return import_module("vllm._h3_flashinfer_C") + except ImportError: + pass + from torch.utils.cpp_extension import load + + root = Path(__file__).resolve().parents[3] + return load( + name="onecat_h3_flashinfer_sm70", + sources=[str(root / "flashinfer-sm70/csrc/h3_noncausal_sm70.cu")], + extra_include_paths=[str(root / "flashinfer-sm70/include")], + extra_cuda_cflags=["-O3", "-gencode=arch=compute_70,code=sm_70"], + verbose=False, + ) + + +@lru_cache(maxsize=1) +def flashattn_extension(): + try: + return import_module("vllm._h3_flashattn_C") + except ImportError: + pass + from torch.utils.cpp_extension import load + + root = Path(__file__).resolve().parents[3] + cutlass_path = os.environ.get("VLLM_CUTLASS_SRC_DIR") + if not cutlass_path: + raise RuntimeError( + "Build the H3 FlashAttention-V100 extension (_h3_flashattn_C), or " + "set VLLM_CUTLASS_SRC_DIR to CUTLASS v4.4.2 for source development" + ) + cutlass = Path(cutlass_path) + return load( + name="onecat_h3_flashattn_sm70", + sources=[str(root / "flash-attention-v100/kernel/h3/forward.cu")], + extra_include_paths=[ + str(cutlass / "include"), + str(cutlass / "examples/41_fused_multi_head_attention"), + ], + extra_cuda_cflags=["-O3", "-gencode=arch=compute_70,code=sm_70"], + verbose=False, + ) + + +def noncausal_attention(q, k, v, *, scale, backend, query_tile=64, key_tile=0): + """Run an explicitly selected SM70 implementation without changing precision. + + CUDA entrypoints validate device, FP16 dtype, D128 heads, index limits and + shapes. FlashAttention supports unequal Q/K lengths and strided inputs; + FlashInfer supports matching lengths and receives contiguous BSHD tensors. + Padding/masking and model-specific sparse selection belong to callers. + The default four-argument call remains compatible with existing wheels. + """ + if not math.isfinite(scale) or not 0 < scale <= 3.4028234663852886e38: + raise ValueError("Attention scale must be positive and finite in FP32") + if backend == "FLASH_ATTN_V100": + if query_tile not in (64, 128) or key_tile not in (0, 64, 128): + raise ValueError("Unsupported SM70 attention tile") + ops = flashattn_extension() + if query_tile == 64 and key_tile == 0: + return ops.forward(q, k, v, scale) + return ops.forward(q, k, v, scale, key_tile, query_tile) + if backend == "FLASHINFER_SM70": + if query_tile != 64 or key_tile != 0: + raise ValueError("Explicit attention tiles require FLASH_ATTN_V100") + return flashinfer_extension().forward( + q.contiguous(), k.contiguous(), v.contiguous(), scale + ) + raise ValueError(f"Unsupported SM70 attention backend: {backend}") diff --git a/vllm/model_executor/layers/sm70_collective_calibration.py b/vllm/model_executor/layers/sm70_collective_calibration.py new file mode 100644 index 0000000000..119840ac36 --- /dev/null +++ b/vllm/model_executor/layers/sm70_collective_calibration.py @@ -0,0 +1,276 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GPU classification of FP32 addition trees for explicit SM70 row plans.""" + +from vllm.triton_utils import tl, triton + + +@triton.jit +def update_mask( + reference, + masks, + expected, + N: tl.constexpr, + OFFSET: tl.constexpr, + BLOCK: tl.constexpr, +): + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + bits = tl.load(reference + OFFSET + i, i < N, other=0).to(tl.uint32, bitcast=True) + mask = tl.load(masks + i, i < N, other=0).to(tl.uint32) + for tree in tl.static_range(15): + correct = bits == tl.load(expected + tree) + mask = mask & tl.where(correct, 0x7FFF, 0x7FFF ^ (1 << tree)).to(tl.uint32) + tl.store(masks + i, mask.to(tl.int16), i < N) + + +@triton.jit +def decode_mask(masks, codes, stats, N: tl.constexpr, BLOCK: tl.constexpr): + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = tl.load(masks + i, i < N, other=0).to(tl.uint32) + unique = (mask != 0) & ((mask & (mask - 1)) == 0) + code = tl.full((BLOCK,), 255, tl.uint8) + for tree in tl.static_range(15): + code = tl.where(mask == (1 << tree), tree, code).to(tl.uint8) + tl.store(codes + i, code, i < N) + tl.atomic_add(stats, tl.sum(((i < N) & (mask == 0)).to(tl.int32))) + tl.atomic_add(stats + 1, tl.sum(((i < N) & (mask != 0) & (~unique)).to(tl.int32))) + + +# Fixed finite probes distinguish all 15 four-input binary addition trees. +PROBES = ( + ( + (-55.133113861083984, -124364.703125, -361660.78125, 0.006582669448107481), + ( + 3370997780, + 3370997780, + 3370997779, + 3370997780, + 3370997779, + 3370997780, + 3370997780, + 3370997780, + 3370997780, + 3370997780, + 3370997780, + 3370997780, + 3370997780, + 3370997780, + 3370997780, + ), + ), + ( + ( + -1.0654968036760692e-06, + 7.75692081451416, + -3.822844155365601e-05, + 0.0002807896235026419, + ), + ( + 1090009773, + 1090009773, + 1090009773, + 1090009773, + 1090009773, + 1090009773, + 1090009773, + 1090009773, + 1090009773, + 1090009773, + 1090009773, + 1090009773, + 1090009772, + 1090009772, + 1090009772, + ), + ), + ( + ( + 2.3551223193862825e-07, + 142.72190856933594, + -489.4043884277344, + 77.70187377929688, + ), + ( + 3280371076, + 3280371077, + 3280371076, + 3280371076, + 3280371076, + 3280371077, + 3280371077, + 3280371077, + 3280371077, + 3280371076, + 3280371076, + 3280371076, + 3280371076, + 3280371076, + 3280371076, + ), + ), + ( + ( + 204.39108276367188, + -0.6318132877349854, + -46.683990478515625, + 0.0033716242760419846, + ), + ( + 1125979170, + 1125979170, + 1125979170, + 1125979170, + 1125979170, + 1125979170, + 1125979171, + 1125979170, + 1125979170, + 1125979170, + 1125979170, + 1125979170, + 1125979170, + 1125979170, + 1125979170, + ), + ), + ( + (-1198.590576171875, -0.01978362910449505, 12.744208335876465, -956837.625), + ( + 3379160183, + 3379160183, + 3379160183, + 3379160184, + 3379160184, + 3379160183, + 3379160184, + 3379160184, + 3379160184, + 3379160183, + 3379160184, + 3379160183, + 3379160183, + 3379160184, + 3379160183, + ), + ), + ( + ( + 4.6832619204906223e-07, + -0.45981907844543457, + -60.8045768737793, + -15.982765197753906, + ), + ( + 3264904843, + 3264904844, + 3264904844, + 3264904843, + 3264904844, + 3264904844, + 3264904844, + 3264904844, + 3264904844, + 3264904844, + 3264904844, + 3264904844, + 3264904843, + 3264904843, + 3264904843, + ), + ), + ( + (-185.40750122070312, -7355930.0, 2.246054172515869, 54945.8671875), + ( + 3403599967, + 3403599967, + 3403599967, + 3403599967, + 3403599966, + 3403599967, + 3403599967, + 3403599967, + 3403599966, + 3403599967, + 3403599967, + 3403599967, + 3403599967, + 3403599967, + 3403599967, + ), + ), + ( + ( + -2.810704131661623e-07, + -1.8851414651521736e-08, + 1.4957002303361833e-10, + 1.1888499784618034e-07, + ), + ( + 3024239082, + 3024239082, + 3024239082, + 3024239081, + 3024239080, + 3024239081, + 3024239082, + 3024239080, + 3024239080, + 3024239081, + 3024239081, + 3024239081, + 3024239081, + 3024239080, + 3024239081, + ), + ), + ( + ( + 125.35652160644531, + -16.128854751586914, + -30.964479446411133, + -1.538109358989459e-06, + ), + ( + 1117554368, + 1117554368, + 1117554368, + 1117554368, + 1117554368, + 1117554368, + 1117554368, + 1117554369, + 1117554369, + 1117554368, + 1117554369, + 1117554369, + 1117554369, + 1117554369, + 1117554369, + ), + ), + ( + ( + -10.238574981689453, + 7.449639797210693, + 0.18453934788703918, + 0.0034961116034537554, + ), + ( + 3223745828, + 3223745828, + 3223745828, + 3223745828, + 3223745826, + 3223745828, + 3223745828, + 3223745828, + 3223745826, + 3223745828, + 3223745828, + 3223745828, + 3223745830, + 3223745826, + 3223745826, + ), + ), +) diff --git a/vllm/model_executor/layers/sm70_collectives.py b/vllm/model_executor/layers/sm70_collectives.py new file mode 100644 index 0000000000..ba47473030 --- /dev/null +++ b/vllm/model_executor/layers/sm70_collectives.py @@ -0,0 +1,270 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Explicit, calibrated FP32 local-row reduction for SM70 TP4 callers. + +This experimental interface has no automatic dispatch. Prepare collectively +outside CUDA graphs, use one bound stream, consume each returned view before +calling again, and close collectively before destroying the process group. +Calibration depends on the communicator and shape, never on model values. +""" + +import sys +from functools import lru_cache +from importlib import import_module +from operator import index +from pathlib import Path +from typing import Any + +import torch +import torch.distributed as dist + + +def _layout(shape): + if len(shape) != 2 or any(isinstance(value, bool) for value in shape): + raise ValueError("Exact row reduction requires two integer dimensions") + rows, columns = map(index, shape) + if rows <= 0 or columns <= 0 or rows % 4: + raise ValueError("Exact row reduction requires positive TP4-aligned rows") + count = rows * columns + if count > (2**63 - 1) // 4: + raise ValueError("Exact row reduction allocation size overflows int64") + local = count // 4 + raw = count * 4 + 2 * 80 * 4 * 4 + resident = raw + local * 5 + calibration_peak = resident + count * 8 + local * 2 + 256 + return (rows, columns), raw, resident, calibration_peak + + +@lru_cache(maxsize=1) +def _extension(): + try: + return import_module("vllm._sm70_exact_reduce_C") + except ImportError: + from torch.utils.cpp_extension import load + + root = Path(__file__).resolve().parents[3] + source = root / "csrc/sm70_turbomind/ops/exact_row_reduce.cu" + if not source.is_file(): + raise RuntimeError( + "SM70 exact row reduction requires the source build" + ) from None + extension = load( + name="onecat_sm70_exact_reduce", + sources=[str(source)], + extra_cuda_cflags=[ + "-O3", + "--fmad=false", + "-gencode=arch=compute_70,code=sm_70", + ], + verbose=False, + ) + sys.modules["onecat_sm70_exact_reduce"] = extension + return extension + + +class SM70ExactRowReductionPlan: + """Own one shape's IPC buffers and calibrated native FP32 addition order. + + ``group`` is a vLLM TP group with four ranks and a CPU process group. + ``memory_budget_bytes`` must cover persistent buffers and calibration + scratch. Callers retain their ordinary collective when this explicit plan + is unsuitable. Returned tensors alias plan storage. CUDA graphs and use + from another stream/device are rejected rather than silently changing + synchronization semantics. + """ + + @staticmethod + def required_memory_bytes(shape): + """Conservative explicit budget for setup scratch and persistent buffers.""" + return _layout(shape)[3] + + def __init__(self, group, shape, *, memory_budget_bytes): + error = None + try: + self.shape, self.raw_ipc_bytes, self.resident_bytes, peak = _layout(shape) + budget = index(memory_budget_bytes) + if isinstance(memory_budget_bytes, bool) or budget < peak: + error = "Exact row reduction exceeds the explicit memory budget" + except (TypeError, ValueError, OverflowError) as exc: + self.shape, peak = None, 0 + error = str(exc) + self.calibration_peak_bytes = peak + self.group = group + self.rank = group.rank_in_group + self.device = torch.accelerator.current_device_index() + self.stream = torch.cuda.current_stream(self.device).cuda_stream + self._closed = False + self._buffers = [] + self._epoch = 0 + self.calls = 0 + properties = torch.cuda.get_device_properties(self.device) + if group.world_size != 4 or not 0 <= self.rank < 4: + error = "Exact row reduction requires TP4" + elif (properties.major, properties.minor) != (7, 0): + error = "Exact row reduction requires SM70" + elif properties.multi_processor_count < 80: + error = "Exact row reduction requires at least 80 SMs" + elif torch.cuda.is_current_stream_capturing(): + error = "Prepare exact row reduction outside CUDA graphs" + metadata: list[Any] = [None] * group.world_size + with torch.inference_mode(False): + dist.all_gather_object( + metadata, + (self.shape, self.device, error, str(properties.uuid)), + group=group.cpu_group, + ) + errors = [item[2] for item in metadata if item[2]] + if errors or any(item[0] != self.shape for item in metadata): + raise ValueError(errors or "Ranks requested different reduction shapes") + if len({item[3] for item in metadata}) != 4 or any( + not 0 <= item[1] < torch.accelerator.device_count() + or str(torch.cuda.get_device_properties(item[1]).uuid) != item[3] + for item in metadata + ): + error = ( + "Exact row reduction requires one host with " + "consistent CUDA device visibility" + ) + elif any( + peer != self.rank + and not torch.cuda.can_device_access_peer(self.device, item[1]) + for peer, item in enumerate(metadata) + ): + error = "Exact row reduction requires peer access to every TP rank" + self._agree(error) + self.ops = _extension() + self.output = None + self.codes = None + try: + with torch.inference_mode(False): + count = self.shape[0] * self.shape[1] + self.pointers = self._shared(count * 4) + self.flags = self._shared(2 * 80 * 4 * 4) + self.output = torch.empty( + (self.shape[0] // 4, self.shape[1]), + device=self.device, + dtype=torch.float32, + ) + self.codes = self._calibrate(count) + dist.barrier(group=group.cpu_group) + except BaseException: + self.close() + raise + + def _agree(self, error): + errors = [None] * self.group.world_size + with torch.inference_mode(False): + dist.all_gather_object(errors, error, group=self.group.cpu_group) + if any(errors): + raise RuntimeError(f"Exact row reduction setup failed: {errors}") + + def _shared(self, size): + pointer, handle, error = 0, None, None + try: + pointer, handle = self.ops.allocate(size) + torch.accelerator.synchronize() + except RuntimeError as exc: + error = str(exc) + handles: list[Any] = [None] * 4 + dist.all_gather_object(handles, (handle, error), group=self.group.cpu_group) + if any(item[1] for item in handles): + if pointer: + self.ops.release(pointer, True) + raise RuntimeError(f"Exact row reduction IPC allocation failed: {handles}") + pointers = [0] * 4 + pointers[self.rank] = pointer + self._buffers.append(pointers) + for peer, (handle, _) in enumerate(handles): + if peer != self.rank: + try: + pointers[peer] = self.ops.open_handle(handle) + except RuntimeError as exc: + error = str(exc) + break + self._agree(error) + return pointers + + def _calibrate(self, count): + from vllm.triton_utils import triton + + from .sm70_collective_calibration import PROBES, decode_mask, update_mask + + n = count // 4 + masks = torch.full((n,), 0x7FFF, device=self.device, dtype=torch.int16) + codes = torch.empty(n, device=self.device, dtype=torch.uint8) + stats = torch.zeros(2, device=self.device, dtype=torch.int64) + value = torch.empty(self.shape, device=self.device, dtype=torch.float32) + for inputs, expected_bits in PROBES: + value.fill_(inputs[self.rank]) + reference = self.group.all_reduce(value) + expected = torch.tensor( + expected_bits, device=self.device, dtype=torch.uint32 + ) + update_mask[(triton.cdiv(n, 256),)]( + reference, masks, expected, n, n * self.rank, 256 + ) + del reference, expected + decode_mask[(triton.cdiv(n, 256),)](masks, codes, stats, n, 256) + self._agree( + None + if stats.tolist() == [0, 0] + else "Native FP32 addition order cannot be classified uniquely" + ) + return codes + + def reduce(self, value): + """Return this rank's local rows; adapters must already be included.""" + if self._closed: + raise RuntimeError("Exact row reduction plan is closed") + if ( + not value.is_cuda + or value.device.index != self.device + or value.dtype != torch.float32 + or tuple(value.shape) != self.shape + or not value.is_contiguous() + or value.requires_grad + ): + raise ValueError("Exact row reduction requires the prepared FP32 layout") + if ( + torch.accelerator.current_device_index() != self.device + or torch.cuda.current_stream(self.device).cuda_stream != self.stream + or torch.cuda.is_current_stream_capturing() + ): + raise RuntimeError( + "Exact row reduction requires its original uncaptured stream" + ) + if self._epoch == 2**32 - 1: + raise RuntimeError( + "Exact row reduction epoch exhausted; prepare a new plan" + ) + self._epoch += 1 + self.calls += 1 + self.ops.run( + value, + self.codes, + self.output, + self.pointers, + self.flags, + self.rank, + self._epoch, + ) + return self.output + + def close(self): + """Collectively release peer handles before freeing their owners.""" + if self._closed: + return + with torch.accelerator.device_index(self.device): + torch.accelerator.synchronize() + dist.barrier(group=self.group.cpu_group) + for pointers in self._buffers: + for peer, pointer in enumerate(pointers): + if pointer and peer != self.rank: + self.ops.release(pointer, False) + dist.barrier(group=self.group.cpu_group) + for pointers in self._buffers: + if pointers[self.rank]: + self.ops.release(pointers[self.rank], True) + self._buffers.clear() + self.output = self.codes = None + self._closed = True diff --git a/vllm/model_executor/layers/sm70_diffusion.py b/vllm/model_executor/layers/sm70_diffusion.py new file mode 100644 index 0000000000..6a6f33db08 --- /dev/null +++ b/vllm/model_executor/layers/sm70_diffusion.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Model-independent SM70 FP16 GEMM with FP32 scaling and accumulation. + +The extension ABI retains its historical H3 name. Dispatch depends on tensor +properties, not checkpoint names, quantization labels or model families. +""" + +from functools import lru_cache +from importlib import import_module +from pathlib import Path + +import torch + + +@lru_cache(maxsize=1) +def sm70_extension(): + try: + return import_module("vllm._h3_w8a16_C") + except ImportError: + pass + from torch.utils.cpp_extension import load + + root = Path(__file__).resolve().parents[3] + source = root / "csrc/sm70_turbomind/ops/h3_w8a16.cu" + if not source.is_file(): + raise RuntimeError("SM70 diffusion operators require the 1Cat source build") + return load( + name="onecat_h3_w8a16", + sources=[str(source)], + extra_cuda_cflags=["-O3", "-gencode=arch=compute_70,code=sm_70"], + extra_ldflags=["-lcublas", "-lcublasLt"], + verbose=False, + ) + + +@lru_cache(maxsize=32) +def _column_major_plan(device, m, n, k, output_fp32): + return sm70_extension().ColumnMajorGemmPlan(device, m, n, k, output_fp32) + + +def fp16_gemm(input, weight, output_fp32=False): + """Use a zero-workspace Volta plan for dense column-major H3 weights. + + Plan entries contain host descriptors only. Unaligned inputs, empty shapes + and library versions without the validated algorithm use the original + row-major GEMM. Warm up each shape before capturing a CUDA graph. + """ + ops = sm70_extension() + if weight.is_contiguous(): + return ops.gemm(input, weight, output_fp32) + if ( + input.is_cuda + and weight.device == input.device + and input.dim() == weight.dim() == 2 + and input.is_contiguous() + and input.dtype == weight.dtype == torch.float16 + and weight.stride() == (1, weight.shape[0]) + and input.shape[1] == weight.shape[1] + and min(*input.shape, weight.shape[0]) > 0 + and input.data_ptr() % 16 == weight.data_ptr() % 16 == 0 + ): + plan = _column_major_plan( + input.device.index, + input.shape[0], + weight.shape[0], + input.shape[1], + bool(output_fp32), + ) + if plan.supported: + return plan.run(input, weight) + return ops.gemm(input, weight.contiguous(), output_fp32) + + +def fp16_gemm_input(x): + """Scale wide-range activations by exact powers of two before FP16 GEMM. + + Leave headroom for the 256-channel rotation's worst-case amplification. + Row scaling is restored in FP32 after the projection. + """ + flat = x.reshape(-1, x.shape[-1]).contiguous() + if flat.dtype == torch.float16: + return flat, None + if flat.dtype != torch.float32: + raise ValueError("SM70 GEMM activations must be FP16 or FP32") + if flat.is_cuda: + return sm70_extension().prepare_fp16(flat) + maximum = flat.abs().amax(-1, keepdim=True) + _, exponent = torch.frexp(maximum) + scale = torch.ldexp(torch.ones_like(maximum), (exponent - 11).clamp_min(0)) + return (flat / scale).half(), scale + + +def fp16_linear_prepared(values, weight, scale=None, *, output_fp32=False): + """Restore per-row scales in FP32 before any distributed reduction.""" + if values.dtype != torch.float16 or weight.dtype != torch.float16: + raise ValueError("Prepared SM70 linear operands must be FP16") + output_fp32 = output_fp32 or scale is not None + flat = values.reshape(-1, values.shape[-1]).contiguous() + if flat.is_cuda: + output = fp16_gemm(flat, weight, output_fp32) + else: + output = torch.nn.functional.linear(flat.float(), weight.float()) + if not output_fp32: + output = output.half() + if scale is not None: + output = output * scale + return output.reshape(*values.shape[:-1], weight.shape[0]) + + +def supports_fused_scaled_add(output): + """Old wheels and unsupported output layouts keep ordinary epilogues.""" + return ( + output.is_cuda + and output.dtype in (torch.float16, torch.float32) + and output.is_contiguous() + and torch.cuda.get_device_capability(output.device) == (7, 0) + and hasattr(sm70_extension(), "scaled_add_") + ) + + +def fp16_linear_add(x, weight, output, *, alpha, offset=0): + """Add a scaled FP16 projection into a contiguous output slice in place. + + GEMM and row-scale restoration retain FP32 boundaries. FP16 output is + rounded after this contribution, so callers combining overlapping deltas + must retain an FP32 accumulation buffer until their last contribution. + """ + if not output.is_contiguous(): + raise ValueError("SM70 projection addition requires contiguous output") + values, scale = fp16_gemm_input(x) + delta = fp16_gemm(values, weight, output_fp32=True) + flat = output.view(-1, output.shape[-1]) + if supports_fused_scaled_add(output): + sm70_extension().scaled_add_(flat, delta, scale, alpha, offset) + else: + if scale is not None: + delta = delta * scale + target = flat[:, offset : offset + weight.shape[0]] + result = target.float().add(delta, alpha=alpha).to(output.dtype) + target.copy_(result) + return output diff --git a/vllm/model_executor/layers/sm70_sparse_attention.py b/vllm/model_executor/layers/sm70_sparse_attention.py new file mode 100644 index 0000000000..b512831213 --- /dev/null +++ b/vllm/model_executor/layers/sm70_sparse_attention.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Explicit SM70 block-sparse attention for pre-tiled DiT operands.""" + +import os +from functools import lru_cache +from importlib import import_module +from pathlib import Path + + +@lru_cache(maxsize=1) +def sparse_extension(): + try: + return import_module("vllm._sm70_sparse_attention_C") + except ImportError: + pass + from torch.utils.cpp_extension import load + + cutlass_path = os.environ.get("VLLM_CUTLASS_SRC_DIR") + if not cutlass_path: + raise RuntimeError("SM70 sparse attention requires CUTLASS v4.4.2 sources") + cutlass = Path(cutlass_path) + root = Path(__file__).resolve().parents[3] + return load( + name="onecat_sm70_sparse_attention", + sources=[str(root / "flash-attention-v100/kernel/h3/forward_sparse.cu")], + extra_include_paths=[ + str(cutlass / "include"), + str(cutlass / "examples/41_fused_multi_head_attention"), + ], + extra_cuda_cflags=["-O3", "-gencode=arch=compute_70,code=sm_70"], + verbose=False, + ) + + +def block_sparse_attention(q, k, v, block_map, block_sizes, *, scale): + """Execute only selected 64-token blocks, excluding every padded edge. + + Operands use FP16 [B,64*N,H,128]; block_map is bool [B,H,N,N] and + block_sizes is int32 [N]. Each selected block is visited in ascending + logical order. Unused output rows remain zero. No dense fallback exists. + """ + return sparse_extension().forward(q, k, v, block_map, block_sizes, scale) + + +def _h3_block_sparse_attention(q, k, v, block_map, block_sizes, *, scale): + """Private route for H3-owned sizes and a nonempty mask built by H3. + + Arbitrary external maps must use block_sparse_attention and its value + checks. Older wheels retain that checked entrypoint until rebuilt. + """ + ops = sparse_extension() + forward = getattr(ops, "_forward_prevalidated", ops.forward) + return forward(q, k, v, block_map, block_sizes, scale) diff --git a/vllm/model_executor/models/minimax_h3/UPSTREAM.md b/vllm/model_executor/models/minimax_h3/UPSTREAM.md index 3873495c29..c29f5386fe 100644 --- a/vllm/model_executor/models/minimax_h3/UPSTREAM.md +++ b/vllm/model_executor/models/minimax_h3/UPSTREAM.md @@ -18,7 +18,11 @@ execution uses the same SM70 staged delta path. The pruned INT8 integration restores original dense AdaLN/time tensors, preserving backbone INT8 data. FastH3 Dense mapping and fusion follow `diffusion/models/minimax_h3/fasth3.py` at the initial pinned Omni revision. Original weights are fused before native -TP loading and staging; native INT8 fusion and VSA remain unimplemented. +TP loading and staging. Native VSA now follows the same revision's +`attention/backends/fastvideo_vsa.py` geometry and learned compression, with a +true SM70 block-sparse CUTLASS kernel. All three VSA adapter identities and +complete gate inventories are validated. Native INT8 fusion remains unsupported; +full VSA sampling quality and performance acceptance are still pending. Approximate diffusion caches, step batching, Ulysses/Ring parallelism and other model families remain outside this implementation. Model weights and checkpoint remote code retain their respective upstream terms; they are not vendored here. diff --git a/vllm/model_executor/models/minimax_h3/attention.py b/vllm/model_executor/models/minimax_h3/attention.py index 4cb91e208e..ba96746ca7 100644 --- a/vllm/model_executor/models/minimax_h3/attention.py +++ b/vllm/model_executor/models/minimax_h3/attention.py @@ -81,6 +81,8 @@ def __init__( self.backend = attention_backend.get() self.scale = softmax_scale self.head_size = head_size + self.vsa_topk = 64 + self.query_tile = 64 @property def attn_backend(self): @@ -104,18 +106,40 @@ def forward(self, q, k, v, metadata): if not 0 < used <= q.shape[1] or k.shape != v.shape: raise ValueError("invalid packed H3 attention lengths") q_valid, k_valid, v_valid = (x[:, :used].contiguous() for x in (q, k, v)) - if self.backend == "FLASH_ATTN_V100": - from .cuda_ops import flashattn_extension - - attended = flashattn_extension().forward( - q_valid, k_valid, v_valid, self.scale + if self.backend in ("FLASH_ATTN_V100", "FLASHINFER_SM70"): + from vllm.model_executor.layers.sm70_attention import noncausal_attention + + attended = noncausal_attention( + q_valid, + k_valid, + v_valid, + scale=self.scale, + backend=self.backend, + query_tile=self.query_tile, ) - elif self.backend == "FLASHINFER_SM70": - from .cuda_ops import flashinfer_extension - - attended = flashinfer_extension().forward( - q_valid, k_valid, v_valid, self.scale + elif self.backend == "FASTVIDEO_VSA": + from .vsa import h3_vsa_attention + + if metadata.video_layout is None or not metadata.video_layout.video_spans: + raise ValueError("VSA requires the complete target video layout") + target = metadata.video_layout.video_spans[-1] + prefix = metadata.extra.get("vsa_h3_prefix_segments", ()) + if target.role != "target" or sum(prefix) != target.start: + raise ValueError("VSA prefix segments disagree with the target video") + gate = metadata.extra.get("gate_compress") + if gate is None or gate.shape[1] < used: + raise ValueError("VSA requires its learned compression gate") + attended, work = h3_vsa_attention( + q_valid, + k_valid, + v_valid, + prefix_segments=prefix, + video_shape=target.latent_grid, + gate_compress=gate[:, :used], + topk=self.vsa_topk, + scale=self.scale, ) + metadata.extra["sparse_work"] = work elif self.backend == "TORCH_SDPA": attended = chunked_attention_reference( q_valid, k_valid, v_valid, scale=self.scale diff --git a/vllm/model_executor/models/minimax_h3/collectives.py b/vllm/model_executor/models/minimax_h3/collectives.py new file mode 100644 index 0000000000..9ecb0d9d63 --- /dev/null +++ b/vllm/model_executor/models/minimax_h3/collectives.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Pipeline ownership and request accounting for explicit residual collectives.""" + +import time + +from vllm.model_executor.layers.sm70_collectives import SM70ExactRowReductionPlan + + +class H3ResidualReduction: + def __init__(self, group, *, memory_budget_bytes): + self.group = group + self.memory_budget_bytes = memory_budget_bytes + self.plan: SM70ExactRowReductionPlan | None = None + self.begin_request() + + def begin_request(self): + self.peer_calls = 0 + self.native_calls = 0 + self.setup_seconds = 0.0 + self.raw_peak_bytes = self.plan.raw_ipc_bytes if self.plan is not None else 0 + self.fallback_reason = None + + def reduce(self, value): + shape = tuple(value.shape) + if self.plan is not None and self.plan.shape != shape: + self.plan.close() + self.plan = None + if self.group.world_size != 4: + self.fallback_reason = "peer execution requires TP4" + elif ( + SM70ExactRowReductionPlan.required_memory_bytes(shape) + > self.memory_budget_bytes + ): + self.fallback_reason = "shape exceeds residual communication budget" + else: + if self.plan is None: + started = time.perf_counter() + self.plan = SM70ExactRowReductionPlan( + self.group, shape, memory_budget_bytes=self.memory_budget_bytes + ) + self.setup_seconds += time.perf_counter() - started + self.raw_peak_bytes = max(self.raw_peak_bytes, self.plan.raw_ipc_bytes) + self.peer_calls += 1 + return self.plan.reduce(value) + self.native_calls += 1 + rows = value.shape[0] // self.group.world_size + return self.group.all_reduce(value).narrow( + 0, self.group.rank_in_group * rows, rows + ) + + def snapshot(self): + return { + "configured_backend": "peer", + "peer_calls": self.peer_calls, + "native_calls": self.native_calls, + "fallback_reason": self.fallback_reason, + "setup_seconds": self.setup_seconds, + "raw_ipc_peak_bytes": self.raw_peak_bytes, + "memory_budget_bytes": self.memory_budget_bytes, + } + + def close(self): + if self.plan is not None: + self.plan.close() + self.plan = None diff --git a/vllm/model_executor/models/minimax_h3/config.py b/vllm/model_executor/models/minimax_h3/config.py index 68016afefd..3708331a81 100644 --- a/vllm/model_executor/models/minimax_h3/config.py +++ b/vllm/model_executor/models/minimax_h3/config.py @@ -38,16 +38,56 @@ class H3Config: transformer_path: str | None = None tensor_parallel_size: int = 4 attention_backend: str = "FLASH_ATTN_V100" + vsa_topk: int = 64 + attention_query_tile: Literal[64, 128] = 64 fp16_weight_cache_gib: float = 0.0 fp16_cache_layers: tuple[str, ...] = () lora_path: str | None = None int8_weight_layout: str = "column" + fp16_weight_layout: Literal["row", "column"] = "row" residual_sequence_parallel: bool = False + residual_reduction: Literal["native", "peer"] = "native" + residual_reduction_memory_gib: float = 4.0 + host_weight_pin_memory: bool = True + share_host_vae_weights: bool = False + weight_offload: Literal["component", "layer"] = "component" video_encoder: Literal["libx264", "h264_nvenc"] = "libx264" host_memory_mode: Literal["auto", "pinned", "mmap"] = "auto" host_memory_directory: str | None = None def __post_init__(self) -> None: + if self.residual_reduction not in ("native", "peer"): + raise H3InputError("residual reduction must be native or peer") + if self.residual_reduction == "peer" and not self.residual_sequence_parallel: + raise H3InputError("peer reduction requires residual sequence parallelism") + if ( + isinstance(self.residual_reduction_memory_gib, bool) + or not math.isfinite(self.residual_reduction_memory_gib) + or not math.isfinite(self.residual_reduction_memory_gib * 2**30) + or self.residual_reduction_memory_gib <= 0 + ): + raise H3InputError("residual communication budget must be finite and > 0") + if self.attention_query_tile not in (64, 128): + raise H3InputError("Attention query tile must be 64 or 128") + if ( + self.attention_query_tile != 64 + and self.attention_backend != "FLASH_ATTN_V100" + ): + raise H3InputError("Explicit query tiling requires FLASH_ATTN_V100") + if self.weight_offload not in ("component", "layer"): + raise H3InputError("weight offload must be component or layer") + if self.weight_offload == "layer" and self.fp16_cache_layers: + raise H3InputError("layer offload cannot retain a fixed GPU weight cache") + if not isinstance(self.host_weight_pin_memory, bool): + raise H3InputError("host weight pinning must be a boolean") + if not isinstance(self.share_host_vae_weights, bool): + raise H3InputError("shared host VAE weights must be a boolean") + if ( + self.share_host_vae_weights + and self.tensor_parallel_size > 1 + and self.host_weight_pin_memory + ): + raise H3InputError("shared host VAE weights require pageable host masters") if self.host_memory_mode not in ("auto", "pinned", "mmap"): raise H3InputError("host memory mode must be auto, pinned or mmap") if self.video_encoder not in ("libx264", "h264_nvenc"): @@ -56,26 +96,25 @@ def __post_init__(self) -> None: raise H3InputError("partition must be fl2va or ref2va") if self.int8_weight_layout not in ("row", "column"): raise H3InputError("H3 INT8 weight layout must be row or column") + if self.fp16_weight_layout not in ("row", "column"): + raise H3InputError("H3 FP16 weight layout must be row or column") if self.tensor_parallel_size not in (1, 2, 4): raise H3InputError("native H3 supports TP1, TP2, or TP4") if self.attention_backend not in ( "FLASH_ATTN_V100", "FLASHINFER_SM70", "TORCH_SDPA", + "FASTVIDEO_VSA", ): raise H3InputError(f"unsupported H3 attention: {self.attention_backend}") - if self.residual_sequence_parallel and ( - self.tensor_parallel_size != 4 - or self.attention_backend not in ("FLASH_ATTN_V100", "FLASHINFER_SM70") - or self.partition != "fl2va" - or not self.transformer_path - or self.lora_path is not None + if ( + isinstance(self.vsa_topk, bool) + or not isinstance(self.vsa_topk, int) + or self.vsa_topk <= 0 ): - raise H3InputError( - "experimental residual sequence parallelism requires TP4, " - "FLASH_ATTN_V100 or FLASHINFER_SM70 and an FL2VA INT8 " - "ConvRot checkpoint without an adapter" - ) + raise H3InputError("VSA topk must be a positive integer") + if self.attention_backend == "FASTVIDEO_VSA" and not self.lora_path: + raise H3InputError("H3 VSA requires an explicit FastH3 VSA artifact") if ( not math.isfinite(self.fp16_weight_cache_gib) or self.fp16_weight_cache_gib < 0 diff --git a/vllm/model_executor/models/minimax_h3/cuda_ops.py b/vllm/model_executor/models/minimax_h3/cuda_ops.py index e27d3a75ec..c99e8dcbac 100644 --- a/vllm/model_executor/models/minimax_h3/cuda_ops.py +++ b/vllm/model_executor/models/minimax_h3/cuda_ops.py @@ -1,120 +1,21 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Lazy build of native H3 SM70 extensions for source development.""" - -import os -from functools import lru_cache -from pathlib import Path - -import torch - - -@lru_cache(maxsize=1) -def w8a16_extension(): - try: - from vllm import _h3_w8a16_C - - return _h3_w8a16_C - except ImportError: - pass - from torch.utils.cpp_extension import load - - root = Path(__file__).resolve().parents[4] - source = root / "csrc/sm70_turbomind/ops/h3_w8a16.cu" - if not source.is_file(): - raise RuntimeError("H3 W8A16 extension requires the 1Cat source build") - return load( - name="onecat_h3_w8a16", - sources=[str(source)], - extra_cuda_cflags=["-O3", "-gencode=arch=compute_70,code=sm_70"], - extra_ldflags=["-lcublas", "-lcublasLt"], - verbose=False, - ) - - -@lru_cache(maxsize=32) -def _column_major_plan(device, m, n, k, output_fp32): - return w8a16_extension().ColumnMajorGemmPlan(device, m, n, k, output_fp32) - - -def fp16_gemm(input, weight, output_fp32=False): - """Use a zero-workspace Volta plan for dense column-major H3 weights. - - Plan entries contain host descriptors only. Unaligned inputs, empty shapes - and library versions without the validated algorithm use the original - row-major GEMM. Warm up each shape before capturing a CUDA graph. - """ - ops = w8a16_extension() - if weight.is_contiguous(): - return ops.gemm(input, weight, output_fp32) - if ( - input.is_cuda - and weight.device == input.device - and input.dim() == weight.dim() == 2 - and input.is_contiguous() - and input.dtype == weight.dtype == torch.float16 - and weight.stride() == (1, weight.shape[0]) - and input.shape[1] == weight.shape[1] - and min(*input.shape, weight.shape[0]) > 0 - and input.data_ptr() % 16 == weight.data_ptr() % 16 == 0 - ): - plan = _column_major_plan( - input.device.index, - input.shape[0], - weight.shape[0], - input.shape[1], - bool(output_fp32), - ) - if plan.supported: - return plan.run(input, weight) - return ops.gemm(input, weight.contiguous(), output_fp32) - - -@lru_cache(maxsize=1) -def flashinfer_extension(): - try: - from vllm import _h3_flashinfer_C - - return _h3_flashinfer_C - except ImportError: - pass - from torch.utils.cpp_extension import load - - root = Path(__file__).resolve().parents[4] - return load( - name="onecat_h3_flashinfer_sm70", - sources=[str(root / "flashinfer-sm70/csrc/h3_noncausal_sm70.cu")], - extra_include_paths=[str(root / "flashinfer-sm70/include")], - extra_cuda_cflags=["-O3", "-gencode=arch=compute_70,code=sm_70"], - verbose=False, - ) - - -@lru_cache(maxsize=1) -def flashattn_extension(): - try: - from vllm import _h3_flashattn_C - - return _h3_flashattn_C - except ImportError: - pass - from torch.utils.cpp_extension import load - - root = Path(__file__).resolve().parents[4] - cutlass_path = os.environ.get("VLLM_CUTLASS_SRC_DIR") - if not cutlass_path: - raise RuntimeError( - "Build the H3 FlashAttention-V100 extension (_h3_flashattn_C), or " - "set VLLM_CUTLASS_SRC_DIR to CUTLASS v4.4.2 for source development" - ) - cutlass = Path(cutlass_path) - return load( - name="onecat_h3_flashattn_sm70", - sources=[str(root / "flash-attention-v100/kernel/h3/forward.cu")], - extra_include_paths=[ - str(cutlass / "include"), - str(cutlass / "examples/41_fused_multi_head_attention"), - ], - extra_cuda_cflags=["-O3", "-gencode=arch=compute_70,code=sm_70"], - verbose=False, - ) +"""Compatibility exports for model-independent SM70 diffusion operators.""" + +from vllm.model_executor.layers.sm70_attention import ( + flashattn_extension as flashattn_extension, +) +from vllm.model_executor.layers.sm70_attention import ( + flashinfer_extension as flashinfer_extension, +) +from vllm.model_executor.layers.sm70_diffusion import ( + _column_major_plan as _column_major_plan, +) +from vllm.model_executor.layers.sm70_diffusion import ( + fp16_gemm as fp16_gemm, +) +from vllm.model_executor.layers.sm70_diffusion import ( + sm70_extension, +) + +w8a16_extension = sm70_extension diff --git a/vllm/model_executor/models/minimax_h3/fasth3.py b/vllm/model_executor/models/minimax_h3/fasth3.py index b45a9419db..4e79eecefa 100644 --- a/vllm/model_executor/models/minimax_h3/fasth3.py +++ b/vllm/model_executor/models/minimax_h3/fasth3.py @@ -22,7 +22,10 @@ logger = init_logger(__name__) FASTH3_FILENAME = "adapter_model.safetensors" -_IDENTITY = "fastvideo/fastvideo-fasth3-dense-4-step-v1" +_DENSE_IDENTITY = "fastvideo/fastvideo-fasth3-dense-4-step-v1" +_VSA_IDENTITIES = frozenset( + f"fastvideo/fastvideo-fasth3-4-step-{version}" for version in ("v1", "v1.1", "v1.2") +) _MODEL_TARGETS = { "proj_in": "video_patch_proj", "proj_out": "final_layer.video_out", @@ -39,6 +42,7 @@ "attn.to_k": ("attn.qkv_proj", "k"), "attn.to_v": ("attn.qkv_proj", "v"), "attn.to_out.0": ("attn.out_proj", "plain"), + "attn.to_gate_compress": ("attn.to_gate_compress", "plain"), "ff.net.0.proj": ("mlp.fc1", "swap"), "ff.net.2": ("mlp.fc2", "plain"), "adaln_proj.linear": ("adaln_proj.linear", "plain"), @@ -54,6 +58,7 @@ ".lora_B.weight": "b", ".diff_b": "bias", ".diff": "diff", + ".set_weight": "set", } @@ -67,6 +72,7 @@ class FastH3Spec: denoise_steps: int = 4 api_steps: int = 4 supported_tasks: frozenset[str] = frozenset({"t2va"}) + requires_vsa: bool = False @dataclass @@ -74,6 +80,7 @@ class _Patch: layout: str pairs: dict[str, dict[str, str]] = field(default_factory=dict) diff: str | None = None + assigned: str | None = None def _native_target(module: str): @@ -99,18 +106,18 @@ def _read_index(path: str | Path, partition: str): counted = {"low_rank_tensors": 0, "diff_tensors": 0, "set_weight_tensors": 0} with safe_open(path, framework="pt", device="cpu") as checkpoint: metadata = checkpoint.metadata() or {} + identity = metadata.get("finetuned_model", "").lower() + requires_vsa = identity in _VSA_IDENTITIES if ( metadata.get("format") != "fastvideo-lora-v2" - or metadata.get("finetuned_model", "").lower() != _IDENTITY + or identity not in (_DENSE_IDENTITY, *_VSA_IDENTITIES) or metadata.get("base_model", "").lower() != "minimaxai/minimax-h3" or metadata.get("rank") != "64" ): raise H3InputError( - "Only the official rank-64 FastH3 Dense release is supported" + "Only the official rank-64 FastH3 Dense/VSA releases are supported" ) for name in checkpoint.keys(): # noqa: SIM118 - if name.endswith(".set_weight"): - raise H3InputError("FastH3 VSA needs a sparse attention implementation") match = next( ( (name[: -len(suffix)], role) @@ -131,6 +138,21 @@ def _read_index(path: str | Path, partition: str): shape = value.get_shape() if value.get_dtype() not in ("F16", "BF16", "F32"): raise H3InputError(f"FastH3 requires floating-point deltas: {name}") + is_gate = native.endswith(".attn.to_gate_compress") + if role == "set" or is_gate: + if ( + not requires_vsa + or role != "set" + or not is_gate + or not native.startswith("blocks.") + or len(shape) != 2 + or min(shape) <= 0 + or patch.assigned is not None + ): + raise H3InputError(f"Invalid FastH3 VSA compression gate: {name}") + patch.assigned = name + counted["set_weight_tensors"] += 1 + continue if role in ("a", "b"): if len(shape) != 2 or shape[0 if role == "a" else 1] != 64: raise H3InputError( @@ -157,6 +179,16 @@ def _read_index(path: str | Path, partition: str): for prefix, (_, count) in _PREFIXES.items(): if coverage[prefix] != set(range(count)): raise H3InputError(f"FastH3 must edit every {prefix} block") + expected_gates = ( + {f"blocks.{i}.attn.to_gate_compress.weight" for i in range(50)} + if requires_vsa + else set() + ) + actual_gates = {key for key, patch in patches.items() if patch.assigned} + if actual_gates != expected_gates: + raise H3InputError( + "FastH3 VSA must assign every main-block compression gate" + ) for param, patch in patches.items(): if any(set(pair) != {"a", "b"} for pair in patch.pairs.values()): raise H3InputError(f"FastH3 has an unpaired factor for {param}") @@ -164,7 +196,7 @@ def _read_index(path: str | Path, partition: str): raise H3InputError( f"FastH3 grouped QKV requires all three projections: {param}" ) - return FastH3Spec(), patches + return FastH3Spec(requires_vsa=requires_vsa), patches def inspect_fasth3_lora(path: str | Path, partition: str) -> FastH3Spec: @@ -186,6 +218,8 @@ def _fuse(self, checkpoint, name, weight): patch = self.patches.get(name) if patch is None: return weight + if patch.assigned is not None: + raise H3InputError("FastH3 gate must be new, not replace a base parameter") if weight.dtype not in (torch.float16, torch.bfloat16, torch.float32): raise H3InputError("FastH3 fusion requires original floating-point weights") @@ -248,6 +282,13 @@ def apply(self, weights): if name in self.applied: raise H3InputError(f"Duplicate FastH3 base parameter: {name}") yield name, self._fuse(checkpoint, name, weight) + for name, patch in self.patches.items(): + if patch.assigned is not None: + value = checkpoint.get_tensor(patch.assigned) + if not torch.isfinite(value).all(): + raise H3InputError(f"FastH3 gate is non-finite: {name}") + self.applied.add(name) + yield name, value self.validate_fully_applied() def validate_fully_applied(self, loaded=None): diff --git a/vllm/model_executor/models/minimax_h3/lora.py b/vllm/model_executor/models/minimax_h3/lora.py index da7f3eadcf..bea798b089 100644 --- a/vllm/model_executor/models/minimax_h3/lora.py +++ b/vllm/model_executor/models/minimax_h3/lora.py @@ -21,6 +21,11 @@ from vllm.logger import init_logger from vllm.model_executor.layers.linear import LinearMethodBase +from vllm.model_executor.layers.sm70_diffusion import ( + fp16_linear_add, + fp16_linear_prepared, + supports_fused_scaled_add, +) from .config import H3InputError from .fasth3 import FASTH3_FILENAME, FastH3Spec @@ -150,6 +155,11 @@ def inspect_deployment_adapter(config): raise H3InputError( "FastH3 Dense fusion requires original weights; omit --transformer-path" ) + sparse = isinstance(spec, FastH3Spec) and spec.requires_vsa + if sparse != (config.attention_backend == "FASTVIDEO_VSA"): + raise H3InputError( + "FastH3 VSA artifacts require the FASTVIDEO_VSA backend together" + ) return spec @@ -272,22 +282,75 @@ def __init__(self, base, parts, alpha_over_rank): self.base = base self.parts = parts self.alpha_over_rank = alpha_over_rank + spans = sorted((offset, offset + width) for _, offset, width in parts) + self.disjoint_parts = all(a[1] <= b[0] for a, b in zip(spans, spans[1:])) def create_weights(self, *args, **kwargs): raise RuntimeError("Turbo is installed after base checkpoint loading") + @property + def supports_prepared_fp16(self): + return bool(getattr(self.base, "supports_prepared_fp16", False)) + + @property + def supports_rotated_input(self): + return bool(getattr(self.base, "supports_rotated_input", False)) + + @property + def requires_original_input(self): + return lora_scale.get() != 0 + + def apply_prepared( + self, layer, values, input_scale, *, input_is_rotated=False, original_input=None + ): + if not self.supports_prepared_fp16: + raise ValueError("LoRA base does not support prepared FP16 operands") + if input_is_rotated and self.requires_original_input and original_input is None: + raise ValueError("Active LoRA requires the original unrotated input") + output = self.base.apply_prepared( + layer, values, input_scale, input_is_rotated=input_is_rotated + ) + scale = lora_scale.get() * self.alpha_over_rank + if scale == 0: + return output + original = original_input if input_is_rotated else values + dtype = output.dtype + fused = self.disjoint_parts and supports_fused_scaled_add(output) + if not fused: + output = output.float() + for index, offset, width in self.parts: + a = getattr(layer, f"h3_lora_a_{index}") + b = getattr(layer, f"h3_lora_b_{index}") + # Restore the first projection's row scale before preparing B's + # input, retaining the established intermediate rounding boundary. + intermediate = fp16_linear_prepared( + original, a, input_scale, output_fp32=True + ) + if fused: + fp16_linear_add(intermediate, b, output, alpha=scale, offset=offset) + else: + delta = _linear_fp32(intermediate, b) + output[..., offset : offset + width].add_(delta, alpha=scale) + return output.to(dtype) + def apply(self, layer, x, bias=None): output = self.base.apply(layer, x, bias) scale = lora_scale.get() * self.alpha_over_rank if scale == 0: return output dtype = output.dtype - output = output.float() + fused = self.disjoint_parts and supports_fused_scaled_add(output) + if not fused: + output = output.float() for index, offset, width in self.parts: a = getattr(layer, f"h3_lora_a_{index}") b = getattr(layer, f"h3_lora_b_{index}") - delta = _linear_fp32(_linear_fp32(x, a), b) - output[..., offset : offset + width].add_(delta, alpha=scale) + intermediate = _linear_fp32(x, a) + if fused: + fp16_linear_add(intermediate, b, output, alpha=scale, offset=offset) + else: + delta = _linear_fp32(intermediate, b) + output[..., offset : offset + width].add_(delta, alpha=scale) return output.to(dtype) diff --git a/vllm/model_executor/models/minimax_h3/pipeline.py b/vllm/model_executor/models/minimax_h3/pipeline.py index dd4b1114a4..211ba2ce73 100644 --- a/vllm/model_executor/models/minimax_h3/pipeline.py +++ b/vllm/model_executor/models/minimax_h3/pipeline.py @@ -29,7 +29,7 @@ from vllm.utils.mem_utils import get_cpu_memory from vllm.video.metrics import DenoiseWorkCounter -from .attention import attention_backend +from .attention import Attention, attention_backend from .comfy_checkpoint import inspect_comfy_checkpoint, resolve_comfy_checkpoint_path from .condition_noise import ( minimax_h3_audio_cond_noise_aug_rows, @@ -84,15 +84,16 @@ validate_reference_audio_files, validate_reference_audio_waveforms, ) -from .residency import MMapHostWeights, PinnedModuleStager +from .residency import LayerwiseModuleStager, MMapHostWeights, PinnedModuleStager from .sigma_schedule import DMD2SigmaSchedule from .time_request import ( MINIMAX_H3_SHAPE_PLANNER, minimax_h3_align_frame_count, minimax_h3_time_shift_sigmas, ) -from .transformer import MiniMaxH3DiTModel +from .transformer import MiniMaxH3DiTBlock, MiniMaxH3DiTModel from .vae import MiniMaxH3AudioVAE, MiniMaxH3VideoVAE +from .vsa import h3_vsa_workspace from .weight_cache import FP16WeightCache from .weights import iter_checkpoint_weights, resolve_model_root @@ -426,8 +427,16 @@ def _broadcast_tensor( class MiniMaxH3Pipeline(nn.Module): - def __init__(self, config: H3Config): + def __init__(self, config: H3Config, *, shared_weights_dir: str | None = None): super().__init__() + if ( + config.share_host_vae_weights + and config.tensor_parallel_size > 1 + and shared_weights_dir is None + ): + raise H3InputError( + "shared host VAE weights require an engine-owned directory" + ) self.config = config self.partition = config.partition from vllm.media.progress import report_loading @@ -529,6 +538,20 @@ def loading(done, component): ) finally: attention_backend.reset(token) + for module in self.transformer.modules(): + if isinstance(module, Attention): + module.query_tile = config.attention_query_tile + self._residual_reduction = None + if config.residual_reduction == "peer": + from .collectives import H3ResidualReduction + + self._residual_reduction = H3ResidualReduction( + get_tp_group(), + memory_budget_bytes=int(config.residual_reduction_memory_gib * 2**30), + ) + for module in self.transformer.modules(): + if isinstance(module, MiniMaxH3DiTBlock): + module.residual_reducer = self._residual_reduction if self._host_backing is not None: PinnedModuleStager.map_cpu_weights( self.transformer, self._host_backing, preserve_parameters=False @@ -538,6 +561,8 @@ def loading(done, component): weights = restore_dense_adaln_weights(weights, path / "transformer") fusion = None if isinstance(adapter_spec, FastH3Spec): + if adapter_spec.requires_vsa: + self.transformer.enable_vsa_gates(config.vsa_topk) fusion = FastH3Fusion( select_adapter_file(config.lora_path), partition=config.partition, @@ -556,6 +581,7 @@ def loading(done, component): for layer in self.transformer.modules(): method = getattr(layer, "quant_method", None) if method is not None: + layer.h3_fp16_weight_layout = config.fp16_weight_layout method.process_weights_after_loading(layer) if self._host_backing is not None: PinnedModuleStager.map_cpu_weights(layer, self._host_backing) @@ -570,7 +596,10 @@ def loading(done, component): self.transformer, config.lora_path, self.partition ) self._dit_stager = PinnedModuleStager( - self.transformer, self.device, host_backing=self._host_backing + self.transformer, + self.device, + pin_memory=config.host_weight_pin_memory, + host_backing=self._host_backing, ) self._weight_cache = FP16WeightCache( self.transformer, @@ -599,13 +628,37 @@ def loading(done, component): ) self.text_encoder.load_weights(iter_checkpoint_weights(shared / "text_encoder")) self._encoder_stager = PinnedModuleStager( - self.text_encoder, self.device, host_backing=self._host_backing + self.text_encoder, + self.device, + pin_memory=config.host_weight_pin_memory, + host_backing=self._host_backing, ) + self._dit_layer_stager: LayerwiseModuleStager | None = None + self._encoder_layer_stager: LayerwiseModuleStager | None = None + if config.weight_offload == "layer": + self._dit_layer_stager = LayerwiseModuleStager( + self._dit_stager, + (*self.transformer.token_refiner.blocks, *self.transformer.blocks), + # Cache decision probes may consume these outside block.forward. + resident_modules=( + self.transformer.blocks[0].norm1, + self.transformer.blocks[0].adaln_proj, + ), + ) + self._encoder_layer_stager = LayerwiseModuleStager( + self._encoder_stager, + ( + *self.text_encoder.vision.blocks, + *self.text_encoder.text_model.layers, + ), + ) loading(2, "video_vae") self.video_vae = MiniMaxH3VideoVAE( str(shared / "video_vae"), device=self.device, load_device=torch.device("cpu"), + pin_memory=config.host_weight_pin_memory, + shared_weights_dir=shared_weights_dir, ) self.video_vae.set_parallel_size(config.tensor_parallel_size) loading(3, "audio_vae") @@ -613,6 +666,8 @@ def loading(done, component): str(shared / "audio_vae"), device=self.device, load_device=torch.device("cpu"), + pin_memory=config.host_weight_pin_memory, + shared_weights_dir=shared_weights_dir, ) self.stage_durations = {} self.actual_dit_calls = 0 @@ -642,6 +697,20 @@ def _encode_text_hidden(self, input_ids, vision_kwargs): @contextmanager def _component_on_device(self, component): + if ( + component is self.text_encoder + and getattr(self, "_encoder_layer_stager", None) is not None + ): + plan = self._encoder_layer_stager + try: + with plan.on_device(): + yield + finally: + self.stage_durations["encoder_layer_weight_staging"] = plan.load_seconds + self.stage_durations["encoder_layer_weight_offload"] = ( + plan.offload_seconds + ) + return stager = self._encoder_stager if component is self.text_encoder else None if stager is not None: stager.load() @@ -658,6 +727,18 @@ def _component_on_device(self, component): @contextmanager def _resident_dit_layers_on_device(self, *, enabled=True): started = time.perf_counter() + if getattr(self, "_dit_layer_stager", None) is not None: + plan = self._dit_layer_stager + try: + with plan.on_device(): + self.stage_durations["dit_staging_and_weight_cache"] = ( + time.perf_counter() - started + ) + yield + finally: + self.stage_durations["dit_layer_weight_staging"] = plan.load_seconds + self.stage_durations["dit_layer_weight_offload"] = plan.offload_seconds + return self._dit_stager.load() try: self._weight_cache.prepare() @@ -673,11 +754,25 @@ def _resident_dit_layers_on_device(self, *, enabled=True): def progress_bar(self, *, total): return tqdm(total=total, desc="H3 denoise", disable=self._dit_rank != 0) + def residual_reduction_stats(self): + reducer = getattr(self, "_residual_reduction", None) + if reducer is None: + return {"configured_backend": "native", "raw_ipc_peak_bytes": 0} + return reducer.snapshot() + + def close(self): + reducer = getattr(self, "_residual_reduction", None) + if reducer is not None: + reducer.close() + @torch.inference_mode() def forward(self, request: H3Request): from vllm.media.progress import report self.stage_durations = {} + reducer = getattr(self, "_residual_reduction", None) + if reducer is not None: + reducer.begin_request() self.actual_dit_calls = 0 report("encoding") started = time.perf_counter() @@ -1575,10 +1670,53 @@ def diffuse( video_outputs=int(branch.update_mask.sum()), audio_outputs=int(branch.audio_update_mask.sum()), ) + self.denoise_workload = { + "work_accounting": ( + "sparse_tp_v1" + if self.config.attention_backend == "FASTVIDEO_VSA" + else "dense_tp_lora_v2" + ), + "partition": self.partition, + "task": task, + "adapter": ( + type(self.turbo_spec).__name__ if self.turbo_spec is not None else None + ), + "video_sigmas": list(inputs["sigmas_video"]), + "audio_sigmas": list(inputs["sigmas_audio"]), + "used_length": branch.used_len, + "blocks_per_call": counter.blocks_per_call, + "attention_algorithm": ( + "vsa" if self.config.attention_backend == "FASTVIDEO_VSA" else "dense" + ), + "cache_algorithm": None, + "actual_backends": sorted( + { + module.backend + for module in transformer.modules() + if isinstance(module, Attention) + } + ), + } + if self.config.attention_backend == "FASTVIDEO_VSA": + layout = branch.static_kwargs["video_token_layout"] + self.denoise_workload["sparse_config"] = { + "topk": self.config.vsa_topk, + "prefix_segments": list( + branch.static_kwargs["packed_seq_params"]["vsa_prefix_segments"] + ), + "video_shape": list(layout.video_spans[-1].latent_grid), + "gated_blocks": len(transformer.blocks), + "heads": transformer.blocks[0].attn.num_heads, + "head_size": transformer.blocks[0].attn.head_dim, + } from vllm.media.progress import report report("staging_model") - with self._resident_dit_layers_on_device(enabled=True): + with ( + counter, + h3_vsa_workspace(), + self._resident_dit_layers_on_device(enabled=True), + ): torch.accelerator.synchronize() dist.barrier() torch.accelerator.synchronize() @@ -1608,15 +1746,21 @@ def on_step(step, video, audio): MINIMAX_H3_AUDIO_REF_COND_TIMESTEP ), on_step=on_step, + step_profiler=counter.step, ) torch.accelerator.synchronize() + counter.finish_sparse() dist.barrier() torch.accelerator.synchronize() self.stage_durations["denoise"] = time.perf_counter() - started self.useful_denoise_flops = counter.flops self.actual_dit_calls = counter.calls self.denoise_flops_by_layer = counter.by_layer - counter.close() + self.redundant_denoise_flops = counter.redundant_flops + self.redundant_flops_by_layer = counter.redundant_by_layer + self.denoise_steps = counter.finish_steps() + self.denoise_executed_blocks = dict(counter.blocks) + self.denoise_sparse_work_by_layer = counter.sparse_by_layer return self._unpack_denoised_rows( branch, diff --git a/vllm/model_executor/models/minimax_h3/quantization.py b/vllm/model_executor/models/minimax_h3/quantization.py index 5b4b31e9ce..a3e09c78ea 100644 --- a/vllm/model_executor/models/minimax_h3/quantization.py +++ b/vllm/model_executor/models/minimax_h3/quantization.py @@ -37,6 +37,12 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( is_layer_skipped, ) +from vllm.model_executor.layers.sm70_diffusion import ( + fp16_gemm_input as fp16_gemm_input, +) +from vllm.model_executor.layers.sm70_diffusion import ( + fp16_linear_prepared, +) from vllm.model_executor.parameter import ( ChannelQuantScaleParameter, ModelWeightParameter, @@ -63,45 +69,48 @@ def create_weight_parameter( _FORMAT = "int8_tensorwise" -def fp16_gemm_input(x): - """Scale wide-range activations by exact powers of two before FP16 GEMM. - - Leave headroom for the 256-channel rotation's worst-case amplification. - Row scaling is restored in FP32 after the projection. - """ - flat = x.reshape(-1, x.shape[-1]).contiguous() - if flat.dtype == torch.float16: - return flat, None - if flat.dtype != torch.float32: - raise ValueError("H3 GEMM activations must be FP16 or FP32") - if flat.is_cuda: - from .cuda_ops import w8a16_extension - - return w8a16_extension().prepare_fp16(flat) - maximum = flat.abs().amax(-1, keepdim=True) - _, exponent = torch.frexp(maximum) - scale = torch.ldexp(torch.ones_like(maximum), (exponent - 11).clamp_min(0)) - return (flat / scale).half(), scale - +class FP16LinearMethod(UnquantizedLinearMethod): + """Dense Tensor Core operands with explicit preparation and output precision.""" -class FP32OutputLinearMethod(UnquantizedLinearMethod): - """Keep wide-range projection outputs in FP32 with FP16 Tensor Core inputs.""" + output_fp32 = False + supports_prepared_fp16 = True + supports_rotated_input = False def process_weights_after_loading(self, layer): - layer.weight.data = layer.weight.data.contiguous() + if getattr(layer, "h3_fp16_weight_layout", "row") == "column": + layer.weight.data = layer.weight.data.t().contiguous().t() + else: + layer.weight.data = layer.weight.data.contiguous() def apply(self, layer, x, bias=None): if x.is_cuda: - from .cuda_ops import w8a16_extension - values, scale = fp16_gemm_input(x) - output = w8a16_extension().gemm(values, layer.weight, True) - if scale is not None: - output = output * scale + output = self.apply_prepared(layer, values, scale) output = output.reshape(*x.shape[:-1], layer.weight.shape[0]) else: output = torch.nn.functional.linear(x.float(), layer.weight.float()) - return output if bias is None else output + bias.float() + if not self.output_fp32: + output = output.to(x.dtype) + return output if bias is None else output + bias.to(output.dtype) + + def apply_prepared( + self, layer, values, scale, *, input_is_rotated=False, original_input=None + ): + if input_is_rotated: + raise ValueError("Dense weights require unrotated activations") + return fp16_linear_prepared( + values, layer.weight, scale, output_fp32=self.output_fp32 + ) + + +class FP32OutputLinearMethod(FP16LinearMethod): + """Keep wide-range projection outputs in FP32 with FP16 Tensor Core inputs.""" + + output_fp32 = True + + +def supports_prepared_fp16(layer): + return bool(getattr(layer.quant_method, "supports_prepared_fp16", False)) def preserve_fp32_output(layer): @@ -396,10 +405,20 @@ def apply( output = torch.nn.functional.linear(x, weight, bias) return output.reshape(*original_shape[:-1], layer.weight.shape[0]) - def apply_prepared(self, layer, values, scale, *, input_is_rotated=False): + supports_prepared_fp16 = True + + @property + def supports_rotated_input(self): + return self.layer_config.convrot and self.layer_config.convrot_groupsize == 256 + + def apply_prepared( + self, layer, values, scale, *, input_is_rotated=False, original_input=None + ): """Project FP16 rows with an explicit scale restored before TP reduction.""" from .cuda_ops import fp16_gemm, w8a16_extension + if input_is_rotated and not self.supports_rotated_input: + raise ValueError("Pre-rotated input requires matching ConvRot weights") ops = w8a16_extension() x = values.reshape(-1, values.shape[-1]) if self.layer_config.convrot and not input_is_rotated: @@ -423,9 +442,8 @@ def rotate_local_fp16(layer, values): if ( values.is_cuda and values.dtype == torch.float16 - and isinstance(method, Int8ConvRotLinearMethod) - and method.layer_config.convrot - and method.layer_config.convrot_groupsize == 256 + and getattr(method, "supports_rotated_input", False) + and not getattr(method, "requires_original_input", False) and layer.bias is None and not layer.gather_output ): @@ -436,7 +454,12 @@ def rotate_local_fp16(layer, values): class _H3RotatedColumnInput(ColumnParallelLinear): - def forward(self, input_, *, input_is_rotated=False): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + if type(self.quant_method) is UnquantizedLinearMethod: + self.quant_method = FP16LinearMethod() + + def forward(self, input_, *, input_is_rotated=False, original_input=None): if not input_is_rotated: return super().forward(input_) method = self.quant_method @@ -445,16 +468,20 @@ def forward(self, input_, *, input_is_rotated=False): or input_.dtype != torch.float16 or self.bias is not None or self.gather_output - or not isinstance(method, Int8ConvRotLinearMethod) - or not method.layer_config.convrot - or method.layer_config.convrot_groupsize != 256 + or not getattr(method, "supports_rotated_input", False) ): raise ValueError( "Pre-rotated H3 columns require bias-free INT8 FP16 inputs" ) # The module still receives every gathered row: ordinary forward hooks # and useful-FLOP accounting retain the original full projection shape. - output = method.apply_prepared(self, input_, None, input_is_rotated=True) + output = method.apply_prepared( + self, + input_, + None, + input_is_rotated=True, + original_input=original_input, + ) return (output, None) if self.return_bias else output @@ -475,10 +502,10 @@ def forward(self, input_, input_scale=None): if ( not self.input_is_parallel or self.bias is not None - or not isinstance(self.quant_method, Int8ConvRotLinearMethod) + or not supports_prepared_fp16(self) ): raise ValueError( - "Prepared H3 rows require a bias-free INT8 local projection" + "Prepared H3 rows require a bias-free local FP16-capable projection" ) output = self.quant_method.apply_prepared(self, input_, input_scale) if self.reduce_results and self.tp_size > 1: diff --git a/vllm/model_executor/models/minimax_h3/residency.py b/vllm/model_executor/models/minimax_h3/residency.py index 5defd62a05..2a73a7f5c0 100644 --- a/vllm/model_executor/models/minimax_h3/residency.py +++ b/vllm/model_executor/models/minimax_h3/residency.py @@ -5,10 +5,15 @@ from __future__ import annotations +import hashlib +import json import os import tempfile +import time import uuid from collections.abc import Iterable +from contextlib import contextmanager +from copy import copy from dataclasses import dataclass from itertools import chain from pathlib import Path @@ -187,6 +192,7 @@ def __init__( self._ready_event = torch.cuda.Event() self.cache_retention = cache_retention self.loaded = False + self._layerwise_active = False self._groups = self._snapshot_groups( modules, pin_memory=pin_memory, host_backing=host_backing ) @@ -289,13 +295,109 @@ def map_cpu_weights( @staticmethod def _view(backing: torch.Tensor, binding: _TensorBinding) -> torch.Tensor: + element_size = binding.dtype.itemsize + backing_offset = backing.storage_offset() * backing.element_size() + if backing_offset % element_size: + raise ValueError("shared weight storage must preserve dtype alignment") return torch.empty(0, dtype=binding.dtype, device=backing.device).set_( backing.untyped_storage(), - binding.storage_offset, + backing_offset // element_size + binding.storage_offset, binding.shape, binding.stride, ) + def share_cpu_storage(self, directory: str | Path) -> None: + """Share a checked immutable replica across this engine's TP workers. + + The engine owns the temporary directory and removes it after workers + stop. Private mappings prevent accidental CPU writes from affecting a + different rank. Only identical complete storage groups may be shared. + """ + import torch.distributed as dist + + if self.loaded or any(group.master.is_pinned() for group in self._groups): + raise ValueError("shared host storage requires unloaded pageable masters") + rank = dist.get_rank() if dist.is_initialized() else 0 + directory = Path(directory) + if rank == 0: + self._write_shared_groups(directory) + if dist.is_initialized(): + dist.barrier() + self._read_shared_groups(directory) + + @staticmethod + def _group_description(group: _StorageGroup) -> dict: + return { + "bytes": group.master.numel(), + "bindings": [ + { + "dtype": str(binding.dtype), + "shape": list(binding.shape), + "stride": list(binding.stride), + "storage_offset": binding.storage_offset, + } + for binding in group.bindings + ], + } + + @staticmethod + def _storage_digest(master: torch.Tensor) -> str: + return hashlib.sha256(memoryview(master.numpy())).hexdigest() + + def _write_shared_groups(self, directory: Path) -> None: + directory.mkdir(mode=0o700) + records, total = [], 0 + for group in self._groups: + total = (total + 255) // 256 * 256 + records.append({**self._group_description(group), "offset": total}) + total += group.master.numel() + data_path = directory / "weights.bin" + fd = os.open(data_path, os.O_RDWR | os.O_CREAT | os.O_EXCL, 0o600) + try: + # Reserve tmpfs space before mapping so a later write cannot SIGBUS + # because unrelated processes consume the remaining free space. + if total: + os.posix_fallocate(fd, 0, total) + finally: + os.close(fd) + shared = torch.from_file( + str(data_path), shared=True, size=total, dtype=torch.uint8 + ) + for group, record in zip(self._groups, records): + target = shared.narrow(0, record["offset"], record["bytes"]) + target.copy_(group.master) + record["sha256"] = self._storage_digest(target) + (directory / "metadata.json").write_text( + json.dumps({"version": 1, "bytes": total, "groups": records}) + ) + + def _read_shared_groups(self, directory: Path) -> None: + metadata = json.loads((directory / "metadata.json").read_text()) + records = metadata["groups"] + if metadata["version"] != 1 or len(records) != len(self._groups): + raise ValueError("shared component storage inventory differs across ranks") + # COW mappings share physical pages while keeping the file immutable. + shared = torch.from_file( + str(directory / "weights.bin"), + shared=False, + size=metadata["bytes"], + dtype=torch.uint8, + ) + for group, record in zip(self._groups, records): + if any( + record[key] != value + for key, value in self._group_description(group).items() + ): + raise ValueError("shared component tensor layouts differ across ranks") + if self._storage_digest(group.master) != record["sha256"]: + raise ValueError("shared component weights differ across ranks") + target = shared.narrow(0, record["offset"], record["bytes"]) + if self._storage_digest(target) != record["sha256"]: + raise ValueError("shared component snapshot failed its checksum") + group.master = target + for binding in group.bindings: + set_tensor_storage(binding.target, self._view(target, binding)) + def _bind(self, storages: list[torch.Tensor]) -> None: for storage, group in zip(storages, self._groups): for binding in group.bindings: @@ -347,6 +449,8 @@ def _cleanup_failed_load(self) -> None: self._release_cache(force=True) def load(self) -> None: + if getattr(self, "_layerwise_active", False): + raise RuntimeError("whole-module load overlaps layerwise weight staging") if self.loaded: return try: @@ -383,4 +487,115 @@ def offload(self) -> None: self._release_cache() +class LayerwiseModuleStager: + """Execute disjoint blocks from the same immutable host snapshot. + + Storage shared across blocks, or between a block and the outer module, + remains resident throughout the context. Other block storage is loaded + immediately before its forward and released afterwards, including errors. + Transfers synchronize at block boundaries; this is a capacity policy. + """ + + def __init__( + self, + snapshot: PinnedModuleStager, + blocks: Iterable[nn.Module], + *, + resident_modules: Iterable[nn.Module] = (), + ): + self.snapshot = snapshot + self.blocks = tuple(blocks) + if len({id(block) for block in self.blocks}) != len(self.blocks): + raise ValueError("layerwise staging blocks must be unique") + block_ids = {id(block) for block in self.blocks} + for block in self.blocks: + if any(id(child) in block_ids for child in tuple(block.modules())[1:]): + raise ValueError("layerwise staging blocks must not be nested") + owners: dict[int, set[int]] = {} + for index, block in enumerate(self.blocks): + for target in chain(block.parameters(), block.buffers()): + owners.setdefault(id(target), set()).add(index) + for module in resident_modules: + for target in chain(module.parameters(), module.buffers()): + owners.setdefault(id(target), set()).add(-1) + grouped: list[list[_StorageGroup]] = [[] for _ in self.blocks] + resident = [] + for group in snapshot._groups: + group_owners: set[int] = set().union( + *(owners.get(id(binding.target), {-1}) for binding in group.bindings) + ) + if len(group_owners) == 1 and -1 not in group_owners: + grouped[next(iter(group_owners))].append(group) + else: + resident.append(group) + + def subset(groups: list[_StorageGroup]) -> PinnedModuleStager: + stager = copy(snapshot) + stager._groups = groups + stager._device_storages = [] + stager.loaded = False + stager.cache_retention = snapshot.cache_retention or BoundedAllocatorCache( + snapshot.device + ) + return stager + + self.resident = subset(resident) + self.stagers = tuple(subset(groups) for groups in grouped) + self.load_seconds = 0.0 + self.offload_seconds = 0.0 + self.loaded_bytes = 0 + + def _load(self, stager): + started = time.perf_counter() + stager.load() + torch.accelerator.synchronize() + self.load_seconds += time.perf_counter() - started + self.loaded_bytes += sum(group.master.numel() for group in stager._groups) + + def _offload(self, stager): + started = time.perf_counter() + stager.offload() + self.offload_seconds += time.perf_counter() - started + + @contextmanager + def on_device(self): + if self.snapshot.loaded or getattr(self.snapshot, "_layerwise_active", False): + raise RuntimeError("layerwise staging requires an idle host snapshot") + self.snapshot._layerwise_active = True + self.load_seconds = self.offload_seconds = 0.0 + self.loaded_bytes = 0 + hooks = [] + try: + self._load(self.resident) + for block, stager in zip(self.blocks, self.stagers): + hooks.append( + block.register_forward_pre_hook( + lambda module, args, stager=stager: self._load(stager) + ) + ) + hooks.append( + block.register_forward_hook( + lambda module, args, result, stager=stager: self._offload( + stager + ), + always_call=True, + ) + ) + yield + finally: + for hook in hooks: + hook.remove() + errors = [] + try: + for stager in (*self.stagers, self.resident): + try: + self._offload(stager) + except Exception as exc: + errors.append(exc) + finally: + self.snapshot._layerwise_active = False + if errors: + raise errors[0] + + __all__ = ["BoundedAllocatorCache", "PinnedModuleStager"] diff --git a/vllm/model_executor/models/minimax_h3/transformer.py b/vllm/model_executor/models/minimax_h3/transformer.py index 7c9c9e8759..ba401d35bd 100644 --- a/vllm/model_executor/models/minimax_h3/transformer.py +++ b/vllm/model_executor/models/minimax_h3/transformer.py @@ -40,12 +40,13 @@ ) from .ops import RMSNorm, RotaryEmbedding, fused_qk_norm_rope from .quantization import ( + FP16LinearMethod, H3MergedColumnParallelLinear, H3QKVParallelLinear, H3RowParallelLinear, - Int8ConvRotLinearMethod, preserve_fp32_output, rotate_local_fp16, + supports_prepared_fp16, ) if TYPE_CHECKING: @@ -53,6 +54,8 @@ QuantizationConfig, ) + from .collectives import H3ResidualReduction + logger = init_logger(__name__) @@ -403,6 +406,9 @@ def __init__( prefix=f"{prefix}.out_proj", ) preserve_fp32_output(self.out_proj) + self.to_gate_compress: ColumnParallelLinear | None = None + self._gate_dimensions = (arch.hidden_size, inner_dim) + self._gate_prefix = f"{prefix}.to_gate_compress" self.attention = Attention( num_heads=self.num_heads, num_kv_heads=self.num_kv_heads, @@ -417,6 +423,18 @@ def __init__( prefix=prefix, ) + def enable_vsa_gate(self, topk: int) -> None: + if self.to_gate_compress is None: + self.to_gate_compress = ColumnParallelLinear( + *self._gate_dimensions, + bias=False, + params_dtype=_COMPUTE_DTYPE, + prefix=self._gate_prefix, + ) + self.to_gate_compress.quant_method = FP16LinearMethod() + nn.init.zeros_(self.to_gate_compress.weight) + self.attention.vsa_topk = topk + def _apply_rope(self, x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: """Rotate the first rot_dim head dims; pass the rest through. @@ -442,6 +460,8 @@ def _run_packed_attention( packed_total: int, num_requests: int = 1, video_layout: VideoTokenLayout | None = None, + vsa_prefix_segments: tuple[int, ...] = (), + gate_compress: torch.Tensor | None = None, ) -> torch.Tensor: """Run packed attention as a small eager island. @@ -527,6 +547,14 @@ def _run_packed_attention( # (see MINIMAX_H3_LASER_INPUT_SCALE). Ignored by every other # backend/path. "laser_input_scale": MINIMAX_H3_LASER_INPUT_SCALE, + **( + { + "gate_compress": gate_compress.unsqueeze(0), + "vsa_h3_prefix_segments": vsa_prefix_segments, + } + if gate_compress is not None + else {} + ), }, video_layout=video_layout, ) @@ -549,6 +577,7 @@ def forward( sp_seq_lens: list[int] | None = None, video_layout: VideoTokenLayout | None = None, input_is_rotated: bool = False, + vsa_prefix_segments: tuple[int, ...] = (), ) -> torch.Tensor: """x: [T, hidden] packed thd rows -> [T, hidden]. @@ -585,6 +614,13 @@ def forward( self.q_norm.variance_epsilon, ) + gate_compress = None + if self.to_gate_compress is not None: + if input_is_rotated: + raise ValueError("VSA gate requires its original unrotated activation") + gate_compress, _ = self.to_gate_compress(x) + gate_compress = gate_compress.view(total, self.num_heads, self.head_dim) + # Each request contributes a document for its rows plus one for any # nonempty alignment padding. Local/Ulysses backends unpad it, while # Ring keeps aligned rows for fixed-size P2P buffers. @@ -600,6 +636,8 @@ def forward( packed_total=packed_total if packed_total is not None else q.shape[0], num_requests=num_requests, video_layout=video_layout, + vsa_prefix_segments=vsa_prefix_segments, + gate_compress=gate_compress, ) out = out.reshape(total, self.num_heads * self.head_dim) out, _ = self.out_proj(out) @@ -647,7 +685,7 @@ def forward(self, x: torch.Tensor, *, input_is_rotated=False) -> torch.Tensor: hidden.is_cuda and hidden.dtype == torch.float16 and 0 < hidden.shape[-1] <= 32768 - and isinstance(self.fc2.quant_method, Int8ConvRotLinearMethod) + and supports_prepared_fp16(self.fc2) ): from .activation import silu_prepare_fp16 @@ -824,12 +862,18 @@ def __init__( quant_config, prefix=f"{prefix}.mlp", ) - self.residual_group = get_tp_group() if residual_sequence_parallel else None + self.residual_reducer: H3ResidualReduction | None = None + self.residual_group = ( + get_tp_group() + if residual_sequence_parallel and get_tensor_model_parallel_world_size() > 1 + else None + ) if self.residual_group is not None: - if self.residual_group.world_size != 4: - raise ValueError("H3 residual sequence parallelism requires TP4") - # These projections return unreduced FP32 partial sums. Reduce-scatter - # keeps that precision and assigns each rank its residual rows. + if self.residual_group.world_size not in (2, 4): + raise ValueError("H3 residual sequence parallelism requires TP2 or TP4") + # Keep the replicated path's FP32 sum order before selecting local + # residual rows. NCCL reduce-scatter uses a different reduction + # order and fails the full four-step latent quality gate. self.attn.out_proj.reduce_results = False self.mlp.fc2.reduce_results = False self.adaln_proj = MiniMaxH3AdalnProj( @@ -854,6 +898,7 @@ def forward( num_requests: int = 1, sp_seq_lens: list[int] | None = None, video_layout: VideoTokenLayout | None = None, + vsa_prefix_segments: tuple[int, ...] = (), ) -> torch.Tensor: """x: [T, H]; t_emb: [M, t_dim]; combined_indices: [T] (= inverse_indices * modality_num + token_tags.clamp(min=0)). @@ -914,9 +959,15 @@ def forward( sp_seq_lens=sp_seq_lens, video_layout=video_layout, input_is_rotated=input_is_rotated, + vsa_prefix_segments=vsa_prefix_segments, ) if group is not None: - h = group.reduce_scatter(h, dim=0) + if self.residual_reducer is not None: + h = self.residual_reducer.reduce(h) + else: + h = group.all_reduce(h).narrow( + 0, group.rank_in_group * residual.shape[0], residual.shape[0] + ) x, h = indexed_gate_rms_norm_scale_shift( residual, gate_msa, @@ -935,7 +986,12 @@ def forward( h = group.all_gather(h, dim=0) h = self.mlp(h, input_is_rotated=input_is_rotated) if group is not None: - h = group.reduce_scatter(h, dim=0) + if self.residual_reducer is not None: + h = self.residual_reducer.reduce(h) + else: + h = group.all_reduce(h).narrow( + 0, group.rank_in_group * residual.shape[0], residual.shape[0] + ) return indexed_gate(residual, gate_mlp, h, combined_indices) @@ -1087,15 +1143,9 @@ def __init__( self._qkv_checkpoint_is_runtime_layout = bool( getattr(quant_config, "is_checkpoint_int8_convrot_serialized", False) ) - self.residual_sequence_parallel = residual_sequence_parallel - if residual_sequence_parallel and ( - get_tensor_model_parallel_world_size() != 4 - or not self._qkv_checkpoint_is_runtime_layout - ): - raise ValueError( - "H3 residual sequence parallelism requires TP4 " - "and serialized INT8 ConvRot" - ) + self.residual_sequence_parallel = ( + residual_sequence_parallel and get_tensor_model_parallel_world_size() > 1 + ) self.hidden_size = arch.hidden_size self.num_attention_heads = arch.num_attention_heads self.num_channels_latents = arch.latents_dim @@ -1178,6 +1228,16 @@ def __init__( if callable(validate_bindings): validate_bindings(self) + def enable_vsa_gates(self, topk: int) -> None: + for block in self.blocks: + if block.attn.attention.backend != "FASTVIDEO_VSA": + raise ValueError("VSA gates require the explicit sparse backend") + block.attn.enable_vsa_gate(topk) + # The official adapter has no token-refiner gates or video tile layout. + for block in self.token_refiner.blocks: + block.attn.attention.backend = "FLASH_ATTN_V100" + self._mark_missing_params_required() + def _mark_missing_params_required(self) -> None: for _, param in self.named_parameters(): param.missing_param_init = "error" @@ -1322,6 +1382,11 @@ def load_weights( weight_loader(param, up, 1) else: weight_loader(param, loaded_weight) + if param.dtype == _COMPUTE_DTYPE and not torch.isfinite(param).all(): + raise ValueError( + f"H3 weight {name} cannot be represented as finite FP16; " + "checkpoint conversion must not silently overflow" + ) loaded.add(name) return loaded @@ -1499,6 +1564,9 @@ def forward(self, **kwargs: Any) -> tuple[torch.Tensor, torch.Tensor]: ) psp = _required_kwarg(kwargs, "packed_seq_params") + vsa_prefix_segments = tuple( + int(n) for n in self._psp_optional(psp, "vsa_prefix_segments", ()) + ) cu_seqlens = self._psp_field(psp, "packed_seq_params", "cu_seqlens_q").to( torch.int32 ) @@ -1609,6 +1677,7 @@ def forward(self, **kwargs: Any) -> tuple[torch.Tensor, torch.Tensor]: packed_total=seq_len, num_requests=num_requests, video_layout=video_layout, + vsa_prefix_segments=vsa_prefix_segments, ) if residual_group is not None: # Final heads and the existing padding boundary consume full FP32 rows. diff --git a/vllm/model_executor/models/minimax_h3/vae.py b/vllm/model_executor/models/minimax_h3/vae.py index dd1f6b4a28..ed4024e62b 100644 --- a/vllm/model_executor/models/minimax_h3/vae.py +++ b/vllm/model_executor/models/minimax_h3/vae.py @@ -136,6 +136,8 @@ def __init__( *, device: torch.device, load_device: torch.device | None = None, + pin_memory: bool = True, + shared_weights_dir: str | None = None, ) -> None: super().__init__() self._device_target = device @@ -155,8 +157,10 @@ def __init__( self._stager = PinnedModuleStager( self.remote, device, - pin_memory=True, + pin_memory=pin_memory, ) + if shared_weights_dir is not None: + self._stager.share_cpu_storage(Path(shared_weights_dir) / "video") self.model = self.remote.model self.use_tiling = True self.use_slicing = False @@ -423,6 +427,8 @@ def __init__( *, device: torch.device, load_device: torch.device | None = None, + pin_memory: bool = True, + shared_weights_dir: str | None = None, ) -> None: super().__init__() self._device_target = device @@ -440,8 +446,10 @@ def __init__( self._stager = PinnedModuleStager( self.remote, device, - pin_memory=True, + pin_memory=pin_memory, ) + if shared_weights_dir is not None: + self._stager.share_cpu_storage(Path(shared_weights_dir) / "audio") self.model = self.remote.model self.sample_rate = int(self.config_dict["sample_rate"]) diff --git a/vllm/model_executor/models/minimax_h3/vsa.py b/vllm/model_executor/models/minimax_h3/vsa.py new file mode 100644 index 0000000000..fe368bea1a --- /dev/null +++ b/vllm/model_executor/models/minimax_h3/vsa.py @@ -0,0 +1,325 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""H3 VSA geometry and compression from the pinned official Omni implementation. + +Source: fastvideo_vsa.py at 7be014bce6374f06c95b703763bdbac4c6198f31. +The selected blocks execute through a native SM70 sparse kernel. Geometry +caches retain indices only; pooled activations, scores and gates are per-call. +""" + +import functools +import math +from contextlib import contextmanager +from contextvars import ContextVar + +import torch + +from vllm.model_executor.layers.sm70_sparse_attention import ( + _h3_block_sparse_attention as block_sparse_attention, +) +from vllm.model_executor.layers.sm70_sparse_attention import ( + sparse_extension, +) + +_layout_buffers: ContextVar[dict[tuple[torch.device, int], torch.Tensor] | None] = ( + ContextVar("h3_vsa_layout_buffers", default=None) +) + + +@contextmanager +def h3_vsa_workspace(): + """Own intermediate buffers for one denoise request, including failures.""" + buffers: dict[tuple[torch.device, int], torch.Tensor] = {} + token = _layout_buffers.set(buffers) + try: + yield + finally: + buffers.clear() + _layout_buffers.reset(token) + + +def _layout_scratch(q: torch.Tensor, rows: int) -> torch.Tensor: + shape = (3, q.shape[0], rows, q.shape[2], q.shape[3]) + buffers = _layout_buffers.get() + if buffers is None: + return torch.empty(shape, device=q.device, dtype=q.dtype) + key = (q.device, torch.cuda.current_stream(q.device).cuda_stream) + scratch = buffers.get(key) + if scratch is None or scratch.shape != shape or scratch.dtype != q.dtype: + scratch = torch.empty(shape, device=q.device, dtype=q.dtype) + buffers[key] = scratch + return scratch + + +def _layout_ops(*operands: torch.Tensor): + # Preserve the generic path for strided or unaligned views and old wheels. + # In particular, a fused inference primitive must not hide autograd inputs. + if any( + not x.is_cuda or not x.is_contiguous() or x.requires_grad or x.data_ptr() % 16 + for x in operands + ): + return None + ops = sparse_extension() + if not all( + hasattr(ops, name) + for name in ("_h3_tile_qkv_prevalidated", "_h3_gate_untile_prevalidated") + ): + return None + return ops + + +@functools.lru_cache(maxsize=32) +def _get_tile_partition_indices( + dit_seq_shape: tuple[int, int, int], + tile_size: tuple[int, int, int], + device: torch.device, +) -> torch.Tensor: + t_size, h_size, w_size = dit_seq_shape + tile_t, tile_h, tile_w = tile_size + indices = torch.arange( + t_size * h_size * w_size, device=device, dtype=torch.long + ).reshape(t_size, h_size, w_size) + tiles = [] + for tile_t_idx in range(math.ceil(t_size / tile_t)): + for tile_h_idx in range(math.ceil(h_size / tile_h)): + for tile_w_idx in range(math.ceil(w_size / tile_w)): + tiles.append( + indices[ + tile_t_idx * tile_t : min((tile_t_idx + 1) * tile_t, t_size), + tile_h_idx * tile_h : min((tile_h_idx + 1) * tile_h, h_size), + tile_w_idx * tile_w : min((tile_w_idx + 1) * tile_w, w_size), + ].flatten() + ) + return torch.cat(tiles, dim=0) + + +@functools.lru_cache(maxsize=32) +def _construct_variable_block_sizes( + dit_seq_shape: tuple[int, int, int], + tile_size: tuple[int, int, int], + device: torch.device, +) -> torch.Tensor: + num_tiles = tuple( + math.ceil(seq_dim / tile_dim) + for seq_dim, tile_dim in zip(dit_seq_shape, tile_size) + ) + + def _sizes(dim_len: int, tile: int, n_tiles: int) -> torch.Tensor: + sizes = torch.full((n_tiles,), tile, dtype=torch.int32, device=device) + remainder = dim_len - (n_tiles - 1) * tile + sizes[-1] = remainder if remainder > 0 else tile + return sizes + + t_sizes = _sizes(dit_seq_shape[0], tile_size[0], num_tiles[0]) + h_sizes = _sizes(dit_seq_shape[1], tile_size[1], num_tiles[1]) + w_sizes = _sizes(dit_seq_shape[2], tile_size[2], num_tiles[2]) + return ( + t_sizes[:, None, None] * h_sizes[None, :, None] * w_sizes[None, None, :] + ).reshape(-1) + + +@functools.lru_cache(maxsize=32) +def _get_non_pad_index( + variable_block_sizes: torch.Tensor, max_block_size: int +) -> torch.Tensor: + num_blocks = variable_block_sizes.shape[0] + device = variable_block_sizes.device + starts = torch.arange(num_blocks, device=device) * max_block_size + padded_index = ( + starts[:, None] + torch.arange(max_block_size, device=device)[None, :] + ) + valid = ( + torch.arange(max_block_size, device=device)[None, :] + < variable_block_sizes[:, None] + ) + return padded_index[valid] + + +@functools.lru_cache(maxsize=32) +def _get_h3_tile_metadata( + prefix_segments: tuple[int, ...], + video_shape: tuple[int, int, int], + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int, int]: + """Official FastVideo H3 geometry: pure prefix chunks + 3-D video tiles.""" + block_size = (4, 4, 4) + block_elements = 64 + prefix_len = sum(prefix_segments) + prefix_sizes: list[int] = [] + for segment in prefix_segments: + full, remainder = divmod(segment, block_elements) + prefix_sizes.extend([block_elements] * full) + if remainder: + prefix_sizes.append(remainder) + + video_indices = ( + _get_tile_partition_indices(video_shape, block_size, device) + prefix_len + ) + video_sizes = _construct_variable_block_sizes(video_shape, block_size, device) + partition = torch.cat( + [torch.arange(prefix_len, device=device, dtype=torch.long), video_indices] + ) + sizes = torch.cat( + [ + torch.tensor(prefix_sizes, device=device, dtype=torch.int32), + video_sizes.to(torch.int32), + ] + ) + non_pad = _get_non_pad_index(sizes, block_elements) + untile = non_pad[torch.argsort(partition)] + total = prefix_len + math.prod(video_shape) + if int(sizes.sum()) != total or untile.numel() != total: + raise ValueError( + f"invalid H3 VSA geometry: prefix={prefix_segments}, video={video_shape}, " + f"sizes_sum={int(sizes.sum())}, total={total}" + ) + return ( + partition, + sizes, + non_pad, + untile, + len(prefix_sizes), + int(video_sizes.numel()), + ) + + +def _pool_h3_tiles(x: torch.Tensor, sizes: torch.Tensor) -> torch.Tensor: + batch, seq_len, heads, dim = x.shape + blocks = seq_len // 64 + pooled = x.view(batch, blocks, 64, heads, dim).sum(dim=2, dtype=torch.float32) + pooled = pooled / sizes.view(1, -1, 1, 1).clamp_min(1) + return pooled.permute(0, 2, 1, 3) + + +@functools.lru_cache(maxsize=32) +def _get_h3_fused_indices(prefix_segments, video_shape, device): + """Cache only indices derived from already validated H3 geometry.""" + partition, sizes, non_pad, untile, _, _ = _get_h3_tile_metadata( + prefix_segments, video_shape, device + ) + source = torch.full((sizes.numel() * 64,), -1, device=device, dtype=torch.int32) + source[non_pad] = partition.to(torch.int32) + return source, untile.to(torch.int32) + + +def _build_h3_block_map( + scores: torch.Tensor, + num_prefix_blocks: int, + num_video_blocks: int, + topk: int, +) -> torch.Tensor: + """Prefix K/V are exempt and prefix queries stay dense, as in FastVideo.""" + keep_video = min(topk, num_video_blocks) + if keep_video == num_video_blocks: + return torch.ones_like(scores, dtype=torch.bool) + block_map = torch.zeros_like(scores, dtype=torch.bool) + indices = ( + scores[..., num_prefix_blocks:].topk(keep_video, dim=-1).indices + + num_prefix_blocks + ) + block_map.scatter_(-1, indices, True) + block_map[..., :num_prefix_blocks] = True + block_map[:, :, :num_prefix_blocks, :] = True + return block_map + + +def h3_vsa_attention( + q, k, v, *, prefix_segments, video_shape, gate_compress, topk, scale +): + """Return sparse output and exact useful pair/block counts. + + Prefix queries stay dense. Video queries select all prefix blocks plus + top-k video blocks, then add the official learned compressed contribution. + """ + if ( + len(video_shape) != 3 + or any(n <= 0 for n in video_shape) + or any(n <= 0 for n in prefix_segments) + or isinstance(topk, bool) + or not isinstance(topk, int) + or topk <= 0 + ): + raise ValueError("VSA needs positive segment/grid dimensions and topk") + expected = sum(prefix_segments) + math.prod(video_shape) + if ( + q.ndim != 4 + or q.shape != k.shape + or q.shape != v.shape + or q.shape[1] != expected + or q.shape[3] != 128 + or q.dtype != torch.float16 + or k.dtype != q.dtype + or v.dtype != q.dtype + or k.device != q.device + or v.device != q.device + ): + raise ValueError("VSA operands must match FP16 [B,valid_rows,H,128] geometry") + if ( + gate_compress is None + or gate_compress.shape != q.shape + or gate_compress.dtype != q.dtype + or gate_compress.device != q.device + ): + raise ValueError("FastH3 VSA requires its matching learned compression gate") + if not math.isclose(scale, 128**-0.5, rel_tol=0, abs_tol=1e-6): + raise ValueError("FastH3 VSA requires the official head-dimension scale") + partition, sizes, non_pad, untile, prefix_blocks, video_blocks = ( + _get_h3_tile_metadata(tuple(prefix_segments), tuple(video_shape), q.device) + ) + blocks = sizes.numel() + shape = (q.shape[0], blocks * 64, q.shape[2], q.shape[3]) + layout = _layout_ops(q, k, v, gate_compress) + if layout is None: + tiled = [] + for operand in (q, k, v): + target = torch.zeros(shape, device=q.device, dtype=q.dtype) + target[:, non_pad] = operand[:, partition] + tiled.append(target) + q_tiled, k_tiled, v_tiled = tiled + else: + source, fused_untile = _get_h3_fused_indices( + tuple(prefix_segments), tuple(video_shape), q.device + ) + scratch = _layout_scratch(q, blocks * 64) + q_tiled, k_tiled, v_tiled = layout._h3_tile_qkv_prevalidated( + q, k, v, source, scratch + ).unbind(0) + q_pool, k_pool = (_pool_h3_tiles(x, sizes) for x in (q_tiled, k_tiled)) + scores = torch.matmul(q_pool, k_pool.transpose(-2, -1)) * scale + block_map = _build_h3_block_map(scores, prefix_blocks, video_blocks, topk) + output = block_sparse_attention( + q_tiled, k_tiled, v_tiled, block_map, sizes, scale=scale + ) + v_pool = _pool_h3_tiles(v_tiled, sizes) + compressed = torch.matmul(torch.softmax(scores, dim=-1), v_pool) + compressed = compressed.permute(0, 2, 1, 3).to(output.dtype) + if layout is None: + gate_tiled = torch.zeros_like(q_tiled) + gate_tiled[:, non_pad] = gate_compress[:, partition] + output = ( + output.view(q.shape[0], blocks, 64, q.shape[2], q.shape[3]) + + compressed.unsqueeze(2) + * gate_tiled.view(q.shape[0], blocks, 64, q.shape[2], q.shape[3]) + ).view_as(output) + output = output[:, untile].contiguous() + else: + # Match the existing separate FP16 multiply and add rounding exactly. + # The returned tensor is fresh; only intermediate QKV storage is reused. + output = layout._h3_gate_untile_prevalidated( + output, compressed.contiguous(), gate_compress, fused_untile + ) + # Keep dynamic counts on the device until complete denoise accounting. + # Padding and unselected blocks never contribute useful model FLOPs. + pair_sizes = sizes.to(torch.int64)[:, None] * sizes.to(torch.int64)[None, :] + work = { + "dense_token_pairs": q.shape[0] * q.shape[2] * q.shape[1] ** 2, + "selected_token_pairs": (block_map * pair_sizes).sum(), + "selected_blocks": block_map.sum(), + "compression_flops": 4 * q.shape[0] * q.shape[2] * blocks**2 * q.shape[3], + "prefix_blocks": prefix_blocks, + "video_blocks": video_blocks, + "heads": q.shape[2], + "head_size": q.shape[3], + } + return output, work diff --git a/vllm/video/benchmark.py b/vllm/video/benchmark.py new file mode 100644 index 0000000000..3978690e46 --- /dev/null +++ b/vllm/video/benchmark.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Complete native H3 warmup and three-run performance measurements. + +Run with ``python -m vllm.video.benchmark --contract contract.json --output DIR``. +The contract contains the same ``config`` and ``request`` as native run.json. +""" + +import argparse +import hashlib +import json +import subprocess +import time +from dataclasses import asdict +from pathlib import Path + +from vllm.model_executor.models.minimax_h3.config import ( + H3Config, + H3Request, + H3SamplingParams, +) +from vllm.video.engine import H3Engine +from vllm.video.metrics import evaluate_performance + + +def source_provenance(): + import torch + + import vllm + + package = Path(vllm.__file__).resolve().parent + root = package.parent + provenance = { + "torch": torch.__version__, + "cuda_runtime": torch.version.cuda, + "vllm": vllm.__version__, + "package_path": str(package), + "sources_sha256": {}, + } + paths = [ + *package.joinpath("video").glob("*.py"), + *package.joinpath("model_executor/models/minimax_h3").glob("*.py"), + package / "model_executor/layers/linear.py", + *package.joinpath("model_executor/layers").glob("sm70_*.py"), + ] + for path in sorted(paths): + if path.is_file(): + provenance["sources_sha256"][str(path.relative_to(package))] = ( + hashlib.sha256(path.read_bytes()).hexdigest() + ) + try: + provenance["git_head"] = subprocess.check_output( + ["git", "rev-parse", "HEAD"], cwd=root, text=True, stderr=subprocess.DEVNULL + ).strip() + provenance["git_status"] = subprocess.check_output( + ["git", "status", "--porcelain"], cwd=root, text=True + ).splitlines() + except (subprocess.CalledProcessError, FileNotFoundError): + provenance["git_head"] = None + return provenance + + +def benchmark(config, request, output): + output = Path(output) + output.mkdir(parents=True, exist_ok=False) + provenance = source_provenance() + contract = { + "config": asdict(config), + "request": asdict(request), + "provenance": provenance, + "state": "running", + "started": time.time(), + } + (output / "contract.json").write_text(json.dumps(contract, indent=2)) + try: + runs = [] + with H3Engine(config) as engine: + for index in range(4): + directory = output / ("warmup" if index == 0 else f"run-{index}") + result = engine.generate(request, directory) + result["measurement"] = { + "warmup": index == 0, + "profiled": False, + "capture": False, + } + result["provenance"] = provenance + (directory / "run.json").write_text(json.dumps(result, indent=2)) + runs.append(result) + report = evaluate_performance(runs[1:], warmup=runs[0]) + (output / "performance.json").write_text(json.dumps(report, indent=2)) + contract["state"] = "completed" + return report + except BaseException as exc: + contract.update(state="failed", error=repr(exc)) + raise + finally: + contract["finished"] = time.time() + (output / "contract.json").write_text(json.dumps(contract, indent=2)) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--contract", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + contract = json.loads(args.contract.read_text()) + config = H3Config(**contract["config"]) + request = dict(contract["request"]) + request["sampling"] = H3SamplingParams(**request["sampling"]) + report = benchmark(config, H3Request(**request), args.output) + print(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/vllm/video/engine.py b/vllm/video/engine.py index a5437d32b7..d58a5a7fe1 100644 --- a/vllm/video/engine.py +++ b/vllm/video/engine.py @@ -8,9 +8,11 @@ import multiprocessing as mp import os import socket +import tempfile import threading import time import traceback +import uuid from dataclasses import asdict from multiprocessing.connection import Connection from pathlib import Path @@ -23,7 +25,7 @@ from .gpu import select_gpu_group as select_gpu_group -def _worker(rank, config, gpu_ids, endpoint, connection): +def _worker(rank, config, gpu_ids, endpoint, connection, shared_weights_dir=None): os.environ["CUDA_VISIBLE_DEVICES"] = worker_device_mask(gpu_ids) from datetime import timedelta @@ -44,6 +46,7 @@ def _worker(rank, config, gpu_ids, endpoint, connection): current = VllmConfig( parallel_config=ParallelConfig(tensor_parallel_size=config.tensor_parallel_size) ) + pipeline = None try: with set_current_vllm_config(current): init_distributed_environment( @@ -56,7 +59,8 @@ def _worker(rank, config, gpu_ids, endpoint, connection): ) initialize_model_parallel(config.tensor_parallel_size) started = time.perf_counter() - pipeline = MiniMaxH3Pipeline(config) + pipeline = MiniMaxH3Pipeline(config, shared_weights_dir=shared_weights_dir) + kernel_provenance = None connection.send( { "ready": True, @@ -80,13 +84,33 @@ def _worker(rank, config, gpu_ids, endpoint, connection): reporting(observer), ): video, audio = pipeline(request) + if kernel_provenance is None: + from .metrics import loaded_kernel_provenance + + kernel_provenance = loaded_kernel_provenance() + communication = pipeline.residual_reduction_stats() + torch_peak = torch.accelerator.max_memory_allocated() + raw_peak = communication["raw_ipc_peak_bytes"] result = { "rank": rank, + "residual_communication": communication, + "torch_peak_allocated_bytes": torch_peak, + "raw_ipc_peak_bytes": raw_peak, + "peak_allocation_is_upper_bound": bool(raw_peak), "stage_seconds": pipeline.stage_durations, "dit_calls": pipeline.actual_dit_calls, "useful_denoise_flops": pipeline.useful_denoise_flops, "denoise_flops_by_layer": pipeline.denoise_flops_by_layer, - "peak_allocated_bytes": torch.accelerator.max_memory_allocated(), + "redundant_denoise_flops": pipeline.redundant_denoise_flops, + "redundant_flops_by_layer": pipeline.redundant_flops_by_layer, + "denoise_workload": pipeline.denoise_workload, + "denoise_steps": pipeline.denoise_steps, + "denoise_executed_blocks": pipeline.denoise_executed_blocks, + "denoise_sparse_work_by_layer": ( + pipeline.denoise_sparse_work_by_layer + ), + "kernel_provenance": kernel_provenance, + "peak_allocated_bytes": torch_peak + raw_peak, } if rank == 0: from .media import export_video @@ -133,14 +157,21 @@ def _worker(rank, config, gpu_ids, endpoint, connection): except BaseException: connection.send({"error": traceback.format_exc(), "rank": rank}) finally: - cleanup_dist_env_and_memory() - connection.close() + try: + if pipeline is not None: + pipeline.close() + finally: + cleanup_dist_env_and_memory() + connection.close() class H3Engine: def __init__(self, config: H3Config): self.config = config + self.session_id = str(uuid.uuid4()) + self.request_index = 0 self._gpu_lease = None + self._shared_weights = None self._lock = threading.Lock() self._closed = False self.workers = [] @@ -150,15 +181,20 @@ def __init__(self, config: H3Config): try: self._gpu_lease = acquire_gpu_group(config.tensor_parallel_size) self.gpu_ids = self._gpu_lease.gpu_ids + if config.share_host_vae_weights and config.tensor_parallel_size > 1: + self._shared_weights = tempfile.TemporaryDirectory( + prefix="vllm-h3-vae-", dir="/dev/shm" + ) with socket.socket() as sock: sock.bind(("127.0.0.1", 0)) port = sock.getsockname()[1] endpoint = f"tcp://127.0.0.1:{port}" for rank in range(config.tensor_parallel_size): parent, child = context.Pipe() - worker = context.Process( - target=_worker, args=(rank, config, self.gpu_ids, endpoint, child) - ) + args: tuple = (rank, config, self.gpu_ids, endpoint, child) + if self._shared_weights is not None: + args = (*args, self._shared_weights.name) + worker = context.Process(target=_worker, args=args) worker.start() child.close() self.workers.append(worker) @@ -229,6 +265,8 @@ def generate( self.close() raise result = { + "engine_session_id": self.session_id, + "request_index": self.request_index, "config": asdict(self.config), "request": asdict(request), "gpus": self.gpu_ids, @@ -239,6 +277,7 @@ def generate( (output_dir / "run.json").write_text( json.dumps(result, indent=2, ensure_ascii=False) ) + self.request_index += 1 return result def close(self): @@ -258,6 +297,9 @@ def close(self): worker.join() for connection in self.connections: connection.close() + if self._shared_weights is not None: + self._shared_weights.cleanup() + self._shared_weights = None if self._gpu_lease is not None: self._gpu_lease.close() self._gpu_lease = None diff --git a/vllm/video/metrics.py b/vllm/video/metrics.py index 824e635b88..b7d93448c0 100644 --- a/vllm/video/metrics.py +++ b/vllm/video/metrics.py @@ -13,6 +13,45 @@ from pathlib import Path +def loaded_kernel_provenance(): + """Identify actual loaded native binaries, including task-owned JIT builds.""" + import hashlib + import sys + + paths = { + str(Path(filename).resolve()) + for name, module in list(sys.modules.items()) + if name.startswith(("vllm._h3_", "onecat_h3_", "vllm._sm70_", "onecat_sm70_")) + and (filename := getattr(module, "__file__", None)) + } + return { + path: hashlib.sha256(Path(path).read_bytes()).hexdigest() + for path in sorted(paths) + } + + +def lora_work(layer, parts, rows, *, replicated_a): + """Separate identical column-A replicas from useful TP partial products. + + Column-parallel A has identical weights and full input rows on every rank. + Attribute each row once, balanced by rank, including non-divisible tails. + Row-parallel A instead consumes distinct input shards; its B products are + distinct partial contributions and are not identical replicas. + """ + a_elements = sum( + getattr(layer, f"h3_lora_a_{index}").numel() for index, _, _ in parts + ) + b_elements = sum( + getattr(layer, f"h3_lora_b_{index}").numel() for index, _, _ in parts + ) + executed = 2 * rows * (a_elements + b_elements) + redundant = 0 + if replicated_a: + owner_rows = rows // layer.tp_size + (layer.tp_rank < rows % layer.tp_size) + redundant = 2 * (rows - owner_rows) * a_elements + return executed - redundant, redundant + + class DenoiseWorkCounter: """Count actual TP-local matrix shapes, excluding structural padding. @@ -21,7 +60,7 @@ class DenoiseWorkCounter: """ def __init__(self, model, *, used_length, video_outputs, audio_outputs): - from vllm.model_executor.layers.linear import LinearBase + from vllm.model_executor.layers.linear import ColumnParallelLinear, LinearBase from vllm.model_executor.models.minimax_h3.attention import Attention from vllm.model_executor.models.minimax_h3.lora import ( TurboLinearMethod, @@ -29,16 +68,43 @@ def __init__(self, model, *, used_length, video_outputs, audio_outputs): ) self.flops = 0 + self.redundant_flops = 0 self.calls = 0 + self.sparse_blocks = 0 + self.sparse_pairs = 0 + self.sparse_compression_flops = 0 + self.attention_avoided_flops = 0 + self.sparse_by_layer: dict[str, dict] = {} + self.blocks: dict[str, int] = {} + self.steps: list[dict] = [] + self._step_events = [] + self._pending_sparse: list[tuple[int | None, str, int, dict]] = [] + self._active_step: int | None = None self.by_layer: dict[str, int] = {} + self.redundant_by_layer: dict[str, int] = {} self.handles = [] def completed_call(module, inputs, output): self.calls += 1 self.handles.append(model.register_forward_hook(completed_call)) + from vllm.model_executor.models.minimax_h3.transformer import ( + MiniMaxH3DiTBlock, + MiniMaxH3TokenRefinerBlock, + ) + + self.blocks_per_call = sum( + isinstance(module, (MiniMaxH3DiTBlock, MiniMaxH3TokenRefinerBlock)) + for module in model.modules() + ) for name, module in model.named_modules(): - if isinstance(module, LinearBase): + if isinstance(module, (MiniMaxH3DiTBlock, MiniMaxH3TokenRefinerBlock)): + + def block_hook(layer, inputs, output, name=name): + self.blocks[name] = self.blocks.get(name, 0) + 1 + + self.handles.append(module.register_forward_hook(block_hook)) + elif isinstance(module, LinearBase): def linear_hook(layer, inputs, output, name=name): rows = inputs[0].numel() // inputs[0].shape[-1] @@ -53,18 +119,19 @@ def linear_hook(layer, inputs, output, name=name): self.by_layer[name] = self.by_layer.get(name, 0) + count method = layer.quant_method if isinstance(method, TurboLinearMethod) and lora_scale.get() != 0: - work = sum( - 2 - * effective - * ( - getattr(layer, f"h3_lora_a_{index}").numel() - + getattr(layer, f"h3_lora_b_{index}").numel() - ) - for index, _, _ in method.parts + work, redundant = lora_work( + layer, + method.parts, + effective, + replicated_a=isinstance(layer, ColumnParallelLinear), ) self.flops += work key = name + ".lora" self.by_layer[key] = self.by_layer.get(key, 0) + work + self.redundant_flops += redundant + self.redundant_by_layer[key] = ( + self.redundant_by_layer.get(key, 0) + redundant + ) self.handles.append(module.register_forward_hook(linear_hook)) elif isinstance(module, Attention): @@ -72,15 +139,143 @@ def linear_hook(layer, inputs, output, name=name): def attention_hook(layer, inputs, output, name=name): q, k, v, metadata = inputs used = metadata.extra.get("valid_kv_length", q.shape[1]) - count = 4 * q.shape[0] * q.shape[2] * used * used * q.shape[3] + if layer.backend == "FASTVIDEO_VSA": + work = metadata.extra["sparse_work"] + # Dynamic counts remain on the GPU while layers enqueue. + self._pending_sparse.append( + (self._active_step, name, q.shape[3], dict(work)) + ) + return + else: + count = 4 * q.shape[0] * q.shape[2] * used * used * q.shape[3] self.flops += count self.by_layer[name] = self.by_layer.get(name, 0) + count self.handles.append(module.register_forward_hook(attention_hook)) + @contextlib.contextmanager + def step(self, index): + """Record stream spans without introducing per-step synchronization. + + GPU events include dependent communication and stream waits. CPU enqueue + time is separate; neither replaces complete synchronized denoise wall time. + """ + import torch + + start, end = (torch.cuda.Event(enable_timing=True) for _ in range(2)) + before = self.flops, self.calls, sum(self.blocks.values()), self.redundant_flops + sparse_before = self.sparse_blocks + pairs_before = self.sparse_pairs + compression_before = self.sparse_compression_flops + avoided_before = self.attention_avoided_flops + if self._active_step is not None: + raise RuntimeError("denoise counter steps cannot be nested") + self._active_step = len(self.steps) + start.record() + started = time.perf_counter() + try: + yield + finally: + self._active_step = None + end.record() + self.steps.append( + { + "index": index, + "cpu_enqueue_seconds": time.perf_counter() - started, + "useful_flops": self.flops - before[0], + "redundant_flops": self.redundant_flops - before[3], + "dit_calls": self.calls - before[1], + "executed_blocks": sum(self.blocks.values()) - before[2], + "sparse_blocks": self.sparse_blocks - sparse_before, + "sparse_token_pairs": self.sparse_pairs - pairs_before, + "sparse_compression_flops": self.sparse_compression_flops + - compression_before, + "attention_avoided_flops": self.attention_avoided_flops + - avoided_before, + "cache_hits": 0, + } + ) + self._step_events.append((start, end)) + + def finish_sparse(self): + """Resolve all dynamic counters once, inside complete-denoise timing.""" + if not self._pending_sparse: + return + import torch + + tensors = [] + for _, _, _, work in self._pending_sparse: + for key in ("selected_token_pairs", "selected_blocks"): + value = work[key] + if isinstance(value, torch.Tensor): + if value.ndim != 0 or value.dtype != torch.int64: + raise ValueError("sparse work counters must be int64 scalars") + tensors.append(value) + values = iter(torch.stack(tensors).cpu().tolist() if tensors else []) + pending, self._pending_sparse = self._pending_sparse, [] + for step_index, name, head_size, work in pending: + for key in ("selected_token_pairs", "selected_blocks"): + if isinstance(work[key], torch.Tensor): + work[key] = next(values) + pairs, blocks = work["selected_token_pairs"], work["selected_blocks"] + compression = work["compression_flops"] + count = 4 * pairs * head_size + avoided = 4 * (work["dense_token_pairs"] - pairs) * head_size + self.flops += count + compression + self.sparse_pairs += pairs + self.sparse_blocks += blocks + self.sparse_compression_flops += compression + self.attention_avoided_flops += avoided + self.by_layer[name] = self.by_layer.get(name, 0) + count + key = name + ".compression" + self.by_layer[key] = self.by_layer.get(key, 0) + compression + record = self.sparse_by_layer.setdefault( + name, + dict( + head_size=head_size, + heads=work["heads"], + selected_blocks=0, + selected_token_pairs=0, + compression_flops=0, + dense_token_pairs=0, + ), + ) + for key in ( + "selected_blocks", + "selected_token_pairs", + "compression_flops", + "dense_token_pairs", + ): + record[key] += work[key] + if step_index is not None: + step = self.steps[step_index] + step["useful_flops"] += count + compression + step["sparse_blocks"] += blocks + step["sparse_token_pairs"] += pairs + step["sparse_compression_flops"] += compression + step["attention_avoided_flops"] += avoided + + def finish_steps(self): + """Read events only after the caller's complete-denoise synchronization.""" + self.finish_sparse() + for record, (start, end) in zip(self.steps, self._step_events): + record["gpu_seconds"] = start.elapsed_time(end) / 1000 + return self.steps + + def __enter__(self): + return self + + def __exit__(self, exc_type, *_): + try: + if exc_type is None: + self.finish_sparse() + finally: + self.close() + def close(self): for handle in self.handles: handle.remove() + self._pending_sparse.clear() class NVMLMonitor: @@ -166,52 +361,220 @@ def __exit__(self, *_): self.thread.join(timeout=5) -def evaluate_performance(runs): - """Evaluate three unprofiled post-warmup runs; never use utilization as FLOPs.""" - if len(runs) != 3 or any(len(run["ranks"]) != 4 for run in runs): - raise ValueError("acceptance requires three four-rank measurements") +def _validate_sparse_work(rank, calls): + workload = rank["denoise_workload"] + config = workload["sparse_config"] + for key in ("topk", "gated_blocks", "heads", "head_size"): + if type(config[key]) is not int or config[key] <= 0: + raise ValueError("invalid sparse execution configuration") + if ( + config["head_size"] != 128 + or config["gated_blocks"] != workload["blocks_per_call"] - 2 + or "FASTVIDEO_VSA" not in workload["actual_backends"] + ): + raise ValueError("sparse execution must cover the actual gated H3 blocks") + prefix, shape = config["prefix_segments"], config["video_shape"] + if len(shape) != 3 or any(type(n) is not int or n <= 0 for n in (*prefix, *shape)): + raise ValueError("invalid sparse geometry") + if sum(prefix) + math.prod(shape) != workload["used_length"]: + raise ValueError("sparse geometry disagrees with valid token count") + prefix_blocks = sum((n + 63) // 64 for n in prefix) + video_blocks = math.prod((n + 3) // 4 for n in shape) + blocks = prefix_blocks + video_blocks + per_layer_blocks = config["heads"] * ( + prefix_blocks * blocks + + video_blocks * (prefix_blocks + min(config["topk"], video_blocks)) + ) + per_layer_dense_pairs = config["heads"] * workload["used_length"] ** 2 + per_layer_compression = 4 * config["heads"] * blocks**2 * config["head_size"] + layers = rank["denoise_sparse_work_by_layer"] + if len(layers) != config["gated_blocks"]: + raise ValueError("missing sparse layer execution records") + sums = dict( + selected_blocks=0, + selected_token_pairs=0, + compression_flops=0, + dense_token_pairs=0, + ) + for name, layer in layers.items(): + if any(type(layer[key]) is not int or layer[key] <= 0 for key in sums): + raise ValueError("sparse work must use positive integer counts") + if ( + layer["head_size"] != config["head_size"] + or layer["heads"] != config["heads"] + or layer["selected_blocks"] != calls * per_layer_blocks + or layer["dense_token_pairs"] != calls * per_layer_dense_pairs + or layer["compression_flops"] != calls * per_layer_compression + or not layer["selected_blocks"] + <= layer["selected_token_pairs"] + <= min(layer["dense_token_pairs"], layer["selected_blocks"] * 64**2) + or rank["denoise_flops_by_layer"][name] + != 4 * layer["selected_token_pairs"] * config["head_size"] + or rank["denoise_flops_by_layer"][name + ".compression"] + != layer["compression_flops"] + ): + raise ValueError("sparse layer pairs, blocks or compression disagree") + for key in sums: + sums[key] += layer[key] + steps = rank["denoise_steps"] + for step in steps: + if ( + step["sparse_blocks"] != per_layer_blocks * config["gated_blocks"] + or type(step["sparse_token_pairs"]) is not int + or not 0 + < step["sparse_token_pairs"] + <= per_layer_dense_pairs * config["gated_blocks"] + or step["sparse_compression_flops"] + != per_layer_compression * config["gated_blocks"] + or step["attention_avoided_flops"] + != 4 + * config["head_size"] + * ( + per_layer_dense_pairs * config["gated_blocks"] + - step["sparse_token_pairs"] + ) + ): + raise ValueError("sparse step counters disagree with executed geometry") + if any( + sum(step[key] for step in steps) != sums[target] + for key, target in ( + ("sparse_blocks", "selected_blocks"), + ("sparse_token_pairs", "selected_token_pairs"), + ("sparse_compression_flops", "compression_flops"), + ) + ): + raise ValueError("sparse step and layer totals disagree") + + +def _validate_workload(rank): + workload = rank["denoise_workload"] + sparse = workload["attention_algorithm"] == "vsa" + expected_version = "sparse_tp_v1" if sparse else "dense_tp_lora_v2" + if workload.get("work_accounting") != expected_version: + raise ValueError("legacy work counts may include replicated LoRA projections") + if ( + workload["attention_algorithm"] not in ("dense", "vsa") + or workload["cache_algorithm"] is not None + ): + raise ValueError( + "sparse/cache workflows require their own measured work accounting" + ) + schedules = [workload[key] for key in ("video_sigmas", "audio_sigmas")] + for schedule in schedules: + if ( + len(schedule) < 2 + or any(not math.isfinite(s) or not 0 <= s <= 1 for s in schedule) + or any(a <= b for a, b in zip(schedule, schedule[1:])) + or schedule[-1] != 0 + ): + raise ValueError("invalid measured sigma schedule") + if len(schedules[0]) != len(schedules[1]): + raise ValueError("video/audio schedules must have equal length") + calls = len(schedules[0]) - 1 + blocks = workload["blocks_per_call"] + if type(blocks) is not int or blocks <= 0 or rank["dit_calls"] != calls: + raise ValueError("incomplete denoiser execution for the measured schedule") + steps = rank["denoise_steps"] + if len(steps) != calls or [step["index"] for step in steps] != list(range(calls)): + raise ValueError("missing, duplicate or unordered denoise steps") + for step in steps: + if ( + step["dit_calls"] != 1 + or step["executed_blocks"] != blocks + or step["cache_hits"] != 0 + or (not sparse and step["sparse_blocks"] != 0) + ): + raise ValueError("dense step work does not match the workflow") + if type(step["useful_flops"]) is not int or step["useful_flops"] <= 0: + raise ValueError("useful FLOPs must be positive integer counts") + if type(step["redundant_flops"]) is not int or step["redundant_flops"] < 0: + raise ValueError("redundant work must be a nonnegative integer count") + for key in ("gpu_seconds", "cpu_enqueue_seconds"): + if not math.isfinite(step[key]) or step[key] <= 0: + raise ValueError("invalid measured step duration") + if ( + sum(step["useful_flops"] for step in steps) != rank["useful_denoise_flops"] + or sum(rank["denoise_flops_by_layer"].values()) != rank["useful_denoise_flops"] + or len(rank["denoise_executed_blocks"]) != blocks + or any(value != calls for value in rank["denoise_executed_blocks"].values()) + or sum(step["redundant_flops"] for step in steps) + != rank["redundant_denoise_flops"] + or sum(rank["redundant_flops_by_layer"].values()) + != rank["redundant_denoise_flops"] + ): + raise ValueError("step, layer and complete-denoise work counts disagree") + if sparse: + _validate_sparse_work(rank, calls) + return workload + + +def evaluate_performance(runs, *, warmup): + """Evaluate three full requests after a completed, same-session warmup. + + The runtime descriptor supplies the actual sigma intervals. API step counts + are deliberately not interpreted here: LightX2V and DMD2 differ. Passing a + shape does not establish coverage of other workloads or their quality. + """ + if len(runs) != 3: + raise ValueError("acceptance requires three post-warmup measurements") baseline = runs[0] - for run in runs: - if sorted(rank["rank"] for rank in run["ranks"]) != list(range(4)): - raise ValueError("each measurement must contain four unique TP ranks") - if any(run[key] != baseline[key] for key in ("config", "request", "gpus")): - raise ValueError("acceptance measurements must use the same configuration") - sampling = run["request"]["sampling"] - expected = { - "width": 1344, - "height": 768, - "num_frames": 243, - "fps": 24, - "seed": 42, - "num_inference_steps": 50, - } - if any(sampling.get(key) != value for key, value in expected.items()): - raise ValueError("acceptance requires the fixed primary workload") - if run.get("measurement", {}).get("profiled") is not False: + tp = baseline["config"]["tensor_parallel_size"] + if tp not in (1, 2, 4): + raise ValueError("invalid H3 tensor parallel size") + if not baseline.get("engine_session_id"): + raise ValueError("completed same-session warmup evidence is required") + indices = [run["request_index"] for run in (warmup, *runs)] + if indices != list(range(indices[0], indices[0] + 4)): + raise ValueError( + "warmup and measurements must be consecutive complete requests" + ) + descriptor = baseline["ranks"][0]["denoise_workload"] + for index, run in enumerate((warmup, *runs)): + if sorted(rank["rank"] for rank in run["ranks"]) != list(range(tp)): + raise ValueError("each measurement must contain every unique TP rank") + for key in ("config", "request", "gpus", "engine_session_id"): + if run[key] != baseline[key]: + raise ValueError( + "acceptance measurements must use the same configuration" + ) + if len(run["gpus"]) != tp or len(set(run["gpus"])) != tp: + raise ValueError("invalid physical GPU group") + measurement = run.get("measurement", {}) + if measurement.get("profiled") is not False: raise ValueError("formal timing must be explicitly recorded as unprofiled") + if measurement.get("warmup") is not (index == 0): + raise ValueError("warmup must be complete and excluded from measurements") + if measurement.get("capture") is not False: + raise ValueError("quality captures must be separate from performance runs") if run.get("timing_valid") is False: raise ValueError("run timing was excluded from performance evidence") - if any(rank["dit_calls"] != 49 for rank in run["ranks"]): - raise ValueError("the primary schedule requires 49 completed DiT calls") - if any( - type(rank["useful_denoise_flops"]) is not int - or rank["useful_denoise_flops"] <= 0 - for rank in run["ranks"] + if ( + not math.isfinite(run["end_to_end_seconds"]) + or run["end_to_end_seconds"] <= 0 ): - raise ValueError("useful FLOPs must be positive integer counts") + raise ValueError("invalid end-to-end duration") + for rank in run["ranks"]: + if _validate_workload(rank) != descriptor: + raise ValueError( + "all ranks and requests must execute the same workflow" + ) + seconds = rank["stage_seconds"]["denoise"] + if not math.isfinite(seconds) or seconds <= 0: + raise ValueError("invalid denoise duration") + memory = rank["peak_allocated_bytes"] + if type(memory) is not int or memory <= 0: + raise ValueError("invalid peak memory measurement") ordered = [sorted(run["ranks"], key=lambda rank: rank["rank"]) for run in runs] seconds = [ - max(rank["stage_seconds"]["denoise"] for rank in run["ranks"]) for run in runs + max(rank["stage_seconds"]["denoise"] for rank in ranks) for ranks in ordered ] - if any(not math.isfinite(value) or value <= 0 for value in seconds): - raise ValueError("invalid denoise duration") - medians = [] - for rank in range(4): - values = [ + medians = [ + statistics.median( ranks[rank]["useful_denoise_flops"] / duration / 1e12 for ranks, duration in zip(ordered, seconds) - ] - medians.append(statistics.median(values)) + ) + for rank in range(tp) + ] cv = statistics.pstdev(seconds) / statistics.mean(seconds) memory_passed = all( rank["peak_allocated_bytes"] <= 30 * 1024**3 @@ -219,9 +582,28 @@ def evaluate_performance(runs): for rank in run["ranks"] ) return { + "workflow": descriptor, + "sampling": baseline["request"]["sampling"], + "tensor_parallel_size": tp, "rank_median_tflops": medians, "denoise_seconds": seconds, "denoise_cv": cv, + "end_to_end_seconds": [run["end_to_end_seconds"] for run in runs], + "peak_allocated_bytes": [ + max(ranks[rank]["peak_allocated_bytes"] for ranks in ordered) + for rank in range(tp) + ], "memory_passed": memory_passed, + "attention_avoided_flops_by_run_and_rank": [ + [ + sum( + step.get("attention_avoided_flops", 0) + for step in rank["denoise_steps"] + ) + for rank in ranks + ] + for ranks in ordered + ], "performance_passed": all(value > 80 for value in medians) and cv <= 0.05, + "quality_status": "requires_separate_numerical_and_human_review", } diff --git a/vllm/video/vsa_acceptance.py b/vllm/video/vsa_acceptance.py new file mode 100644 index 0000000000..c2a008dcc5 --- /dev/null +++ b/vllm/video/vsa_acceptance.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Evaluate the primary VSA timing stage independently of the future 80-TF goal. + +This checks measured timing, work and matched configuration. Weight identity, +numerical quality, human review and other shapes remain separate evidence. +Inputs are native benchmark directories or JSON with ``warmup`` and three +``runs``. No profiler or quality-capture run is eligible for formal timing. +""" + +import argparse +import json +import statistics +from pathlib import Path + +from vllm.model_executor.models.minimax_h3.fasth3 import FastH3Spec +from vllm.model_executor.models.minimax_h3.sigma_schedule import DMD2SigmaSchedule +from vllm.video.metrics import evaluate_performance + + +def evaluate_vsa_stage(vsa, dense): + reports = { + name: evaluate_performance(data["runs"], warmup=data["warmup"]) + for name, data in (("vsa", vsa), ("dense", dense)) + } + for report in reports.values(): + report["future_80_tflops_passed"] = report.pop("performance_passed") + candidate = vsa["runs"][0] + control = dense["runs"][0] + allowed_differences = { + "attention_backend", + "attention_query_tile", + "vsa_topk", + "lora_path", + } + configs = [ + { + key: value + for key, value in run["config"].items() + if key not in allowed_differences + } + for run in (candidate, control) + ] + if configs[0] != configs[1] or any( + candidate[key] != control[key] for key in ("request", "gpus") + ): + raise ValueError( + "VSA and Dense require matched host, VAE, communication, " + "export and request settings" + ) + sampling = candidate["request"]["sampling"] + if ( + candidate["config"]["tensor_parallel_size"] != 4 + or tuple(sampling[key] for key in ("width", "height", "num_frames", "fps")) + != (1280, 736, 120, 24) + or sampling["num_inference_steps"] != 4 + or sampling["num_outputs_per_prompt"] != 1 + ): + raise ValueError( + "31.3-second acceptance applies to the TP4 primary four-step request only" + ) + sparse_work = reports["vsa"]["workflow"] + dense_work = reports["dense"]["workflow"] + spec = FastH3Spec() + schedule = DMD2SigmaSchedule(spec.base_schedule) + if ( + sparse_work["attention_algorithm"] != "vsa" + or dense_work["attention_algorithm"] != "dense" + or sparse_work["sparse_config"]["topk"] != 64 + or sparse_work["sparse_config"]["video_shape"] != [37, 23, 40] + or sparse_work["sparse_config"]["prefix_segments"] != [97, 414] + or sparse_work["sparse_config"]["gated_blocks"] != 50 + or candidate["config"]["attention_backend"] != "FASTVIDEO_VSA" + or candidate["config"]["vsa_topk"] != 64 + or control["config"]["attention_backend"] != "FLASH_ATTN_V100" + or sparse_work["adapter"] != "FastH3Spec" + or dense_work["adapter"] != "FastH3Spec" + or sparse_work["video_sigmas"] != schedule.shifted_sigmas(spec.video_shift) + or sparse_work["audio_sigmas"] != schedule.shifted_sigmas(spec.audio_shift) + or sparse_work["task"] != "t2va" + or any( + sparse_work[key] != dense_work[key] + for key in ( + "partition", + "task", + "video_sigmas", + "audio_sigmas", + "used_length", + "blocks_per_call", + "cache_algorithm", + ) + ) + or len(sparse_work["video_sigmas"]) != 5 + ): + raise ValueError( + "primary acceptance requires matching FastH3 four-step Dense " + "and top-k64 VSA algorithms" + ) + denoise = statistics.median(reports["vsa"]["denoise_seconds"]) + requests = { + name: statistics.median(report["end_to_end_seconds"]) + for name, report in reports.items() + } + checks = { + "denoise_median": denoise <= 31.3, + "denoise_cv": reports["vsa"]["denoise_cv"] <= 0.05, + "whole_request_beats_dense": requests["vsa"] < requests["dense"], + } + return { + "stage": "primary_vsa_31_3_seconds", + "checks": checks, + "stage_performance_passed": all(checks.values()), + "denoise_median_seconds": denoise, + "request_median_seconds": requests, + "measurements": reports, + "quality_status": "requires_independent_numerical_and_human_review", + "weight_identity_status": "requires_separate_frozen_base_and_adapter_identity", + "official_hardware_status": "deferred", + "other_shapes_status": "not_established_by_primary_timing", + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--vsa", type=Path, required=True) + parser.add_argument("--dense", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + + def read(path): + if path.is_file(): + return json.loads(path.read_text()) + return { + "warmup": json.loads((path / "warmup/run.json").read_text()), + "runs": [ + json.loads((path / f"run-{i}/run.json").read_text()) + for i in range(1, 4) + ], + } + + result = evaluate_vsa_stage(read(args.vsa), read(args.dense)) + args.output.write_text(json.dumps(result, indent=2)) + print(json.dumps(result, indent=2)) + + +if __name__ == "__main__": + main()