diff --git a/benchmarks/compare_sm70_dflash2_natural_audit.py b/benchmarks/compare_sm70_dflash2_natural_audit.py index fd3a83111d..946938df77 100644 --- a/benchmarks/compare_sm70_dflash2_natural_audit.py +++ b/benchmarks/compare_sm70_dflash2_natural_audit.py @@ -19,6 +19,10 @@ sampling_difference, tensor_difference, ) +from benchmarks.sm70_dflash2_state_layout import ( + check_slot_mapping, + explain_state_difference, +) def _load(directory: Path) -> dict: @@ -104,21 +108,36 @@ def _proposal_tensors(row: dict) -> dict[str, torch.Tensor]: return values -def compare_natural(left_dir: Path, right_dir: Path) -> dict: +def compare_natural( + left_dir: Path, right_dir: Path, *, conv_width: int | None = None +) -> dict: left, right = _load(left_dir), _load(right_dir) cases = {k[0] for k in left} if cases != {k[0] for k in right}: raise ValueError("Case coverage differs") - result = {"left": str(left_dir), "right": str(right_dir), "cases": []} + result = { + "left": str(left_dir), + "right": str(right_dir), + "conv_width": conv_width, + "cases": [], + } for case in sorted(cases): lengths = [len({k[1] for k in arm if k[0] == case}) for arm in (left, right)] first = None mappings = [] + state_mappings = {rank: ({}, {}) for rank in range(4)} + explained_storage = [] for step in range(min(lengths)): for phase_index, phase in enumerate(("target", "proposal")): differences = [] for rank in range(4): rows = [arm[case, step, rank][phase_index] for arm in (left, right)] + if phase == "target" and conv_width is not None: + check_slot_mapping( + rows[0]["states"], + rows[1]["states"], + *state_mappings[rank], + ) values = [ _target_tensors(row) if phase == "target" @@ -152,7 +171,20 @@ def compare_natural(left_dir: Path, right_dir: Path) -> dict: continue if name == "native_logits": diff.update(sampling_difference(a, b)) - differences.append({"rank": rank, "name": name, **diff}) + observation = {"rank": rank, "name": name, **diff} + if name.startswith("states/") and conv_width is not None: + reason = explain_state_difference( + name.removeprefix("states/"), + rows[0]["states"], + rows[1]["states"], + conv_width, + ) + if reason is not None: + explained_storage.append( + {"step": step, "reason": reason, **observation} + ) + continue + differences.append(observation) if differences: first = {"step": step, "phase": phase, "differences": differences} break @@ -164,6 +196,7 @@ def compare_natural(left_dir: Path, right_dir: Path) -> dict: "steps_per_arm": lengths, "first_observed_difference": first, "different_request_slot_mappings": mappings, + "explained_storage_differences": explained_storage, "all_logical_tensors_equal": first is None and lengths[0] == lengths[1], } ) @@ -175,9 +208,10 @@ def main() -> None: parser.add_argument("left", type=Path) parser.add_argument("right", type=Path) parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--conv-width", type=int, choices=range(2, 7)) args = parser.parse_args() torch.set_num_threads(4) - result = compare_natural(args.left, args.right) + result = compare_natural(args.left, args.right, conv_width=args.conv_width) args.output.write_text(json.dumps(result, indent=2) + "\n") for case in result["cases"]: first = case["first_observed_difference"] diff --git a/benchmarks/compare_sm70_dflash2_state_audit.py b/benchmarks/compare_sm70_dflash2_state_audit.py index aafe8ab860..226c79dcdc 100644 --- a/benchmarks/compare_sm70_dflash2_state_audit.py +++ b/benchmarks/compare_sm70_dflash2_state_audit.py @@ -10,6 +10,11 @@ import torch +from benchmarks.sm70_dflash2_state_layout import ( + check_slot_mapping, + explain_state_difference, +) + def tensor_difference(left: torch.Tensor, right: torch.Tensor) -> dict: if left.shape != right.shape or left.dtype != right.dtype: @@ -37,7 +42,9 @@ def raw_bytes(tensor): } -def sampling_difference(left: torch.Tensor, right: torch.Tensor) -> dict: +def sampling_difference( + left: torch.Tensor, right: torch.Tensor, eos_token_ids: tuple[int, ...] = () +) -> dict: from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p_pytorch if left.shape != right.shape: @@ -52,17 +59,37 @@ def probabilities(logits): ).softmax(-1) p, q = probabilities(left), probabilities(right) - full_tv = (left.float().softmax(-1) - right.float().softmax(-1)).abs().sum(-1) / 2 - return { + full_p, full_q = left.float().softmax(-1), right.float().softmax(-1) + full_tv = (full_p - full_q).abs().sum(-1) / 2 + result = { "full_softmax_tv": full_tv.tolist(), "sampling_tv": ((p - q).abs().sum(-1) / 2).tolist(), "support_changed": ((p > 0) != (q > 0)).any(-1).tolist(), "top1_changed": (left.argmax(-1) != right.argmax(-1)).tolist(), } + if eos_token_ids: + if any(token < 0 or token >= left.shape[-1] for token in eos_token_ids): + raise ValueError("EOS token ID outside the captured vocabulary") + ids = list(eos_token_ids) + result["eos_token_ids"] = ids + result["full_eos_probabilities_left"] = full_p[:, ids].tolist() + result["full_eos_probabilities_right"] = full_q[:, ids].tolist() + result["sampling_eos_probabilities_left"] = p[:, ids].tolist() + result["sampling_eos_probabilities_right"] = q[:, ids].tolist() + result["max_eos_probability_abs_difference"] = max( + (full_p[:, ids] - full_q[:, ids]).abs().max().item(), + (p[:, ids] - q[:, ids]).abs().max().item(), + ) + return result def compare( - left_dir: Path, right_dir: Path, *, right_verifier_route: str | None = None + left_dir: Path, + right_dir: Path, + *, + right_verifier_route: str | None = None, + eos_token_ids: tuple[int, ...] = (), + conv_width: int | None = None, ) -> dict: left_files = {p.name: p for p in left_dir.glob("*-rank*-step*.pt")} right_files = {p.name: p for p in right_dir.glob("*-rank*-step*.pt")} @@ -73,7 +100,10 @@ def compare( "right": str(right_dir), "comparisons": [], "logits": [], + "explained_storage_differences": [], + "conv_width": conv_width, } + mappings: dict[tuple[str, int], tuple[dict[int, int], dict[int, int]]] = {} coverage: dict[tuple[str, int], set[int]] = {} seen_states: dict[tuple[str, int, str], set[str]] = {} for name in sorted(left_files): @@ -83,6 +113,11 @@ def compare( if left[key] != right[key]: raise ValueError(f"{name}: {key} differs") identity = {key: left[key] for key in ("case", "rank", "step", "phase")} + if conv_width is not None: + mapping, reverse = mappings.setdefault( + (left["case"], left["rank"]), ({}, {}) + ) + check_slot_mapping(left["states"], right["states"], mapping, reverse) if right_verifier_route is not None and right["phase"] == "verify": expected = { f"route/verify/layer{layer}/{right_verifier_route}" @@ -134,6 +169,14 @@ def compare( result["comparisons"].append( {**identity, "group": group, "label": label, **difference} ) + if group == "state" and conv_width is not None: + reason = explain_state_difference( + label, left["states"], right["states"], conv_width + ) + if reason is not None: + result["explained_storage_differences"].append( + {**identity, "label": label, "reason": reason} + ) if left["rank"] == 0: if not all(torch.isfinite(d["native_logits"]).all() for d in (left, right)): raise ValueError(f"{name}: nonfinite native logits") @@ -143,7 +186,7 @@ def compare( "positions": left["positions"].tolist(), **tensor_difference(left["native_logits"], right["native_logits"]), **sampling_difference( - left["native_logits"], right["native_logits"] + left["native_logits"], right["native_logits"], eos_token_ids ), } ) @@ -171,6 +214,8 @@ def compare( result["summary"] = { "files_per_arm": len(left_files), "differing_intermediates": len(result["comparisons"]), + "unexplained_intermediates": len(result["comparisons"]) + - len(result["explained_storage_differences"]), "max_sampling_tv": max(max(row["sampling_tv"]) for row in result["logits"]), "support_changed_rows": sum( sum(row["support_changed"]) for row in result["logits"] @@ -179,6 +224,11 @@ def compare( "all_logits_bitwise_equal": all( row["bitwise_equal"] for row in result["logits"] ), + "max_eos_probability_abs_difference": max( + row["max_eos_probability_abs_difference"] for row in result["logits"] + ) + if eos_token_ids + else None, } return result @@ -189,10 +239,21 @@ def main() -> None: parser.add_argument("right", type=Path) parser.add_argument("--output", required=True, type=Path) parser.add_argument("--right-verifier-route", choices=("split", "packed")) + parser.add_argument("--eos-token-ids", type=int, nargs="+", default=[]) + parser.add_argument( + "--conv-width", + type=int, + choices=range(2, 7), + help="Frozen model convolution width; explain raw storage differences", + ) args = parser.parse_args() torch.set_num_threads(4) result = compare( - args.left, args.right, right_verifier_route=args.right_verifier_route + args.left, + args.right, + right_verifier_route=args.right_verifier_route, + eos_token_ids=tuple(args.eos_token_ids), + conv_width=args.conv_width, ) args.output.write_text(json.dumps(result, indent=2) + "\n") print(json.dumps(result["summary"], indent=2)) diff --git a/benchmarks/kernels/benchmark_sm70_grouped_attention_long.py b/benchmarks/kernels/benchmark_sm70_grouped_attention_long.py new file mode 100644 index 0000000000..2da32f76cd --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_grouped_attention_long.py @@ -0,0 +1,309 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Check compensated q8 schedules and time a sixteen-layer KV working set.""" + +import argparse +import hashlib +import importlib.util +import json +import statistics +from pathlib import Path + +import torch + + +def load_operator(manifest_path): + manifest = json.loads(manifest_path.read_text()) + path = Path(manifest["library"]) + assert hashlib.sha256(path.read_bytes()).hexdigest() == manifest["library_sha256"] + spec = importlib.util.spec_from_file_location(path.name.split(".")[0], path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + if Path(module.__file__).resolve() != path.resolve(): + raise RuntimeError( + f"Extension module alias: requested {path}, loaded {module.__file__}. " + "Build candidates with distinct native module names." + ) + entrypoint = manifest.get("entrypoint", "run") + if entrypoint not in ("run", "grouped_e4m3_fp32_paged_fwd"): + raise ValueError(f"Unsupported grouped attention entrypoint: {entrypoint}") + if ( + entrypoint == "grouped_e4m3_fp32_paged_fwd" + and int(module.grouped_e4m3_fp32_precision_version()) < 4 + ): + raise ValueError("The frozen native reference requires precision revision 4") + return getattr(module, entrypoint), manifest + + +def make_case(rows, page, length, num_arms, stride_padding=0, split_counts=None): + pages = (length + page - 1) // page + raw = torch.randn((pages, 2, page, 1, 256), device="cuda", dtype=torch.float16) + raw[:, 1].add_(torch.linspace(-4, 4, 256, device="cuda")) + encoded = raw.to(torch.float8_e4m3fn).view(torch.uint8) + backing = torch.empty( + (pages, 2, page, 1, 256 + stride_padding), device="cuda", dtype=torch.uint8 + ) + backing[..., :256].copy_(encoded) + k, v = backing[..., :256].unbind(1) + table = torch.randperm(pages, device="cuda").int()[None].contiguous() + q = torch.randn((rows, 6, 256), device="cuda", dtype=torch.float16) * 0.5 + lengths = torch.arange( + length - rows + 1, length + 1, device="cuda", dtype=torch.int32 + ) + arms = [] + split_counts = split_counts or [80] * num_arms + assert len(split_counts) == num_arms + for splits in split_counts: + guard = torch.full( + (rows + 2, 6, 256), -777.0, device="cuda", dtype=torch.float16 + ) + arms.append( + { + "guard": guard, + "out": guard[1:-1], + "partial": torch.full( + (splits, 8, 6, 256), -777.0, device="cuda", dtype=torch.float32 + ), + "lse": torch.full( + (splits, 8, 6, 2), -777.0, device="cuda", dtype=torch.float32 + ), + } + ) + return { + "q": q, + "k": k, + "v": v, + "table": table, + "lengths": lengths, + "initial_lengths": lengths.clone(), + "arms": arms, + } + + +def call(case, arm, operator): + buffers = case["arms"][arm] + operator( + case["q"], + case["k"], + case["v"], + buffers["out"], + case["table"], + case["lengths"], + buffers["partial"], + buffers["lse"], + 0.0625, + 0.5, + 1.25, + ) + + +def byte_equal(a, b): + return torch.equal( + a.contiguous().view(torch.uint8), b.contiguous().view(torch.uint8) + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--baseline", required=True, type=Path) + parser.add_argument("--candidate", action="append", type=Path, default=[]) + parser.add_argument("--output", required=True, type=Path) + parser.add_argument("--correctness-only", action="store_true") + parser.add_argument("--short", action="store_true") + parser.add_argument("--sanitizer", action="store_true") + parser.add_argument("--profile-one", action="store_true") + parser.add_argument("--performance-page", type=int, default=1648) + parser.add_argument( + "--performance-query-rows", type=int, choices=range(2, 9), default=8 + ) + parser.add_argument( + "--tail-queries", + action="store_true", + help="Audit each q2..q7 shape on the traced 3296-token page layout", + ) + parser.add_argument( + "--performance-contexts", + type=int, + nargs="+", + default=[1024, 32768, 65536, 131072], + ) + parser.add_argument( + "--extended-boundary", + type=int, + choices=(262144, 262152), + help="Additional operator-only 256K boundary; does not enable serving", + ) + args = parser.parse_args() + if not args.candidate and not args.profile_one: + parser.error("At least one --candidate is required for a comparison") + if args.performance_page <= 0: + parser.error("--performance-page must be positive") + if any(n < 8 or n > 262152 for n in args.performance_contexts): + parser.error("Performance contexts must be between 8 and 262152 tokens") + assert not args.output.exists(), args.output + assert torch.cuda.get_device_capability() == (7, 0) + torch.manual_seed(20260909) + loaded = [load_operator(p) for p in [args.baseline, *args.candidate]] + operators = [item[0] for item in loaded] + if len({id(operator) for operator in operators}) != len(operators): + raise RuntimeError("Candidate operators share a Python function binding") + report = { + "manifests": [item[1] for item in loaded], + "device": torch.cuda.get_device_name(), + "scope": "Operator screen, not model or complete-round admission", + "checks": [], + "performance": [], + "complete": False, + } + + def save(): + args.output.write_text(json.dumps(report, indent=2) + "\n") + + if args.profile_one: + case = make_case(8, 1648, 131072, len(operators)) + call(case, 0, operators[0]) + torch.cuda.synchronize() + report["profile_input"] = {"rows": 8, "page": 1648, "length": 131072} + report["complete"] = True + save() + return + + cases = [ + (8, 3296, 63, 0), + (8, 3296, 65, 0), + (8, 3296, 129, 0), + (8, 1648, 1649, 0), + (8, 3296, 3297, 0), + (2, 848, 8197, 8), + (8, 1648, 3297, 8), + ] + if not args.short: + cases += [ + (5, 1648, 32768, 0), + (8, 1648, 65536, 0), + (8, 3296, 131072, 0), + (8, 3296, 132096, 0), # 128K prompt plus bounded generation headroom. + ] + if args.sanitizer: + cases = [(8, 1648, 1649, 0), (8, 1648, 3297, 8), (8, 3296, 6593, 0)] + if args.extended_boundary is not None: + cases += [ + (8, 3296, args.extended_boundary, 0), + (8, 1648, args.extended_boundary, 8), + ] + if args.tail_queries: + tail_boundary = args.extended_boundary or 262144 + cases = [(q, 3296, 3297, 8) for q in range(2, 8)] + if args.sanitizer: + cases.append((6, 3296, tail_boundary, 0)) + else: + cases += [(q, 3296, tail_boundary, 0) for q in range(2, 8)] + for rows, page, length, padding in cases: + case = make_case(rows, page, length, len(operators), padding) + graphs = [] + for arm, operator in enumerate(operators): + call(case, arm, operator) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + call(case, arm, operator) + graphs.append(graph) + for state in ("live", "all_zero", "tail_zero", "first_zero", "restore"): + case["lengths"].copy_(case["initial_lengths"]) + if state == "all_zero": + case["lengths"].zero_() + elif state == "tail_zero": + case["lengths"][-1] = 0 + elif state == "first_zero": + case["lengths"][0] = 0 + for buffers, graph in zip(case["arms"], graphs): + buffers["partial"].fill_(-777.0) + buffers["lse"].fill_(-777.0) + graph.replay() + for arm in range(1, len(operators)): + checks = { + key: byte_equal(case["arms"][0][key], case["arms"][arm][key]) + for key in ("out", "partial", "lse") + } + guards = all( + bool((buffers["guard"][[0, -1]] == -777).all()) + for buffers in case["arms"] + ) + row = { + "rows": rows, + "page": page, + "length": length, + "stride_padding": padding, + "state": state, + "arm": arm, + "byte_exact": checks, + "guards": guards, + } + report["checks"].append(row) + if not all(checks.values()) or not guards: + save() + raise AssertionError(row) + print(json.dumps({"exact_shape": [rows, page, length, padding]}), flush=True) + del case, graphs + save() + + if not args.correctness_only: + for length in args.performance_contexts: + # Sixteen distinct KV allocations represent the target's real + # layer working set and prevent single-layer hot-cache claims. + cases = [ + make_case( + args.performance_query_rows, + args.performance_page, + length, + len(operators), + ) + for _ in range(16) + ] + graphs = [] + for arm, operator in enumerate(operators): + for case in cases: + call(case, arm, operator) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for case in cases: + call(case, arm, operator) + graphs.append(graph) + for _ in range(5): + for graph in graphs: + graph.replay() + samples = [[] for _ in operators] + for trial in range(5): + order = list(range(len(operators))) + if trial % 2: + order.reverse() + for arm in order: + start, end = [ + torch.cuda.Event(enable_timing=True) for _ in range(2) + ] + start.record() + for _ in range(8): + graphs[arm].replay() + end.record() + end.synchronize() + samples[arm].append(start.elapsed_time(end) / 8) + row = { + "context": length, + "page": args.performance_page, + "query_rows": args.performance_query_rows, + "layers": 16, + "samples_ms": samples, + "median_ms": [statistics.median(values) for values in samples], + "complete_round_performance": False, + } + report["performance"].append(row) + save() + print(json.dumps(row), flush=True) + del cases, graphs + report["complete"] = True + save() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_sm70_grouped_attention_precision.py b/benchmarks/kernels/benchmark_sm70_grouped_attention_precision.py new file mode 100644 index 0000000000..957da81f67 --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_grouped_attention_precision.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Reject arithmetic candidates before timing when FP64-reference errors grow. + +Independent native FP32-output builds expose the actual final accumulator. +Their partial workspaces and converted output must match the original builds; +an inaccurate or aliased diagnostic cannot authorize an arithmetic change. +This operator screen never provides model, recursion or acceptance admission. +""" + +import argparse +import json +from pathlib import Path + +import torch + +from benchmarks.kernels.benchmark_sm70_grouped_attention_long import ( + byte_equal, + call, + load_operator, + make_case, +) +from benchmarks.kernels.benchmark_sm70_grouped_attention_splits import ( + error, + fp64_reference, +) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + for name in ("baseline", "candidate", "baseline-precast", "candidate-precast"): + parser.add_argument("--" + name, required=True, type=Path) + parser.add_argument("--output", required=True, type=Path) + args = parser.parse_args() + assert not args.output.exists() + assert torch.cuda.get_device_capability() == (7, 0) + paths = [ + args.baseline, + args.candidate, + args.baseline_precast, + args.candidate_precast, + ] + operators, manifests = zip(*(load_operator(p) for p in paths)) + assert len({id(op) for op in operators}) == 4 + assert all(m.get("splits", 80) == 80 for m in manifests) + assert all(m["diagnostic_output_fp32"] for m in manifests[2:]) + report = { + "manifests": manifests, + "reference": "Independent FP64 QK, softmax and PV on decoded E4M3 operands", + "scope": "Operator reference screen; no model or arithmetic admission", + "arithmetic_admitted": False, + "complete": False, + "checks": [], + } + + def save() -> None: + args.output.write_text(json.dumps(report, indent=2) + "\n") + + for seed in (20260909, 20260910): + torch.manual_seed(seed) + for length in (129, 3297, 32768, 131072, 262144): + case = make_case(8, 3296, length, 4, stride_padding=8 if seed % 2 else 0) + if seed % 2: + case["lengths"][-1] = 0 + for arm in (2, 3): + guard = torch.full( + (10, 6, 256), -777.0, dtype=torch.float32, device="cuda" + ) + case["arms"][arm]["guard"] = guard + case["arms"][arm]["out"] = guard[1:-1] + reference = fp64_reference(case, length) + for arm, operator in enumerate(operators): + call(case, arm, operator) + assert torch.isfinite(case["arms"][arm]["out"]).all() + assert (case["arms"][arm]["guard"][[0, -1]] == -777).all() + for ordinary, precast in ((0, 2), (1, 3)): + left, right = case["arms"][ordinary], case["arms"][precast] + for key in ("partial", "lse"): + assert byte_equal(left[key], right[key]), ( + "Invalid pre-cast diagnostic", + seed, + length, + ordinary, + key, + ) + assert byte_equal(left["out"], right["out"].half()) + metrics = [error(case["arms"][i]["out"], reference) for i in range(4)] + nonexpansion = { + name: all(metrics[b][key] <= metrics[a][key] for key in metrics[a]) + for name, a, b in (("fp16_output", 0, 1), ("precast_fp32", 2, 3)) + } + row = { + "seed": seed, + "length": length, + "padding_row": bool(seed % 2), + "stride_padding": 8 if seed % 2 else 0, + "precast_diagnostics_validated": True, + "errors": dict( + zip(("baseline", "candidate", "base32", "new32"), metrics) + ), + "reference_error_nonexpansion": nonexpansion, + "byte_equal": { + key: byte_equal(case["arms"][0][key], case["arms"][1][key]) + for key in ("out", "partial", "lse") + }, + "precast_byte_equal": byte_equal( + case["arms"][2]["out"], case["arms"][3]["out"] + ), + } + report["checks"].append(row) + save() + print(json.dumps(row), flush=True) + del case, reference + report["all_reference_errors_nonexpanding"] = all( + all(r["reference_error_nonexpansion"].values()) for r in report["checks"] + ) + report["complete"] = True + save() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_sm70_grouped_attention_splits.py b/benchmarks/kernels/benchmark_sm70_grouped_attention_splits.py new file mode 100644 index 0000000000..29caddfd5f --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_grouped_attention_splits.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Screen split schedules against FP64 and a sixteen-layer KV working set. + +This records final FP16-output error only. It cannot admit arithmetic changes: +native pre-cast FP32 error, sanitizers and model quality remain separate gates. +""" + +import argparse +import json +import statistics +from pathlib import Path + +import torch + +from benchmarks.kernels.benchmark_sm70_grouped_attention_long import ( + call, + load_operator, + make_case, +) + + +def fp64_reference(case, length): + order = case["table"][0].long() + key = case["k"].index_select(0, order).reshape(-1, 256)[:length] + value = case["v"].index_select(0, order).reshape(-1, 256)[:length] + key = key.view(torch.float8_e4m3fn).double() + value = value.view(torch.float8_e4m3fn).double() + query = case["q"].reshape(-1, 256).double() + scores = (query @ key.T) * (0.0625 * 0.5) + lengths = case["lengths"].repeat_interleave(6) + visible = torch.arange(length, device="cuda")[None, :] < lengths[:, None] + scores.masked_fill_(~visible, -torch.inf) + probabilities = scores.softmax(-1) + probabilities.masked_fill_(lengths[:, None] == 0, 0) + return ((probabilities @ value) * 1.25).reshape_as(case["q"]) + + +def error(output, reference): + difference = output.double() - reference + absolute = difference.abs().reshape(-1) + return { + "max_abs": absolute.max().item(), + "p99_abs": torch.quantile(absolute, 0.99).item(), + "relative_l2": (difference.norm() / reference.norm().clamp_min(1e-30)).item(), + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--baseline", type=Path, required=True) + parser.add_argument("--candidate", type=Path, action="append", required=True) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + assert not args.output.exists() + assert torch.cuda.get_device_capability() == (7, 0) + torch.manual_seed(20260910) + loaded = [load_operator(p) for p in [args.baseline, *args.candidate]] + operators, manifests = zip(*loaded) + assert len({id(op) for op in operators}) == len(operators) + splits = [m.get("splits", 80) for m in manifests] + report = { + "manifests": manifests, + "splits": splits, + "reference": "Independent FP64 QK, softmax and PV; decoded E4M3 operands", + "error_scope": "Final FP16 output; native pre-cast FP32 not yet measured", + "arithmetic_admitted": False, + "complete_round_performance": False, + "checks": [], + "performance": [], + "complete": False, + } + + def save(): + args.output.write_text(json.dumps(report, indent=2) + "\n") + + for length in (129, 3297, 32768, 131072): + case = make_case(8, 3296, length, len(operators), split_counts=splits) + case["lengths"][-1] = 0 + reference = fp64_reference(case, length) + for arm, operator in enumerate(operators): + call(case, arm, operator) + output = case["arms"][arm]["out"] + assert torch.isfinite(output).all() + assert (output[-1] == 0).all() + assert (case["arms"][arm]["guard"][[0, -1]] == -777).all() + report["checks"].append( + {"length": length, "splits": splits[arm], **error(output, reference)} + ) + save() + print(json.dumps({"reference_length": length}), flush=True) + del case, reference + + for length in (1024, 32768, 65536, 131072): + cases = [ + make_case(8, 3296, length, len(operators), split_counts=splits) + for _ in range(16) + ] + graphs = [] + for arm, operator in enumerate(operators): + for case in cases: + call(case, arm, operator) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for case in cases: + call(case, arm, operator) + graphs.append(graph) + for _ in range(5): + for graph in graphs: + graph.replay() + samples = [[] for _ in operators] + for trial in range(5): + order = list(range(len(operators))) + if trial % 2: + order.reverse() + for arm in order: + start, end = [torch.cuda.Event(enable_timing=True) for _ in range(2)] + start.record() + for _ in range(8): + graphs[arm].replay() + end.record() + end.synchronize() + samples[arm].append(start.elapsed_time(end) / 8) + row = { + "context": length, + "layers": 16, + "page": 3296, + "samples_ms": samples, + "median_ms": [statistics.median(x) for x in samples], + } + report["performance"].append(row) + save() + print(json.dumps(row), flush=True) + del cases, graphs + report["complete"] = True + save() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_sm70_scalar_attention.py b/benchmarks/kernels/benchmark_sm70_scalar_attention.py new file mode 100644 index 0000000000..c1c96d8d67 --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_scalar_attention.py @@ -0,0 +1,182 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compare scalar q1 native workspaces and a sixteen-layer KV working set.""" + +import argparse +import hashlib +import importlib.util +import json +import statistics +from pathlib import Path + +import torch + +from benchmarks.kernels.benchmark_sm70_grouped_attention_long import ( + byte_equal, + load_operator, +) + +parser = argparse.ArgumentParser(description=__doc__) +parser.add_argument("--reference-library", type=Path, required=True) +parser.add_argument("--reference-sha256", required=True) +parser.add_argument("--candidate", type=Path, action="append", required=True) +parser.add_argument("--output", type=Path, required=True) +parser.add_argument("--correctness-only", action="store_true") +parser.add_argument("--sanitizer", action="store_true") +args = parser.parse_args() +library = args.reference_library.resolve() +assert hashlib.sha256(library.read_bytes()).hexdigest() == args.reference_sha256 +assert torch.cuda.get_device_capability() == (7, 0) +spec = importlib.util.spec_from_file_location("flash_attn_v100_cuda", library) +native = importlib.util.module_from_spec(spec) +spec.loader.exec_module(native) +assert Path(native.__file__).resolve() == library +loaded = [load_operator(p) for p in args.candidate] +operators = [None] + [r[0] for r in loaded] +report = dict( + scope="Scalar q1 operator screen; no full-model admission", + frozen_library=dict(library=str(library), sha256=args.reference_sha256), + manifests=[r[1] for r in loaded], + checks=[], + performance=[], + complete=False, +) +out = args.output +assert not out.exists() + + +def make_case(length, padding=0): + page = 3296 + pages = (length + page - 1) // page + raw = torch.randn((pages, 2, page, 1, 256), dtype=torch.float16, device="cuda") + backing = torch.empty( + (pages, 2, page, 1, 256 + padding), dtype=torch.uint8, device="cuda" + ) + backing[..., :256].copy_(raw.to(torch.float8_e4m3fn).view(torch.uint8)) + k, v = backing[..., :256].unbind(1) + q = torch.randn((1, 6, 256), dtype=torch.float16, device="cuda") * 0.5 + table = torch.randperm(pages, device="cuda").int()[None].contiguous() + lengths = torch.tensor([length], dtype=torch.int32, device="cuda") + active = torch.full((1,), 256, dtype=torch.int32, device="cuda") + arms = [] + for _ in operators: + guard = torch.full((3, 6, 256), -777.0, dtype=torch.float16, device="cuda") + arms.append( + dict( + guard=guard, + out=guard[1:2], + partial=torch.full((1, 6, 256, 256), -777.0, device="cuda"), + maximum=torch.full((1, 6, 256), -777.0, device="cuda"), + sums=torch.full((1, 6, 256), -777.0, device="cuda"), + ) + ) + return dict(q=q, k=k, v=v, table=table, lengths=lengths, active=active, arms=arms) + + +def run(case, arm): + r = case["arms"][arm] + args = [ + case["q"], + case["k"], + case["v"], + r["out"], + case["table"], + case["lengths"], + r["partial"], + r["maximum"], + r["sums"], + case["active"], + ] + if arm == 0: + native.decode_paged_fwd( + *args, 0.0625, 1024, 256, "fp8_e4m3", 0.5, 1.25, -1, -1, None, 0 + ) + else: + operators[arm](*args, 0.0625, 0.5, 1.25) + + +torch.manual_seed(20260910) +cases_to_check = [(1024, 0), (3295, 0), (3297, 8), (131072, 0), (262144, 8)] +if args.sanitizer: + cases_to_check = [(3297, 8), (262144, 0)] +for length, padding in cases_to_check: + case = make_case(length, padding) + graphs = [] + for arm in range(len(operators)): + run(case, arm) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run(case, arm) + graphs.append(graph) + for state in [length, 0, length]: + case["lengths"].fill_(state) + for graph in graphs: + graph.replay() + torch.cuda.synchronize() + for arm in range(1, len(operators)): + exact = { + key: byte_equal(case["arms"][0][key], case["arms"][arm][key]) + for key in ["out", "partial", "maximum", "sums"] + } + guards = bool((case["arms"][arm]["guard"][[0, 2]] == -777).all()) + row = dict( + length=length, + padding=padding, + state=state, + arm=arm, + byte_exact=exact, + guards=guards, + ) + report["checks"].append(row) + out.write_text(json.dumps(report, indent=2) + "\n") + assert all(exact.values()) and guards, row + print("SCALAR_Q1_EXACT", length, padding, flush=True) + del graphs, case +for length in ( + [] + if args.correctness_only or args.sanitizer + else [1024, 32768, 65536, 131072, 261888] +): + cases = [make_case(length) for _ in range(16)] + graphs = [] + for arm in range(len(operators)): + for case in cases: + run(case, arm) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for case in cases: + run(case, arm) + graphs.append(graph) + for _ in range(3): + for graph in graphs: + graph.replay() + torch.cuda.synchronize() + samples = [[] for _ in operators] + for repeat in range(5): + for arm in ( + range(len(operators)) + if repeat % 2 == 0 + else reversed(range(len(operators))) + ): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(4): + graphs[arm].replay() + end.record() + end.synchronize() + samples[arm].append(start.elapsed_time(end) / 4) + row = dict( + length=length, + layers=16, + samples_ms=samples, + median_ms=[statistics.median(x) for x in samples], + ) + report["performance"].append(row) + out.write_text(json.dumps(report, indent=2) + "\n") + print("SCALAR_Q1_PERFORMANCE", row, flush=True) + del graphs, cases +report["complete"] = True +out.write_text(json.dumps(report, indent=2) + "\n") diff --git a/benchmarks/kernels/build_sm70_grouped_attention_candidate.py b/benchmarks/kernels/build_sm70_grouped_attention_candidate.py index 42dda02c61..19799f903d 100644 --- a/benchmarks/kernels/build_sm70_grouped_attention_candidate.py +++ b/benchmarks/kernels/build_sm70_grouped_attention_candidate.py @@ -3,8 +3,8 @@ """Build a private E4M3 grouped-attention scheduling candidate. The original per-head arithmetic, FP32 numerator/max/sum workspace and -native input validation remain intact. Only the number of heads per CTA -changes. This builder installs no serving route or default. +native input validation remain intact. Candidates change CTA grouping, +loop scheduling or address-equivalent loads. The builder installs no route. """ import argparse @@ -20,25 +20,741 @@ def replace_once(source: str, old: str, new: str) -> str: return source.replace(old, new) +def prefetch_values(source: str) -> str: + """Fill a disjoint V panel with idle QK warps before the existing barrier.""" + start = source.index("__launch_bounds__(kGroupedVerifyThreads, 1) void ") + end = source.index( + "void flash_attention_grouped_verify_e5m2_combine_kernel(", start + ) + partial = source[start:end] + value_argument = partial.index(" shared_kv, v_cache, page_ids,") + load_start = partial.rfind(" load_xqa_tc_kv_panel<", 0, value_argument) + barrier = " __syncthreads();" + load_end = partial.index(barrier, value_argument) + len(barrier) + load = partial[load_start:load_end] + load = load[: load.rfind(barrier)] + load = load.replace("shared_kv", "shared_values") + load = load.replace("kGroupedVerifyThreads", "kValueLoadThreads") + load = load.replace("idx = tid +", "idx = value_load_tid +") + load = replace_once( + load, + "v_block_stride, v_token_stride, v_head_stride, 0);", + "v_block_stride, v_token_stride, v_head_stride, 0, value_load_tid);", + ) + partial = partial[:load_start] + partial[load_end:] + marker = " constexpr int kResidualStride = kGroupedVerifyProbStride;" + partial = replace_once( + partial, + marker, + marker + "\n __half* shared_values = shared_prob_residual + " + "kGroupedVerifyRows * kResidualStride;", + ) + qk = ( + " grouped_verify_qk(shared_q, shared_kv, shared_scores,\n" + " qk_scale, active_m_tiles);" + ) + partial = replace_once( + partial, + qk, + qk + "\n if (warp_id >= kGroupedVerifyQKWarps) {\n" + " constexpr int kValueLoadThreads = kGroupedVerifyThreads - " + "kGroupedVerifyQKWarps * kWarpSize;\n" + " const int value_load_tid = tid - kGroupedVerifyQKWarps * kWarpSize;\n" + + load + + " }\n", + ) + partial = replace_once( + partial, + "shared_kv + k_offset * kGroupedVerifyKVStride + d_tile * 16,", + "shared_values + k_offset * kGroupedVerifyKVStride + d_tile * 16,", + ) + source = source[:start] + partial + source[end:] + source = replace_once( + source, + "kGroupedVerifyRows * kGroupedVerifyProbStride * sizeof(__half);", + "kGroupedVerifyRows * kGroupedVerifyProbStride * sizeof(__half) +\n" + " kGroupedVerifyBlockN * kGroupedVerifyKVStride * sizeof(__half);", + ) + source = replace_once( + source, + "static_assert(kCompensatedSmemBytes <= 64 * 1024,", + "static_assert(kCompensatedSmemBytes <= 96 * 1024,", + ) + source = replace_once( + source, + '"compensated P must fit the SM70 shared-memory budget");', + '"compensated P and prefetched V must fit the SM70 budget");\n' + " TORCH_CHECK(properties->sharedMemPerBlockOptin >= kCompensatedSmemBytes,\n" + ' "V prefetch exceeds device opt-in shared memory");', + ) + return source + + +def prefetch_keys(source: str, key_load_warps: int) -> str: + """Overlap the next K panel with softmax using dedicated load warps.""" + start = source.index("__launch_bounds__(kGroupedVerifyThreads, 1) void ") + end = source.index( + "void flash_attention_grouped_verify_e5m2_combine_kernel(", start + ) + partial = source[start:end] + loop_start = partial.rindex(" for (int tile_start = split_start;") + loop = partial[loop_start:] + barrier = " __syncthreads();" + load_start = loop.index(" load_xqa_tc_kv_panel<") + load_end = loop.index(barrier, load_start) + len(barrier) + load = loop[load_start:load_end] + # The first panel is loaded by the whole CTA. Subsequent panels are ready + # at the previous tile's softmax barrier, before PV consumes disjoint V. + loop = ( + loop[:load_start] + + " if (tile_start == split_start) {\n" + + load + + "\n }\n" + + loop[load_end:] + ) + load = load[: load.rfind(barrier)] + load = load.replace("kGroupedVerifyThreads", "kKeyLoadThreads") + load = load.replace("idx = tid +", "idx = key_load_tid +") + load = load.replace("valid_k_rows", "next_k_rows") + load = load.replace("tile_page_offset", "next_page_offset") + load = replace_once( + load, + "k_block_stride, k_token_stride, k_head_stride, 0);", + "k_block_stride, k_token_stride, k_head_stride, 0, key_load_tid);", + ) + softmax_start = loop.index( + "#pragma unroll\n for (int row = warp_id; row < kGroupedVerifyRows;" + ) + softmax_end = loop.index(" __syncthreads();", softmax_start) + softmax = loop[softmax_start:softmax_end] + softmax = replace_once( + softmax, "row += kGroupedVerifyWarps)", "row += kSoftmaxWarps)" + ) + replacement = ( + f" constexpr int kSoftmaxWarps = {16 - key_load_warps};\n" + " if (warp_id < kSoftmaxWarps) {\n" + + softmax + + " } else if (tile_start + kGroupedVerifyBlockN < split_end) {\n" + " constexpr int kKeyLoadThreads = " + "kGroupedVerifyThreads - kSoftmaxWarps * kWarpSize;\n" + " const int key_load_tid = tid - kSoftmaxWarps * kWarpSize;\n" + " const int next_k_rows = min(kGroupedVerifyBlockN, " + "split_end - tile_start - kGroupedVerifyBlockN);\n" + " const int next_page_offset = tile_page_offset + " + "kGroupedVerifyBlockN;\n" + load + " }\n" + ) + loop = loop[:softmax_start] + replacement + loop[softmax_end:] + partial = partial[:loop_start] + loop + return source[:start] + partial + source[end:] + + +def reuse_pv_values(source: str) -> str: + """Interchange independent M tiles to reuse raw/scaled V fragments. + + Each accumulator still receives main0, residual0, main16, residual16 in that + order, followed by one N32 online-state update. The six-head/16-warp layout + assigns the same D tile to a warp for all three M tiles. + """ + start = source.index("__launch_bounds__(kGroupedVerifyThreads, 1) void ") + end = source.index( + "void flash_attention_grouped_verify_e5m2_combine_kernel(", start + ) + partial = source[start:end] + begin = partial.index( + "#pragma unroll\n for (int fragment_idx = 0; " + "fragment_idx < kGroupedVerifyOutputTilesPerWarp;" + ) + finish = partial.index(" __syncthreads();\n }", begin) + replacement = r""" static_assert(COMPENSATE_P && kGroupedVerifyWarps == 16, + "PV reuse is isolated to six-head compensated E4M3"); + volta::fragment + tile_fragments[kGroupedVerifyOutputTilesPerWarp]; +#pragma unroll + for (int i = 0; i < kGroupedVerifyOutputTilesPerWarp; ++i) + volta::fill_fragment(tile_fragments[i], 0.0f); + const int d_tile = warp_id; +#pragma unroll + for (int k_offset = 0; k_offset < kGroupedVerifyBlockN; k_offset += 16) { + volta::fragment + value_fragment; + volta::load_matrix_sync( + value_fragment, + shared_values + k_offset * kGroupedVerifyKVStride + d_tile * 16, + kGroupedVerifyKVStride); + auto residual_value_fragment = value_fragment; +#pragma unroll + for (int i = 0; i < value_fragment.num_elements / 2; ++i) { + union { + uint32_t bits; + __half2 pair; + } packed_value; + packed_value.bits = value_fragment.x[i]; + packed_value.pair = + __hmul2(packed_value.pair, __float2half2_rn(1.0f / 2048.0f)); + residual_value_fragment.x[i] = packed_value.bits; + } +#pragma unroll + for (int fragment_idx = 0; + fragment_idx < kGroupedVerifyOutputTilesPerWarp; ++fragment_idx) { + const int m_tile = fragment_idx; + if ((active_m_tiles & (1 << m_tile)) == 0) continue; + volta::fragment + probability_fragment; + volta::load_matrix_sync( + probability_fragment, + shared_probs + m_tile * 16 * kGroupedVerifyProbStride + k_offset, + kGroupedVerifyProbStride); + volta::mma_sync(tile_fragments[fragment_idx], probability_fragment, + value_fragment, tile_fragments[fragment_idx]); + volta::load_matrix_sync( + probability_fragment, + shared_prob_residual + m_tile * 16 * kResidualStride + k_offset, + kResidualStride); + volta::mma_sync(tile_fragments[fragment_idx], probability_fragment, + residual_value_fragment, tile_fragments[fragment_idx]); + } + } +#pragma unroll + for (int fragment_idx = 0; + fragment_idx < kGroupedVerifyOutputTilesPerWarp; ++fragment_idx) { + const int m_tile = fragment_idx; + if ((active_m_tiles & (1 << m_tile)) == 0) continue; + grouped_verify_add_output_tile(output_fragments[fragment_idx], + tile_fragments[fragment_idx], + smem.row_scale, m_tile * 16); + } +""" + partial = partial[:begin] + replacement + partial[finish:] + return source[:start] + partial + source[end:] + + +def pair_qk_products(source: str) -> str: + """Produce two independent K16 products before their ordered corrections. + + The dot products retain their own zero-initialized accumulators. Correction + consumes the first product and then the second, in the original K16 order. + This isolates instruction scheduling from a change in reduction arithmetic. + """ + start = source.index("__device__ __forceinline__ void grouped_verify_qk(") + end = source.index( + "__device__ __forceinline__ void grouped_verify_scale_output_fragment(", + start, + ) + qk = source[start:end] + qk = replace_once( + qk, + "k_offset < kGroupedVerifyHeadDim; k_offset += 16)", + "k_offset < kGroupedVerifyHeadDim; k_offset += (COMPENSATE ? 32 : 16))", + ) + product = ( + " volta::mma_sync(tile_fragment, q_fragment, k_fragment, tile_fragment);" + ) + qk = replace_once( + qk, + product, + product + + r""" + volta::fragment next_tile_fragment; + volta::fill_fragment(next_tile_fragment, 0.0f); + volta::load_matrix_sync( + q_fragment, + shared_q + m_tile * 16 * kGroupedVerifyQStride + k_offset + 16, + kGroupedVerifyQStride); + volta::load_matrix_sync( + k_fragment, + shared_k + n_tile * 16 * kGroupedVerifyKVStride + k_offset + 16, + kGroupedVerifyKVStride); + volta::mma_sync(next_tile_fragment, q_fragment, k_fragment, + next_tile_fragment); +""", + ) + correction = r"""#pragma unroll + for (int i = 0; i < score_fragment.num_elements; ++i) { + const float y = __fsub_rn(tile_fragment.x[i], correction[i]); + const float sum = __fadd_rn(score_fragment.x[i], y); + correction[i] = __fsub_rn(__fsub_rn(sum, score_fragment.x[i]), y); + score_fragment.x[i] = sum; + }""" + qk = replace_once( + qk, + correction, + correction + + "\n" + + correction.replace("tile_fragment.x", "next_tile_fragment.x"), + ) + return source[:start] + qk + source[end:] + + +def retain_fp32_output(source: str) -> str: + """Expose the final FP32 accumulator for an independent numerical audit.""" + symbol = source.index("void flash_attention_grouped_verify_e5m2_combine_kernel(") + start = source.rfind("template <", 0, symbol) + end = source.index("\ntemplate <", symbol) + combine = source[start:end] + for old, new in ( + ("__half* __restrict__ out,", "float* __restrict__ out,"), + ("__float2half_rn(0.0f)", "0.0f"), + ("__float2half_rn(accumulator)", "accumulator"), + ): + combine = replace_once(combine, old, new) + source = source[:start] + combine + source[end:] + start = source.index("at::Tensor flash_attention_grouped_e4m3_fp32_paged(") + host = source[start:] + for old, new in ( + ("out.scalar_type() == at::kHalf", "out.scalar_type() == at::kFloat"), + ("output must be contiguous FP16", "audit output must be contiguous FP32"), + ( + "reinterpret_cast<__half*>(out.data_ptr())", + "reinterpret_cast(out.data_ptr())", + ), + ): + host = replace_once(host, old, new) + return source[:start] + host + + +def accumulate_qk_fp64(source: str) -> str: + """Arithmetic experiment: sum the unchanged K16 products with FP64 adds. + + This is not an exact scheduling optimization. Even if a sampled output + agrees, independent pre-cast reference and model audits remain mandatory. + """ + start = source.index("template ") + end = source.index("void grouped_verify_scale_output_fragment(", start) + qk = source[start:end] + qk = replace_once(qk, "float correction[8] = {};", "double wide_sum[8] = {};") + qk = replace_once( + qk, + """ // E4M3 x FP16 products fit comfortably in FP32, but a D256 Tensor + // Core accumulation can still lose low bits. Sum short K16 products + // with compensated FP32 additions; explicit RN operations preserve + // the correction under the standard fast-math build.""", + """ // Arithmetic probe: retain each original FP32 K16 Tensor Core + // product, but sum those products in FP64 before the final FP32 cast. + // This needs an independent reference audit; it is not bit-exact by + // construction and is not admitted by an operator timing result.""", + ) + qk = replace_once( + qk, + """ const float y = __fsub_rn(tile_fragment.x[i], correction[i]); + const float sum = __fadd_rn(score_fragment.x[i], y); + correction[i] = __fsub_rn(__fsub_rn(sum, score_fragment.x[i]), y); + score_fragment.x[i] = sum;""", + """ wide_sum[i] = __dadd_rn( + wide_sum[i], static_cast(tile_fragment.x[i]));""", + ) + qk = replace_once( + qk, + " score_fragment.x[i] *= qk_scale;", + """ if constexpr (COMPENSATE) { + score_fragment.x[i] = __double2float_rn(wide_sum[i]) * qk_scale; + } else { + score_fragment.x[i] *= qk_scale; + }""", + ) + return source[:start] + qk + source[end:] + + +def register_softmax_state(partial: str) -> str: + """Keep each warp's three online rows private until the final publication. + + Every lane consumes the same broadcast tile maximum/sum and runs the + original N32 update. Only row_scale is shared with PV between tiles. + """ + marker = " // Recompute QK for the conservative path" + partial = replace_once( + partial, + marker, + " static_assert(!TWO_PASS && kGroupedVerifyWarps == 16 &&\n" + ' kGroupedVerifyRows == 48, "Three fixed online rows per warp");\n' + " float online_max[3] = {kXQANegInf, kXQANegInf, kXQANegInf};\n" + " float online_sum[3] = {};\n" + marker, + ) + begin = partial.index(marker) + end = partial.index(" // The compute buffers are dead.", begin) + loop = partial[begin:end] + loop = replace_once( + loop, + "const float old_max = smem.row_max[row];", + "const float old_max = online_max[row / kGroupedVerifyWarps];", + ) + loop = replace_once( + loop, + " // Finish every lane's shared-state reads before lane 0 " + "overwrites the\n" + """ // online maximum. Shuffle synchronization does not order memory. + __syncwarp(); + if (lane_id == 0) { + if (tile_sum > 0.0f) { + smem.row_sum[row] = smem.row_sum[row] * exp_diff + tile_sum; + smem.row_max[row] = new_max; + } + smem.row_scale[row] = exp_diff; + }""", + " // Each lane owns an identical copy; no shared maximum " + "is overwritten.\n" + """ const int local_row = row / kGroupedVerifyWarps; + if (tile_sum > 0.0f) { + online_sum[local_row] = online_sum[local_row] * exp_diff + tile_sum; + online_max[local_row] = new_max; + } + if (lane_id == 0) smem.row_scale[row] = exp_diff;""", + ) + return ( + partial[:begin] + + loop + + """ // Publish final row statistics before the existing output barrier. + if (lane_id == 0) { +#pragma unroll + for (int i = 0; i < 3; ++i) { + smem.row_max[warp_id + i * kGroupedVerifyWarps] = online_max[i]; + smem.row_sum[warp_id + i * kGroupedVerifyWarps] = online_sum[i]; + } + } + +""" + + partial[end:] + ) + + +def specialize_full_q8( + source: str, visible_tiles: bool, register_state: bool = False +) -> str: + """Remove variable-Q branches only when all eight query rows exist. + + Per-row GPU lengths remain authoritative. The optional all-visible branch + requires the complete N32 tile to precede the minimum of all eight lengths; + padding, rejected rows and the causal tail use the original visibility code. + """ + old_name = "flash_attention_grouped_verify_e5m2_partial_kernel" + new_name = "flash_attention_grouped_verify_e4m3_full_q8_kernel" + start = source.index( + "template ;", + " static_assert(MAX_QUERY_TOKENS == 8 && COMPENSATE_P &&\n" + ' ROW_SEQLENS && !SPARSE_PAGE4, "Full q8 compensated contract");\n' + " if (runtime_query_len != 8) return;\n" + " constexpr int query_len = 8;\n" + " using Traits = GroupedVerifyTraits;", + ) + if register_state: + partial = register_softmax_state(partial) + if visible_tiles: + loop_start = partial.index(" // Recompute QK for the conservative path") + partial = ( + partial[:loop_start] + " int minimum_visible_length = row_lengths[0];\n" + "#pragma unroll\n" + " for (int i = 1; i < 8; ++i)\n" + " minimum_visible_length =\n" + " min(minimum_visible_length, row_lengths[i]);\n" + + partial[loop_start:] + ) + begin = partial.index( + "#pragma unroll\n for (int row = warp_id;", loop_start + ) + finish = partial.index(" __syncthreads();", begin) + original = partial[begin:finish] + visible = original + a = visible.index(" const bool visible =") + b = visible.index(" const float score =", a) + visible = visible[:a] + " constexpr bool visible = true;\n" + visible[b:] + partial = ( + partial[:begin] + " if (tile_start + kGroupedVerifyBlockN <=\n" + " minimum_visible_length) {\n" + + visible + + " } else {\n" + + original + + " }\n" + + partial[finish:] + ) + source = source[:end] + partial + source[end:] + host_start = source.index(" auto kernel =", source.index("at::Tensor ")) + host_end = source.index(" constexpr int kCompensatedSmemBytes", host_start) + selection = source[host_start:host_end].replace("auto kernel =", "kernel =", 1) + selection = selection.replace(old_name, new_name) + return ( + source[:host_end] + + " if (q.size(0) == 8) {\n" + + selection + + " }\n" + + source[host_end:] + ) + + +def qk_head_rows(source: str) -> str: + """Audit M8/N32 per head; WMMA shape equality is not assumed. + + Q is staged as six groups of eight rows. Scores are stored back to the + parent's token/head order, so softmax, PV and merge retain their layout. + """ + start = source.index("template ") + end = source.index( + "__device__ __forceinline__ void grouped_verify_scale_output_fragment(", + start, + ) + qk = source[start:end].replace("grouped_verify_qk(", "grouped_verify_qk_head(") + qk = qk.replace("16, 16, 16", "8, 32, 16") + a = qk.index(" const int m_tile =") + b = qk.index(" volta::fragment<", a) + qk = qk[:a] + " const int head = warp_id;\n" + qk[b:] + qk = replace_once( + qk, + "shared_q + m_tile * 16 * kGroupedVerifyQStride + k_offset", + "shared_q + head * 8 * kGroupedVerifyQStride + k_offset", + ) + qk = replace_once( + qk, + "shared_k + n_tile * 16 * kGroupedVerifyKVStride + k_offset", + "shared_k + k_offset", + ) + qk = replace_once( + qk, + """ shared_scores + m_tile * 16 * kGroupedVerifyScoreStride + n_tile * 16, + score_fragment, kGroupedVerifyScoreStride, volta::mem_row_major);""", + """ shared_scores + head * kGroupedVerifyScoreStride, + score_fragment, 6 * kGroupedVerifyScoreStride, volta::mem_row_major);""", + ) + source = source[:end] + qk + source[end:] + start = source.index("void flash_attention_grouped_verify_e4m3_full_q8_kernel(") + end = source.index("template (", "grouped_verify_qk_head(" + ) + return source[:start] + partial + source[end:] + + +def pipeline_qk_operands(source: str) -> str: + """Rotate Q/K fragment loads before the preceding K16 correction. + + Current operands are dead after MMA. Their registers can hold the next + operands while the original FP32 correction consumes the current product. + """ + start = source.index("template ") + end = source.index( + "__device__ __forceinline__ void grouped_verify_scale_output_fragment(", + start, + ) + qk = source[start:end] + begin = qk.index(" volta::load_matrix_sync(") + finish = qk.index(" if constexpr (COMPENSATE)", begin) + loads = qk[begin:finish] + qk = qk[:begin] + qk[finish:] + before_loop = qk.index("#pragma unroll") + qk = qk[:before_loop] + loads.replace("k_offset", "0") + qk[before_loop:] + marker = ( + " volta::mma_sync(tile_fragment, q_fragment, k_fragment, tile_fragment);" + ) + qk = replace_once( + qk, + marker, + marker + + "\n if (k_offset + 16 < kGroupedVerifyHeadDim) {\n" + + loads.replace("k_offset", "(k_offset + 16)") + + " }", + ) + marker = ( + " volta::mma_sync(score_fragment, q_fragment, k_fragment, score_fragment);" + ) + qk = replace_once( + qk, + marker, + marker + + "\n if (k_offset + 16 < kGroupedVerifyHeadDim) {\n" + + loads.replace("k_offset", "(k_offset + 16)") + + " }", + ) + return source[:start] + qk + source[end:] + + def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--output-dir", required=True, type=Path) - parser.add_argument("--head-groups", type=int, choices=(1, 3), default=3) + parser.add_argument("--head-groups", type=int, choices=(1, 2, 3), default=3) + parser.add_argument("--qk-unroll", type=int, choices=(1, 2, 4, 8, 16)) + parser.add_argument("--vector-load", action="store_true") + parser.add_argument("--page-specialize", action="store_true") + parser.add_argument("--prefetch-v", action="store_true") + parser.add_argument("--prefetch-k", action="store_true") + parser.add_argument("--prefetch-k-warps", type=int, choices=(4, 8), default=8) + parser.add_argument("--reuse-pv-values", action="store_true") + parser.add_argument("--qk-paired-products", action="store_true") + parser.add_argument("--qk-fp64-sum", action="store_true") + parser.add_argument("--specialize-full-q8", action="store_true") + parser.add_argument("--all-visible-tiles", action="store_true") + parser.add_argument("--register-softmax-state", action="store_true") + parser.add_argument("--qk-head-rows", action="store_true") + parser.add_argument("--qk-operand-pipeline", action="store_true") + parser.add_argument("--diagnostic-output-fp32", action="store_true") + parser.add_argument("--splits", type=int, choices=(80, 160, 320), default=80) + parser.add_argument("--grouped-only", action="store_true") parser.add_argument("--build", action="store_true") args = parser.parse_args() + if args.prefetch_v and args.head_groups != 1: + parser.error("--prefetch-v currently requires --head-groups 1") + if args.prefetch_k and not (args.prefetch_v and args.grouped_only): + parser.error("--prefetch-k requires --prefetch-v and --grouped-only") + if args.reuse_pv_values and not ( + args.head_groups == 1 + and args.prefetch_v + and args.grouped_only + and not args.prefetch_k + ): + parser.error("--reuse-pv-values needs six heads, V prefetch and grouped-only") + if args.specialize_full_q8 and not ( + args.grouped_only and args.head_groups == 1 and args.reuse_pv_values + ): + parser.error("Full-q8 specialization requires the six-head PV-reuse path") + if args.all_visible_tiles and not args.specialize_full_q8: + parser.error("Visible-tile specialization requires --specialize-full-q8") + if args.register_softmax_state and not args.specialize_full_q8: + parser.error("Register softmax state requires --specialize-full-q8") + if args.qk_head_rows and not ( + args.specialize_full_q8 + and not args.register_softmax_state + and not args.qk_paired_products + and not args.qk_fp64_sum + ): + parser.error("Audit the M8/N32 QK shape independently on fixed q8") + if args.qk_operand_pipeline and not ( + args.head_groups == 1 + and args.grouped_only + and not args.qk_paired_products + and not args.qk_fp64_sum + and not args.qk_head_rows + and not args.register_softmax_state + ): + parser.error("Audit QK operand rotation independently on six-head groups") + if args.qk_fp64_sum and not ( + args.grouped_only and args.head_groups == 1 and not args.qk_paired_products + ): + parser.error("FP64 K16 accumulation needs an isolated six-head grouped build") + if args.diagnostic_output_fp32 and not args.grouped_only: + parser.error("Pre-cast FP32 output is an isolated grouped-operator audit") root = Path(__file__).resolve().parents[2] / "flash-attention-v100" original = root / "kernel/flash_decode_paged.cu" source = original.read_text() + if args.grouped_only: + prefix_end = source.index( + "at::Tensor flash_attention_grouped_sparse_page4_plan(" + ) + entry_start = source.index( + "at::Tensor flash_attention_grouped_e4m3_fp32_paged(" + ) + entry_end = source.index( + "int64_t flash_attention_grouped_e4m3_fp32_precision_version()" + ) + source = source[:prefix_end] + source[entry_start:entry_end] + if args.qk_unroll is not None: + start = source.index("__device__ __forceinline__ void grouped_verify_qk(") + end = source.index( + "__device__ __forceinline__ void grouped_verify_scale_", start + ) + qk = source[start:end] + qk = replace_once( + qk, + "#pragma unroll\n for (int k_offset = 0; " + "k_offset < kGroupedVerifyHeadDim; k_offset += 16)", + f"#pragma unroll {args.qk_unroll}\n for (int k_offset = 0; " + "k_offset < kGroupedVerifyHeadDim; k_offset += 16)", + ) + source = source[:start] + qk + source[end:] + if args.vector_load: + source = replace_once( + source, + "bool ROW_SEQLENS = false, bool COMPENSATE_P = false>", + "bool ROW_SEQLENS = false, bool COMPENSATE_P = false, " + "bool PAIR_E4M3 = false>", + ) + start = source.index("__launch_bounds__(kGroupedVerifyThreads, 1) void ") + end = source.index( + "void flash_attention_grouped_verify_e5m2_combine_kernel(", start + ) + partial = source[start:end] + old = "!SPARSE_PAGE4 && !ROW_SEQLENS" + if partial.count(old) != 3: + raise ValueError("Expected three grouped KV panel loads") + partial = partial.replace(old, "!SPARSE_PAGE4 && (!ROW_SEQLENS || PAIR_E4M3)") + source = source[:start] + partial + source[end:] + source = replace_once( + source, + " auto kernel = flash_attention_grouped_verify_e5m2_partial_kernel<\n" + " 8, false, 0, false, false, false, " + "flash_v100::KV_CACHE_DTYPE_FP8_E4M3,\n" + " false, float, true, true>;", + " bool paired = true;\n" + " for (const auto* tensor : {&k, &v}) {\n" + " for (int dim = 0; dim < 3; ++dim)\n" + " paired = paired && tensor->stride(dim) % 16 == 0;\n" + " }\n" + " auto kernel = paired\n" + " ? flash_attention_grouped_verify_e5m2_partial_kernel<\n" + " 8, false, 0, false, false, false, " + "flash_v100::KV_CACHE_DTYPE_FP8_E4M3,\n" + " false, float, true, true, true>\n" + " : flash_attention_grouped_verify_e5m2_partial_kernel<\n" + " 8, false, 0, false, false, false, " + "flash_v100::KV_CACHE_DTYPE_FP8_E4M3,\n" + " false, float, true, true, false>;", + ) + if args.page_specialize: + statements = [] + for page in (1648, 3296): + prefix = ( + "flash_attention_grouped_verify_e5m2_partial_kernel<" + f"8, false, {page}, false, false, false, " + "flash_v100::KV_CACHE_DTYPE_FP8_E4M3, false, float, true, true" + ) + expression = ( + f"paired ? {prefix}, true> : {prefix}, false>" + if args.vector_load + else prefix + ">" + ) + statements.append(f" if (k.size(1) == {page}) kernel = {expression};\n") + marker = " constexpr int kCompensatedSmemBytes =" + source = replace_once(source, marker, "".join(statements) + marker) + if args.prefetch_v: + source = prefetch_values(source) + if args.prefetch_k: + source = prefetch_keys(source, args.prefetch_k_warps) + if args.reuse_pv_values: + source = reuse_pv_values(source) + if args.qk_paired_products: + if not args.grouped_only or args.head_groups != 1 or args.qk_unroll != 1: + parser.error( + "QK paired products require grouped-only, six heads and unroll1" + ) + source = pair_qk_products(source) + if args.qk_fp64_sum: + source = accumulate_qk_fp64(source) barrier = """ __syncwarp(); if (lane_id == 0) { if (tile_sum > 0.0f) {""" if source.count(barrier) != 1: raise ValueError("The grouped online-softmax warp-state fix is required") - if args.head_groups == 3: + if args.head_groups in (2, 3): source = replace_once( source, "constexpr int kGroupedVerifyRows = 48;", - "constexpr int kGroupedVerifyRows = 16;", + "constexpr int kGroupedVerifyRows = " + f"{32 if args.head_groups == 2 else 16};", ) source = replace_once( source, @@ -49,9 +765,50 @@ def main() -> None: source, "kernel<<>>", - "kernel<<>>", ) + if args.head_groups == 2: + source = replace_once( + source, + "static constexpr int kHeadsPerCta = " + "kGroupedVerifyRows / MAX_QUERY_TOKENS;", + "static constexpr int kHeadsPerCta = " + "MAX_QUERY_TOKENS == kGroupedVerifyQ8MaxQ " + "? 3 : kGroupedVerifyRows / MAX_QUERY_TOKENS;", + ) + source = replace_once( + source, + "MAX_QUERY_TOKENS * kHeadsPerCta == kGroupedVerifyRows,", + "MAX_QUERY_TOKENS * kHeadsPerCta <= kGroupedVerifyRows,", + ) + if args.splits != 80: + source = replace_once( + source, + "constexpr int kGroupedVerifyQ8Splits = 80;", + f"constexpr int kGroupedVerifyQ8Splits = {args.splits};", + ) + for old, new in ( + ("{80, 8, 6, 256}", f"{{{args.splits}, 8, 6, 256}}"), + ("{80, 8, 6, 2}", f"{{{args.splits}, 8, 6, 2}}"), + ("[80,8,6,256]", f"[{args.splits},8,6,256]"), + ("[80,8,6,2]", f"[{args.splits},8,6,2]"), + ( + f"kernel<< None: shutil.copy2(root / "LICENSE", sources) path = sources / "kernel/grouped-attention.cu" path.write_text(source) + source_sha256 = hashlib.sha256(path.read_bytes()).hexdigest() + module_name = "sm70_grouped_attention_" + source_sha256[:12] # Retain Flash-V100's existing math flags; this is a scheduling candidate. flags = [ "-O3", @@ -89,13 +848,32 @@ def main() -> None: ] manifest = { "input_source_sha256": hashlib.sha256(original.read_bytes()).hexdigest(), - "source_sha256": hashlib.sha256(path.read_bytes()).hexdigest(), + "source_sha256": source_sha256, + "module_name": module_name, "source_files": { str(p.relative_to(sources)): hashlib.sha256(p.read_bytes()).hexdigest() for p in sorted(sources.rglob("*")) if p.is_file() }, "head_groups": args.head_groups, + "qk_unroll": args.qk_unroll, + "vector_load": args.vector_load, + "page_specialize": args.page_specialize, + "prefetch_v": args.prefetch_v, + "prefetch_k": args.prefetch_k, + "prefetch_k_warps": args.prefetch_k_warps if args.prefetch_k else None, + "specialize_full_q8": args.specialize_full_q8, + "register_softmax_state": args.register_softmax_state, + "qk_head_rows": args.qk_head_rows, + "qk_operand_pipeline": args.qk_operand_pipeline, + "all_visible_tiles": args.all_visible_tiles, + "qk_fp64_sum": args.qk_fp64_sum, + "arithmetic_change": args.qk_fp64_sum or args.splits != 80, + "diagnostic_output_fp32": args.diagnostic_output_fp32, + "reuse_pv_values": args.reuse_pv_values, + "qk_paired_products": args.qk_paired_products, + "splits": args.splits, + "grouped_only": args.grouped_only, "extra_cuda_cflags": flags, "scope": "Private operator candidate; full-model admission required", } @@ -106,7 +884,7 @@ def main() -> None: build.mkdir(exist_ok=True) library = Path( load( - name="sm70_grouped_attention_candidate", + name=module_name, sources=[str(path)], build_directory=str(build), extra_cuda_cflags=flags, diff --git a/benchmarks/kernels/build_sm70_grouped_attention_n64.py b/benchmarks/kernels/build_sm70_grouped_attention_n64.py new file mode 100644 index 0000000000..bc0bcaeb72 --- /dev/null +++ b/benchmarks/kernels/build_sm70_grouped_attention_n64.py @@ -0,0 +1,308 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Batch two independent QK tiles while retaining both ordered N32 updates. + +Twelve QK warps and four raw-V prefetch warps share a 512-thread CTA. The +candidate retains 80 logical partitions and the original K16 compensation, +softmax and PV operations. Only aligned q8 inputs enter this private prototype. +""" + +import argparse +import hashlib +import json +import shutil +from pathlib import Path + +import regex as re + +from benchmarks.kernels.build_sm70_grouped_attention_candidate import replace_once + +STORAGE = r""" +struct alignas(256) GroupedN64Smem { + union { + struct { + alignas(16) __half q[48 * 264]; + alignas(16) __half kv[64 * 264]; + alignas(16) float scores[48 * 64]; + alignas(16) __half probs[48 * 40]; + alignas(16) __half residual[48 * 40]; + } compute; + alignas(16) float output[48 * 256]; + } storage; + alignas(16) float row_max[48]; + alignas(16) float row_sum[48]; + alignas(16) float row_scale[48]; + alignas(16) int page_ids[kGroupedVerifyPageIdsCapacity]; + alignas(16) uint32_t sparse_token_masks[8]; +}; +static_assert(sizeof(GroupedN64Smem) <= 96 * 1024, "N64 SM70 storage budget"); +""" + +PREFETCH = r""" + uint4 prefetched_v[8]; + if (warp_id < 12) { + grouped_verify_qk_n64(shared_q, shared_kv, shared_scores, + qk_scale, active_m_tiles); + } else { + const int copy_tid = tid - 384; +#pragma unroll + for (int i = 0; i < 8; ++i) { + const int index = copy_tid + i * 128; + const int row = index / 16; + const int vector = index % 16; + uint4 raw = make_uint4(0, 0, 0, 0); + if (row < valid_batch_rows) { + const int token = batch_start + row; + const int page = PAGE_BLOCK_SIZE > 0 + ? token / PAGE_BLOCK_SIZE : token / page_block_size; + const int offset = token - page * + (PAGE_BLOCK_SIZE > 0 ? PAGE_BLOCK_SIZE : page_block_size); + const int64_t base = static_cast(page_ids[page]) * + v_block_stride + static_cast(offset) * v_token_stride; + raw = __ldg(reinterpret_cast(v_cache) + base / 16 + vector); + } + prefetched_v[i] = raw; + } + } + __syncthreads(); + // K is dead after QK. The prefetch warps may now publish decoded V in + // the same panel; all consumers wait until publication is complete. + if (warp_id >= 12) { + const int copy_tid = tid - 384; +#pragma unroll + for (int i = 0; i < 8; ++i) { + const int index = copy_tid + i * 128; + const int shared_offset = (index / 16) * 33 + (index % 16) * 2; + const uint4 raw = prefetched_v[i]; + const uint64_t lo = static_cast(raw.x) | + (static_cast(raw.y) << 32); + const uint64_t hi = static_cast(raw.z) | + (static_cast(raw.w) << 32); + reinterpret_cast(shared_values)[shared_offset] = + fp8_e4m3fn_vector_to_half8_fast(lo); + reinterpret_cast(shared_values)[shared_offset + 1] = + fp8_e4m3fn_vector_to_half8_fast(hi); + } + } + __syncthreads(); +""" + + +def batched_source(source: str, v_loading: str = "prefetch") -> str: + begin = source.index("template ") + end = source.index("__device__ __forceinline__ void grouped_verify_scale", begin) + qk = source[begin:end] + qk = qk.replace("grouped_verify_qk(", "grouped_verify_qk_n64(") + qk = qk.replace("kGroupedVerifyQKWarps", "12") + qk = qk.replace("(kGroupedVerifyBlockN / 16)", "4") + qk = qk.replace("kGroupedVerifyScoreStride", "64") + source = source[:end] + STORAGE + qk + source[end:] + + begin = source.index( + "template ", a) + prefix = partial[a:b] + prefix = re.sub(r"\btile_start\b", "batch_start", prefix) + prefix = re.sub(r"\bvalid_k_rows\b", "valid_batch_rows", prefix) + prefix = prefix.replace("kGroupedVerifyBlockN", "64") + c = partial.index(" if constexpr (TWO_PASS)", b) + d = partial.index(" __syncthreads();\n }", c) + d += len(" __syncthreads();") + ordered = partial[c:d] + ordered = ordered.replace( + "shared_scores[row * kGroupedVerifyScoreStride + lane_id]", + "shared_scores[row * 64 + subtile_offset + lane_id]", + ) + ordered = replace_once( + ordered, + "shared_values + k_offset * kGroupedVerifyKVStride + d_tile * 16,", + "shared_values + (subtile_offset + k_offset) *\n" + " kGroupedVerifyKVStride + d_tile * 16,", + ) + value_loading = PREFETCH + publish_first = "" + if v_loading == "after-qk": + value_loading = r""" + grouped_verify_qk_n64(shared_q, shared_kv, shared_scores, + qk_scale, active_m_tiles); + __syncthreads(); + load_xqa_tc_kv_panel( + shared_values, v_cache, page_ids, valid_batch_rows, 32, 33, + batch_start, 0, page_block_size, 0, v_block_stride, v_token_stride, + v_head_stride, 0); + for (int i = tid + valid_batch_rows * 33; i < 64 * 33; i += 512) + reinterpret_cast(shared_values)[i] = make_uint4(0, 0, 0, 0); + __syncthreads(); +""" + elif v_loading == "softmax": + # Hold only the first half in registers. Publishing it and reading the + # other half are independent of the first N32 softmax's score/state. + prefetch = PREFETCH.replace("prefetched_v[8]", "prefetched_v[4]") + prefetch = prefetch.replace("i < 8", "i < 4") + at = prefetch.index(" // K is dead after QK.") + value_loading = prefetch[:at] + publish = prefetch[at:] + ending = " }\n __syncthreads();\n" + assert publish.endswith(ending) + publish = ( + publish[: -len(ending)] + + r""" + if (valid_batch_rows > 32) { + load_xqa_tc_kv_panel( + shared_values + 32 * 264, v_cache, page_ids, + valid_batch_rows - 32, 32, 33, batch_start + 32, 0, + page_block_size, 0, v_block_stride, v_token_stride, + v_head_stride, 0, copy_tid); + for (int i = copy_tid + (valid_batch_rows - 32) * 33; + i < 32 * 33; i += 128) + reinterpret_cast(shared_values + 32 * 264)[i] = + make_uint4(0, 0, 0, 0); + } + } +""" + ) + publish_first = " if (subtile_offset == 0) {\n" + publish + " }\n" + ordered = replace_once( + ordered, + """#pragma unroll + for (int row = warp_id; row < kGroupedVerifyRows; + row += kGroupedVerifyWarps) {""", + """ const int softmax_warps = subtile_offset == 0 ? 12 : 16; +#pragma unroll + for (int row = warp_id < softmax_warps ? warp_id : kGroupedVerifyRows; + row < kGroupedVerifyRows; row += softmax_warps) {""", + ) + else: + assert v_loading == "prefetch" + partial = ( + partial[:a] + + prefix + + value_loading + + " for (int subtile_offset = 0; subtile_offset < valid_batch_rows;\n" + " subtile_offset += 32) {\n" + " const int tile_start = batch_start + subtile_offset;\n" + " const int valid_k_rows = min(32, valid_batch_rows - subtile_offset);\n" + + publish_first + + ordered + + "\n }" + + partial[d:] + ) + source = source[:end] + partial + source[end:] + a = source.index(" auto kernel =", source.index("private_grouped_e4m3_fp32_paged")) + b = source.index(" constexpr int kCompensatedSmemBytes", a) + selected = source[a:b] + # Only paired aligned q8 is admitted; all other layouts use the parent. + selected = selected.replace("auto kernel = paired", "kernel = true", 1) + selected = selected.replace("paired ?", "true ?").replace(old, new) + # Both sides of a constant conditional must instantiate a valid template. + selected = selected.replace("true, true, false>", "true, true, true>") + at = source.index(" C10_CUDA_CHECK(\n cudaFuncSetAttribute(kernel", b) + source = ( + source[:at] + " int shared_bytes = kCompensatedSmemBytes;\n" + " if (q.size(0) == 8 && paired) {\n" + + selected + + " shared_bytes = sizeof(GroupedN64Smem);\n" + " }\n" + " TORCH_CHECK(properties->sharedMemPerBlockOptin >= shared_bytes,\n" + ' "N64 shared memory exceeds device limit");\n' + source[at:] + ) + tail = source[at:] + tail = tail.replace( + " kCompensatedSmemBytes));", + " shared_bytes));", + ) + tail = tail.replace("kCompensatedSmemBytes, stream>>>", "shared_bytes, stream>>>") + return source[:at] + tail + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-manifest", required=True, type=Path) + parser.add_argument("--output-dir", required=True, type=Path) + parser.add_argument("--build", action="store_true") + parser.add_argument( + "--v-loading", choices=("prefetch", "after-qk", "softmax"), default="prefetch" + ) + args = parser.parse_args() + base = json.loads(args.base_manifest.read_text()) + assert base["head_groups"] == 1 and base["splits"] == 80 + assert base["prefetch_v"] and base["reuse_pv_values"] + original = args.base_manifest.parent / "sources" + for name, digest in base["source_files"].items(): + assert hashlib.sha256((original / name).read_bytes()).hexdigest() == digest + directory = args.output_dir.resolve() + sources = directory / "sources" + shutil.copytree(original, sources) + path = sources / "kernel/grouped-attention.cu" + path.write_text(batched_source(path.read_text(), args.v_loading)) + digest = hashlib.sha256(path.read_bytes()).hexdigest() + module_name = "sm70_grouped_n64_" + digest[:12] + manifest = { + **{k: v for k, v in base.items() if k not in ("library", "library_sha256")}, + "input_source_sha256": base["source_sha256"], + "source_sha256": digest, + "module_name": module_name, + "scope": "Private physical N64 schedule; both logical N32 updates retained", + "physical_tile_n": 64, + "logical_update_n": 32, + "prefetch_v": args.v_loading != "after-qk", + "raw_v_prefetch": args.v_loading != "after-qk", + "v_loading": args.v_loading, + "source_files": { + str(p.relative_to(sources)): hashlib.sha256(p.read_bytes()).hexdigest() + for p in sorted(sources.rglob("*")) + if p.is_file() + }, + } + if args.build: + from torch.utils.cpp_extension import load + + build = directory / "build" + build.mkdir() + module = load( + name=module_name, + sources=[str(path)], + extra_include_paths=[str(sources / "include"), str(sources / "kernel")], + extra_cuda_cflags=base["extra_cuda_cflags"], + build_directory=str(build), + verbose=True, + ) + manifest["library"] = str(Path(module.__file__).resolve()) + manifest["library_sha256"] = hashlib.sha256( + Path(module.__file__).read_bytes() + ).hexdigest() + (directory / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n") + print(json.dumps(manifest, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/build_sm70_grouped_attention_phase_probe.py b/benchmarks/kernels/build_sm70_grouped_attention_phase_probe.py new file mode 100644 index 0000000000..01283c7a5c --- /dev/null +++ b/benchmarks/kernels/build_sm70_grouped_attention_phase_probe.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Build an operator-only clock64 probe of the ordered N32 attention loop. + +The explicit ``profile`` entrypoint writes five timestamps per visible tile. +Its timings include probe overhead and synchronization, and are neither achieved +occupancy nor uninstrumented latency. No serving route imports this module. +""" + +import argparse +import hashlib +import json +import shutil +from pathlib import Path + +from benchmarks.kernels.build_sm70_grouped_attention_candidate import replace_once + + +def phase_source(source: str) -> str: + start = source.index( + "template str: + return ( + " if (clock_stamps != nullptr && tid == 0)\n" + f" clock_stamps[(tile_start / 32) * 5 + {index}] = clock64();\n" + ) + + loop = replace_once( + loop, + " const int valid_k_rows =", + stamp(0) + " const int valid_k_rows =", + ) + loop = replace_once( + loop, + " __syncthreads();\n\n int active_m_tiles", + " __syncthreads();\n" + stamp(1) + "\n int active_m_tiles", + ) + loop = replace_once( + loop, + " __syncthreads();\n\n if constexpr (TWO_PASS)", + " __syncthreads();\n" + stamp(2) + "\n if constexpr (TWO_PASS)", + ) + loop = replace_once( + loop, + " static_assert(COMPENSATE_P && kGroupedVerifyWarps == 16,", + stamp(3) + " static_assert(COMPENSATE_P && kGroupedVerifyWarps == 16,", + ) + loop = replace_once( + loop, + " __syncthreads();\n }", + " __syncthreads();\n" + stamp(4) + " }", + ) + partial = partial[:loop_start] + loop + partial[loop_end:] + source = source[:start] + partial + source[end:] + host_start = source.index("at::Tensor private_grouped_e4m3_fp32_paged(") + host_end = source.index("PYBIND11_MODULE(", host_start) + host = source[host_start:host_end] + host = replace_once( + host, + " 1, row_lengths.data_ptr());", + " 1, row_lengths.data_ptr(), nullptr);", + ) + profile = host.replace( + "private_grouped_e4m3_fp32_paged", "profile_grouped_e4m3_fp32_paged" + ) + profile = replace_once( + profile, + "float scale, float k_scale, float v_scale) {", + "float scale, float k_scale, float v_scale, at::Tensor& stamps) {\n" + " TORCH_CHECK(stamps.device() == q.device() &&\n" + " stamps.scalar_type() == at::kLong && stamps.is_contiguous() &&\n" + " stamps.dim() == 2 && stamps.size(1) == 5 &&\n" + " stamps.size(0) >= (block_table.numel() * k.size(1) + 31) / 32,\n" + ' "The diagnostic needs CUDA int64 timestamps [capacity_tiles,5]");', + ) + profile = replace_once( + profile, + " 1, row_lengths.data_ptr(), nullptr);", + " 1, row_lengths.data_ptr(),\n" + " reinterpret_cast(stamps.data_ptr()));", + ) + source = source[:host_start] + host + profile + source[host_end:] + return replace_once( + source, + ' m.def("run", &private_grouped_e4m3_fp32_paged);', + ' m.def("run", &private_grouped_e4m3_fp32_paged);\n' + ' m.def("profile", &profile_grouped_e4m3_fp32_paged);', + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-manifest", required=True, type=Path) + parser.add_argument("--output-dir", required=True, type=Path) + parser.add_argument("--build", action="store_true") + args = parser.parse_args() + base = json.loads(args.base_manifest.read_text()) + assert base["head_groups"] == 1 and base["splits"] == 80 + assert base["prefetch_v"] and base["reuse_pv_values"] + assert base.get("physical_tile_n", 32) == 32 + source_dir = args.base_manifest.parent / "sources" + for relative, digest in base["source_files"].items(): + assert ( + hashlib.sha256((source_dir / relative).read_bytes()).hexdigest() == digest + ) + directory = args.output_dir.resolve() + sources = directory / "sources" + shutil.copytree(source_dir, sources) + path = sources / "kernel/grouped-attention.cu" + path.write_text(phase_source(path.read_text())) + digest = hashlib.sha256(path.read_bytes()).hexdigest() + module_name = "sm70_grouped_phase_" + digest[:12] + manifest = { + **{k: v for k, v in base.items() if k not in ("library", "library_sha256")}, + "input_source_sha256": base["source_sha256"], + "source_sha256": digest, + "module_name": module_name, + "scope": "Instrumented phase attribution only, never performance admission", + "phase_names": ["K load", "QK and V load", "online softmax", "ordered PV"], + "source_files": { + str(p.relative_to(sources)): hashlib.sha256(p.read_bytes()).hexdigest() + for p in sorted(sources.rglob("*")) + if p.is_file() + }, + } + if args.build: + from torch.utils.cpp_extension import load + + build = directory / "build" + build.mkdir() + module = load( + name=module_name, + sources=[str(path)], + extra_include_paths=[str(sources / "include"), str(sources / "kernel")], + extra_cuda_cflags=base["extra_cuda_cflags"], + build_directory=str(build), + verbose=True, + ) + manifest["library"] = str(Path(module.__file__).resolve()) + manifest["library_sha256"] = hashlib.sha256( + Path(module.__file__).read_bytes() + ).hexdigest() + (directory / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n") + print(json.dumps(manifest, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/build_sm70_grouped_attention_pipeline.py b/benchmarks/kernels/build_sm70_grouped_attention_pipeline.py new file mode 100644 index 0000000000..963bf95f3c --- /dev/null +++ b/benchmarks/kernels/build_sm70_grouped_attention_pipeline.py @@ -0,0 +1,380 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Private SM70 QK/PV warp pipeline; no serving route is installed. + +Four to eight producer warps and sixteen consumers share two score/value panels. +Named ready/free barriers protect each panel. K16 compensation, N32 updates, +80 logical partitions and the final merge are copied from the audited source. +""" + +import argparse +import hashlib +import json +import shutil +from pathlib import Path + +from benchmarks.kernels.build_sm70_grouped_attention_candidate import replace_once + +PREFIX = r""" +constexpr int kPipelineProducerThreads = 256; +constexpr int kPipelineConsumerThreads = 512; +constexpr int kPipelineThreads = 768; +struct alignas(256) GroupedPipelineSmem { + union { + struct { + alignas(16) __half q[48 * 264]; + alignas(16) __half k[32 * 264]; + alignas(16) __half values[2][32 * 264]; + alignas(16) float scores[2][48 * 32]; + alignas(16) __half probs[48 * 40]; + alignas(16) __half residual[48 * 40]; + } compute; + alignas(16) float output[48 * 256]; + } storage; + alignas(16) float row_max[48]; + alignas(16) float row_sum[48]; + alignas(16) float row_scale[48]; +}; +static_assert(sizeof(GroupedPipelineSmem) <= 96 * 1024, "pipeline storage budget"); + +// Deliberately use unaligned barrier instructions: producer and consumer +// warps follow different control flow. A panel cannot be reused until every +// consumer has arrived at its free barrier. No warp spins on shared flags. +template +__device__ __forceinline__ void pipeline_sync(int id) { + asm volatile("barrier.sync %0, %1;" :: "r"(id), "n"(COUNT) : "memory"); +} +template +__device__ __forceinline__ void pipeline_arrive(int id) { + asm volatile("barrier.arrive %0, %1;" :: "r"(id), "n"(COUNT) : "memory"); +} + +template +__global__ __launch_bounds__(kPipelineThreads, 1) +void grouped_pipeline_partial_kernel( + const __half* q, const void* k_cache, const void* v_cache, + const int* page_ids, const int* row_lengths, + float* partial_out, float* partial_lse, int query_len, + int page_block_size, int64_t k_block_stride, int64_t k_token_stride, + int64_t k_head_stride, int64_t v_block_stride, int64_t v_token_stride, + int64_t v_head_stride, float qk_scale, float v_scale) { + using Traits = GroupedVerifyTraits<8>; + using PARTIAL_T = float; + constexpr int MAX_QUERY_TOKENS = 8; + constexpr bool SPARSE_PAGE4 = false; + constexpr bool ROW_SEQLENS = true; + constexpr bool COMPENSATE_P = true; + constexpr int group_idx = 0; + constexpr int head_start = 0; + constexpr int active_m_tiles = 7; + constexpr int kResidualStride = 40; + const int block_tid = threadIdx.x; + const int split_id = blockIdx.x; + int total_kv = 0; + for (int i = 0; i < query_len; ++i) + total_kv = max(total_kv, row_lengths[i]); + if (total_kv <= 0) return; + const int active_splits = grouped_verify_active_splits<8, false>(total_kv); + if (split_id >= active_splits) return; + const int total_tiles = (total_kv + 31) / 32; + const int base_tiles = total_tiles / active_splits; + const int extra_tiles = total_tiles - base_tiles * active_splits; + const int split_tile_start = split_id * base_tiles + min(split_id, extra_tiles); + const int split_tiles = base_tiles + (split_id < extra_tiles ? 1 : 0); + const int split_start = split_tile_start * 32; + const int prefix_kv_len = max(0, total_kv - query_len); + + extern __shared__ char pipeline_smem_raw[]; + auto& smem = *reinterpret_cast(pipeline_smem_raw); + __half* shared_q = smem.storage.compute.q; + __half* shared_k = smem.storage.compute.k; + auto* q_vec = reinterpret_cast(q); + for (int i = block_tid; i < 48 * 32; i += kPipelineThreads) { + reinterpret_cast(shared_q)[(i / 32) * 33 + i % 32] = __ldg(q_vec + i); + } + if (block_tid < 48) { + smem.row_max[block_tid] = kXQANegInf; + smem.row_sum[block_tid] = 0.0f; + smem.row_scale[block_tid] = 1.0f; + } + __syncthreads(); + + if (block_tid < kPipelineProducerThreads) { + for (int tile = 0; tile < split_tiles; ++tile) { + const int panel = tile & 1; + if (tile >= 2) pipeline_sync(3 + panel); + const int tile_start = split_start + tile * 32; + const int valid_k_rows = min(32, total_kv - tile_start); + load_xqa_tc_kv_panel( + shared_k, k_cache, page_ids, valid_k_rows, 32, 33, + tile_start, 0, page_block_size, 0, k_block_stride, k_token_stride, + k_head_stride, 0, block_tid); + for (int i = block_tid + valid_k_rows * 33; i < 32 * 33; + i += kPipelineProducerThreads) + reinterpret_cast(shared_k)[i] = make_uint4(0, 0, 0, 0); + pipeline_sync(5); + grouped_verify_qk(shared_q, shared_k, + smem.storage.compute.scores[panel], qk_scale, 7); + __half* values = smem.storage.compute.values[panel]; + load_xqa_tc_kv_panel( + values, v_cache, page_ids, valid_k_rows, 32, 33, + tile_start, 0, page_block_size, 0, v_block_stride, v_token_stride, + v_head_stride, 0, block_tid); + for (int i = block_tid + valid_k_rows * 33; i < 32 * 33; + i += kPipelineProducerThreads) + reinterpret_cast(values)[i] = make_uint4(0, 0, 0, 0); + // Every producer must finish reading K before any starts the next load. + pipeline_sync(5); + pipeline_arrive(1 + panel); + } + // Consumers overlay compute storage only after production has finished. + pipeline_sync(7); + pipeline_sync(8); + return; + } + + const int tid = block_tid - kPipelineProducerThreads; + const int warp_id = tid / 32; + const int lane_id = tid % 32; + __half* shared_probs = smem.storage.compute.probs; + __half* shared_prob_residual = smem.storage.compute.residual; + volta::fragment output_fragments[3]; +#pragma unroll + for (int i = 0; i < 3; ++i) volta::fill_fragment(output_fragments[i], 0.0f); + for (int tile = 0; tile < split_tiles; ++tile) { + const int panel = tile & 1; + pipeline_sync(1 + panel); + const int tile_start = split_start + tile * 32; + const int valid_k_rows = min(32, total_kv - tile_start); + float* shared_scores = smem.storage.compute.scores[panel]; + __half* shared_values = smem.storage.compute.values[panel]; +""" + + +def pipeline_source( + source: str, + serialize: bool = False, + debug_first_tile: bool = False, + producer_warps: int = 8, +) -> str: + assert producer_warps in (4, 6, 8) + start = source.index("__launch_bounds__(kGroupedVerifyThreads, 1) void ") + end = source.index( + "template (7);", 1 + ) + tail = tail.replace( + " __syncthreads();", " pipeline_sync(8);", 1 + ) + assert "__syncthreads" not in tail + prefix = replace_once( + PREFIX, + "constexpr int kPipelineProducerThreads = 256;", + f"constexpr int kPipelineProducerThreads = {producer_warps * 32};", + ) + prefix = replace_once( + prefix, + "constexpr int kPipelineThreads = 768;", + f"constexpr int kPipelineThreads = {producer_warps * 32 + 512};", + ) + if producer_warps == 4: + # Six/eight producers both still spill at the resource limit. Reuse + # four producer warps across the six independent QK output tiles; + # each tile retains the original K16 products and corrections. + qk_start = source.index( + "template \n" + "__device__ __forceinline__ void grouped_verify_qk(" + ) + qk_end = source.index( + "__device__ __forceinline__ void grouped_verify_scale_output_fragment(", + qk_start, + ) + qk = source[qk_start:qk_end].rstrip() + qk = replace_once(qk, "void grouped_verify_qk(", "void pipeline_qk(") + qk = replace_once(qk, "warp_id >= kGroupedVerifyQKWarps", "warp_id >= 4") + qk = replace_once( + qk, + " if ((active_m_tiles & (1 << m_tile)) == 0) {\n return;\n }", + " if ((active_m_tiles & (1 << m_tile)) == 0) {\n continue;\n }", + ) + qk = qk.replace("m_tile = warp_id /", "m_tile = qk_tile /") + qk = qk.replace("n_tile = warp_id %", "n_tile = qk_tile %") + loop = qk.index(" const int m_tile =") + assert qk.endswith("}") + qk = ( + qk[:loop] + + " for (int qk_tile = warp_id; qk_tile < kGroupedVerifyQKWarps; " + "qk_tile += 4) {\n" + qk[loop:-1] + " }\n}\n" + ) + prefix = qk + replace_once( + prefix, "grouped_verify_qk", "pipeline_qk" + ) + kernel = ( + prefix + + softmax + + " pipeline_sync(6);\n" + + pv + + " if (tile + 2 < split_tiles) " + "pipeline_arrive(3 + panel);\n }\n" + tail + ) + if serialize: + kernel = replace_once( + kernel, + " pipeline_arrive(1 + panel);", + " pipeline_arrive(1 + panel);\n" + " pipeline_sync(9);", + ) + kernel = replace_once( + kernel, + "pipeline_arrive(3 + panel);\n }", + "pipeline_arrive(3 + panel);\n" + " pipeline_sync(9);\n }", + ) + source = source[:end] + kernel + source[end:] + a = source.index( + " auto kernel = paired", source.index("private_grouped_e4m3_fp32_paged(") + ) + b = source.index(" flash_attention_grouped_verify_e5m2_combine_kernel<", a) + fallback = source[a:b] + launch = r""" if (q.size(0) == 8 && block_table.size(1) * k.size(1) > 32768) { + auto pipeline = paired ? grouped_pipeline_partial_kernel<0, true> + : grouped_pipeline_partial_kernel<0, false>; + if (k.size(1) == 1648) + pipeline = paired ? grouped_pipeline_partial_kernel<1648, true> + : grouped_pipeline_partial_kernel<1648, false>; + if (k.size(1) == 3296) + pipeline = paired ? grouped_pipeline_partial_kernel<3296, true> + : grouped_pipeline_partial_kernel<3296, false>; + TORCH_CHECK(properties->sharedMemPerBlockOptin >= sizeof(GroupedPipelineSmem), + "warp pipeline exceeds the opt-in shared memory budget"); + C10_CUDA_CHECK(cudaFuncSetAttribute(pipeline, + cudaFuncAttributeMaxDynamicSharedMemorySize, sizeof(GroupedPipelineSmem))); + C10_CUDA_CHECK(cudaFuncSetAttribute(pipeline, + cudaFuncAttributePreferredSharedMemoryCarveout, 100)); + pipeline<<<80, kPipelineThreads, sizeof(GroupedPipelineSmem), stream>>>( + reinterpret_cast(aligned_q.data_ptr()), k.data_ptr(), + v.data_ptr(), block_table.data_ptr(), row_lengths.data_ptr(), + partial.data_ptr(), lse.data_ptr(), q.size(0), k.size(1), + k.stride(0), k.stride(1), k.stride(2), + v.stride(0), v.stride(1), v.stride(2), scale * k_scale, v_scale); + } else { +""" + return source[:a] + launch + fallback + " }\n" + source[b:] + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-manifest", required=True, type=Path) + parser.add_argument("--output-dir", required=True, type=Path) + parser.add_argument("--build", action="store_true") + parser.add_argument("--serialize", action="store_true") + parser.add_argument("--debug-first-tile", action="store_true") + parser.add_argument("--producer-warps", type=int, choices=(4, 6, 8), default=8) + args = parser.parse_args() + base = json.loads(args.base_manifest.read_text()) + if not ( + base["head_groups"] == 1 + and base.get("reuse_pv_values") + and base["splits"] == 80 + ): + parser.error("The pipeline requires an 80-split six-head PV-reuse source") + source_dir = args.base_manifest.parent / "sources" + source_file = source_dir / "kernel/grouped-attention.cu" + assert hashlib.sha256(source_file.read_bytes()).hexdigest() == base["source_sha256"] + source = pipeline_source( + source_file.read_text(), + args.serialize, + args.debug_first_tile, + args.producer_warps, + ) + directory = args.output_dir.resolve() + if directory.exists(): + parser.error("Use a new output directory for every prototype") + sources = directory / "sources" + shutil.copytree(source_dir, sources) + path = sources / "kernel/grouped-attention.cu" + path.write_text(source) + digest = hashlib.sha256(path.read_bytes()).hexdigest() + module_name = "sm70_grouped_attention_" + digest[:12] + manifest = dict(base) + manifest.pop("library", None) + manifest.pop("library_sha256", None) + manifest.update( + input_source_sha256=base["source_sha256"], + source_sha256=digest, + module_name=module_name, + warp_pipeline=True, + serialized_diagnostic=args.serialize, + debug_first_tile=args.debug_first_tile, + producer_warps=args.producer_warps, + source_files={ + str(p.relative_to(sources)): hashlib.sha256(p.read_bytes()).hexdigest() + for p in sorted(sources.rglob("*")) + if p.is_file() + }, + scope="Private warp-pipeline operator; no serving admission", + synchronization_reference="https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-bar", + ) + if args.build: + from torch.utils.cpp_extension import load + + build = directory / "build" + build.mkdir() + library = Path( + load( + name=module_name, + sources=[str(path)], + build_directory=str(build), + extra_cuda_cflags=base["extra_cuda_cflags"], + extra_include_paths=[str(sources / "kernel"), str(sources / "include")], + verbose=True, + ).__file__ + ) + manifest["library"] = str(library) + manifest["library_sha256"] = hashlib.sha256(library.read_bytes()).hexdigest() + (directory / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n") + print(json.dumps(manifest, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/build_sm70_grouped_attention_staged.py b/benchmarks/kernels/build_sm70_grouped_attention_staged.py new file mode 100644 index 0000000000..dc3e3f19b3 --- /dev/null +++ b/benchmarks/kernels/build_sm70_grouped_attention_staged.py @@ -0,0 +1,367 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Private two-stage q8 experiment; does not install or admit a serving route. + +QK tiles are independent. A separate producer computes the same compensated +K16 scores; two PV column partitions retain all 80 logical context partitions +and each partition's N32 online-update order. The prototype allocates its score +scratch during warmup/capture. Serving integration and scratch admission are +deliberately separate from this feasibility experiment. +""" + +import argparse +import hashlib +import json +import shutil +from pathlib import Path + +from benchmarks.kernels.build_sm70_grouped_attention_candidate import replace_once + +PRODUCER = r""" +template +__global__ __launch_bounds__(256, 2) void grouped_staged_qk_kernel( + const __half* q, const void* k_cache, const int* page_ids, + const int* row_lengths, float* scores, int page_block_size, + int64_t k_block_stride, int64_t k_token_stride, int64_t k_head_stride, + float qk_scale) { + int total_kv = 0; +#pragma unroll + for (int i = 0; i < 8; ++i) total_kv = max(total_kv, row_lengths[i]); + const int tile_start = blockIdx.x * 32; + if (tile_start >= total_kv) return; + const int valid_k_rows = min(32, total_kv - tile_start); + const int tid = threadIdx.x; + __shared__ __align__(16) __half shared_q[48 * 264]; + __shared__ __align__(16) __half shared_k[32 * 264]; + __shared__ __align__(16) float shared_scores[48 * 32]; + const uint4* q_vec = reinterpret_cast(q); + uint4* shared_q_vec = reinterpret_cast(shared_q); + for (int i = tid; i < 48 * 32; i += 256) { + shared_q_vec[(i / 32) * 33 + i % 32] = __ldg(q_vec + i); + } + load_xqa_tc_kv_panel( + shared_k, k_cache, page_ids, valid_k_rows, 32, 33, + tile_start, 0, page_block_size, 0, k_block_stride, k_token_stride, + k_head_stride, 0); + for (int i = tid + valid_k_rows * 33; i < 32 * 33; i += 256) { + reinterpret_cast(shared_k)[i] = make_uint4(0, 0, 0, 0); + } + __syncthreads(); + grouped_verify_qk(shared_q, shared_k, shared_scores, qk_scale, 0x7); + __syncthreads(); + for (int i = tid; i < 48 * 32; i += 256) { + scores[static_cast(blockIdx.x) * 48 * 32 + i] = shared_scores[i]; + } +} + +constexpr int kStagedPVThreads = 256; +constexpr int kStagedPVWarps = 8; +constexpr int kStagedPVHeadDim = 128; +constexpr int kStagedPVStride = 136; +constexpr int kStagedPVOutputTilesPerWarp = 3; +struct StagedPVSmem { + union { + struct { + alignas(16) __half kv[32 * 136]; + alignas(16) float scores[48 * 32]; + alignas(16) __half probs[48 * 40]; + alignas(16) __half residual[48 * 40]; + } compute; + alignas(16) float output[48 * 128]; + } storage; + alignas(16) float row_max[48]; + alignas(16) float row_sum[48]; + alignas(16) float row_scale[48]; + alignas(16) int page_ids[16]; + alignas(16) uint32_t sparse_token_masks[8]; +}; +static_assert(sizeof(StagedPVSmem) <= 48 * 1024, "PV storage budget"); +""" + + +def staged_source(source: str, pv_columns: int = 2) -> str: + assert pv_columns in (1, 2) + start = source.index( + "template (v_cache) + column_partition * 128; +""" + + partial[b:] + ) + # Q is consumed only by the independent producer. Do not reserve or fill a + # query panel in the PV block. The TWO_PASS template is never instantiated. + a = partial.index(" constexpr int kVecsPerRow =") + b = partial.index(" if (tid < kGroupedVerifyRows)", a) + partial = partial[:a] + partial[b:] + a = partial.index(" // The conservative baseline computes") + b = partial.index(" // Recompute QK for the conservative path", a) + partial = partial[:a] + partial[b:] + a = partial.index(" load_xqa_tc_kv_panel<") + b = partial.index(" int active_m_tiles =", a) + partial = ( + partial[:a] + " for (int i = tid; i < 48 * 32; i += kStagedPVThreads) {\n" + " shared_scores[i] = staged_scores[" + "static_cast(tile_start / 32) * 48 * 32 + i];\n" + " }\n" + partial[b:] + ) + a = partial.index(" grouped_verify_qk") + b = partial.index(" load_xqa_tc_kv_panel<", a) + partial = ( + partial[:a] + + """ constexpr int kValueLoadThreads = kStagedPVThreads; + const int value_load_tid = tid; +""" + + partial[b:] + ) + partial = replace_once( + partial, " }\n }\n\n __syncthreads();", " }\n\n __syncthreads();" + ) + # Both D halves retain the same N32 updates. Only their disjoint numerator + # columns are stored; one partition alone publishes the common max/sum. + a = partial.index(" int64_t output_idx;") + b = partial.index(" if constexpr (std::is_same_v)", a) + partial = ( + partial[:a] + + """ const int64_t output_idx = + (((static_cast(split_id) * MAX_QUERY_TOKENS + token_idx) * + kGroupedVerifyHeads + head_idx) * 256 + column_partition * 128 + d); +""" + + partial[b:] + ) + a = partial.rindex(" if (tid < kGroupedVerifyRows)") + partial = partial[:a] + partial[a:].replace( + "if (tid < kGroupedVerifyRows)", + "if (column_partition == 0 && tid < kGroupedVerifyRows)", + 1, + ) + producer = PRODUCER + if pv_columns == 1: + for old, new in ( + ("kStagedPVThreads = 256", "kStagedPVThreads = 512"), + ("kStagedPVWarps = 8", "kStagedPVWarps = 16"), + ("kStagedPVHeadDim = 128", "kStagedPVHeadDim = 256"), + ("kStagedPVStride = 136", "kStagedPVStride = 264"), + ("kv[32 * 136]", "kv[32 * 264]"), + ("output[48 * 128]", "output[48 * 256]"), + ("<= 48 * 1024", "<= 64 * 1024"), + ): + producer = replace_once(producer, old, new) + partial = replace_once( + partial, + "__launch_bounds__(kStagedPVThreads, 2)", + "__launch_bounds__(kStagedPVThreads, 1)", + ) + elif "PV reuse is isolated" in partial: + # Each D128 block has eight warps, one D tile per warp. The same + # three independent M accumulators can still share both V fragments. + partial = replace_once( + partial, + "COMPENSATE_P && kStagedPVWarps == 16", + "COMPENSATE_P && kStagedPVWarps == 8", + ) + source = source[:end] + producer + partial + source[end:] + # Keep the complete original host validation and fallback for untested + # small/other shapes. Prototype scratch is capture-owned, not a cache. + a = source.index( + " auto kernel = paired", source.index("private_grouped_e4m3_fp32_paged(") + ) + b = source.index(" flash_attention_grouped_verify_e5m2_combine_kernel<", a) + original = source[a:b] + declarations = [] + for page in (0, 1648, 3296): + prefix = "auto" if page == 0 else f"if (k.size(1) == {page})" + for op, name in ( + ("grouped_staged_qk_kernel", "producer"), + ("grouped_staged_pv_kernel", "consumer"), + ): + suffix = ( + "" + if name == "producer" + else ", false, false, false, flash_v100::KV_CACHE_DTYPE_FP8_E4M3, " + "false, float, true, true" + ) + args = ( + str(page) if name == "producer" else "8, false, " + str(page) + suffix + ) + assign = f"auto {name} =" if page == 0 else f"{prefix} {name} =" + declarations.append( + f" {assign} paired ? {op}<{args}, true> : {op}<{args}, false>;" + ) + replacement = ( + """ if (q.size(0) == 8 && block_table.size(1) * k.size(1) > 32768) { + const int tiles = (block_table.size(1) * k.size(1) + 31) / 32; + auto scores = at::empty({tiles, 48, 32}, q.options().dtype(at::kFloat)); +""" + + "\n".join(declarations) + + """ + C10_CUDA_CHECK(cudaFuncSetAttribute(consumer, + cudaFuncAttributeMaxDynamicSharedMemorySize, sizeof(StagedPVSmem))); + producer<<>>( + reinterpret_cast(aligned_q.data_ptr()), k.data_ptr(), + block_table.data_ptr(), row_lengths.data_ptr(), + scores.data_ptr(), k.size(1), k.stride(0), k.stride(1), + k.stride(2), scale * k_scale); + consumer<<>>( + reinterpret_cast(aligned_q.data_ptr()), k.data_ptr(), + v.data_ptr(), block_table.data_ptr(), row_lengths.data_ptr(), + partial.data_ptr(), lse.data_ptr(), q.size(0), + block_table.size(1), k.size(1), k.stride(0), k.stride(1), k.stride(2), + v.stride(0), v.stride(1), v.stride(2), scale * k_scale, v_scale, + nullptr, 1, row_lengths.data_ptr(), scores.data_ptr()); + } else { +""" + + original + + " }\n" + ) + if pv_columns == 1: + replacement = replace_once( + replacement, "consumer<<, 256, 0); + result["pv"] = describe(grouped_staged_pv_kernel<8, false, 3296, + false, false, false, flash_v100::KV_CACHE_DTYPE_FP8_E4M3, + false, float, true, true, true>, kStagedPVThreads, sizeof(StagedPVSmem)); + return result; +} +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("resource_report", &staged_resource_report); +""", + ) + path.write_text(source) + digest = hashlib.sha256(path.read_bytes()).hexdigest() + module_name = "sm70_staged_attention_" + digest[:12] + manifest = { + **base, + "input_source_sha256": base["source_sha256"], + "source_sha256": digest, + "module_name": module_name, + "staged_qk": True, + "pv_column_partitions": args.pv_columns, + "reuse_pv_values": base.get("reuse_pv_values", False), + "capture_owned_score_scratch": True, + "shared_carveout": args.shared_carveout, + "source_files": { + str(p.relative_to(sources)): hashlib.sha256(p.read_bytes()).hexdigest() + for p in sorted(sources.rglob("*")) + if p.is_file() + }, + "scope": "Private feasibility candidate, not admitted for serving", + } + manifest.pop("library", None) + manifest.pop("library_sha256", None) + if args.build: + from torch.utils.cpp_extension import load + + build = directory / "build" + build.mkdir() + module = load( + name=module_name, + sources=[str(path)], + build_directory=str(build), + extra_cuda_cflags=base["extra_cuda_cflags"], + extra_include_paths=[str(sources / "kernel"), str(sources / "include")], + verbose=True, + ) + library = Path(module.__file__) + manifest["library"] = str(library) + manifest["library_sha256"] = hashlib.sha256(library.read_bytes()).hexdigest() + (directory / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n") + print(json.dumps(manifest, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/build_sm70_scalar_attention_candidate.py b/benchmarks/kernels/build_sm70_scalar_attention_candidate.py new file mode 100644 index 0000000000..aa55866e92 --- /dev/null +++ b/benchmarks/kernels/build_sm70_scalar_attention_candidate.py @@ -0,0 +1,445 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Isolate the traced FP32 scalar q1 fallback and its ordered PV loop. + +This builder installs no backend. Candidates keep each head's scalar FP32 +operations and the frozen partition/merge order. Independent load scheduling +and six-head KV reuse are screened separately. +""" + +import argparse +import hashlib +import json +import shutil +from pathlib import Path + +from benchmarks.kernels.build_sm70_grouped_attention_candidate import replace_once + +SHARED_HEADS_BODY = r"""{ + static_assert(D == 256 && PARTITION_SIZE == 1024 && + KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP8_E4M3 && + SEQ_LEN_ROUTE == 0 && !ANCHORED_SWA && + std::is_same_v, "Private scalar six-head contract"); + const int partition_idx = blockIdx.z; + const int seq_len = seq_lens[0]; + const int start_token_idx = partition_idx * PARTITION_SIZE; + if (seq_len <= 0 || start_token_idx >= seq_len || + partition_idx >= max_num_partitions) return; + const int effective_num_partitions = min(max_num_partitions, + max(active_num_partitions[0], (seq_len + PARTITION_SIZE - 1) / PARTITION_SIZE)); + if (partition_idx >= effective_num_partitions) return; + const int part_tokens = min(PARTITION_SIZE, seq_len - start_token_idx); + const int lane = threadIdx.x % kWarpSize; + const int warp_idx = threadIdx.x / kWarpSize; + const float score_scale = softmax_scale * k_scale; + __shared__ __half q_shared[6][D]; + __shared__ float scores_shared[6][PARTITION_SIZE]; + __shared__ int block_idx_shared[PARTITION_SIZE]; + __shared__ int block_offset_shared[PARTITION_SIZE]; + for (int i = threadIdx.x; i < 6 * D; i += blockDim.x) + q_shared[i / D][i % D] = q[(i / D) * q_stride1 + i % D]; + for (int i = threadIdx.x; i < part_tokens; i += blockDim.x) { + const int token_idx = start_token_idx + i; + const int logical_block = token_idx / block_size; + block_idx_shared[i] = block_table[logical_block]; + block_offset_shared[i] = token_idx - logical_block * block_size; + } + __syncthreads(); + float local_max[6]; +#pragma unroll + for (int head = 0; head < 6; ++head) local_max[head] = -1.0e20f; + for (int token_local = warp_idx; token_local < part_tokens; + token_local += kWarpsPerBlock) { + const int64_t k_index = + static_cast(block_idx_shared[token_local]) * k_block_stride + + static_cast(block_offset_shared[token_local]) * k_token_stride; + float score[6] = {}; + // Each head retains d=lane,lane+32,... and the original warp reduction. + // Only independent heads share the decoded K value. +#pragma unroll + for (int d = lane; d < D; d += kWarpSize) { + const float kv = flash_v100::load_kv_cache_float_unscaled( + k_cache, k_index + d); +#pragma unroll + for (int head = 0; head < 6; ++head) + score[head] = fmaf(__half2float(q_shared[head][d]), kv, score[head]); + } +#pragma unroll + for (int head = 0; head < 6; ++head) { + const float reduced = warp_reduce_sum(score[head]); + if (lane == 0) { + const float value = reduced * score_scale; + scores_shared[head][token_local] = value; + local_max[head] = fmaxf(local_max[head], value); + } + } + } + float inv_sum[6]; +#pragma unroll + for (int head = 0; head < 6; ++head) { + const float part_max = block_reduce_max(local_max[head]); + float local_sum = 0.f; + for (int i = threadIdx.x; i < part_tokens; i += blockDim.x) { + const float p = __expf(scores_shared[head][i] - part_max); + scores_shared[head][i] = p; + local_sum += p; + } + const float part_sum = block_reduce_sum(local_sum); + inv_sum[head] = part_sum > 0.f ? 1.f / part_sum : 0.f; + if (threadIdx.x == 0) { + const int64_t stats_index = head * stats_stride1 + partition_idx; + max_logits[stats_index] = part_max; + exp_sums[stats_index] = part_sum; + } + // Complete reads of the reduction helper's shared result before the + // next head reuses it; all 256 threads participate in every reduction. + __syncthreads(); + } + for (int d = threadIdx.x; d < D; d += blockDim.x) { + float acc[6] = {}; + for (int i = 0; i < part_tokens; ++i) { + const int64_t v_index = + static_cast(block_idx_shared[i]) * v_block_stride + + static_cast(block_offset_shared[i]) * v_token_stride + d; + const float vv = + flash_v100::load_kv_cache_float_unscaled(v_cache, v_index); +#pragma unroll + for (int head = 0; head < 6; ++head) + acc[head] = fmaf(scores_shared[head][i], vv, acc[head]); + } +#pragma unroll + for (int head = 0; head < 6; ++head) { + const float out_scale = inv_sum[head] * v_scale; + const int64_t tmp_out_base = head * tmp_out_stride1 + + static_cast(partition_idx) * tmp_out_stride2; + tmp_out[tmp_out_base + d] = acc[head] * out_scale; + } + } +} +""" + + +def compact_page_map(partition: str) -> str: + """Keep two page IDs instead of per-token IDs/offsets; preserve FMA order. + + The host validates page3296 and partition1024, so a partition spans at most + two physical pages. Both PV segments visit exactly the original token order. + """ + partition = replace_once( + partition, + " __shared__ int block_idx_shared[PARTITION_SIZE];\n" + " __shared__ int block_offset_shared[PARTITION_SIZE];", + " __shared__ int block_idx_shared[2];\n" + " const int first_block = start_token_idx / block_size;\n" + " const int first_offset = start_token_idx - first_block * block_size;\n" + " const int first_page_tokens = min(part_tokens, block_size - first_offset);", + ) + a = partition.index(" for (int i = threadIdx.x; i < part_tokens;") + b = partition.index(" __syncthreads();", a) + partition = ( + partition[:a] + + """ if (threadIdx.x == 0) { + block_idx_shared[0] = block_table[first_block]; + block_idx_shared[1] = first_page_tokens < part_tokens + ? block_table[first_block + 1] : block_idx_shared[0]; + } +""" + + partition[b:] + ) + partition = replace_once( + partition, + """ const int64_t k_index = + static_cast(block_idx_shared[token_local]) * k_block_stride + + static_cast(block_offset_shared[token_local]) * k_token_stride;""", + """ const bool second_page = token_local >= first_page_tokens; + const int token_offset = second_page ? token_local - first_page_tokens + : first_offset + token_local; + const int64_t k_index = + static_cast(block_idx_shared[second_page]) * k_block_stride + + static_cast(token_offset) * k_token_stride;""", + ) + a = partition.index(" for (int i = 0; i < part_tokens; ++i) {") + b = partition.index("#pragma unroll\n for (int head = 0;", a) + partition = ( + partition[:a] + + """ for (int page = 0; page < 2; ++page) { + const int begin = page == 0 ? 0 : first_page_tokens; + const int end = page == 0 ? first_page_tokens : part_tokens; + int64_t v_index = static_cast(block_idx_shared[page]) * + v_block_stride + (page == 0 ? first_offset * v_token_stride : 0) + d; + for (int i = begin; i < end; ++i, v_index += v_token_stride) { + const float vv = + flash_v100::load_kv_cache_float_unscaled(v_cache, v_index); +#pragma unroll + for (int head = 0; head < 6; ++head) + acc[head] = fmaf(scores_shared[head][i], vv, acc[head]); + } + } +""" + + partition[b:] + ) + return partition + + +def template_function(source: str, name: str) -> tuple[int, str]: + symbol = source.index("void " + name + "(") + start = source.rfind("template <", 0, symbol) + brace = source.index("{", symbol) + depth = 1 + end = brace + 1 + while depth: + depth += (source[end] == "{") - (source[end] == "}") + end += 1 + return start, source[start:end] + + +HOST = r""" +} // namespace +at::Tensor scalar_attention_candidate( + const at::Tensor& q, const at::Tensor& k, const at::Tensor& v, + at::Tensor& out, const at::Tensor& table, const at::Tensor& lengths, + at::Tensor& partial, at::Tensor& maximum, at::Tensor& sums, + const at::Tensor& active, float scale, float k_scale, float v_scale) { + TORCH_CHECK(q.is_cuda() && q.scalar_type() == at::kHalf && + q.sizes() == at::IntArrayRef({1, 6, 256}) && q.is_contiguous(), + "Private scalar q1 probe requires contiguous FP16 [1,6,256]"); + TORCH_CHECK(k.dim() == 4 && k.size(2) == 1 && k.size(3) == 256 && + k.size(1) > 0 && k.scalar_type() == at::kByte && + v.sizes() == k.sizes() && v.scalar_type() == at::kByte, + "Private scalar q1 probe requires E4M3 [pages,page,1,256]"); + for (const auto* t : {&k, &v}) + TORCH_CHECK(t->stride(3) == 1, "Head dimension must be contiguous"); + TORCH_CHECK(out.sizes() == q.sizes() && out.scalar_type() == at::kHalf && + out.is_contiguous(), "Output shape/dtype mismatch"); + TORCH_CHECK(table.dim() == 2 && table.size(0) == 1 && table.size(1) > 0 && + table.scalar_type() == at::kInt && table.is_contiguous() && + table.size(1) * k.size(1) <= 266240 && + lengths.sizes() == at::IntArrayRef({1}) && + lengths.scalar_type() == at::kInt && lengths.is_contiguous() && + active.sizes() == at::IntArrayRef({1}) && + active.scalar_type() == at::kInt && active.is_contiguous(), + "Invalid scalar q1 page/length metadata"); + TORCH_CHECK(partial.sizes() == at::IntArrayRef({1, 6, 256, 256}) && + partial.scalar_type() == at::kFloat && partial.is_contiguous() && + maximum.sizes() == at::IntArrayRef({1, 6, 256}) && + sums.sizes() == maximum.sizes() && + maximum.scalar_type() == at::kFloat && + sums.scalar_type() == at::kFloat && + maximum.is_contiguous() && sums.is_contiguous(), + "Scalar q1 requires unchanged FP32 partition workspaces"); + for (const auto* t : {&k, &v, static_cast(&out), &table, + &lengths, static_cast(&partial), + static_cast(&maximum), + static_cast(&sums), &active}) + TORCH_CHECK(t->device() == q.device(), "Device mismatch"); + TORCH_CHECK(std::isfinite(scale) && std::isfinite(k_scale) && + std::isfinite(v_scale) && k_scale > 0 && v_scale > 0, + "Finite positive KV scales required"); + c10::cuda::CUDAGuard guard(q.device()); + const auto* properties = at::cuda::getCurrentDeviceProperties(); + TORCH_CHECK(properties->major == 7 && properties->minor == 0, "SM70 only"); + const auto stream = at::cuda::getCurrentCUDAStream().stream(); + flash_attention_decode_partition_kernel<256, 1024, + flash_v100::KV_CACHE_DTYPE_FP8_E4M3, 0, false, float> + <<>>( + reinterpret_cast(q.data_ptr()), k.data_ptr(), v.data_ptr(), + partial.data_ptr(), maximum.data_ptr(), sums.data_ptr(), + table.data_ptr(), lengths.data_ptr(), active.data_ptr(), + 1, table.size(1), 256, 6, 1, k.size(1), q.stride(0), q.stride(1), + partial.stride(0), partial.stride(1), partial.stride(2), + maximum.stride(0), maximum.stride(1), + k.stride(0), k.stride(1), k.stride(2), + v.stride(0), v.stride(1), v.stride(2), + scale, k_scale, v_scale, -1, -1, 0, 0, 0, nullptr, 0); + flash_attention_decode_reduce_kernel<256, 1024, 0, float> + <<>>( + partial.data_ptr(), maximum.data_ptr(), sums.data_ptr(), + lengths.data_ptr(), active.data_ptr(), + reinterpret_cast<__half*>(out.data_ptr()), 1, 256, 6, + partial.stride(0), partial.stride(1), partial.stride(2), + maximum.stride(0), maximum.stride(1), out.stride(0), out.stride(1), + 0, 0, 0, 0); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("run", &scalar_attention_candidate); +} +""" + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--source", required=True, type=Path) + parser.add_argument("--output-dir", required=True, type=Path) + parser.add_argument("--pv-unroll", type=int, choices=(1, 2, 4, 8), default=1) + parser.add_argument("--pv-prefetch", action="store_true") + parser.add_argument("--share-kv-six-heads", action="store_true") + parser.add_argument("--e4m3-lut", action="store_true") + parser.add_argument("--compact-page-map", action="store_true") + parser.add_argument( + "--dynamic-shared-bytes", type=int, choices=(0, 4096), default=0 + ) + parser.add_argument("--build", action="store_true") + args = parser.parse_args() + original = args.source.read_text() + start, partition = template_function( + original, "flash_attention_decode_partition_kernel" + ) + _, reduce = template_function(original, "flash_attention_decode_reduce_kernel") + if args.pv_unroll != 1: + partition = replace_once( + partition, + " for (int i = 0; i < part_tokens; ++i) {", + f"#pragma unroll {args.pv_unroll}\n" + " for (int i = 0; i < part_tokens; ++i) {", + ) + if args.pv_prefetch: + if args.pv_unroll != 1: + parser.error("Isolate explicit prefetch from the unroll-hint experiment") + a = partition.index(" for (int i = 0; i < part_tokens; ++i) {") + b = partition.index(" const float out_scale =", a) + partition = ( + partition[:a] + + r""" for (int base = 0; base < part_tokens; base += 4) { + float values[4], weights[4]; +#pragma unroll + for (int stage = 0; stage < 4; ++stage) { + const int i = base + stage; + if (i < part_tokens) { + const int physical_block = block_idx_shared[i]; + const int block_offset = block_offset_shared[i]; + const int64_t v_index = + static_cast(physical_block) * v_block_stride + + static_cast(block_offset) * v_token_stride + + static_cast(kv_head_idx) * v_head_stride + d; + values[stage] = + flash_v100::load_kv_cache_float_unscaled(v_cache, v_index); + weights[stage] = scores_shared[i]; + } + } +#pragma unroll + for (int stage = 0; stage < 4; ++stage) { + if (base + stage < part_tokens) + acc = fmaf(weights[stage], values[stage], acc); + } + } +""" + + partition[b:] + ) + host = HOST + if args.share_kv_six_heads: + if args.pv_prefetch or args.pv_unroll != 1: + parser.error("Screen shared-head KV reuse independently") + partition = partition[: partition.index("{")] + SHARED_HEADS_BODY + host = replace_once(host, "<<(threadIdx.x));\n" + " __shared__ __half q_shared[6][D];", + ) + partition = replace_once( + partition, + "flash_v100::load_kv_cache_float_unscaled(\n" + " k_cache, k_index + d)", + "kv_lut[static_cast(k_cache)[k_index + d]]", + ) + partition = replace_once( + partition, + "flash_v100::load_kv_cache_float_unscaled(v_cache, v_index)", + "kv_lut[static_cast(v_cache)[v_index]]", + ) + if args.dynamic_shared_bytes: + if not (args.compact_page_map and args.e4m3_lut): + parser.error("Shared-memory occupancy tuning requires the compact LUT") + # No data is stored here. The extra reservation tests the two-block + # resource limit against the compact layout's three-block limit. + host = replace_once( + host, + "<<>>", + f"<<>>", + ) + source = original[:start] + partition + "\n" + reduce + "\n" + host + directory = args.output_dir.resolve() + if directory.exists(): + parser.error("Use a new directory for each native variant") + sources = directory / "sources" + root = args.source.parent.parent + for sub in ("kernel", "include"): + target = sources / sub + target.mkdir(parents=True) + for pattern in ("*.h", "*.cuh"): + for header in (root / sub).glob(pattern): + shutil.copy2(header, target) + shutil.copy2(root / "LICENSE", sources) + path = sources / "kernel/scalar-attention.cu" + path.write_text(source) + digest = hashlib.sha256(path.read_bytes()).hexdigest() + name = "sm70_scalar_attention_" + digest[:12] + flags = [ + "-O3", + "-std=c++17", + "-gencode=arch=compute_70,code=sm_70", + "-U__CUDA_NO_HALF_OPERATORS__", + "-U__CUDA_NO_HALF_CONVERSIONS__", + "-U__CUDA_NO_HALF2_OPERATORS__", + "--expt-relaxed-constexpr", + "--expt-extended-lambda", + "--use_fast_math", + "-lineinfo", + "-Xptxas=-v", + ] + manifest = dict( + input_source_sha256=hashlib.sha256(args.source.read_bytes()).hexdigest(), + source_sha256=digest, + module_name=name, + pv_unroll=args.pv_unroll, + pv_prefetch=args.pv_prefetch, + share_kv_six_heads=args.share_kv_six_heads, + e4m3_lut=args.e4m3_lut, + compact_page_map=args.compact_page_map, + dynamic_shared_bytes=args.dynamic_shared_bytes, + max_context=262144, + source_files={ + str(p.relative_to(sources)): hashlib.sha256(p.read_bytes()).hexdigest() + for p in sorted(sources.rglob("*")) + if p.is_file() + }, + extra_cuda_cflags=flags, + scope="Private scalar q1 fallback screen; no serving route", + ) + if args.build: + from torch.utils.cpp_extension import load + + build = directory / "build" + build.mkdir() + library = Path( + load( + name=name, + sources=[str(path)], + build_directory=str(build), + extra_cuda_cflags=flags, + extra_include_paths=[str(sources / "kernel"), str(sources / "include")], + verbose=True, + ).__file__ + ) + manifest.update( + library=str(library), + library_sha256=hashlib.sha256(library.read_bytes()).hexdigest(), + ) + (directory / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n") + print(json.dumps(manifest, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/sm70_dflash2_scalar_q1_probe.py b/benchmarks/sm70_dflash2_scalar_q1_probe.py new file mode 100644 index 0000000000..fa631420a7 --- /dev/null +++ b/benchmarks/sm70_dflash2_scalar_q1_probe.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Explicit worker-extension probe for an independently built long q1 operator. + +Import only with VLLM_SM70_SCALAR_Q1_PROBE_MANIFEST set. Both A/B arms retain +this wrapper, and only eligible eager calls select the candidate. It installs +no default serving route and never reads a GPU sequence length on the host. +""" + +import functools +import hashlib +import os +from collections import Counter +from pathlib import Path + +import flash_attn_v100_cuda as native +import torch + +from benchmarks.kernels.benchmark_sm70_grouped_attention_long import load_operator +from vllm.forward_context import get_forward_context, is_forward_context_available +from vllm.v1.attention.backends.flash_attn_v100 import FlashAttnV100Impl + +operator, manifest = load_operator( + Path(os.environ["VLLM_SM70_SCALAR_Q1_PROBE_MANIFEST"]) +) +assert manifest["share_kv_six_heads"] and manifest["max_context"] == 262144 +MODE = "control" +_ELIGIBLE = False +SEEN = Counter() +HITS = Counter() +_original_call = FlashAttnV100Impl._call_flash_attn_decode_paged +_original_native = native.decode_paged_fwd + + +def cpu_context_bound(kwargs): + bound = kwargs.get("max_seq_len_hint") + if type(bound) is int: + return bound + if is_forward_context_available(): + metadata = get_forward_context().attn_metadata + if isinstance(metadata, dict): + bounds = [getattr(m, "max_seq_len", None) for m in metadata.values()] + bounds = [b for b in bounds if type(b) is int and b > 0] + if bounds: + return max(bounds) + return None + + +@functools.wraps(_original_call) +def call(self, query, key_cache, value_cache, *args, **kwargs): + global _ELIGIBLE + previous = _ELIGIBLE + _ELIGIBLE = False + if ( + query.shape == (1, 6, 256) + and key_cache.dtype == torch.uint8 + and not torch.cuda.is_current_stream_capturing() + ): + bound = cpu_context_bound(kwargs) + SEEN[str(bound)] += 1 + if type(bound) is int and 131072 <= bound <= 262144: + _ELIGIBLE = MODE == "candidate" + try: + return _original_call(self, query, key_cache, value_cache, *args, **kwargs) + finally: + _ELIGIBLE = previous + + +def run(*args, **kwargs): + if _ELIGIBLE and not kwargs and len(args) == 20: + q, k, v, out, table, lengths, partial, maximum, sums, active = args[:10] + ( + scale, + partition, + count, + dtype, + k_scale, + v_scale, + left, + right, + anchor, + window, + ) = args[10:] + if ( + q.shape == (1, 6, 256) + and k.shape[1] == 3296 + and k.shape[2:] == (1, 256) + and k.dtype == v.dtype == torch.uint8 + and partition == 1024 + and count == 256 + and dtype in ("fp8", "fp8_e4m3") + and left == right == -1 + and anchor is None + and window == 0 + and partial.shape == (1, 6, 256, 256) + and partial.dtype == torch.float32 + and maximum.shape == sums.shape == (1, 6, 256) + ): + HITS["eager_long_q1"] += 1 + return operator(*args[:10], scale, k_scale, v_scale) + return _original_native(*args, **kwargs) + + +FlashAttnV100Impl._call_flash_attn_decode_paged = call +native.decode_paged_fwd = run + + +class ScalarQ1ProbeExtension: + def scalar_attention_switch(self, mode): + global MODE + assert mode in ("control", "candidate") + MODE = mode + return self.scalar_attention_snapshot() + + def scalar_attention_snapshot(self): + result = dict( + rank=torch.distributed.get_rank(), + mode=MODE, + seen_cpu_context_bounds=dict(SEEN), + hits=dict(HITS), + manifest=manifest, + native_library=str(Path(native.__file__).resolve()), + native_library_sha256=hashlib.sha256( + Path(native.__file__).read_bytes() + ).hexdigest(), + scope="Explicit eager long q1 probe; other calls retain the original entry", + ) + return result diff --git a/benchmarks/sm70_dflash2_state_audit.py b/benchmarks/sm70_dflash2_state_audit.py index d796483c19..9347d2071c 100644 --- a/benchmarks/sm70_dflash2_state_audit.py +++ b/benchmarks/sm70_dflash2_state_audit.py @@ -49,6 +49,24 @@ def cpu_request_slots(values: torch.Tensor, indices: torch.Tensor) -> torch.Tens return values.index_select(0, indices.to(torch.int64)).detach().cpu().clone() +def target_auxiliary_states(owner, batch) -> list[torch.Tensor]: + """Find the matching runner frame through opt-in sampling wrappers.""" + frame = sys._getframe(1) + try: + for _ in range(12): + values = frame.f_locals + if values.get("self") is owner and values.get("input_batch") is batch: + auxiliary = values.get("aux_hidden_states") + if auxiliary is not None: + return auxiliary + frame = frame.f_back + if frame is None: + break + finally: + del frame + raise RuntimeError("Natural audit requires matching target auxiliary states") + + def install() -> None: from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn as gd from vllm.model_executor.models import qwen3_next as qn @@ -358,10 +376,7 @@ def cpu(tensor): # token IDs or acceptance decisions. Still diagnostic-only: CPU # snapshots synchronize execution and cannot measure performance. result["control"] = "natural_sampling" - frame = sys._getframe(1) - aux = frame.f_locals.get("aux_hidden_states") - if aux is None: - raise RuntimeError("Natural audit requires target auxiliary states") + aux = target_auxiliary_states(self, batch) result["aux_hidden_states"] = [cpu(t) for t in aux] if batch.num_draft_tokens: result["draft_logits"] = cpu( diff --git a/benchmarks/sm70_dflash2_state_layout.py b/benchmarks/sm70_dflash2_state_layout.py new file mode 100644 index 0000000000..d01ba8e6e2 --- /dev/null +++ b/benchmarks/sm70_dflash2_state_layout.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Explain physical-slot and unused conv-storage differences in raw captures. + +This never edits captures or replaces the raw byte comparison. The caller must +provide the convolution width from the frozen model. Unknown layouts fail +closed. All values consumed by convolution, and all verifier output storage, +remain subject to the exact comparison. +""" + +from __future__ import annotations + +import torch + + +def _field(states: dict, name: str) -> torch.Tensor: + matches = [v for k, v in states.items() if k.startswith(name + ":")] + if len(matches) != 1: + raise ValueError(f"Missing or ambiguous state metadata: {name}") + return matches[0] + + +def check_slot_mapping( + left: dict, right: dict, mapping: dict[int, int], reverse: dict[int, int] +) -> None: + """Require one consistent bijection across every observed state access.""" + if left.keys() != right.keys(): + raise ValueError("State snapshot coverage differs") + for label, a in left.items(): + name = label.split(":")[0] + if not name.endswith(("/indices", "/slot_table")): + continue + b = right[label] + if a.shape != b.shape or a.dtype != b.dtype or a.is_floating_point(): + raise ValueError(f"Invalid slot metadata: {label}") + for x, y in zip(a.reshape(-1).tolist(), b.reshape(-1).tolist()): + if x < 0 or y < 0: + if x != y: + raise ValueError("Padding slot changed into a live slot") + continue + if mapping.get(x, y) != y or reverse.get(y, x) != x: + raise ValueError("Inconsistent or aliased physical-slot mapping") + mapping[x], reverse[y] = y, x + + +def explain_state_difference( + label: str, left: dict, right: dict, conv_width: int +) -> str | None: + """Classify only source-defined unused storage; retain every raw mismatch.""" + if not 2 <= conv_width <= 6: + raise ValueError("Unsupported convolution width") + name = label.split(":")[0] + if name.endswith(("/indices", "/slot_table")): + # check_slot_mapping must already have validated the complete captures. + return "bijective_physical_slot_renaming" + if "/conv/" not in name or not name.endswith("/values"): + return None + prefix, suffix = name.split("/conv/", 1) + a, b = left[label], right[label] + if a.shape != b.shape or a.dtype != b.dtype or a.ndim != 3: + raise ValueError("Unsupported convolution state layout") + n, _, storage = a.shape + history = conv_width - 1 + if storage < history: + raise ValueError("Convolution state smaller than its history") + validity = [ + _field(side, name.rsplit("/", 1)[0] + "/valid").reshape(-1) + for side in (left, right) + ] + if not torch.equal(*validity) or validity[0].numel() != n: + raise ValueError("Convolution slot validity differs") + columns = torch.arange(storage).reshape(1, -1) + active = columns < history + if prefix.startswith("prefill/"): + if suffix == "input_state/values": + initial = [ + _field(side, prefix + "/conv/has_initial_state").reshape(-1) + for side in (left, right) + ] + if not torch.equal(*initial) or initial[0].numel() != n: + raise ValueError("Convolution initial-state contract differs") + active = active & initial[0].reshape(-1, 1) + elif suffix != "output_state/values": + return None + # causal_conv1d_fn fixes state_len = KERNEL_WIDTH - 1. It does not + # initialize the extra speculative history columns in the allocation. + elif prefix.startswith("verify/") and suffix == "input_state/values": + selectors = [ + _field(side, prefix + "/conv/num_accepted_tokens").reshape(-1)[:n] + for side in (left, right) + ] + if not torch.equal(*selectors) or selectors[0].numel() != n: + raise ValueError("Convolution acceptance selectors differ") + offset = selectors[0] - 1 + valid = validity[0] + if ((offset < 0) | (offset + history > storage))[valid].any(): + raise ValueError("Convolution selector outside stored history") + # Both the convolution and the rolling-state copy read within this + # window; output storage is always compared in full, including q8. + active = (columns >= offset[:, None]) & (columns < offset[:, None] + history) + else: + return None + mask = (active & validity[0].reshape(-1, 1))[:, None, :].expand_as(a) + aa, bb = a[mask].contiguous(), b[mask].contiguous() + if not torch.equal(aa.view(torch.uint8), bb.view(torch.uint8)): + return None + return "unused_convolution_storage" diff --git a/benchmarks/summarize_sm70_dflash2_long_curve.py b/benchmarks/summarize_sm70_dflash2_long_curve.py new file mode 100644 index 0000000000..073d82d65f --- /dev/null +++ b/benchmarks/summarize_sm70_dflash2_long_curve.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Freeze or evaluate a repeated, unprofiled DFlash2 context-cost curve.""" + +import argparse +import hashlib +import json +import statistics +from pathlib import Path + +import numpy as np + +LENGTHS = (1024, 32768, 65536, 131072) +INTERVALS = ((32768, 65536), (65536, 131072), (32768, 131072)) + + +def summarize(paths): + if len(paths) != 3 or len(set(p.resolve() for p in paths)) != 3: + raise ValueError("Three distinct startup reports are required") + reports = [json.loads(p.read_text()) for p in paths] + startup_ids = { + tuple( + row["pid"] for row in sorted(r["initial_routes"], key=lambda x: x["rank"]) + ) + for r in reports + } + if len(startup_ids) != 3: + raise ValueError("Reports do not identify three independent worker startups") + contract_keys = ("sampling", "corpus_sha256", "context_capacity") + contract = {key: reports[0][key] for key in contract_keys} + for report in reports: + assert report["complete"] and not report["profiler"] + assert report["require_native_prefill"] + assert report["require_original_gdn_prefill"] + assert {key: report[key] for key in contract_keys} == contract + assert {r["prompt_tokens"] for r in report["cases"]} == set(LENGTHS) + result = { + "contract": contract, + "source_reports": { + str(p): hashlib.sha256(p.read_bytes()).hexdigest() for p in paths + }, + "rows": {}, + } + for length in LENGTHS: + cold, measured, startup_rows = [], [], [] + for report in reports: + rows = [r for r in report["cases"] if r["prompt_tokens"] == length] + warmup = [r for r in rows if r["warmup"]] + repeats = [r for r in rows if not r["warmup"]] + assert len(warmup) == 1 and len(repeats) == 5 + assert warmup[0]["prefill_computed_tokens"] == length + assert sorted(r["repeat"] for r in repeats) == list(range(5)) + assert all(r["prompt_sha256"] == rows[0]["prompt_sha256"] for r in rows) + cold.append(warmup[0]["engine_prefill_s"]) + measured.extend(repeats) + startup_rows.append(rows) + row = { + "cold_prefill_seconds": cold, + "cold_prefill_tps": length / statistics.median(cold), + "startup_count": 3, + "measured_request_count": len(measured), + "within_startup_same_tokens": all( + all(r["token_ids"] == rows[0]["token_ids"] for r in rows) + for rows in startup_rows + ), + "between_startup_same_tokens": all( + rows[0]["token_ids"] == startup_rows[0][0]["token_ids"] + for rows in startup_rows + ), + "finish_reasons": sorted({r["finish_reason"] for r in measured}), + } + assert len({rows[0]["prompt_sha256"] for rows in startup_rows}) == 1 + for metric in ( + "complete_round_ms", + "pure_decode_tps", + "accepted_drafts_per_round", + "emitted_tokens_per_round", + "accepted_over_proposed", + "ttft_s", + ): + values = [r[metric] for r in measured] + row[metric] = float(statistics.median(values)) + row[metric + "_request_distribution"] = dict( + zip(("p50", "p90", "p99"), np.percentile(values, [50, 90, 99]).tolist()) + ) + result["rows"][str(length)] = row + result["percentile_scope"] = ( + "Distributions across request averages; not individual GPU-round percentiles" + ) + return result + + +def context_increments(rows): + result = {} + for a, b in INTERVALS: + start = rows[str(a)]["complete_round_ms"] + end = rows[str(b)]["complete_round_ms"] + result[f"{a}:{b}"] = { + "additional_round_ms": end - start, + "round_ms_per_1024_context_tokens": (end - start) / ((b - a) / 1024), + "round_growth_ratio": end / start, + } + return result + + +def compare_performance(result, baseline): + """Require absolute/increment improvement; prefill ratios are context only.""" + assert result["contract"] == baseline["contract"] + checks = {} + for length in LENGTHS: + current = result["rows"][str(length)]["complete_round_ms"] + previous = baseline["rows"][str(length)]["complete_round_ms"] + checks[f"absolute_{length}"] = ( + current <= previous if length == 1024 else current < previous + ) + previous_increments = context_increments(baseline["rows"]) + current_increments = context_increments(result["rows"]) + for interval, current in current_increments.items(): + previous = previous_increments[interval] + current["baseline_additional_round_ms"] = previous["additional_round_ms"] + current["increment_reduction_ms"] = ( + previous["additional_round_ms"] - current["additional_round_ms"] + ) + checks[f"increment_{interval}"] = current["increment_reduction_ms"] >= 0 + result["context_increments"] = current_increments + result["curve_checks"] = checks + result["performance_curve_passed"] = all(checks.values()) + result["prefill_growth_reference"] = { + f"{a}:{b}": baseline["rows"][str(a)]["cold_prefill_tps"] + / baseline["rows"][str(b)]["cold_prefill_tps"] + for a, b in INTERVALS + } + result["prefill_growth_is_admission_gate"] = False + result["quality_and_acceptance_admission"] = "Requires separate paired evidence" + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--reports", nargs=3, type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--baseline", type=Path) + args = parser.parse_args() + if args.output.exists(): + raise FileExistsError("Refusing to overwrite a frozen curve") + result = summarize(args.reports) + result["objective"] = "Reduce absolute round cost and long-context increments" + result["context_increments"] = context_increments(result["rows"]) + if args.baseline is None: + result["prefill_growth_reference"] = { + f"{a}:{b}": result["rows"][str(a)]["cold_prefill_tps"] + / result["rows"][str(b)]["cold_prefill_tps"] + for a, b in INTERVALS + } + else: + baseline = json.loads(args.baseline.read_text()) + compare_performance(result, baseline) + result["baseline_sha256"] = hashlib.sha256( + args.baseline.read_bytes() + ).hexdigest() + args.output.write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(result, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/docs/design/sm70_dflash2_long_verify_curve_20260909.md b/docs/design/sm70_dflash2_long_verify_curve_20260909.md new file mode 100644 index 0000000000..80940523db --- /dev/null +++ b/docs/design/sm70_dflash2_long_verify_curve_20260909.md @@ -0,0 +1,928 @@ +# DFlash2 long-context verification curve + +## Frozen scope + +Integration base: `80545c010bbf6f5ed06458d992c189d75d0eff8f`. The task includes +the diagnostic-only changes from PR #586 at +`720457ff4a8f6361e8160bd1cc9210d327d9b8de`. Runtime numerical/performance +comparisons use a frozen source and native-library manifest, physical GPUs +4–7, TP4/B1/q8, QUASAR target revision +`d8e6fbfa3e3a78899b440222b827430045a05b44`, DFlash2 draft revision +`dedf8df68adfb1afeaf7b7480c0a0243108177b4`, CUDA 12.8, Torch 2.10.0+cu128, +E4M3 target KV, FP32 logits/state and the existing compensated attention. +Original FlashQLA GDN prefill and the verified FA2 sidecar stay enabled. +Capacity remains 262144. The frozen first campaign stops at 131072 input +tokens; the subsequent target revision below explicitly adds the 256K tier. + +The user revised the objective on September 10: continuously reduce absolute +complete-round cost and the incremental cost of longer contexts. Prefill +growth ratios are reference observations, not admission limits or a stopping +condition. Every long-context absolute round cost must improve, 1K must not +regress, and context increments must not increase. Report both additional +milliseconds and milliseconds per 1024 additional context tokens. Never +improve a ratio by slowing the shorter point, prefill or acceptance. + +The subsequent September 10 target revision sets explicit complete-round +latency goals: + +| Context tier | Complete-round target | +| --- | ---: | +| 32K | <= 17 ms | +| 64K | <= 18 ms | +| 128K | < 20 ms | +| 256K | < 22 ms | + +1K must not regress and the <15 ms short-context goal remains. Output quality, +compensation and acceptance requirements are unchanged. The 256K target +supersedes the previous instruction to stop all measurements at 128K; it does +not qualify any untested kernel range. Keep the current 132096-token serving +gate until the longer operator and model checks pass. Existing frozen reports +remain immutable. Within 262144 service capacity, a boundary-window performance +probe constructs a separate 261888-token prompt and reserves 256 output tokens; +report that exact input length and the actual sampled context range. Do not +truncate an existing prompt, label this as a full 262144-token cold prefill, or +silently raise model capacity. Operator checks separately include 262144 and +the speculative q8 boundary headroom. + +## Implementation order + +1. Repeat the restored-path baseline, one verified cold request and five warm + requests per length/startup. Retain source/library hashes and all-rank route + evidence. Supplement the existing 64K q8 trace with a 128K trace. +2. Independently screen aligned E4M3 vector loads and QK live-range/unrolling + changes, preserving the 80-split, N32, K16 compensation and FP32 partial + contract. Require byte-exact output and full partial/max/sum buffers. +3. Develop overlapping loads only after these measurements. If necessary, + evaluate 80/160/320 splits, grouped-head KV reuse and versioned workspaces. + Any graph specialization belongs in the actual MRV2 graph manager. + +Model gates retain fixed-prefix distributions, acceptance and natural-output +checks. Arithmetic variants require independent FP64 error measurements before +the final FP16 cast. Never remove precision compensation as an optimization. +No experimental route is enabled by importing its builder or benchmark. + +## Progress and artifacts + +Three baseline startups completed 72 requests, with identical token IDs and +acceptance across and within startups at every length. The baseline median +complete rounds at 1K/32K/64K/128K are 16.350/25.821/35.047/53.267 ms. +Cold-prefill throughput is 3525/4063/3640/3009 tokens/s. The frozen growth +reference ratios are 1.1163476857, 1.2095236995 and 1.3502489828 for 32→64, +64→128 and 32→128. The original frozen report remains immutable even though +these ratios no longer gate admission. First-use prefill overhead affects the +first 1K request; medians use +the three independent cold observations, and 1K is not the curve denominator. + +The curve reporter verifies distinct worker startups, the full computed-token +count for cold requests and five measured requests per length. Its percentiles +describe request-average round costs, not individual GPU-round latency. It +refuses to overwrite a frozen curve and checks absolute costs and context +increments. A candidate can now pass despite exceeding the prefill ratio; +slowing an anchor still fails the absolute-cost gate. + +The separate 128K trace measures 36.849 ms target grouped attention per +rank/round, versus about 7.317 ms of QPN2 projections, 0.878 ms of draft +attention and 0.678 ms GDN recurrence. Critical-rank interval is 54.871 ms and +GPU union is 52.899 ms. The NCU probe returns `ERR_NVGPUCTRPERM`; no hardware +counter claim is made. The separate 32K trace is complete: its critical-rank +interval/GPU union is 27.615/25.584 ms. Target attention accounts for almost +all additional GPU time at 128K; draft attention and GDN recurrence remain +approximately 0.88 and 0.68 ms respectively. These are profiler observations, +not the unprofiled acceptance costs. + +Initial operator candidates preserve the 80-split/N32/K16 compensation +contract. With sixteen distinct layer KV allocations, paired GPU-graph +measurements give: + +| Candidate | 32K attention, ms | 64K attention, ms | 128K attention, ms | Byte checks | +| --- | ---: | ---: | ---: | --- | +| Frozen three-group control | 9.232 | 17.985 | 35.453 | Reference | +| Guarded 16-byte KV loads | 6.012 | 11.540 | 22.744 | 50/50 | +| QK unroll1 | 7.529 | 14.379 | 28.264 | 50/50 | +| QK unroll4 | 9.319 | 17.990 | 35.437 | 50/50; no stable gain | + +The combined screen independently compares unroll1, vector loads plus +unroll1, two padded three-head groups plus unroll1, and one six-head group +plus unroll1 with the frozen control. All 200 output/full-workspace/canary +cases match. At 128K the control costs 35.458 ms; the respective candidates +cost 28.258/18.474/17.383/12.766 ms. The six-head candidate's 32K/64K costs +are 3.497/6.592 ms. These remain operator measurements, not complete-round +gains or model admission. In particular, some unroll/grouping variants slow +the 1K operator and cannot replace the short path without further evidence. + +The initial multi-candidate loader exposed a native-module alias: two DSOs +used the same module name, and CPython returned the first module for both. +The second candidate's initial results are withdrawn and the report marked +invalid. Source-derived native module names, actual loaded-file verification +and distinct-callable checks now prevent this failure. The table above uses +the corrected independent bindings. A CPU regression check rejects the old +aliased pair before any GPU work. + +QK unroll1 reduces the compiled register count from 234 to 94 without spills; +unroll4 uses 140 and does not gain stable speed. These are compiler resource +observations, not achieved occupancy. Follow-ups evaluate unroll2 and constant +1648/3296-page addressing with the same arithmetic. Other page sizes and +8-byte-only strides retain their existing address/load implementations. + +Combining six-head reuse, vector loads and unroll1 reduces the sixteen-layer +128K attention cost to 9.959 ms. A disjoint V panel lets otherwise idle QK +warps load and convert values while the six QK warps compute; it reduces that +cost further to 9.092 ms (32K/64K: 2.557/4.736 ms). Both candidates pass all +50 byte-exact output, FP32 partial/max/sum and canary cases. V prefetch uses +73728 bytes of dynamic shared memory with an explicit device opt-in limit +check. It preserves the existing CTA barriers and online-softmax warp barrier. +Page-specialized addressing and unroll2 together reach 9.218 ms without V +prefetch; these independent results determine which combinations to test. + +The system sanitizer executable failed before testing because its injection +library was absent. Reruns use the previously validated, complete CUDA 12.8 +sanitizer bundle. The vector-only path passes memcheck, racecheck and +synccheck with zero errors/warnings. Its first unprofiled startup gives +16.133/22.068/27.824/38.765 ms complete rounds at 1K/32K/64K/128K, with identical +token IDs, finish reasons and acceptance to the frozen control. All four +ranks capture the candidate in 16 actual target attention calls. This is an +independent screen, not three-startup acceptance. It failed the superseded +prefill-ratio gate. Subsequent V-prefetch and qk2/page/V-prefetch candidates +each passed their own memory/race/synchronization checks before model screening. + +The actual model capture uses a 3296-token KV page with strides +`(1687552, 256, 256, 1)`. The initial sixteen-layer performance screen used +1648-token pages; both are correctness cases, and `--performance-page` allows +timing the exact captured page geometry. The loader also retains the +8-byte-only stride fallback. No context result is extrapolated to 256K. + +Private launch/build manifests and raw results are retained in the task +artifact archive; generated libraries and private cache paths are excluded +from Git. The new serving route remains explicitly opt-in and has not been +enabled by default or merged. + +## September 10 implementation and rejected candidates + +The qk2/page/V-prefetch operator takes 2.401/4.417/8.444 ms for sixteen +independent layer allocations at 32K/64K/128K with actual 3296-token pages. +The corresponding initial model startup gives 16.143/18.521/20.484/24.452 ms +at 1K/32K/64K/128K. Tokens and acceptance match the frozen baseline. This +still requires integrated-route quality and repeated-startup admission. + +Next-K prefetch preserves byte-exact output and full FP32 workspace, but +both tested load/softmax warp partitions lose performance. Eight load warps +cost 8.815 ms at 128K versus the 8.439 ms paired control; four load warps +cost 9.394 versus 8.444 ms. The extra warp-role work outweighs the overlap. +Neither is selected for serving. The extended byte gate covers 132096 visible +tokens, providing bounded generation headroom after a 128K input. + +Increasing to 160/320 splits also loses: 128K costs 9.002/9.897 ms versus +8.469 ms for 80. The independent FP64 screen records final-FP16 max error +0.001952 for all three at 128K; this does not replace a native pre-cast FP32 +audit or model admission. Arithmetic variants remain rejected and disabled. + +`VLLM_SM70_E4M3_LONG_ATTENTION_MANIFEST` enables the experimental loader and +an additional MRV2 B1/q8 graph. The loader verifies the actual DSO SHA and +module identity, accepts the 80-split six-head workspace, and retains native +input validation. Eligibility depends on q8/GQA6/D256 E4M3 tensors and +validated 1648/3296 page layouts, not target weight quantization or model name. +The descriptor carries a 132096-token upper bound; replay chooses it from +the existing CPU sequence-length upper bound. Larger bounds and other shapes +retain the full-context graph. No device-to-host length read is introduced. +Workspaces are fixed per operator source SHA, capacity, device and CUDA stream. +CPU tests cover the boundary, fallback, switching back, missing captures and +refusal to read a device hint. Native graph-switch/model gates remain pending. + +The integrated serving source `239d71c7100b3bce5526268be2cafb4cff8ba8f2` +uses candidate source SHA +`8459d57c6b72993ba47f5c3fe3953bd8343c4174974a05f329984e1ef070f738` and DSO SHA +`dac8262f3d023ce35f1618bbd6fe0f569993f1e9e1f2a60e8291880f76a97339`. +The first unprofiled integrated startup gives 16.078/18.382/20.356/24.326 ms +at 1K/32K/64K/128K, with pure decode 293.7/210.2/237.9/177.7 tokens/s. +All 24 request token sequences, acceptance counts and finish reasons match +the frozen control. This is still a single-startup screen. Three independent +paired startups and the frozen seeds 0/1/2 natural-output campaign follow. +The selected DSO also passes expanded memcheck, racecheck and synccheck +coverage, including the actual 3296-token pages, with zero errors/warnings. + +The first fixed-prefix diagnostic completes its 1K control/candidate/control +captures, then exhausts GPU memory during the 32K warmup. Snapshot buffers +grow after the initial memory profile; this instrumented failure is not a +performance result. The retry reserves additional diagnostic memory by using +GPU memory utilization 0.6; capacity stays 262144 and uninstrumented performance +runs keep 0.8. The partial captures and failed report remain in the archive. + +For the completed 1K captures, all full-vocabulary logits, distributions, +top-p support, top-1 and EOS probabilities are exact in both A/B and A/A. +There are 320 raw intermediate mismatches in each pair. Every mismatch is +either a bijective physical-slot renaming or unused convolution storage: +prefill writes only `kernel_width - 1` history columns; the verifier reads the +window beginning at `num_accepted_tokens - 1`. With no initial prefill state, +the old convolution allocation is not read. All verifier output storage is +compared in full. The offline comparer retains raw differences and separately +reports their explanations; it rejects changed live history, invalid selectors, +padding-to-live changes and inconsistent or aliased slot mappings. It must not +use repeated-run TV as a numerical tolerance. EOS IDs come from the frozen +generation configuration, not another tokenizer's constants. + +## Repeated integrated results and the expanded target + +Three independent paired startups complete 144 requests (72 per arm). Each +startup includes one cold warmup and five measured requests per context/arm, +with reversed arm order in the second startup. All paired token sequences, +finish reasons and acceptance records match. The unprofiled request-median +results are: + +| Context | Paired control, ms | Candidate, ms | Candidate pure decode, tokens/s | Accepted drafts/round | Emitted tokens/round | +| --- | ---: | ---: | ---: | ---: | ---: | +| 1K | 16.210 | 16.130 | 292.76 | 3.778 | 4.741 | +| 32K | 25.553 | 18.371 | 210.32 | 2.894 | 3.879 | +| 64K | 34.701 | 20.418 | 237.20 | 3.863 | 4.863 | +| 128K | 52.826 | 24.451 | 176.76 | 3.339 | 4.339 | + +The candidate passes absolute-cost and incremental-cost checks against both +the original frozen curve and the new paired controls. The 32K-to-64K increment +is 2.047 ms, or 0.06397 ms per additional 1024 tokens; 64K-to-128K adds 4.033 ms, +or 0.06302 ms per 1024. These results have not reached the new 17/18/20 ms +targets. The complete natural-output campaign remains a separate admission. + +The diagnostic retry completes all six fixed tapes: 1K, 32K, 64K, 128K, MBPP28 +and MBPP3. Each arm has 96 all-rank target snapshots. All native logits, +full/sampled distributions, support sets, top-1 and EOS probabilities are +exact in control/candidate and repeated-control comparisons. The 1768/1752 +raw state differences are fully explained by the validated storage layout; +no live-state or other unexplained differences remain. Diagnostic GPU memory +utilization 0.6 provides 451076 KV token slots, exceeding the unchanged +262144 service capacity. These dumps do not contribute performance samples. + +The integrated 128K trace now attributes 8.615 ms to target attention and +7.387 ms to the three QPN2 projection categories, averaged across ranks. Draft +proposal GPU service is 3.919 ms. The same fixed rank 0 has a 26.419 ms round +interval and 24.423 ms GPU union. Across critical ranks, actual profiled round +p50/p90/p99 are 26.524/26.631/26.754 ms. Those are individual **profiled** +intervals and must not replace the unprofiled request-average distribution. + +The first expanded operator screen passes 45 byte-exact output/full-FP32- +workspace/canary cases, including 262152-token physical-page and stride +boundaries. At 261888 tokens the sixteen-layer attention working set takes +70.344 ms for the frozen control and 16.335 ms for the selected one-stage +candidate. This is not a 256K model result. The restored serving control's +single cold/warm screen gives 2265.5 cold-prefill tokens/s and a 95.694 ms +warmed complete round; the bounded experimental graph is deliberately not +selected beyond its current admitted domain. Repeated 256K model acceptance +still needs a separately validated extended serving route. + +A new private feasibility builder separates compensated QK production from +two disjoint PV column partitions. It retains 80 logical context partitions, +K16 compensation, N32 online updates and probability residual products. The +QK producer uses 80 registers and 48384 shared-memory bytes; the PV consumer +uses 102 registers and a 25248-byte shared layout, without spills. These are +compiler resource observations, not achieved occupancy. Extra score storage, +kernel boundaries and repeated softmax work may erase the benefit, so byte +checks and complete-working-set timing decide whether to continue. Prototype +score scratch is capture-owned; no serving route is installed by this builder. +Its first screen preserves output and complete partial/max/sum bytes in all +65 cases, including the expanded boundary (130 checks across the one-stage +and staged candidates). It is slower: 32K/64K/128K/261888-token attention costs +3.162/5.849/11.194/21.714 ms versus 2.401/4.417/8.451/16.341 ms for the paired +one-stage candidate. The staged source SHA is +`d8493061867f9d044ce7a70e2298306816d0f283aaa84984c69031b668aee825`; DSO SHA is +`02b2a3080c99ad1c837e98fda1222eeb25fefca71ca97e23b939122f60b32f2d`. +It is rejected for serving. A bounded operator trace separates producer and +consumer costs before any follow-up; extra parallelism alone is not a gain. +The selected one-stage DSO separately passes extended-boundary memcheck, +racecheck and synccheck with zero errors. + +Prior rejected experiments remain recorded in the context-cost and long-verify +worklogs. Historical E5M2 and FP16-partial Pack-GQA timings are design references, +not quality/performance evidence for this E4M3 FP32 path. + +## Expanded-domain model checks and terminal scheduling + +A private process-only extension raises the experimental graph bound to 262152 +before capture. The public source remains bounded at 132096, and service +capacity remains 262144. The original selected DSO completes one startup with +one cold warmup and five measured requests per arm at 261888 input tokens plus +256 output tokens. Median complete cost is 95.419 ms for the control and +38.602 ms for the candidate. All token IDs, acceptance and finish reasons +match. Accepted drafts/round are 4.02 and emitted tokens/round are 5.12. +This is a capacity-bound performance probe, whose responses reach the output +limit; it is not a natural-EOS quality case or three-startup acceptance. + +The 256K fixed-prefix control/candidate/control diagnostic also has exact +logits, TV, support and EOS probabilities. Its 272 raw differences per +comparison are explained storage differences, with zero unexplained changes. +The separate natural-output campaign retains exact paired 10107-token +HumanEval and 75828-token LiveCodeBench responses, both ending naturally. +All twelve structured pairs (JSON, schema, one tool, parallel tools; seeds +0/1/2) have exact tokens/acceptance, valid structure and natural termination. +The other sixteen code pairs are separate resumable jobs. The first of those +also completes an exact 51562-token LiveCodeBench pair with natural EOS; it +does not complete the rest of the campaign. + +The expanded 256K middle-window trace attributes 16.517 ms of rank-average +GPU service to target attention and about 7.387 ms to QPN2 projections. +Its profiled critical interval p50/p90/p99 is 33.834/34.201/34.822 ms. +Attention agrees with the actual-page independent working set; these data do +not establish an extra allocation/TLB bottleneck. + +A second trace observes actual B1 scheduling through the end of the request. +At computed position 262139, after 252 emitted tokens, the final four steps +have one scheduled token and zero proposals. They leave the FULL q8 graph and +run eager target forward. The scalar E4M3 FP32 partition attention kernel has +grid `(1,6,256)`, 256 threads, 40 registers and 12880 static shared bytes; +its sixteen target-layer calls cost 59.377 ms per captured q1 step/rank. +The draft phase still runs. This explains a substantial terminal penalty and +identifies a separate optimization scope. Full request accounting must retain +those steps even though the speculative-round counter does not count them. +Adjacent profiled scheduling intervals overlap GPU work; their service sums +must not be presented as a closed wall-clock decomposition. + +## PV reuse and quality localization + +The PV-reuse builder interchanges independent M fragments so a raw V fragment +and its residual-scaled copy are loaded/formed once per N16 panel. Each output +accumulator still consumes main0, residual0, main16 and residual16 before its +N32 online update. Source SHA is +`c70e6046c374d18f1c51ad126622e03ecfc34a46a656de225ae2376d384968c0`; DSO SHA is +`30c468456c6e5bfb8d97819a3cad1d6e598db20d0ca18e9ca2d681b12ef32961`. +It uses 112 registers without spills and the same 73728 shared-memory bytes. +All 130 paired operator checks pass, as do its own extended memcheck, +racecheck and synccheck. The actual-page sixteen-layer costs are: + +| Input tokens | Prior selected attention, ms | PV reuse attention, ms | +| --- | ---: | ---: | +| 1024 | 0.580 | 0.552 | +| 32768 | 2.401 | 2.244 | +| 65536 | 4.428 | 4.111 | +| 131072 | 8.471 | 7.841 | +| 261888 | 16.384 | 15.150 | + +Its first standalone model startup is **not admissible**. The 1K responses +match retained controls, but longer free generations differ, including +acceptance and a 64K finish reason. The 32K/64K/128K request medians of +18.280/19.993/23.727 ms therefore cannot establish a paired speedup; 256K is +39.875 ms and does not improve the previous screen. The experiment also uses +a fresh compiler-cache namespace and the expanded graph domain. Fifteen +shared native-library hashes still match the frozen run. Do not attribute +the trajectory difference to the new operator without an exact-input check. + +The next 32K/128K fixed-prefix A/B/A captures have exact native logits and +zero TV, support changes and unexplained state differences. Comparing their +control against the earlier startup's matching fixed tapes is also exact. +These checks do not replace the failed free-generation gate. A diagnostic +native shadow and natural proposal/state captures are used to locate the +first difference; the new candidate remains disabled for promotion. + +Paired QK products are another independent scheduling experiment: produce +two separate zero-initialized K16 products, then consume their compensation +updates in the original order. All 130 operator checks pass. The 128K working +set is 7.769 ms versus the paired PV-reuse parent's 7.818 ms; at 261888 it is +15.001 versus 15.095 ms. This small local benefit has no model admission. +Its source SHA is +`0f1d3703382b94ccacee3958eb9f89e0b4f13cdb25ca48269ef7e31006fa731f`; DSO SHA is +`14377f48ce548858c8d2af830e892cb68abe3abe1da29ee2d60713698117a30a`. + +The staged QK/PV follow-up with preferred shared carveout 100 does not help: +128K is 11.208 ms versus 11.193 ms without the preference and 8.446 ms for +the earlier one-stage selection. The resource occupancy API permits two PV +blocks/SM with either preference; this is a resource bound, not achieved +occupancy. The producer/PV trace costs 3.580/8.232 ms, motivating reuse and +consumer scheduling work instead of claiming that extra parallelism suffices. + +An independent warp-pipeline prototype uses eight producer and sixteen +consumer warps, two score/V panels and named ready/free barriers. Its initial +binary (source `416d076ee5962e1587b3c0c9854eb8a1837ba808bc1bcf82a804645216b143cf`) +fails the expanded byte-equality check despite zero reported synccheck errors. +It also spills under the 768-thread register limit. It has no performance or +serving acceptance. A serialized variant isolates overlap from other causes; +both preserve failed artifacts. The serialized variant still produces NaNs +on independent random inputs, so disabling overlap alone does not fix the +prototype. Constant-input, first-tile checks follow before any timing claim. +Named-barrier synchronization follows the +[PTX producer/consumer memory-ordering contract](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-bar), +including explicit participating-thread counts and unaligned instructions +for divergent producer/consumer warps. + +The structured-client diagnostic also found a prompt-counting defect: +canonical JSON serialization reordered tool/schema keys before rendering, +whereas generation preserved their insertion order. Actual prompts differ +in token count (339 versus 338 for one tool and 452 versus 454 for parallel +tools). The private client now uses the chat renderer with generation's wire +order, retaining the original counts as evidence and requiring equality to +the generated request's reported prompt count. Sampling, prompts, natural +EOS and total context capacity are unchanged. + +The natural-state diagnostic also needs to find auxiliary target states +through the context-probe sampling wrapper. Its earlier immediate-caller +lookup fails there before a scored observation. The revised lookup requires +both the same model-runner object and the same input-batch object, rejecting +unrelated frames. The affected CPU suite passes 19 tests with one GPU skip; +this change observes existing tensors and does not alter inference arithmetic. + +An isolated q1 builder reproduces the traced scalar E4M3/FP32 partition and +1024-token merge using the frozen serving source. The extracted reference and +PV-unroll4 candidate pass all 24 output/full-workspace/canary checks against +the frozen production DSO, including stride padding, zero-length replay and +262144 visible tokens. The sixteen-layer 261888-token workset is +60.565/60.574/60.580 ms for production/extracted/unroll4 respectively; the hint +has no useful performance benefit. Source/DSO hashes are: + +| Scalar q1 variant | Source SHA256 | DSO SHA256 | +| --- | --- | --- | +| Extracted reference | `903401903455e91fda49685b170b7bd932b1975d291f0c022b7a951ad99d0501` | `c521b0f96fa17bd8f474e5104ed2b35d1c82074d73af2e447172c90a05b0d584` | +| PV unroll4 | `e81f28ba2bf8215e7a6e037c2f599697e84cef582d1be0148c71f6f2c082a6f5` | `e9336ac7bd888bd3a9b43559423c4661de74b8a5d46d6ee1c7591afc845c06cc` | + +Explicit four-value prefetch applies the original FP32 FMAs in token order +and passes all 24 checks, but regresses the sixteen-layer 261888-token workset +from 60.564 to 63.124 ms. It is rejected without a model trial. Its source/DSO +SHA256 are `9723b252d9c54ebd96d1aa99c4ea3544c8cf66696e58e5fafa218d161e0619a8` / +`e822ade9e58e14510be2c40818327c4a608e4978d9970e5871d7fc3b231bf69d`. +A separate six-head scalar prototype shares each decoded K/V value across +independent per-head FP32 chains, retaining the original 1024-token partitions +and merge. It requires its own full-workspace checks and performance screen. +The scalar builder installs no serving route and does not change q8 arithmetic. + +## September 10 online-softmax correction and natural diagnostics + +The first warp-pipeline failure is localized to code extraction: the builder +copied the TWO_PASS statistics-only row loop instead of the online softmax +branch. It never wrote P, the compensated probability residual, or row rescale +before PV read them. A constant-input first-tile probe confirms an incorrect +P/residual with correct V and max/sum, including with overlap disabled. The +builder now anchors the online branch and requires all three publications and +its warp memory barrier. Corrected source +`15b44fd2fefe692000d18a112cb3f278802c6e5bc7448a62c6a885dc02014385` has DSO SHA256 +`2fdf18b3d0d71d84a9797a046554fe56e15a70182151b15206cbe06080b73717`. +This correction is not a quality or speed pass; the original failed artifacts +remain retained. The current 768-thread prototype still has register spills. + +The no-forced-token 32K control/candidate/control audit has 30 observed steps +and all four ranks. After checking consistent physical-slot renaming and +source-defined unused convolution storage, all captured target/proposal +logical tensors agree. Each comparison retains 2352 raw storage differences. +The natural comparator now supports the same explicit convolution-width +explanation as the fixed-prefix comparator; tests reject changed live values +and inconsistent mappings. The focused suite passes 22 tests with one GPU skip. + +This audit cannot explain away the PV candidate's free-generation failure. +Its 128-token output differs from the retained uninstrumented selected result +at token 97 (zero based), while the failed PV result first differs at token 59. +In addition, graph-captured persistent shadow counters contain non-count +values after control replays. They are invalid evidence. A revised diagnostic +allocates persistent counters and reference workspaces before graph capture, +matching initial workspace contents before each native comparison. A separate +uninstrumented same-startup route A/B is required to distinguish candidate +behavior from startup/layout or diagnostic effects. No claim of a model-wide +PV equality pass is made from these observations. + +The original selected combination has completed 17 of 30 frozen natural-EOS +pairs at this checkpoint. This includes all 12 JSON/schema/tool/parallel-tool +pairs over seeds 0/1/2 and five seed-0 code pairs. Completed long-code outputs +include 75828, 51562, 83004 and 24318 tokens with exact output IDs, acceptance +and finish reason. These are paired non-regression observations, not new +benchmark scores. Remaining natural code pairs continue at complete-pair +boundaries alongside the optimization queue. + +## Matched PV evidence and scalar q1 whole-round result + +The subsequent uninstrumented same-startup A/B uses the original three-group +route as control and PV reuse as candidate. All 12 pairs at 32K/128K have +identical output IDs, acceptance and finish reason. Both arms reproduce the +previously observed historical token differences at positions 59/21, +respectively. The prompt corpus SHA, prompt IDs and sampling contract match +the retained runs. Thus the observed cross-startup drift is not specific to +PV reuse. It remains unresolved and is not converted into an allowed quality +tolerance. Complete-round medians are 25.584/18.401 ms at 32K and +53.061/23.817 ms at 128K for this single control/candidate startup. + +Moving persistent shadow buffers outside the shared graph pool repairs the +counter corruption: the actual-input 32K diagnostic observes exactly 1856 +native comparisons (29 q8 replays × 16 layers × four ranks), with zero changed +output, numerator or max/sum elements. Control replays leave the counters at +zero. All captured target/proposal logical tensors agree across the 30-step +control/candidate/control runs, retaining the explained raw storage differences. +The original selected natural-output campaign is now 19/30 complete; seed-1 +HumanEval-10 adds an exact 16503-token natural-EOS pair. + +The scalar six-head candidate passes memcheck, racecheck and synccheck with +zero errors and six byte/canary checks in each run. Its first uninstrumented +model A/B keeps the prior selected q8 kernel in both arms and uses the scalar +candidate only for eligible eager q1 calls. All four ranks observe 384 real +candidate calls over six boundary requests, using CPU context bounds +262140–262143; no GPU length is read on the CPU. All 12 request pairs have +identical output IDs, acceptance and finish reason: + +| Input | q1 control complete round | q1 candidate complete round | Pure decode control/candidate | +| --- | ---: | ---: | ---: | +| 1K | 15.959 ms | 15.893 ms | 295.901 / 297.120 tokens/s | +| 261888 | 38.563 ms | 36.428 ms | 132.252 / 140.000 tokens/s | + +At the boundary, accepted drafts/round and emitted tokens/round remain +4.02 and 5.12. This is one paired startup, not the required three. The complete +round includes the terminal q1 overhead. The scalar optimization reduces it +by 2.134 ms, or 5.53%; the 22 ms goal is still unmet. The initial service-client +attempt failed to unpack the collective RPC results envelope before any +benchmark request; the corrected client validates all four rank identities. +The explicit worker probe and operator harness are now included for review. +No default route is enabled. + +The revised QK/PV warp pipeline is rejected on performance. At 128K/261888, +the PV-reuse baseline is 7.812/15.093 ms, versus 8.834/17.123 ms for eight +producers and 9.933/19.340 ms for six. Both retain the same 80-register limit +and spills. Four producer warps remove spills at 96 registers, while preserving +all 65 byte checks, but are slower again at 11.301/22.210 ms. Removing spills +alone does not provide a useful pipeline; none of these variants enters a +model trial. + +A separate feasibility probe predecodes the exact E4M3 values into FP16, +without restoring precision lost by E4M3 encoding. It passes 65 byte checks but +only changes the 261888-token attention workset from 15.090 to 15.034 ms, +excluding population/invalidation and the additional mirror memory. It is +rejected: the gain does not justify a mirror cache. The recorded failed build +attempt caught use of the FP8 paired-loader option with FP16 data; the tested +probe uses the existing FP16 vector loader. No serving KV representation changes. + +Available Triton caches share 189 compilation identities; 91 cubin hashes +differ. For those 91 entries, PTX agrees after excluding debug location/file +sections and assert-filename strings. This rules out a PTX arithmetic change +in those shared artifacts, not a change in actual dispatch, arguments or +machine-code behavior. All raw hashes and the excluded differences are retained. + +Native manifests for this stage: + +| Candidate | Source SHA256 | DSO SHA256 | +| --- | --- | --- | +| q1 six-head KV reuse | `e624c2f2c2eaa0d770b46aef7d2b4f84a710bc895bf0141b4670fd21ba498f69` | `8d6ede73f56b9edc270eab507d23d56f437567db362263c96d60b2a5ae05f98a` | +| four-producer QK/PV | `b005452e565468012b25a6f53c297ebc9a3482adde545ac041bfd170b264cde4` | `7dab6ff2d14e60a6f2c9803b9039bceeefc0dd879005f9333beef20d8ede8b99` | +| six-producer QK/PV | `a79561fc9c1a8a0a06590455e8e3d64807efb70a6aed412a50abdf229a1c12f7` | `1e9da088a750af9892b8c50a12187c652e3061dbd645091152cba7ed01a8b74b` | +| lossless decoded mirror probe | `eca466df3e1a64c945627e70447470250b83a1683cb6a2b32fd23724d06071e5` | `c31c0a3c136f453ea3f2b27f0fdd1275d3c29963ec3aa0ded889100651b0aeba` | + +## Three-startup q1 result and the next operator screens + +The six-head scalar q1 comparison has now completed three independent paired +service starts, each with a cold request and five measured requests per length +and arm. All 36 request pairs (72 requests) have identical output IDs, +acceptance and finish reason. The summary checks distinct server PIDs, frozen +source/library hashes, the prompt/sampling contract and 384 candidate calls on +each rank in each startup. Both arms retain the previously selected q8 kernel. + +| Input | Control complete round | Shared-q1 complete round | Pure decode control/candidate | +| --- | ---: | ---: | ---: | +| 1K | 15.951 ms | 15.893 ms | 296.070 / 296.663 tokens/s | +| 261888 | 38.563 ms | 36.436 ms | 132.268 / 139.854 tokens/s | + +Complete-round values are the median of the three startup medians; pure +decode values summarize the measured requests. Terminal q1 time is included. +Boundary accepted drafts/round and emitted tokens/round remain 4.02 and 5.12. +The second startup's 1K candidate is 0.184 ms slower; it remains in the +aggregate and makes no new scalar calls. The aggregate short-context median +does not regress. The request-average p50/p90/p99 distributions are retained +separately from actual GPU-step timings in +`scalar-q1-three-startup-summary.json`. + +An independent actual-input boundary shadow also passes: 64 q1 comparisons +per rank, 256 in total, have byte-identical output, full partial numerator and +max/sum workspaces against the frozen production scalar operator. The paired +requests retain identical output and acceptance. Shadow timing is diagnostic +only. The original selected q8 natural-EOS campaign has completed 20/30 pairs, +including a new exact 62396-token seed-1 LiveCodeBench-21 output. + +The next q1 candidate constructs a 256-entry FP32 shared-memory E4M3 lookup +table using the original decoder. A completed initialization barrier precedes +all reads. It reuses the same decoded values across six independent head +chains, retaining the original dimension/token order, partition size, FP32 +state and merge. Thirty output/full-workspace byte checks pass. Its sixteen-layer +261888-token operator workset is 15.965 ms, versus 33.326 ms for scalar sharing +alone and 60.577 ms for the frozen production operator. This is an operator +screen; its own sanitizer and full-service gates are required. It does not +broaden the public worker probe's context eligibility or enable a default. + +Revisiting staged QK/PV with PV-value reuse still fails the performance screen. +One D256 column uses 512 threads; two D128 columns use 256 threads each. Both +retain all 80 logical partitions and N32 updates and pass 130 byte checks in +total. Sixteen-layer 128K/261888 worksets are 9.276/17.899 ms for one column and +10.537/20.420 ms for two, versus 7.814/15.088 ms for the one-stage PV parent. +Neither staged variant advances to a model trial. The one-column diagnostic +trace attributes 3.931 ms to QK, 6.762 ms to PV and 0.229 ms to merge per +sixteen layers at 128K. These instrumented service sums cannot be used as +unprofiled latency or assumed overlap savings. The PV kernel uses 108 +registers/thread and 49824 bytes of dynamic shared memory; these are static +resources, not achieved occupancy. + +A separate explicit `profile` entrypoint records intra-CTA `clock64` +boundaries around K loading, QK with V loading, online softmax, and ordered +PV. It exists to distinguish phase dependencies before another scheduling +rewrite. Timestamp deltas include probe overhead and waits and do not report +hardware utilization. No serving route imports the probe. + +| Candidate | Source SHA256 | DSO SHA256 | +| --- | --- | --- | +| q1 shared E4M3 lookup | `82e3e0486f549125b94f5b38555e03e792fea51e349d3d5f972cf269871f632f` | `576b7dc765dd6850a6760a7b4dd248e803570d5c5d193ee3af83fc0d68d89a7a` | +| staged PV reuse, one column | `9aed322a8a098e9bb51f7113a0774bab3ca66766051281c20f54bfd9443c3b20` | `f01d7ea03de216445a325a95e5213fdb9f18b0ae8ff1299acabfa3aac7ffafce` | +| staged PV reuse, two columns | `fe00ab5a5e39e177f22bc3ecfda24d155bc8ae25ca78784c61f01d2150aa006f` | `16e5d1271f11e3201c5db9f61120cb6aea51c72d9b502007ddcfaec07f4ddfba` | +| intra-CTA phase probe | `aac2d83f53b52e738e9b71903ddc4f077d4ecb00ab95bcd354860067226f5534` | `abdea58968713b5017bf2d6f5898735a9cb741016264fa79c5df458d3d147906` | + +The lookup candidate's own memcheck, racecheck and synccheck each complete six +byte/canary checks with zero reported errors. Its first boundary shadow uses +261888 input tokens and 256 output tokens: the two outputs and acceptance +records agree, but this trajectory exercises no eligible scalar q1 calls. +The client rejects the missing coverage. This is retained as a route miss, +not a scalar quality pass. A separate 262136-input/eight-output boundary +diagnostic is used to leave no full q8 window; its timing is excluded. + +A subsequent uninstrumented paired startup does exercise the lookup operator: +384 calls per rank over six boundary requests. All twelve request pairs are +exact. Complete-round medians at 1K are 16.028/15.929 ms for control/lookup, +and at 261888 are 38.822/36.796 ms. Boundary pure decode is +119.426/126.003 tokens/s, accepted drafts/round 3.563636 and emitted +tokens/round 4.654545 in both arms. These acceptance values differ from the +older three-startup scalar campaign, so its 36.436 ms result cannot rank the +two scalar implementations. The unchanged control first differs from that +older output at tokens 114/62 for 1K/261888. All fifteen checked shared native +libraries retain their hashes. This startup repeatability issue is retained; +it is not assigned to the lookup kernel, which is disabled in the control. +An explicit same-startup control/shared/lookup comparison follows before +claiming incremental lookup benefit. + +The clock probe completes sixteen-layer working sets at 128K and 261888 with +byte-identical output and full FP32 workspaces. It records 65536/130944 N32 +tiles respectively. Aggregated CTA cycles at 261888 divide into 14.86% K load, +34.19% QK with V load, 20.16% online softmax and 30.79% ordered PV. The 128K +fractions closely agree. These include waits and timestamp overhead, and are +not kernel wall-time fractions or achieved utilization. They motivate removing +common-case branches and auditing the QK dependency chain. The original +selected natural campaign is now 22/30 exact natural-EOS pairs; the new seed-1 +LiveCodeBench-64/93 outputs contain 94767/47005 tokens respectively. + +The explicit short boundary-generation shadow completes 48 actual scalar +comparisons per rank (192 total), with byte-identical output, partial numerator +and max/sum storage; both requests retain identical output and acceptance. +The timestamp probe also passes its own memcheck, racecheck and synccheck, +including full 262144 visibility, page crossing, graph replay and timestamp +canaries. These checks validate the diagnostic; they do not remove its timing +overhead or grant model admission to another operator. + +## Fixed-q8 tile specialization and a separate arithmetic screen + +The fixed-q8 prototype clones the existing partial kernel with a constant +query-row count, selected only when the actual query shape has eight rows. +Queries with two through seven rows retain the original kernel. A second +variant removes visibility checks only when the complete N32 tile precedes +the minimum of all eight GPU row lengths. Zero-length and rejected rows, +partial tiles and the causal tail retain the original masking. No CPU copy of +GPU lengths is introduced. All 80 logical partitions, K16 compensation, +probability residual products and N32 numerator/max/sum updates are unchanged. + +The two variants pass 130 byte-equality checks in total, including full FP32 +workspaces. Sixteen-layer worksets show that fixed-q8 specialization alone +regresses 128K/261888 from 7.816/15.092 to 7.932/15.325 ms and is rejected. +Adding the complete-visible-tile specialization reaches 7.703/14.872 ms, +approximately 1.4% faster than the PV-reuse parent. The smaller gain needs a +same-startup complete-round A/B against that parent, using separate MRV2 +graphs; operator timing does not establish service benefit. Its native +sanitizer checks precede that service job. + +QK remains a substantial phase. A separate arithmetic prototype sums the +unchanged FP32 K16 products with explicit FP64 additions, then rounds to FP32 +before the original scale and softmax. It preserves N32 state updates and PV +compensation, but changes the QK summation arithmetic and has no admission. +The explicit nearest-even operation follows the +[CUDA double-precision intrinsic contract](https://docs.nvidia.com/cuda/cuda-math-api/cuda_math_api/group__CUDA__MATH__INTRINSIC__DOUBLE.html). +This is a feasibility experiment, not a claim that FP64 improves either +reference error or latency for this workload. + +The builder can additionally expose the actual FP32 combine accumulator before +FP16 conversion. The independent numerical harness first proves that each +diagnostic build retains the regular build's full partial/max/sum workspaces +and produces exactly the same final FP16 values after conversion. It then +compares both ordinary and pre-cast outputs against independent FP64 QK, +softmax and PV, reporting maximum absolute error, p99 absolute error and +relative L2 for two seeds at five lengths through 262144. An increase in any +registered error metric rejects the arithmetic candidate before performance +testing. Passing this operator screen still does not establish recursive or +model quality, logits/distribution equality, or acceptance non-inferiority. + +The new options are private builder switches: `--specialize-full-q8`, +`--all-visible-tiles`, `--qk-fp64-sum`, and `--diagnostic-output-fp32`. None +changes the frozen serving source or enables a production default. + +The arithmetic screen completes and rejects the FP64-sum candidate before +timing or any model trial. All ten pre-cast diagnostics agree with their +corresponding regular builds' workspaces and converted outputs. All final +FP16 reference-error metrics are nonexpanding, but four of ten cases expand +at least one pre-cast FP32 metric. At length 3297, seed 20260910, maximum +absolute error grows from `2.779396474e-6` to `3.216708786e-6` (about 15.73%); +p99 grows from `1.495418274e-6` to `1.531674442e-6`. Two cases also change +final FP16 output values. These are operator differences, not measured model +token flips. The original K16 compensation remains selected. + +The exact visible-tile candidate completes 25 byte/canary checks under each +of memcheck, racecheck and synccheck, with zero reported errors or races. +The initial racecheck invocation accidentally included timing worksets and +was stopped; the corrected invocation uses `--correctness-only` and reruns +the complete required checks. Timing collected under a sanitizer is excluded +from all performance claims. The subsequent service A/B compares PV reuse +against visible-tile specialization at 1K/32K/64K/128K/261888, with scalar q1 +unchanged in both arms. + +| Candidate | Source SHA256 | DSO SHA256 | +| --- | --- | --- | +| fixed q8, rejected speed | `eff0d41b399cbdb343f3b5543cd4577af2571b033fa9a3904d2fd9ad96cd07d3` | `6bfb9f20bc062f91731519faa05254bb53c4342052c6b0b733431f3a1ac35134` | +| fixed q8 and complete visible tiles | `3b0c9688ce17e1870408ef81fb5cd9b63a677b7cfd7d4777b8df77dd0fc24132` | `7dc632c2ff110cc751bceeb5fb683ede6066e443dad673d06c9e05429e01d0a2` | +| FP64 sum of K16 products | `e477e605fb94b9003cf71308da371b9051d7c72eb841e4e6e7fb4dbd6a8484da` | `e86ee04a6c7e949f9d3ded8ca7613b208bd89f4d3b642804d57e3f51f37daa1d` | +| PV parent, pre-cast diagnostic | `f30d85ef59aa86ca99d48aeba96eaa4a63dca8fbe4aaa1c375295d42f6ed284e` | `fc08222e8d768691d052817d89aadb2c7e5a7aa0207397da04fe4d70739ed0af` | +| FP64 K16 sum, pre-cast diagnostic | `9ba8011dbdc11103fb4746e07e3b798953bfa17d54557f973cdcb7638084e353` | `23ae5cb0b011b797e9279c6f2377863e2b337d301423b0ca38f4b28f59254699` | + +## Physical N64 and scalar page-map screens + +The physical N64 prototype computes two independent QK tiles together but +retains two consecutive N32 softmax/PV updates, the 80 logical partitions, +K16 compensation and probability residual products. It admits only aligned +q8; other query shapes and unaligned strides use its PV-reuse parent. Three +V-loading schedules are screened independently. All fail the working-set +speed gate despite byte-identical output and full FP32 workspaces: the first +two complete 130 checks in total, and the softmax-overlap variant completes +65. None advances to a service or sanitizer campaign. + +| Sixteen-layer working set | PV parent | Raw V prefetch | V after QK | V during softmax | +| --- | ---: | ---: | ---: | ---: | +| 131072 tokens | 7.811 ms | 9.757 ms | 7.922 ms | 8.977 ms | +| 261888 tokens | 15.084 ms | 18.927 ms | 15.330 ms | 17.378 ms | + +The softmax-overlap run has its own paired parent at 7.815/15.095 ms; it is +not compared by subtracting measurements from the earlier launch. Raw V +prefetch uses 128 registers and spills; V-after-QK uses 116 without spills. +Holding only the first N32 V tile across QK and loading the second during +softmax still uses 128 registers with spills. The wider tile does not produce +a useful local gain on this workload. + +The scalar q1 page-map prototype exploits the validated 3296-token page and +1024-token partition contract: each partition spans at most two physical +pages. It replaces per-token page/offset arrays with two page IDs and two +consecutive PV segments, preserving each head's original FP32 FMA order. +Together with and without the existing E4M3 lookup, the candidates complete +45 byte checks. The compact lookup variant improves the 128K sixteen-layer +workset from 8.011 to 7.554 ms but regresses 261888 from 15.965 to 17.436 ms; +neither compact variant is admitted. A separate unused 4096-byte dynamic +shared-memory reservation tests whether the changed resource limit explains +the long-context regression. Resource bounds are not achieved occupancy. + +The original selected natural campaign reaches 24/30 exact natural-EOS +pairs, including all twelve structured/tool cases. The final seed-1 +LiveCodeBench-131/162 outputs contain 49357/61290 tokens. Six seed-2 code +pairs remain. These comparisons do not establish new benchmark scores or +admit a later untested combination. + +The same-startup scalar control/shared/lookup experiment completes all saved +requests with matching output and acceptance, but records no eligible scalar +q1 calls in any arm. Its postprocessing failed because of a missing `Path` +import; offline recovery validates the retained requests and preserves the +original failure. Neither its timings nor its zero-hit counter establishes +incremental scalar speed or scalar operator quality. + +| Candidate | Source SHA256 | DSO SHA256 | +| --- | --- | --- | +| physical N64, raw V prefetch | `cf4ac891a5c3e7d38354ae5ec7af8d3e8800f23a6d56113828827bfd9b3e088f` | `1350da745048a53137347715c1d4155ce4770ddd5a87b4745ad186af639206be` | +| physical N64, V after QK | `cf41ce8b2d8d171c900a3d943dc9d8f0698b5493a00f46a816939058340ad60d` | `4f342b5cde1189e6b15bb78e87472f7081354f7e722b390e844151aee265a0c3` | +| physical N64, V during softmax | `1d44e0ae130d906877c9ed24308a6d079cbe5191d72db4a9a85d0737945cd5f2` | `9fecce8b8d1626a6ac3bd3de52b736219e3bcfdf03e9c78a0a731b342fa84bac` | +| scalar compact pages | `46980a019914f5a80861f9e7173aada8f292fe2c0bc2a17d1e1b0c6bddf5252d` | `00c602f201323071c5790443e08d8294febae2dbca53f44701e70bce177815b7` | +| scalar compact pages with lookup | `3cf205e3f3c2744b9a059663a6c5ff3db84cef53fd37ad5b663700154d7d2a97` | `eb54d60463e9811b1c116da25e7f56655925f9bfe0d61f691f48f06e2c5c537c` | +| scalar compact lookup, 4096-byte reservation | `5e68578c0252c7525496abab98aef5ed5f774528ac1dab6c5e1c587912cfa632` | `6d2b2b1ec5e0501de9abf49a81670cca2fb2ee700be66365f98db118a62fee21` | + +The next exact q8 candidate keeps each warp's three online max/sum rows in +registers and publishes them at the final output barrier. PV still consumes +the original shared row scales, and every row retains its ordered N32 +update. The builder switch `--register-softmax-state` requires the fixed-q8 +specialization; it changes no serving default. The generated visible-tile +source without this switch retains its original SHA. Native byte checks, +working-set timing and, if faster, its own sanitizers precede any service +trial. New 32K/261888 complete-round traces of the current visible-tile path +are collected separately from uninstrumented admission results. + +## Visible-tile repeated service results and next resource controls + +The visible-tile/PV-parent comparison completes three independent paired +startups (PIDs 1094304, 1101605 and 1103057), 180 requests and 90 exact pairs. +All contexts retain prompt, tokens, finish reason, sampling and acceptance. +The scalar q1 operator remains the original implementation in both arms. +The values below are medians of the three startup medians after one cold +request and five measured requests per context/arm. No profiler or tensor +dump is enabled. + +| Input tokens | PV-parent round | Visible-tile round | Visible pure decode | Accepted drafts/round | Emitted tokens/round | +| --- | ---: | ---: | ---: | ---: | ---: | +| 1024 | 16.008 ms | 15.982 ms | 295.464 tokens/s | 3.777778 | 4.740741 | +| 32768 | 18.222 ms | 18.185 ms | 212.458 tokens/s | 2.893939 | 3.878788 | +| 65536 | 20.084 ms | 20.040 ms | 196.015 tokens/s | 3.030769 | 3.938462 | +| 131072 | 23.786 ms | 23.687 ms | 165.712 tokens/s | 2.984615 | 3.938462 | +| 261888 | 37.547 ms | 37.298 ms | 124.306 tokens/s | 3.563636 | 4.654545 | + +These are small gains; the third startup's 32K candidate is 0.005 ms slower +and remains included. The aggregate short-context median does not regress. +All four revised long-context targets and the short-context <15 ms target +remain unmet. Do not rank this trajectory against earlier scalar-q1 trials +with different accepted outputs. + +| Visible-tile input | Request-average p50/p90/p99 | Cold TTFT median | Cold prefill median | +| --- | --- | ---: | ---: | +| 1024 | 15.982 / 16.061 / 16.149 ms | 0.273 s | 3989.826 tokens/s | +| 32768 | 18.185 / 18.235 / 18.243 ms | 8.288 s | 3977.766 tokens/s | +| 65536 | 20.014 / 20.058 / 20.065 ms | 18.417 s | 3573.957 tokens/s | +| 131072 | 23.674 / 23.706 / 23.711 ms | 44.248 s | 2971.068 tokens/s | +| 261888 | 37.298 / 37.494 / 37.530 ms | 117.995 s | 2223.944 tokens/s | + +Request-average quantiles are not actual GPU-round quantiles. Cold requests +verify the full computed-token count; cached repeat prefill is not used for +the cold throughput. Complete-round increments are 1.855, 3.647 and 13.611 ms +for 32K→64K, 64K→128K and 128K→261888. Their marginal costs are 0.057968, +0.056977 and 0.106546 ms per additional 1024 tokens. The last interval spans +127.75 such units and includes terminal q1 work. See the retained +`full-q8-visible-three-startup-summary.json` and its six hashed arm reports. + +The compact scalar lookup with the extra 4096-byte reservation passes all +45 paired operator checks. At 261888, sixteen-layer working-set cost is +15.026 ms versus 17.444 ms without the reservation and 15.967 ms for the +previous lookup. The 128K pair is 7.556/8.013 ms versus the previous lookup. +Changing this resource reservation removes the observed compact-layout +regression, but it does not measure achieved occupancy or isolate all cache +effects. Its own sanitizer campaign and an actual service comparison follow; +this local result is not a complete-round gain. + +Register-held q8 max/sum passes 65 full-workspace byte checks but regresses +every workset. Its 128K/261888 costs are 7.787/15.079 ms versus the paired +visible-tile parent at 7.695/14.862 ms. The relevant build uses 125 registers, +a 24-byte stack frame and no reported spills. Reject it before sanitizer or +service trials. Source SHA is +`3fb342d6737c1fe81220b2f18d805c3f394d1e728cda4e41683a7aaa4a15a905`; +the DSO SHA is +`153c8521e3dc173de9806fd5e1aef40e9768be8735c89ab448b1d126ed7aea12`. + +A separate `--qk-head-rows` prototype stages Q by head and uses one M8/N32 +QK tile per head, storing scores back in the original token/head order. +K16 compensation and all softmax/PV updates stay in place. Matrix-shape +equivalence is explicitly not assumed: the +[PTX WMMA contract](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-wmma-mma) +does not specify accumulation order or rounding for FP16 operations. The +byte gate must pass before timing; any difference instead requires the +independent reference audit. The normal source remains unchanged without +this experiment's switch. + +The M8/N32 candidate passes 65 byte checks but is slower at every measured +context. At 128K/261888 it takes 7.801/15.076 ms versus its paired parent at +7.699/14.876 ms. Rotating the next Q/K fragment load ahead of the original +current K16 correction also passes 65 byte checks but loses: +7.883/15.279 ms versus 7.696/14.893 ms. Neither advances. Their source/DSO +hashes are respectively: + +| Candidate | Source SHA256 | DSO SHA256 | +| --- | --- | --- | +| M8/N32 QK by head | `b99968cac226fc4443b097f71a4dd0a5ad3882b739eee697ff4f6ca28e04c2a7` | `bafa8a3a5fccb2fc07f94f45a5100b2ee58a6b68d2346a7d3be63d348a2cf06d` | +| QK operand rotation | `cda4d682a262bf859b54885f3044ca8b490fb55100d8277943114c1f43543e35` | `08c18843465cb7ee8e44b69e6ad8842cc9adb051ef33079cb3cae9a46cf149a9` | + +The compact lookup/reservation candidate's own three sanitizers each pass +six cases with zero errors or race hazards. Its first uninstrumented service +pair uses visible-tile q8 in both arms and hits 384 actual scalar calls per +rank. All twelve request pairs are exact. At 1K the round median changes +15.995→15.968 ms; at 261888 it changes 37.160→35.602 ms and pure decode +137.245→143.250 tokens/s. Accepted drafts/round remain 4.02 and emitted +tokens/round 5.12. This is one startup, not completed repeated-start admission. + +An actual-q1 diagnostic compares original, shared, lookup and compact lookup +using saved live operands and isolated output/workspaces. The retained subset +contains eight unique K/V pointer pairs per rank, with 96 exact candidate +comparisons across four ranks. Its final client assertion incorrectly expected +sixteen unique pairs and the job exits 1. Preserve the original failed report +and the separate subset analysis; this is not evidence that all sixteen +attention layers were sampled. Median per-operand graph times across ranks +are approximately 3.69–3.85 ms original, 2.02–2.10 ms shared, +0.96–1.00 ms lookup and 0.90–0.94 ms compact lookup. These operator events +and sums over eight samples are not complete-round costs. Additional route +metadata is required to close the smaller observed complete-round gain. + +## Latest whole-round attribution and small-Q coverage + +The new visible-tile 32K trace contains 57 analyzed q8 intervals after edge +exclusion. Its critical-rank mean interval is 20.121 ms, GPU event union +18.499 ms and uncovered time 1.623 ms. A 39.244-ms outlier remains included. +QPN2 service is 7.370 ms, draft service 3.933 ms and target grouped attention +2.379 ms on those same critical ranks. These instrumented values do not +replace the 18.185-ms unprofiled endpoint result. + +The 261888 trace observes 54 q8 steps, one q6 step and four q1 steps in the +whole diagnostic request. In the analyzed inner intervals, q8 target +attention takes 15.027 ms and draft 3.889 ms. The q6 interval uses eager +target execution and target attention takes 24.435 ms. The q1 intervals use +the original scalar implementation. A partial verifier therefore loses both +the q8 graph and the q8-specific scheduling wrapper. All such costs remain +in the complete-round denominator; do not remove them to claim a target. + +Earlier query-shape fixtures exercised q2 and q5 in addition to q8; they +did not qualify every q2–q7 shape. The benchmark now exposes an explicit +`--tail-queries` suite, a `--performance-query-rows` selector, and a hashed +reference binding to the actual frozen `grouped_e4m3_fp32_paged_fwd` entry. +The reference requires precision revision 4 and retains DSO SHA +`a751fed902279b0de23537c4aad2dc4fee360146d7fce7ef0c4f255a77f48b02`. +Default benchmark queries and candidate entrypoints are unchanged. + +All six q2–q7 shapes pass 60 byte checks against that production reference, +including page crossing, padded strides, zero/rejected rows, restored lengths +and graph replay. Sixteen-layer q6 worksets improve from 12.176 to 7.796 ms +at 128K and from 23.717 to 15.066 ms at 261888. Its own small-Q sanitizer +suite precedes any eager service-route extension. No new shape is enabled +merely by adding the benchmark selector. diff --git a/tests/kernels/core/test_sm70_dflash2_state_audit.py b/tests/kernels/core/test_sm70_dflash2_state_audit.py index f6698e0db6..a45c427ce0 100644 --- a/tests/kernels/core/test_sm70_dflash2_state_audit.py +++ b/tests/kernels/core/test_sm70_dflash2_state_audit.py @@ -10,7 +10,38 @@ cpu_request_slots, gather_state, selected_ssm_slots, + target_auxiliary_states, ) +from benchmarks.sm70_dflash2_state_layout import ( + check_slot_mapping, + explain_state_difference, +) + + +def test_natural_audit_observes_auxiliary_states_through_sampling_wrapper(): + class Runner: + def run(self, input_batch, aux_hidden_states): + return self.wrapper(input_batch) + + def wrapper(self, batch): + # A wrapper may contain unrelated state; only a matching batch + # from the same runner is a valid source for the observation. + input_batch = object() + aux_hidden_states = [torch.tensor([-1.0])] + assert input_batch is not batch and aux_hidden_states + return target_auxiliary_states(self, batch) + + expected = [torch.tensor([1.0])] + assert Runner().run(object(), expected) is expected + + +def test_natural_audit_rejects_unrelated_auxiliary_states(): + class Runner: + def run(self, input_batch, aux_hidden_states): + return target_auxiliary_states(object(), input_batch) + + with pytest.raises(RuntimeError, match="matching target auxiliary states"): + Runner().run(object(), [torch.tensor([1.0])]) @pytest.fixture @@ -96,6 +127,88 @@ def test_audit_comparator_rejects_nonfinite_logits(captures): compare(left, right) +def test_audit_tracks_eos_changes_without_a_top1_flip(captures): + left, right = captures + eos = (0, 30) + assert ( + compare(left, right, eos_token_ids=eos)["summary"][ + "max_eos_probability_abs_difference" + ] + == 0 + ) + path = right / "test-rank0-step1.pt" + data = torch.load(path, weights_only=True) + data["native_logits"][0, 30] += 0.5 + torch.save(data, path) + result = compare(left, right, eos_token_ids=eos) + assert result["summary"]["top1_changed_rows"] == 0 + assert result["summary"]["max_eos_probability_abs_difference"] > 0 + row = result["logits"][-1] + assert row["full_eos_probabilities_left"] != row["full_eos_probabilities_right"] + assert ( + row["sampling_eos_probabilities_left"] + != row["sampling_eos_probabilities_right"] + ) + with pytest.raises(ValueError, match="outside the captured vocabulary"): + compare(left, right, eos_token_ids=(32,)) + + +def test_state_layout_requires_bijective_slots_and_preserves_padding(): + key = "verify/layer0/recurrent/slot_table:(1, 3)" + left = {key: torch.tensor([[52, 53, -1]])} + right = {key: torch.tensor([[154, 155, -1]])} + mapping: dict[int, int] = {} + reverse: dict[int, int] = {} + check_slot_mapping(left, right, mapping, reverse) + chosen = "verify/layer0/recurrent/input_state/indices:(1,)" + # An accepted selector reading the wrong logical slot cannot be explained + # as another physical renaming after the full table has been observed. + with pytest.raises(ValueError, match="Inconsistent or aliased"): + check_slot_mapping( + {chosen: torch.tensor([52])}, + {chosen: torch.tensor([155])}, + mapping, + reverse, + ) + with pytest.raises(ValueError, match="Inconsistent or aliased"): + check_slot_mapping(left, {key: torch.tensor([[154, 154, -1]])}, {}, {}) + with pytest.raises(ValueError, match="Padding slot"): + check_slot_mapping(left, {key: torch.tensor([[154, 155, 0]])}, {}, {}) + + +@pytest.mark.parametrize("phase", ["prefill", "verify"]) +def test_state_layout_compares_only_proven_conv_input_window(phase): + prefix = f"{phase}/layer0/conv" + key = prefix + "/input_state/values:(1, 2, 10)" + left = { + key: torch.zeros(1, 2, 10), + prefix + "/input_state/valid:(1,)": torch.tensor([True]), + prefix + "/has_initial_state:(1,)": torch.tensor([True]), + prefix + "/num_accepted_tokens:(1,)": torch.tensor([4]), + } + right = {k: v.clone() for k, v in left.items()} + # Prefill reads columns 0..2. With selector 4, verify reads columns 3..5. + right[key][..., 8] = 123 + assert explain_state_difference(key, left, right, 4) == ( + "unused_convolution_storage" + ) + active_column = 1 if phase == "prefill" else 4 + right[key][..., active_column] = 1 + assert explain_state_difference(key, left, right, 4) is None + del right[prefix + "/input_state/valid:(1,)"] + with pytest.raises(ValueError, match="Missing or ambiguous"): + explain_state_difference(key, left, right, 4) + + +def test_state_layout_keeps_all_verifier_output_bytes(): + prefix = "verify/layer0/conv/output_state" + key = prefix + "/values:(1, 2, 10)" + left = {key: torch.zeros(1, 2, 10), prefix + "/valid:(1,)": torch.tensor([True])} + right = {k: v.clone() for k, v in left.items()} + right[key][..., 9] = 123 + assert explain_state_difference(key, left, right, 4) is None + + @pytest.fixture def natural_captures(captures): for directory in captures: @@ -174,6 +287,37 @@ def test_natural_audit_aligns_request_slots(natural_captures): assert actual.tolist() == [0] +@pytest.mark.parametrize("mutation", [None, "live_value", "slot_alias"]) +def test_natural_audit_explains_only_consistent_state_slots(natural_captures, mutation): + left, right = natural_captures + for side, directory in enumerate(natural_captures): + for path in directory.glob("*-rank*-step*.pt"): + row = torch.load(path, weights_only=True) + label = f"{row['phase']}/layer0/conv/input_state/indices:(1,)" + row["states"][label] = torch.tensor([5 + side * 100], dtype=torch.int32) + torch.save(row, path) + # The optional explanation preserves the raw byte mismatch. + assert not compare_natural(left, right)["cases"][0]["all_logical_tensors_equal"] + path = right / "test-rank2-step1.pt" + row = torch.load(path, weights_only=True) + if mutation == "live_value": + row["states"]["verify/layer0/recurrent/input_state:(1, 2)"][0, 0] += 1 + elif mutation == "slot_alias": + row["states"]["verify/layer0/conv/input_state/indices:(1,)"][0] = 106 + torch.save(row, path) + if mutation == "slot_alias": + with pytest.raises(ValueError, match="physical-slot mapping"): + compare_natural(left, right, conv_width=4) + else: + result = compare_natural(left, right, conv_width=4)["cases"][0] + assert result["explained_storage_differences"] + assert result["all_logical_tensors_equal"] == (mutation is None) + if mutation == "live_value": + first = result["first_observed_difference"] + assert (first["step"], first["phase"]) == (1, "target") + assert first["differences"][0]["name"].endswith("input_state:(1, 2)") + + def test_state_audit_reads_accepted_slot_and_preserves_invalid_selectors(): table = torch.tensor([[7, 0, 9, 2], [3, 5, 4, 8]], dtype=torch.int32) assert selected_ssm_slots(table, torch.tensor([2, 4])).tolist() == [0, 8] diff --git a/tests/v1/worker/test_sm70_long_attention_graphs.py b/tests/v1/worker/test_sm70_long_attention_graphs.py new file mode 100644 index 0000000000..c8e7781d9d --- /dev/null +++ b/tests/v1/worker/test_sm70_long_attention_graphs.py @@ -0,0 +1,69 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU checks for conservative MRV2 attention graph selection.""" + +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch + +from vllm.config.compilation import CUDAGraphMode +from vllm.v1.attention.ops.sm70_e4m3_long import MAX_CONTEXT +from vllm.v1.worker.gpu.cudagraph_utils import ( + BatchExecutionDescriptor, + ModelCudaGraphManager, +) + + +@pytest.fixture +def graph_pair(): + manager = ModelCudaGraphManager.__new__(ModelCudaGraphManager) + ordinary = BatchExecutionDescriptor(CUDAGraphMode.FULL, 8, 1, 8) + bounded = replace(ordinary, attention_context_bucket=MAX_CONTEXT) + manager._long_attention_graphs = {ordinary: bounded} + manager.graphs = {ordinary: object(), bounded: object()} + return manager, ordinary, bounded + + +def test_context_boundary_and_switch_back(graph_pair): + manager, ordinary, bounded = graph_pair + for upper, expected in ( + (1024, bounded), + (MAX_CONTEXT, bounded), + (MAX_CONTEXT + 1, ordinary), + (262144, ordinary), + (32768, bounded), + (0, ordinary), + ): + assert ( + manager.select_attention_graph(ordinary, torch.tensor([upper])) == expected + ) + + +def test_device_hint_never_copied_to_host(graph_pair): + manager, ordinary, _ = graph_pair + # Accessing a device hint's values would fail: selection must inspect the + # device first and use the full-context graph without requesting a copy. + device_hint = SimpleNamespace(device=torch.device("cuda")) + assert manager.select_attention_graph(ordinary, device_hint) == ordinary + + +def test_other_batch_shapes_and_missing_capture_fall_back(graph_pair): + manager, ordinary, bounded = graph_pair + other = replace(ordinary, num_tokens=16, num_reqs=2) + assert manager.select_attention_graph(other, torch.tensor([1024, 1024])) == other + assert manager.select_attention_graph(ordinary, torch.tensor([1024, 0])) == ordinary + del manager.graphs[bounded] + assert manager.select_attention_graph(ordinary, torch.tensor([1024])) == ordinary + + +def test_disabled_operator_preserves_original_binding(monkeypatch): + from vllm.v1.attention.ops.sm70_e4m3_long import MANIFEST_ENV, wrap_long_attention + + monkeypatch.delenv(MANIFEST_ENV, raising=False) + + def original(*args, **kwargs): + raise AssertionError("Binding inspection must not launch an operator") + + assert wrap_long_attention(original) is original diff --git a/vllm/envs.py b/vllm/envs.py index 84327507f0..154774ad4c 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -187,6 +187,7 @@ VLLM_SM70_FP8_QPN8_LIBRARY: str | None = None VLLM_SM70_SAMPLER_LIBRARY: str | None = None VLLM_SM70_FA2_D256_LIBRARY: str | None = None + VLLM_SM70_E4M3_LONG_ATTENTION_MANIFEST: str | None = None VLLM_SM70_FP8_PREFILL_VISIBLE_DENSE_MM: bool = False VLLM_SM70_NVFP4_QPN2: bool = False VLLM_SM70_NVFP4_QPN2_M16_NATIVE: bool = True @@ -1875,6 +1876,10 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_FP8_QPN8_LIBRARY": lambda: os.getenv("VLLM_SM70_FP8_QPN8_LIBRARY", None), "VLLM_SM70_SAMPLER_LIBRARY": lambda: os.getenv("VLLM_SM70_SAMPLER_LIBRARY", None), "VLLM_SM70_FA2_D256_LIBRARY": lambda: os.getenv("VLLM_SM70_FA2_D256_LIBRARY", None), + # Experimental q8 long attention; unset preserves the full-context route. + "VLLM_SM70_E4M3_LONG_ATTENTION_MANIFEST": lambda: os.getenv( + "VLLM_SM70_E4M3_LONG_ATTENTION_MANIFEST", None + ), "VLLM_SM70_FP8_PREFILL_CUTLASS": lambda: bool( int(os.getenv("VLLM_SM70_FP8_PREFILL_CUTLASS", "1")) ), diff --git a/vllm/v1/attention/ops/sm70_e4m3_grouped.py b/vllm/v1/attention/ops/sm70_e4m3_grouped.py index b4dea81606..0f6935857b 100644 --- a/vllm/v1/attention/ops/sm70_e4m3_grouped.py +++ b/vllm/v1/attention/ops/sm70_e4m3_grouped.py @@ -17,7 +17,9 @@ def load_grouped_e4m3_fp32(): return None if not flash_attn_grouped_e4m3_fp32_available(): return None - return flash_attn_grouped_e4m3_fp32_paged + from vllm.v1.attention.ops.sm70_e4m3_long import wrap_long_attention + + return wrap_long_attention(flash_attn_grouped_e4m3_fp32_paged) def grouped_e4m3_fp32_allowed( diff --git a/vllm/v1/attention/ops/sm70_e4m3_long.py b/vllm/v1/attention/ops/sm70_e4m3_long.py new file mode 100644 index 0000000000..429e8fbbef --- /dev/null +++ b/vllm/v1/attention/ops/sm70_e4m3_long.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Opt-in compensated q8 scheduling with an explicit native build manifest.""" + +import hashlib +import importlib.util +import json +import os +from functools import lru_cache +from pathlib import Path + +import torch + +from vllm.forward_context import get_forward_context, is_forward_context_available +from vllm.logger import init_logger + +logger = init_logger(__name__) + +# Include generation headroom after a 128K prompt. Larger CPU upper bounds use +# the existing full-context graph; device row lengths remain authoritative. +MAX_CONTEXT = 132096 +MANIFEST_ENV = "VLLM_SM70_E4M3_LONG_ATTENTION_MANIFEST" +_WORKSPACES: dict[tuple, tuple[torch.Tensor, torch.Tensor]] = {} + + +def long_attention_enabled() -> bool: + return bool(os.environ.get(MANIFEST_ENV)) + + +@lru_cache(maxsize=1) +def load_long_attention(manifest_name: str): + manifest_path = Path(manifest_name).resolve() + manifest = json.loads(manifest_path.read_text()) + # Different split counts change arithmetic and workspace geometry. They + # must complete their own admission before extending this serving route. + if manifest.get("splits", 80) != 80 or manifest["head_groups"] != 1: + raise ValueError("The long-attention serving route requires 80 six-head splits") + library = Path(manifest["library"]) + if not library.is_absolute(): + library = manifest_path.parent / library + library = library.resolve() + if hashlib.sha256(library.read_bytes()).hexdigest() != manifest["library_sha256"]: + raise ValueError( + "Long-attention native library SHA does not match its manifest" + ) + name = library.name.split(".")[0] + if name != manifest["module_name"]: + raise ValueError( + "Long-attention native module name does not match its manifest" + ) + spec = importlib.util.spec_from_file_location(name, library) + if spec is None or spec.loader is None: + raise ImportError(f"Cannot load long-attention extension {library}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + loaded_file = module.__file__ + if loaded_file is None or Path(loaded_file).resolve() != library: + raise RuntimeError("Long-attention native extension module alias") + logger.info_once( + "Loaded experimental SM70 E4M3 q8 attention: module=%s SHA256=%s " + "max_context=%d; 80 splits, compensated FP32 state.", + name, + manifest["library_sha256"], + MAX_CONTEXT, + scope="process", + ) + return module.run, manifest + + +def wrap_long_attention(fallback): + manifest_name = os.environ.get(MANIFEST_ENV) + if not manifest_name: + return fallback + operator, manifest = load_long_attention(manifest_name) + + def run( + q, k, v, table, row_lengths, *, out, softmax_scale, k_scale=1.0, v_scale=1.0 + ): + descriptor = ( + get_forward_context().batch_descriptor + if is_forward_context_available() + else None + ) + if not ( + descriptor is not None + and descriptor.attention_context_bucket == MAX_CONTEXT + and q.shape == (8, 6, 256) + and k.ndim == 4 + and k.shape[1] in (1648, 3296) + and k.shape[2:] == (1, 256) + and v.shape == k.shape + ): + return fallback( + q, + k, + v, + table, + row_lengths, + out=out, + softmax_scale=softmax_scale, + k_scale=k_scale, + v_scale=v_scale, + ) + # Allocate a fixed workspace once for each warmup/capture stream. Graph + # replay never allocates. Layers reuse it in stream order; different + # streams and versions never share the legacy 80-split buffers. + stream = torch.cuda.current_stream(q.device).cuda_stream + key = (manifest["source_sha256"], MAX_CONTEXT, 80, q.device, stream) + if key not in _WORKSPACES: + _WORKSPACES[key] = ( + torch.empty((80, 8, 6, 256), dtype=torch.float32, device=q.device), + torch.empty((80, 8, 6, 2), dtype=torch.float32, device=q.device), + ) + partial, lse = _WORKSPACES[key] + return operator( + q, + k, + v, + out, + table, + row_lengths, + partial, + lse, + float(softmax_scale), + float(k_scale), + float(v_scale), + ) + + return run diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 4dbe0bce1e..d55d747c10 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections import defaultdict from collections.abc import Callable -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import Any, NamedTuple import torch @@ -100,6 +100,7 @@ class BatchExecutionDescriptor: num_tokens: int num_reqs: int | None # None means no request padding is needed (PIECEWISE graphs) uniform_token_count: int | None = None + attention_context_bucket: int | None = None def _is_compatible( @@ -405,6 +406,47 @@ def __init__( self.aux_hidden_states: list[torch.Tensor] = [] self.use_aux_hidden_state_outputs = False self.intermediate_tensors: IntermediateTensors | None = None + self._long_attention_graphs: dict[ + BatchExecutionDescriptor, BatchExecutionDescriptor + ] = {} + from vllm.v1.attention.ops.sm70_e4m3_long import ( + MAX_CONTEXT, + long_attention_enabled, + ) + + if ( + long_attention_enabled() + and current_platform.is_cuda() + and current_platform.is_device_capability((7, 0)) + and self.dp_size == 1 + and vllm_config.parallel_config.pipeline_parallel_size == 1 + ): + descs = self._capture_descs.get(CUDAGraphMode.FULL, []) + for desc in list(descs): + if (desc.num_tokens, desc.num_reqs, desc.uniform_token_count) == ( + 8, + 1, + 8, + ): + variant = replace(desc, attention_context_bucket=MAX_CONTEXT) + self._long_attention_graphs[desc] = variant + descs.append(variant) + + def select_attention_graph( + self, desc: BatchExecutionDescriptor, cpu_upper_bounds: torch.Tensor + ) -> BatchExecutionDescriptor: + variant = self._long_attention_graphs.get(desc) + if variant is None or variant not in self.graphs: + return desc + # Never materialize device lengths on the host. A missing or oversized + # CPU hint conservatively selects the existing full-context graph. + if cpu_upper_bounds.device.type != "cpu" or cpu_upper_bounds.numel() != 1: + return desc + upper = int(cpu_upper_bounds[0]) + limit = variant.attention_context_bucket + if limit is not None and 0 < upper <= limit: + return variant + return desc def capture( self, @@ -461,10 +503,16 @@ def create_forward_fn( def forward_fn(cg_mode: CUDAGraphMode) -> None: batch_descriptor = None - if cg_mode == CUDAGraphMode.PIECEWISE: - assert attn_metadata is None + if ( + cg_mode == CUDAGraphMode.PIECEWISE + or desc.attention_context_bucket is not None + ): + if cg_mode == CUDAGraphMode.PIECEWISE: + assert attn_metadata is None batch_descriptor = BatchDescriptor( - num_tokens=num_tokens, has_lora=has_lora + num_tokens=num_tokens, + has_lora=has_lora, + attention_context_bucket=desc.attention_context_bucket, ) with ( sm70_decode_graph_compilation(desc.cg_mode == CUDAGraphMode.FULL), diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 8ac18bd9b1..30846ce599 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1633,6 +1633,9 @@ def execute_model( # NOTE(woosuk): Here, we don't need to pass the input tensors, # because they are already copied to the CUDA graph input buffers. assert self.cudagraph_manager is not None + batch_desc = self.cudagraph_manager.select_attention_graph( + batch_desc, input_batch.seq_lens_cpu_upper_bound + ) self.kv_connector.pre_forward(scheduler_output) model_output = self.cudagraph_manager.run_fullgraph(batch_desc) else: