diff --git a/benchmarks/compare_sm70_dflash2_natural_audit.py b/benchmarks/compare_sm70_dflash2_natural_audit.py new file mode 100644 index 0000000000..fd3a83111d --- /dev/null +++ b/benchmarks/compare_sm70_dflash2_natural_audit.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Locate the first observed target/proposal difference with real acceptance. + +This is a diagnostic comparison, not a quality or acceptance noninferiority +gate. Inputs after the first differing proposal need not be the same. A first +observed state difference still requires an operator-level causality check. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import torch + +from benchmarks.compare_sm70_dflash2_state_audit import ( + sampling_difference, + tensor_difference, +) + + +def _load(directory: Path) -> dict: + records = {} + for path in directory.glob("*-rank*-step*.pt"): + row = torch.load(path, map_location="cpu", weights_only=True, mmap=True) + case, rank, step = (row[k] for k in ("case", "rank", "step")) + key = case, step, rank + if key in records or row.get("control") != "natural_sampling": + raise ValueError(f"{path}: duplicate or non-natural observation") + if row["phase"] != ("prefill" if step == 0 else "verify"): + raise ValueError(f"{path}: incorrect phase") + if not row.get("expected_layers") or "capture_epoch" not in row: + raise ValueError(f"{path}: missing snapshot provenance") + if rank == 0 and not torch.isfinite(row["native_logits"]).all(): + raise ValueError(f"{path}: nonfinite native logits") + for field in ("aux_hidden_states", "sampling", "states"): + if not row.get(field): + raise ValueError(f"{path}: missing {field}") + if row["num_sampled"].numel() != 1 or row["num_rejected"].numel() != 1: + raise ValueError(f"{path}: expected B1 counts") + n = int(row["num_sampled"].item()) + rejected = int(row["num_rejected"].item()) + if not 1 <= n <= 8 or not 0 <= rejected <= 7: + raise ValueError(f"{path}: invalid acceptance counts") + if n > row["sampled_token_ids"].shape[1]: + raise ValueError(f"{path}: missing accepted output IDs") + proposal = directory / f"proposal-{case}-tp{rank}-forward{step}.pt" + if not proposal.exists(): + raise ValueError(f"{path}: missing proposal observation") + draft = torch.load(proposal, map_location="cpu", weights_only=True, mmap=True) + if (draft["case"], draft["step"], draft["rank"]) != key: + raise ValueError(f"{proposal}: proposal identity differs") + if "draft_tokens" not in draft or "projected_context" not in draft: + raise ValueError(f"{proposal}: incomplete proposal") + records[key] = row, draft + if not records: + raise ValueError(f"{directory}: empty captures") + for case in {k[0] for k in records}: + steps = {k[1] for k in records if k[0] == case} + if len(steps) < 2 or steps != set(range(max(steps) + 1)): + raise ValueError(f"{case}: missing prefill or verifier steps") + for step in steps: + if {k[2] for k in records if k[:2] == (case, step)} != {0, 1, 2, 3}: + raise ValueError(f"{case}/{step}: every step requires four TP ranks") + return records + + +def _target_tensors(row: dict) -> dict[str, torch.Tensor]: + result = { + k: row[k] + for k in ("positions", "input_ids", "hidden", "num_sampled", "num_rejected") + } + for k in ("native_logits", "draft_logits"): + if row.get(k) is not None: + result[k] = row[k] + for k in ("sampling", "states"): + result.update({f"{k}/{label}": t for label, t in row[k].items()}) + result.update({f"aux/{i}": t for i, t in enumerate(row["aux_hidden_states"])}) + result.update( + { + f"layer/{v['layer_idx']}/{v['label']}": v["tensor"] + for v in row["tensors"].values() + } + ) + # Unwritten output padding is not an emitted or accepted token. + result["accepted_output"] = row["sampled_token_ids"][:, : row["num_sampled"].item()] + return result + + +def _proposal_tensors(row: dict) -> dict[str, torch.Tensor]: + values = { + k: v + for k, v in row.items() + if isinstance(v, torch.Tensor) and k != "idx_mapping" + } + if "idx_mapping" in row and row.get("sampling_layout") != "request_gathered_v1": + # Early captures retained complete arrays indexed by request slot. + indices = row["idx_mapping"].to(torch.int64) + for name in ("temperature", "seeds"): + if name in values: + values[name] = values[name].index_select(0, indices) + return values + + +def compare_natural(left_dir: Path, right_dir: Path) -> 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": []} + 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 = [] + 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)] + values = [ + _target_tensors(row) + if phase == "target" + else _proposal_tensors(row) + for row in rows + ] + if phase == "proposal" and "idx_mapping" in rows[0]: + slots = [row["idx_mapping"].tolist() for row in rows] + if slots[0] != slots[1]: + mappings.append( + {"step": step, "rank": rank, "slots": slots} + ) + if values[0].keys() != values[1].keys(): + raise ValueError( + f"{case}/{step}/{rank}: tensor coverage differs" + ) + for name in sorted(values[0]): + a, b = (v[name] for v in values) + if a.shape != b.shape or a.dtype != b.dtype: + diff = { + "contract_changed": [ + str(a.shape), + str(b.shape), + str(a.dtype), + str(b.dtype), + ] + } + else: + diff = tensor_difference(a, b) + if diff["bitwise_equal"]: + continue + if name == "native_logits": + diff.update(sampling_difference(a, b)) + differences.append({"rank": rank, "name": name, **diff}) + if differences: + first = {"step": step, "phase": phase, "differences": differences} + break + if first is not None: + break + result["cases"].append( + { + "case": case, + "steps_per_arm": lengths, + "first_observed_difference": first, + "different_request_slot_mappings": mappings, + "all_logical_tensors_equal": first is None and lengths[0] == lengths[1], + } + ) + return result + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("left", type=Path) + parser.add_argument("right", type=Path) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + torch.set_num_threads(4) + result = compare_natural(args.left, args.right) + args.output.write_text(json.dumps(result, indent=2) + "\n") + for case in result["cases"]: + first = case["first_observed_difference"] + print( + case["case"], "equal" if first is None else (first["step"], first["phase"]) + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/compare_sm70_dflash2_state_audit.py b/benchmarks/compare_sm70_dflash2_state_audit.py new file mode 100644 index 0000000000..aafe8ab860 --- /dev/null +++ b/benchmarks/compare_sm70_dflash2_state_audit.py @@ -0,0 +1,202 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compare complete four-rank StateAuditExtension captures, failing on gaps.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import torch + + +def tensor_difference(left: torch.Tensor, right: torch.Tensor) -> dict: + if left.shape != right.shape or left.dtype != right.dtype: + raise ValueError( + f"Tensor contract differs: {left.shape}/{left.dtype}, " + f"{right.shape}/{right.dtype}" + ) + + def raw_bytes(tensor): + flat = tensor.contiguous().reshape(-1) + # A size-one tensor can be "contiguous" with stride 8. Reinterpret + # only its logical storage, not padding or neighboring metadata rows. + return flat.as_strided((flat.numel(),), (1,)).view(torch.uint8) + + byte_equal = torch.equal(raw_bytes(left), raw_bytes(right)) + if byte_equal: + return {"bitwise_equal": True} + difference = (left.double() - right.double()).abs() + return { + "bitwise_equal": False, + "different_elements": int((left != right).sum()), + "max_abs": float(difference.max()) if difference.numel() else 0.0, + "left_nonfinite": int((~torch.isfinite(left)).sum()), + "right_nonfinite": int((~torch.isfinite(right)).sum()), + } + + +def sampling_difference(left: torch.Tensor, right: torch.Tensor) -> dict: + from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p_pytorch + + if left.shape != right.shape: + raise ValueError("Logit shapes differ") + count = left.shape[0] + + def probabilities(logits): + return apply_top_k_top_p_pytorch( + logits.float().clone(), + torch.full((count,), 20, dtype=torch.int32), + torch.full((count,), 0.95), + ).softmax(-1) + + p, q = probabilities(left), probabilities(right) + full_tv = (left.float().softmax(-1) - right.float().softmax(-1)).abs().sum(-1) / 2 + return { + "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(), + } + + +def compare( + left_dir: Path, right_dir: Path, *, right_verifier_route: str | 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")} + if not left_files or left_files.keys() != right_files.keys(): + raise ValueError("Missing or mismatched capture files") + result = { + "left": str(left_dir), + "right": str(right_dir), + "comparisons": [], + "logits": [], + } + coverage: dict[tuple[str, int], set[int]] = {} + seen_states: dict[tuple[str, int, str], set[str]] = {} + for name in sorted(left_files): + left = torch.load(left_files[name], map_location="cpu", weights_only=True) + right = torch.load(right_files[name], map_location="cpu", weights_only=True) + for key in ("case", "rank", "step", "phase", "num_draft_tokens"): + if left[key] != right[key]: + raise ValueError(f"{name}: {key} differs") + identity = {key: left[key] for key in ("case", "rank", "step", "phase")} + if right_verifier_route is not None and right["phase"] == "verify": + expected = { + f"route/verify/layer{layer}/{right_verifier_route}" + for layer in right["expected_layers"] + } + if set(right.get("verifier_routes", ())) != expected: + raise ValueError(f"{name}: missing {right_verifier_route} route hit") + if left.get("expected_layers") != right.get("expected_layers"): + raise ValueError(f"{name}: requested layers differ") + if not left.get("expected_layers") or "capture_epoch" not in left: + raise ValueError(f"{name}: missing current-forward snapshot provenance") + for side in (left, right): + for layer in side["expected_layers"]: + prefix = f"{side['phase']}/layer{layer}" + for required in ( + "/conv/input:", + "/conv/output:", + "/recurrent/q:", + "/recurrent/input_state", + "/recurrent/output:", + ): + if not any(k.startswith(prefix + required) for k in side["states"]): + raise ValueError(f"{name}: missing {prefix + required}") + coverage.setdefault((left["case"], left["step"]), set()).add(left["rank"]) + for key in ("positions", "input_ids"): + if not torch.equal(left[key], right[key]): + raise ValueError(f"{name}: forced {key} differs") + groups = { + "boundary": ({"hidden": left["hidden"]}, {"hidden": right["hidden"]}), + "state": (left["states"], right["states"]), + "layer": tuple( + { + f"layer{v['layer_idx']}/{v['label']}": v["tensor"] + for v in d["tensors"].values() + } + for d in (left, right) + ), + "sampling": (left["sampling"], right["sampling"]), + } + seen_states.setdefault( + (left["case"], left["rank"], left["phase"]), set() + ).update(left["states"]) + for group, (lvalues, rvalues) in groups.items(): + if lvalues.keys() != rvalues.keys(): + raise ValueError(f"{name}: {group} snapshot coverage differs") + for label in sorted(lvalues): + difference = tensor_difference(lvalues[label], rvalues[label]) + if not difference["bitwise_equal"]: + result["comparisons"].append( + {**identity, "group": group, "label": label, **difference} + ) + 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") + result["logits"].append( + { + **identity, + "positions": left["positions"].tolist(), + **tensor_difference(left["native_logits"], right["native_logits"]), + **sampling_difference( + left["native_logits"], right["native_logits"] + ), + } + ) + if any(ranks != {0, 1, 2, 3} for ranks in coverage.values()): + raise ValueError("Every step requires four TP ranks") + for key, labels in seen_states.items(): + for required in ( + "/conv/input", + "/conv/output", + "/recurrent/input_state", + "/recurrent/output", + ): + if not any(required in label for label in labels): + raise ValueError(f"{key}: missing {required} evidence") + for case in {key[0] for key in coverage}: + steps = {key[1] for key in coverage if key[0] == case} + if steps != set(range(max(steps) + 1)) or len(steps) < 2: + raise ValueError(f"{case}: missing prefill or verifier steps") + for rank in range(4): + if any( + (case, rank, phase) not in seen_states + for phase in ("prefill", "verify") + ): + raise ValueError(f"{case}/rank{rank}: missing prefill/verifier states") + result["summary"] = { + "files_per_arm": len(left_files), + "differing_intermediates": len(result["comparisons"]), + "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"] + ), + "top1_changed_rows": sum(sum(row["top1_changed"]) for row in result["logits"]), + "all_logits_bitwise_equal": all( + row["bitwise_equal"] for row in result["logits"] + ), + } + return result + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("left", type=Path) + parser.add_argument("right", type=Path) + parser.add_argument("--output", required=True, type=Path) + parser.add_argument("--right-verifier-route", choices=("split", "packed")) + args = parser.parse_args() + torch.set_num_threads(4) + result = compare( + args.left, args.right, right_verifier_route=args.right_verifier_route + ) + args.output.write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(result["summary"], indent=2)) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_sm70_draft_f16_layout.py b/benchmarks/kernels/benchmark_sm70_draft_f16_layout.py new file mode 100644 index 0000000000..8c1a0f0672 --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_draft_f16_layout.py @@ -0,0 +1,182 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Screen cuBLAS operand layouts using retained actual TP4 draft projections. + +Numerical checks replay all four ranks' real weights and inputs on GPU 4. +The separate timing uses twenty consecutive-layer weights from rank zero. +Changed arithmetic requires numerical and model gates; this screen alone does +not admit a serving route. +""" + +import argparse +import hashlib +import json +import os +import statistics +from pathlib import Path + +import torch + + +def errors(value, reference): + difference = (value.double() - reference).abs().flatten() + return dict( + max_abs=difference.max().item(), + p99_abs=torch.quantile(difference, 0.99).item(), + relative_l2=(difference.norm() / reference.norm()).item(), + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--root", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--strict-candidate-reduction", action="store_true") + args = parser.parse_args() + assert os.environ.get("CUDA_VISIBLE_DEVICES") == "4" + assert torch.cuda.get_device_capability() == (7, 0) + modes = ["row_weight", "column_weight", "row_weight_m16", "column_weight_m16"] + numerical, working_set, provenance = [], [], [] + + def run(item, mode): + x = item["padded"] if mode >= 2 else item["x"] + weight = item["column_weight"] if mode % 2 else item["weight"] + previous = torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction + try: + if args.strict_candidate_reduction and mode != 0: + torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = ( + False + ) + torch.mm(x, weight.T, out=item["outputs"][mode]) + finally: + torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = previous + + for rank in range(4): + root = args.root / f"rank{rank}" + sample_paths = sorted(root.glob("inputs-step*.pt")) + weight_paths = sorted(root.glob("*-weight.pt")) + assert len(sample_paths) == 5 and len(weight_paths) == 20 + samples = [ + torch.load(p, map_location="cpu", weights_only=True) for p in sample_paths + ] + for path in sample_paths + weight_paths: + provenance.append( + dict( + file=str(path), sha256=hashlib.sha256(path.read_bytes()).hexdigest() + ) + ) + for path in weight_paths: + saved = torch.load(path, map_location="cpu", weights_only=True) + name, cpu_weight = saved["name"], saved["weight"] + weight = cpu_weight.cuda() + n, k = weight.shape + item = dict( + weight=weight, + column_weight=weight.T.contiguous().T, + x=torch.empty((8, k), device="cuda", dtype=torch.float16), + padded=torch.zeros((16, k), device="cuda", dtype=torch.float16), + outputs=[ + torch.empty( + (8 if m < 2 else 16, n), device="cuda", dtype=torch.float16 + ) + for m in range(4) + ], + ) + assert torch.equal(weight, item["column_weight"]) + weight64 = weight.double() + for sample_path, snapshot in zip(sample_paths, samples): + values = snapshot[name] + item["x"].copy_(values["input"]) + item["padded"][:8].copy_(item["x"]) + oracle = item["x"].double() @ weight64.T + for mode in range(4): + run(item, mode) + output = item["outputs"][mode][:8] + if mode == 0: + control_errors = errors(output, oracle) + saved_equal = torch.equal( + output.cpu().view(torch.uint8), + values["control"].view(torch.uint8), + ) + metrics = errors(output, oracle) + numerical.append( + dict( + rank=rank, + name=name, + snapshot=sample_path.name, + mode=modes[mode], + shape=[8, n, k], + saved_control_equal=saved_equal, + byte_equal=torch.equal( + output.view(torch.uint8), + item["outputs"][0].view(torch.uint8), + ), + finite=bool(torch.isfinite(output).all()), + control=control_errors, + candidate=metrics, + reference_error_not_expanded=all( + metrics[key] <= control_errors[key] for key in metrics + ), + ) + ) + if rank == 0: + working_set.append(item) + print(f"Checked actual rank {rank} inputs", flush=True) + assert len(working_set) == 20 + graphs = [] + for mode in range(4): + for _ in range(3): + for item in working_set: + run(item, mode) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for item in working_set: + run(item, mode) + graphs.append(graph) + timings = [[] for _ in modes] + for trial in range(7): + for mode in range(4) if trial % 2 == 0 else reversed(range(4)): + for _ in range(20): + graphs[mode].replay() + start, end = [torch.cuda.Event(enable_timing=True) for _ in range(2)] + start.record() + for _ in range(50): + graphs[mode].replay() + end.record() + end.synchronize() + timings[mode].append(start.elapsed_time(end) / 50) + summary = {} + for mode in modes: + rows = [r for r in numerical if r["mode"] == mode] + summary[mode] = dict( + comparisons=len(rows), + differing=sum(not r["byte_equal"] for r in rows), + expanded_reference_error=sum( + not r["reference_error_not_expanded"] for r in rows + ), + mismatched_saved_control=sum(not r["saved_control_equal"] for r in rows), + ) + report = dict( + modes=modes, + summary=summary, + numerical=numerical, + provenance=provenance, + samples_ms=timings, + medians_ms=[statistics.median(t) for t in timings], + complete_round_performance=False, + strict_candidate_reduction=args.strict_candidate_reduction, + torch_version=torch.__version__, + allow_fp16_reduced_precision_reduction=torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction, + ) + args.output.write_text(json.dumps(report, indent=2) + "\n") + print( + json.dumps( + {k: v for k, v in report.items() if k not in ("numerical", "provenance")}, + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_sm70_qpn2_chunked.py b/benchmarks/kernels/benchmark_sm70_qpn2_chunked.py new file mode 100644 index 0000000000..bb536a232f --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_qpn2_chunked.py @@ -0,0 +1,269 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compare frozen publication, serial chunks, and event-ordered chunk overlap. + +Uses four ranks of real consecutive-layer weights with changing synthetic +activations. Each chunk owns a separate two-epoch channel; consumers wait for +their local producer before entering the original peer packet protocol. +GPU 4--7 only. No service route is installed by this benchmark. +""" + +import argparse +import hashlib +import json +import os +import statistics +from pathlib import Path + +import torch +import torch.distributed as dist + +from vllm import _custom_ops as custom_ops +from vllm import _sm70_ops as ops +from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--root", type=Path, required=True) + p.add_argument("--library", type=Path, required=True) + p.add_argument("--control-library", type=Path, required=True) + p.add_argument("--output", type=Path, required=True) + p.add_argument("--correctness-only", action="store_true") + p.add_argument("--cycles", type=int, default=5) + p.add_argument("--coordination-backend", choices=["nccl", "gloo"], default="nccl") + args = p.parse_args() + if args.cycles < 1: + p.error("--cycles must be positive") + args.output.parent.mkdir(parents=True, exist_ok=True) + assert os.environ.get("CUDA_VISIBLE_DEVICES") == "4,5,6,7" + rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(rank) + assert torch.cuda.get_device_capability() == (7, 0) + torch.ops.load_library(str(args.control_library)) + torch.ops.load_library(str(args.library)) + dist.init_process_group(args.coordination_backend) + assert dist.get_world_size() == 4 + group = dist.new_group(backend="gloo") + ca = CustomAllreduce(group, rank, max_size=128 * 1024) + assert not ca.disabled and ca.sm70_tp4_push_buffer_ptrs is not None + peers = ca.sm70_tp4_push_buffer_ptrs + channels = [] + comm_stream = torch.cuda.Stream(device=rank) + generator = torch.Generator(device="cuda").manual_seed(20260908 + rank) + items, provenance, guards = [], [], [] + + def output(shape): + storage = torch.full( + (shape[0] * shape[1] + 16,), -37, device="cuda", dtype=torch.float16 + ) + guards.append(storage) + return storage[8:-8].view(shape) + + try: + for path in sorted((args.root / f"rank{rank}").glob("*.pt")): + item = torch.load(path, map_location="cpu", weights_only=True) + assert item["rank"] == rank + provenance.append( + dict( + file=path.name, sha256=hashlib.sha256(path.read_bytes()).hexdigest() + ) + ) + item["codes"] = item["codes"].cuda() + item["scales"] = item["scales"].cuda() + item["input"] = torch.randn( + (8, item["k"]), dtype=torch.float16, device="cuda", generator=generator + ) + width = item["n"] // 2 if item["gated"] else item["n"] + item["output"] = [output((8, width)) for _ in range(3)] + item["row_parallel"] = not item["gated"] and item["n"] == 5120 + item["reduced"] = ( + [output((8, width)) for _ in range(3)] if item["row_parallel"] else None + ) + item["ready_events"] = [torch.cuda.Event() for _ in range(2)] + items.append(item) + assert len(items) == 16 and sum(i["row_parallel"] for i in items) == 8 + tail_input = torch.randn( + (8, 5120), dtype=torch.float16, device="cuda", generator=generator + ) + tail_outputs = [output((8, 5120)) for _ in range(3)] + + for _ in range(2): + channel = ca.create_shared_buffer( + custom_ops.sm70_tp4_push_allreduce_buffer_size(), group=group + ) + channels.append(channel) + torch.ops._qpn2_chunked.initialize(tail_input, channel[rank]) + torch.cuda.synchronize() + dist.barrier() + + def run(arm): + for item in items: + arguments = [ + item["output"][arm], + item["input"], + item["codes"], + item["scales"], + item["global_scale"], + item["split_k"], + item["nacc"], + ] + if item["row_parallel"]: + if arm == 0: + torch.ops._qpn2_candidate.publish(*arguments, peers, rank) + torch.ops._qpn2_candidate.consume( + item["output"][arm], item["reduced"][arm], peers, rank + ) + else: + main_stream = torch.cuda.current_stream() + for chunk, channel in enumerate(channels): + torch.ops._qpn2_chunked.publish( + *arguments, channel, rank, chunk * 2560 + ) + if arm == 1: + torch.ops._qpn2_chunked.consume( + item["reduced"][arm], channel, rank, chunk * 2560 + ) + else: + ready = item["ready_events"][chunk] + ready.record(main_stream) + with torch.cuda.stream(comm_stream): + comm_stream.wait_event(ready) + torch.ops._qpn2_chunked.consume( + item["reduced"][arm], + channel, + rank, + chunk * 2560, + ) + if arm == 2: + main_stream.wait_stream(comm_stream) + else: + op = ( + ops.nvfp4_qpn2_gated_sm70_out + if item["gated"] + else ops.nvfp4_qpn2_gemm_sm70_out + ) + op(*arguments) + # Odd collective count plus an ordinary push call exercises epoch + # transitions across graphs and the two publication mechanisms. + ca.all_reduce(tail_input, out=tail_outputs[arm], registered=True) + + graphs = [] + for arm in range(3): + torch.cuda.synchronize() + dist.barrier() + graph = torch.cuda.CUDAGraph() + with ca.capture(), torch.cuda.graph(graph): + run(arm) + graphs.append(graph) + + def check(): + pairs = ( + [i["output"] for i in items] + + [i["reduced"] for i in items if i["row_parallel"]] + + [tail_outputs] + ) + for index, outputs in enumerate(pairs): + a = outputs[0] + for arm, b in enumerate(outputs[1:], 1): + assert torch.isfinite(a).all() and torch.isfinite(b).all(), ( + "nonfinite", + rank, + index, + arm, + ) + assert torch.equal(a.view(torch.uint8), b.view(torch.uint8)), ( + "mismatch", + rank, + index, + arm, + int(torch.count_nonzero(a != b)), + ) + for g in guards: + assert (g[:8] == -37).all() and (g[-8:] == -37).all() + + for cycle in range(args.cycles): + for item in items: + item["input"].normal_(generator=generator) + tail_input.normal_(generator=generator) + for g in guards: + g[8:-8].fill_(float("nan")) + dist.barrier() + for arm in (0, 1, 2) if cycle % 2 == 0 else (2, 1, 0): + if rank == cycle % 4: + torch.cuda._sleep(20000) + graphs[arm].replay() + torch.cuda.synchronize() + check() + samples = [[], [], []] + for trial in range(0 if args.correctness_only else 7): + for arm in (0, 1, 2) if trial % 2 == 0 else (2, 1, 0): + for _ in range(20): + graphs[arm].replay() + torch.cuda.synchronize() + dist.barrier() + start, end = [torch.cuda.Event(enable_timing=True) for _ in range(2)] + start.record() + for _ in range(50): + graphs[arm].replay() + end.record() + end.synchronize() + ranks = [None] * 4 + dist.all_gather_object(ranks, start.elapsed_time(end) / 50, group=group) + samples[arm].append(max(ranks)) + check() + ranks = [None] * 4 + dist.all_gather_object( + ranks, dict(rank=rank, snapshots=provenance), group=group + ) + if rank == 0: + result = dict( + rank_weights=ranks, + shape="TP4 B1/q8 four consecutive layers", + projections=16, + row_parallel=8, + extra_ordinary_push=1, + activation_kind="synthetic_fp16_normal", + all_outputs_bitwise_equal=True, + canaries_intact=True, + changing_input_cycles=args.cycles, + rank_skew_cycles=20000, + medians_ms=[statistics.median(x) for x in samples] + if not args.correctness_only + else None, + samples_ms=samples, + arms=["frozen_publisher", "two_chunks_serial", "two_chunks_overlap"], + paired_saved_ms=[ + [a - b for a, b in zip(samples[0], candidate)] + for candidate in samples[1:] + ], + control_library_sha256=hashlib.sha256( + args.control_library.read_bytes() + ).hexdigest(), + library_sha256=hashlib.sha256(args.library.read_bytes()).hexdigest(), + communicator_sha256=hashlib.sha256( + Path(os.environ["VLLM_SM70_CUSTOM_AR_LIBRARY"]).read_bytes() + ).hexdigest(), + complete_round_performance=False, + ) + args.output.write_text(json.dumps(result, indent=2)) + print( + json.dumps( + {k: v for k, v in result.items() if k != "rank_weights"}, indent=2 + ), + flush=True, + ) + finally: + # All streams and ranks have finished reading peer allocations before + # the owning rank releases either channel. + torch.cuda.synchronize() + dist.barrier(group=group) + for channel in channels: + ca.free_shared_buffer(channel, rank=rank) + ca.close() + dist.destroy_process_group(group) + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_sm70_qpn2_cooperative_mlp.py b/benchmarks/kernels/benchmark_sm70_qpn2_cooperative_mlp.py new file mode 100644 index 0000000000..33d83bf690 --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_qpn2_cooperative_mlp.py @@ -0,0 +1,276 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Validate cooperative gate/up and down publication on four ranks of real weights. + +This screen uses synthetic activations and sixteen projections from consecutive +layers, with eight dependent all-reduces and an extra ordinary push call to +exercise mixed-protocol epoch transitions. It is not a complete model round. +Set VLLM_SM70_CUSTOM_AR_LIBRARY to a communicator built from the same header +as the candidate. Use Gloo process coordination for focused sanitizer runs. +""" + +import argparse +import hashlib +import json +import os +import statistics +from pathlib import Path + +import torch +import torch.distributed as dist + +from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--root", type=Path, required=True) + p.add_argument("--library", type=Path, required=True) + p.add_argument("--output", type=Path, required=True) + p.add_argument("--publisher-library", type=Path, required=True) + p.add_argument("--capped-library", type=Path, required=True) + p.add_argument("--correctness-only", action="store_true") + p.add_argument("--cycles", type=int, default=5) + p.add_argument("--coordination-backend", choices=["nccl", "gloo"], default="nccl") + args = p.parse_args() + if args.cycles < 1: + p.error("--cycles must be positive") + args.output.parent.mkdir(parents=True, exist_ok=True) + rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(rank) + assert torch.cuda.get_device_capability() == (7, 0) + torch.ops.load_library(str(args.library)) + torch.ops.load_library(str(args.publisher_library)) + torch.ops.load_library(str(args.capped_library)) + resources = torch.ops._qpn2_coop_mlp.resources() + assert resources[3] >= 2, resources + dist.init_process_group(args.coordination_backend) + assert dist.get_world_size() == 4 + group = dist.new_group(backend="gloo") + ca = CustomAllreduce(group, rank, max_size=128 * 1024) + assert not ca.disabled and ca.sm70_tp4_push_buffer_ptrs is not None + peers = ca.sm70_tp4_push_buffer_ptrs + generator = torch.Generator(device="cuda").manual_seed(20260908 + rank) + items, provenance, guards = [], [], [] + + def output(shape): + storage = torch.full( + (shape[0] * shape[1] + 16,), -37, device="cuda", dtype=torch.float16 + ) + guards.append(storage) + return storage[8:-8].view(shape) + + try: + for path in sorted((args.root / f"rank{rank}").glob("*.pt")): + item = torch.load(path, map_location="cpu", weights_only=True) + assert item["rank"] == rank + provenance.append( + dict( + file=path.name, sha256=hashlib.sha256(path.read_bytes()).hexdigest() + ) + ) + item["codes"] = item["codes"].cuda() + item["scales"] = item["scales"].cuda() + item["input"] = torch.randn( + (8, item["k"]), dtype=torch.float16, device="cuda", generator=generator + ) + width = item["n"] // 2 if item["gated"] else item["n"] + item["output"] = [output((8, width)) for _ in range(2)] + item["row_parallel"] = not item["gated"] and item["n"] == 5120 + item["reduced"] = ( + [output((8, width)) for _ in range(2)] if item["row_parallel"] else None + ) + items.append(item) + assert len(items) == 16 and sum(i["row_parallel"] for i in items) == 8 + tail_input = torch.randn( + (8, 5120), dtype=torch.float16, device="cuda", generator=generator + ) + tail_outputs = [output((8, 5120)) for _ in range(2)] + + gates = {x["layer"]: x for x in items if x["gated"]} + downs = {x["layer"]: x for x in items if x["row_parallel"] and x["k"] == 4352} + for layer in range(4): + assert ( + next(i for i, x in enumerate(items) if x is downs[layer]) + == next(i for i, x in enumerate(items) if x is gates[layer]) + 1 + ) + + rejected_rows = [] + gate, down = gates[0], downs[0] + for rows in (1, 2, 4, 7, 9, 16): + try: + torch.ops._qpn2_coop_mlp.mlp( + gate["output"][0].new_empty((rows, 4352)), + down["output"][0].new_empty((rows, 5120)), + gate["input"].new_empty((rows, 5120)), + gate["codes"], + gate["scales"], + down["codes"], + down["scales"], + gate["global_scale"], + down["global_scale"], + peers, + rank, + ) + except RuntimeError as error: + assert "exact TP4 q8 MLP required" in str(error), error + rejected_rows.append(rows) + else: + raise AssertionError(("unexpected non-q8 launch", rows)) + + def run(arm): + for item in items: + down = item["row_parallel"] and item["k"] == 4352 + input_ = gates[item["layer"]]["output"][arm] if down else item["input"] + if arm and item["gated"]: + paired = downs[item["layer"]] + torch.ops._qpn2_coop_mlp.mlp( + item["output"][arm], + paired["output"][arm], + input_, + item["codes"], + item["scales"], + paired["codes"], + paired["scales"], + item["global_scale"], + paired["global_scale"], + peers, + rank, + ) + continue + arguments = [ + item["output"][arm], + input_, + item["codes"], + item["scales"], + item["global_scale"], + item["split_k"], + item["nacc"], + ] + if item["row_parallel"]: + if not (arm and down): + torch.ops._qpn2_candidate.publish(*arguments, peers, rank) + torch.ops._qpn2_candidate.consume( + item["output"][arm], item["reduced"][arm], peers, rank + ) + else: + op = ( + torch.ops._qpn2_capped.gated + if item["gated"] + else torch.ops._qpn2_capped.gemm + ) + op(*arguments) + ca.all_reduce(tail_input, out=tail_outputs[arm], registered=True) + + graphs = [] + for arm in range(2): + torch.cuda.synchronize() + dist.barrier() + graph = torch.cuda.CUDAGraph() + with ca.capture(), torch.cuda.graph(graph): + run(arm) + graphs.append(graph) + + def check(): + pairs = ( + [i["output"] for i in items] + + [i["reduced"] for i in items if i["row_parallel"]] + + [tail_outputs] + ) + for index, (a, b) in enumerate(pairs): + assert torch.isfinite(a).all() and torch.isfinite(b).all(), ( + "nonfinite", + rank, + index, + ) + assert torch.equal(a.view(torch.uint8), b.view(torch.uint8)), ( + "mismatch", + rank, + index, + int(torch.count_nonzero(a != b)), + ) + for g in guards: + assert (g[:8] == -37).all() and (g[-8:] == -37).all() + + for cycle in range(args.cycles): + for item in items: + item["input"].normal_(generator=generator) + tail_input.normal_(generator=generator) + for g in guards: + g[8:-8].fill_(float("nan")) + dist.barrier() + for arm in (0, 1) if cycle % 2 == 0 else (1, 0): + if rank == cycle % 4: + torch.cuda._sleep(20000) + graphs[arm].replay() + torch.cuda.synchronize() + check() + samples = [[], []] + for trial in range(0 if args.correctness_only else 7): + for arm in (0, 1) if trial % 2 == 0 else (1, 0): + for _ in range(20): + graphs[arm].replay() + torch.cuda.synchronize() + dist.barrier() + start, end = [torch.cuda.Event(enable_timing=True) for _ in range(2)] + start.record() + for _ in range(50): + graphs[arm].replay() + end.record() + end.synchronize() + ranks = [None] * 4 + dist.all_gather_object(ranks, start.elapsed_time(end) / 50, group=group) + samples[arm].append(max(ranks)) + check() + ranks = [None] * 4 + dist.all_gather_object( + ranks, dict(rank=rank, snapshots=provenance), group=group + ) + if rank == 0: + result = dict( + rank_weights=ranks, + shape="TP4 B1/q8 four consecutive layers", + projections=16, + cooperative_resources=resources, + gate_down_dependency=True, + rejected_rows=rejected_rows, + publisher_sha256=hashlib.sha256( + args.publisher_library.read_bytes() + ).hexdigest(), + capped_sha256=hashlib.sha256( + args.capped_library.read_bytes() + ).hexdigest(), + script_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(), + row_parallel=8, + extra_ordinary_push=1, + activation_kind="synthetic_fp16_normal", + all_outputs_bitwise_equal=True, + canaries_intact=True, + changing_input_cycles=args.cycles, + rank_skew_cycles=20000, + medians_ms=[statistics.median(x) for x in samples] + if not args.correctness_only + else None, + samples_ms=samples, + paired_saved_ms=[a - b for a, b in zip(*samples)], + library_sha256=hashlib.sha256(args.library.read_bytes()).hexdigest(), + communicator_sha256=hashlib.sha256( + Path(os.environ["VLLM_SM70_CUSTOM_AR_LIBRARY"]).read_bytes() + ).hexdigest(), + complete_round_performance=False, + ) + args.output.write_text(json.dumps(result, indent=2)) + print( + json.dumps( + {k: v for k, v in result.items() if k != "rank_weights"}, indent=2 + ), + flush=True, + ) + finally: + ca.close() + dist.destroy_process_group(group) + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_sm70_qpn2_graph_layout.py b/benchmarks/kernels/benchmark_sm70_qpn2_graph_layout.py new file mode 100644 index 0000000000..bb441cfbd5 --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_qpn2_graph_layout.py @@ -0,0 +1,359 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Check same-graph q8 norm/packed-column switches against frozen libraries. + +Inputs are four consecutive layers of real prepared weights on each of four +ranks. No model route is installed. --timing measures the column/norm working +set only, excluding row projections, communication and the complete round. +""" + +import argparse +import hashlib +import importlib.util +import json +import os +from pathlib import Path + +import torch +import torch.distributed as dist + +from benchmarks.kernels.sm70_qpn2_graph_layout import LayoutTemplates +from benchmarks.kernels.sm70_qpn2_graph_nodes import nodes +from vllm.model_executor.layers import layernorm as norm + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--output", type=Path, required=True) + p.add_argument("--root", type=Path, required=True) + p.add_argument("--column-library", type=Path, required=True) + p.add_argument("--packed-library", type=Path, required=True) + p.add_argument("--dual-norm", type=Path, required=True) + p.add_argument("--timing", action="store_true") + p.add_argument("--cycles", type=int, default=9) + args = p.parse_args() + if args.cycles < 1: + p.error("--cycles must be positive") + args.output.parent.mkdir(parents=True, exist_ok=True) + rank = int(os.environ.get("LOCAL_RANK", "0")) + torch.cuda.set_device(rank) + assert torch.cuda.get_device_capability() == (7, 0) + libraries = {"column": args.column_library, "packed": args.packed_library} + for path in libraries.values(): + torch.ops.load_library(str(path)) + spec = importlib.util.spec_from_file_location("benchmark_dual_norm", args.dual_norm) + dual = importlib.util.module_from_spec(spec) + spec.loader.exec_module(dual) + dist.init_process_group("gloo") + assert dist.get_world_size() == 4 + generator = torch.Generator(device="cuda").manual_seed(20260909 + rank) + items = [] + provenance = [] + for path in sorted((args.root / f"rank{rank}").glob("*.pt")): + item = torch.load(path, map_location="cpu", weights_only=True) + if item["k"] != 5120: + continue + provenance.append( + dict(file=path.name, sha256=hashlib.sha256(path.read_bytes()).hexdigest()) + ) + item["codes"] = item["codes"].cuda() + item["scales"] = item["scales"].cuda() + items.append(item) + assert len(items) == 8 + templates = LayoutTemplates(items, dual) + report = dict( + rank=rank, + templates=templates.describe(), + provenance=provenance, + libraries={ + k: hashlib.sha256(v.read_bytes()).hexdigest() for k, v in libraries.items() + }, + cases=[], + ) + print(json.dumps(dict(rank=rank, templates=templates.describe())), flush=True) + + def make(rows, kind): + records = [] + for item in items: + x = torch.randn( + (rows, 5120), device="cuda", dtype=torch.float16, generator=generator + ) + weight = torch.randn( + (5120,), device="cuda", dtype=torch.float16, generator=generator + ) + residual = ( + None + if kind == "none" + else torch.randn( + (rows, 5120), + device="cuda", + dtype=torch.float16 if kind == "half" else torch.float32, + generator=generator, + ) + ) + out = torch.empty_like(x) + resout = ( + None if residual is None else torch.empty_like(x, dtype=torch.float32) + ) + result = torch.empty( + (rows, item["n"] // (2 if item["gated"] else 1)), + device="cuda", + dtype=torch.float16, + ) + records.append((item, x, weight, residual, out, resout, result)) + + def run(): + for item, x, weight, residual, out, resout, result in records: + if kind == "float": + norm._sm70_dflash2_gemma_fused_add_rms_kernel[(rows,)]( + x, + residual, + weight, + out, + resout, + hidden_size=5120, + BLOCK_SIZE=8192, + epsilon=1e-6, + num_warps=8, + num_stages=1, + ) + else: + norm._sm70_dflash2_fixed_gemma_rms_kernel[(rows,)]( + x, + residual, + weight, + out, + resout, + HAS_RESIDUAL=residual is not None, + epsilon=1e-6, + num_warps=16, + num_stages=1, + enable_fp_fusion=True, + ) + if rows > 8: + continue # Frozen raw QPN2 rejects larger rows; norm-only fallback. + namespace = torch.ops._qpn2_capped if rows == 8 else torch.ops._C + if rows == 8: + op = namespace.gated if item["gated"] else namespace.gemm + else: + op = ( + namespace.nvfp4_qpn2_gated_sm70_out + if item["gated"] + else namespace.nvfp4_qpn2_gemm_sm70_out + ) + op( + result, + out, + item["codes"], + item["scales"], + item["global_scale"], + item["split_k"], + item["nacc"], + ) + + run() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph(keep_graph=True) + with torch.cuda.graph(graph): + run() + graph.instantiate() + return graph, records + + def fingerprint(graph): + kernels, parents = nodes(graph) + data = [ + ( + key, + n.name, + n.params.func, + n.grid, + [a.hex() for a in n.args], + parents[key], + ) + for key, n in sorted(kernels.items()) + ] + return hashlib.sha256(json.dumps(data).encode()).hexdigest() + + count = 0 + for kind in ("none", "half", "float"): + graph, records = make(8, kind) + edit = templates.patch(graph) + assert len(edit.edits) == 16 and len(edit.norms) == 8 + raw_sha = fingerprint(graph) + for cycle in range(args.cycles): + scale = (0.01, 1.0, 8.0)[cycle % 3] + for item, x, weight, residual, out, resout, result in records: + x.copy_( + torch.randn( + x.shape, device="cuda", dtype=x.dtype, generator=generator + ) + * scale + ) + if residual is not None: + residual.copy_( + torch.randn( + residual.shape, + device="cuda", + dtype=residual.dtype, + generator=generator, + ) + * scale + ) + saved = [] + for mode in ("control", "candidate", "control"): + edit.switch(mode) + graph.replay() + torch.cuda.synchronize() + outputs = [t for rec in records for t in rec[4:] if t is not None] + if not saved: + saved = [t.clone() for t in outputs] + else: + for index, (reference, actual) in enumerate(zip(saved, outputs)): + assert torch.equal( + reference.view(torch.uint8), actual.view(torch.uint8) + ), ( + rank, + kind, + cycle, + mode, + index, + int( + ( + reference.view(torch.uint8) + != actual.view(torch.uint8) + ).sum() + ), + ) + assert all(bool(torch.isfinite(t).all()) for t in outputs) + edit.check_canaries() + assert fingerprint(graph) == raw_sha, ( + "Executable edit changed raw graph" + ) + if mode == "candidate": + by_ptr = {r[4].data_ptr(): r[4] for r in records} + for original, _ in edit.edits: + if original.handle not in edit.norms: + continue + key = (original.name, tuple(map(len, original.args))) + idx = templates.norms[key][3] + logical = by_ptr[original.integer(idx)] + packed = ( + edit.norms[original.handle][0] + .view(320, 8, 16) + .permute(1, 0, 2) + .contiguous() + .view_as(logical) + ) + assert torch.equal( + logical.view(torch.uint8), packed.view(torch.uint8) + ), "Packed norm mismatch" + count += len(records) + report["cases"].append( + dict( + kind=kind, + cycle=cycle, + scale=scale, + outputs_state_bits_equal=True, + restore_equal=True, + canaries=True, + ) + ) + timing = None + if args.timing: + samples = {"control": [], "candidate": []} + for trial in range(7): + for mode in ( + ("control", "candidate") if trial % 2 == 0 else ("candidate", "control") + ): + edit.switch(mode) + for _ in range(8): + graph.replay() + start, end = ( + torch.cuda.Event(enable_timing=True), + torch.cuda.Event(enable_timing=True), + ) + start.record() + for _ in range(64): + graph.replay() + end.record() + end.synchronize() + samples[mode].append(start.elapsed_time(end) / 64) + timing = dict( + samples_ms=samples, + scope=( + "eight real column weights plus FP32-residual norms; " + "no row projections, communication or model timing" + ), + ) + fallbacks = [] + for rows in (1, 7, 9, 32): + graph, records = make(rows, "float") + edit = templates.patch(graph) + assert len(edit.edits) == 0 + fallbacks.append( + dict( + rows=rows, + scope="norm+raw QPN2" + if rows <= 8 + else "norm only; raw QPN2 correctly rejects M>8", + ) + ) + matrix_a = torch.randn( + (8, 128), device="cuda", dtype=torch.float16, generator=generator + ) + matrix_b = torch.randn( + (128, 256), device="cuda", dtype=torch.float16, generator=generator + ) + matrix_out = torch.empty((8, 256), device="cuda", dtype=torch.float16) + torch.mm(matrix_a, matrix_b, out=matrix_out) + torch.cuda.synchronize() + unrelated = torch.cuda.CUDAGraph(keep_graph=True) + with torch.cuda.graph(unrelated): + torch.mm(matrix_a, matrix_b, out=matrix_out) + unrelated.instantiate() + unrelated.replay() + torch.cuda.synchronize() + expected_matrix = matrix_out.clone() + untouched = templates.patch(unrelated) + assert len(untouched.edits) == 0 + untouched.switch("candidate") + unrelated.replay() + torch.cuda.synchronize() + assert torch.equal(matrix_out.view(torch.uint8), expected_matrix.view(torch.uint8)) + report["unrelated_cublas_graph_unmodified"] = True + report.update( + timing=timing, + source_hashes={ + path.name: hashlib.sha256(path.read_bytes()).hexdigest() + for path in ( + Path(__file__), + Path(__file__).with_name("sm70_qpn2_graph_layout.py"), + Path(__file__).with_name("sm70_qpn2_graph_nodes.py"), + args.dual_norm, + ) + }, + passed=True, + projection_cases=count, + switch_replays=len(report["cases"]) * 3, + unmodified_rows=fallbacks, + ) + gathered = [None] * dist.get_world_size() + dist.all_gather_object(gathered, report) + if rank == 0: + args.output.write_text( + json.dumps( + dict( + passed=True, + ranks=gathered, + scope="same captured operator graph; not model or timing evidence", + ), + indent=2, + ) + + "\n" + ) + dist.destroy_process_group() + print(json.dumps(dict(rank=rank, passed=True, projection_cases=count)), flush=True) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_sm70_qpn2_packed_mlp.py b/benchmarks/kernels/benchmark_sm70_qpn2_packed_mlp.py new file mode 100644 index 0000000000..374a4002ec --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_qpn2_packed_mlp.py @@ -0,0 +1,306 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Check private packed gate/up outputs feeding packed QPN2 row publishers. + +Both arms use packed column inputs and the same rank-ordered publication +protocol. Four gate/up outputs feed their respective down projections. The +candidate changes only that boundary layout, with no runtime transpose. Other +projections, eight reductions and a ninth ordinary push remain in the working +set. Input packing shared by both arms is outside timing; this is not a model +round. Use libraries produced by the accompanying private candidate builders. +""" + +import argparse +import hashlib +import json +import os +import statistics +from pathlib import Path + +import torch +import torch.distributed as dist + +from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--root", type=Path, required=True) + p.add_argument("--library", type=Path, required=True) + p.add_argument("--output", type=Path, required=True) + p.add_argument("--input-library", type=Path, required=True) + p.add_argument("--gated-library", type=Path, required=True) + p.add_argument("--control-publisher", type=Path, required=True) + p.add_argument("--correctness-only", action="store_true") + p.add_argument("--cycles", type=int, default=5) + p.add_argument("--coordination-backend", choices=["nccl", "gloo"], default="nccl") + args = p.parse_args() + if args.cycles < 1: + p.error("--cycles must be positive") + args.output.parent.mkdir(parents=True, exist_ok=True) + rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(rank) + assert torch.cuda.get_device_capability() == (7, 0) + torch.ops.load_library(str(args.library)) + auxiliary_libraries = [ + args.control_publisher, + args.input_library, + args.gated_library, + ] + for library in auxiliary_libraries: + torch.ops.load_library(str(library)) + dist.init_process_group(args.coordination_backend) + assert dist.get_world_size() == 4 + group = dist.new_group(backend="gloo") + ca = CustomAllreduce(group, rank, max_size=128 * 1024) + assert not ca.disabled and ca.sm70_tp4_push_buffer_ptrs is not None + peers = ca.sm70_tp4_push_buffer_ptrs + generator = torch.Generator(device="cuda").manual_seed(20260908 + rank) + items, provenance, guards = [], [], [] + + def output(shape): + storage = torch.full( + (shape[0] * shape[1] + 16,), -37, device="cuda", dtype=torch.float16 + ) + guards.append(storage) + return storage[8:-8].view(shape) + + try: + for path in sorted((args.root / f"rank{rank}").glob("*.pt")): + item = torch.load(path, map_location="cpu", weights_only=True) + assert item["rank"] == rank + provenance.append( + dict( + file=path.name, sha256=hashlib.sha256(path.read_bytes()).hexdigest() + ) + ) + item["codes"] = item["codes"].cuda() + item["scales"] = item["scales"].cuda() + item["input"] = torch.randn( + (8, item["k"]), dtype=torch.float16, device="cuda", generator=generator + ) + width = item["n"] // 2 if item["gated"] else item["n"] + item["output"] = [output((8, width)) for _ in range(2)] + item["row_parallel"] = not item["gated"] and item["n"] == 5120 + item["reduced"] = ( + [output((8, width)) for _ in range(2)] if item["row_parallel"] else None + ) + item["packed_input"] = ( + item["input"] + .view(8, item["k"] // 16, 16) + .permute(1, 0, 2) + .contiguous() + .view_as(item["input"]) + ) + items.append(item) + assert len(items) == 16 and sum(i["row_parallel"] for i in items) == 8 + rejected_shapes = 0 + for gated in (False, True): + item = next( + i for i in items if i["gated"] == gated and not i["row_parallel"] + ) + op = ( + torch.ops._qpn2_packed_mlp.gated + if gated + else torch.ops._qpn2_packed_input.gemm + ) + width = item["output"][0].shape[1] + for rows in (0, 1, 7, 9, 16, 32): + try: + op( + item["output"][0].new_empty((rows, width)), + item["input"].new_empty((rows, item["k"])), + item["codes"], + item["scales"], + item["global_scale"], + item["split_k"], + item["nacc"], + ) + except RuntimeError as error: + assert "requires M=8" in str(error), str(error) + rejected_shapes += 1 + else: + raise AssertionError(f"Packed operator accepted M={rows}") + gates = {i["layer"]: i for i in items if i["gated"]} + tail_input = torch.randn( + (8, 5120), dtype=torch.float16, device="cuda", generator=generator + ) + tail_outputs = [output((8, 5120)) for _ in range(2)] + + def run(arm): + for item in items: + row = item["row_parallel"] + down = row and item["k"] == 4352 + input_ = ( + gates[item["layer"]]["output"][arm] + if down + else (item["input"] if row else item["packed_input"]) + ) + arguments = [ + item["output"][arm], + input_, + item["codes"], + item["scales"], + item["global_scale"], + item["split_k"], + item["nacc"], + ] + if row: + producer = ( + torch.ops._qpn2_packed_row.publish + if arm and down + else torch.ops._qpn2_candidate.publish + ) + producer(*arguments, peers, rank) + torch.ops._qpn2_candidate.consume( + item["output"][arm], item["reduced"][arm], peers, rank + ) + else: + op = ( + ( + torch.ops._qpn2_packed_mlp.gated + if arm + else torch.ops._qpn2_packed_input.gated + ) + if item["gated"] + else torch.ops._qpn2_packed_input.gemm + ) + op(*arguments) + # Odd collective count plus an ordinary push call exercises epoch + # transitions across graphs and the two publication mechanisms. + ca.all_reduce(tail_input, out=tail_outputs[arm], registered=True) + + graphs = [] + for arm in range(2): + torch.cuda.synchronize() + dist.barrier() + graph = torch.cuda.CUDAGraph() + with ca.capture(), torch.cuda.graph(graph): + run(arm) + graphs.append(graph) + + def check(): + pairs = ( + [ + ( + i["output"][0], + i["output"][1] + .view(i["output"][1].shape[1] // 16, 8, 16) + .permute(1, 0, 2) + .contiguous() + .view_as(i["output"][0]), + ) + if i["gated"] + else i["output"] + for i in items + ] + + [i["reduced"] for i in items if i["row_parallel"]] + + [tail_outputs] + ) + for index, (a, b) in enumerate(pairs): + assert torch.isfinite(a).all() and torch.isfinite(b).all(), ( + "nonfinite", + rank, + index, + ) + assert torch.equal(a.view(torch.uint8), b.view(torch.uint8)), ( + "mismatch", + rank, + index, + int(torch.count_nonzero(a != b)), + ) + for g in guards: + assert (g[:8] == -37).all() and (g[-8:] == -37).all() + + for cycle in range(args.cycles): + for item in items: + item["input"].normal_(generator=generator) + item["packed_input"].copy_( + item["input"] + .view(8, item["k"] // 16, 16) + .permute(1, 0, 2) + .contiguous() + .view_as(item["input"]) + ) + tail_input.normal_(generator=generator) + for g in guards: + g[8:-8].fill_(float("nan")) + dist.barrier() + for arm in (0, 1) if cycle % 2 == 0 else (1, 0): + if rank == cycle % 4: + torch.cuda._sleep(20000) + graphs[arm].replay() + torch.cuda.synchronize() + check() + samples = [[], []] + for trial in range(0 if args.correctness_only else 7): + for arm in (0, 1) if trial % 2 == 0 else (1, 0): + for _ in range(20): + graphs[arm].replay() + torch.cuda.synchronize() + dist.barrier() + start, end = [torch.cuda.Event(enable_timing=True) for _ in range(2)] + start.record() + for _ in range(50): + graphs[arm].replay() + end.record() + end.synchronize() + ranks = [None] * 4 + dist.all_gather_object(ranks, start.elapsed_time(end) / 50, group=group) + samples[arm].append(max(ranks)) + check() + ranks = [None] * 4 + dist.all_gather_object( + ranks, dict(rank=rank, snapshots=provenance), group=group + ) + if rank == 0: + result = dict( + rank_weights=ranks, + shape="TP4 B1/q8 four consecutive layers", + projections=16, + row_parallel=8, + extra_ordinary_push=1, + activation_kind="synthetic_fp16_normal", + all_outputs_bitwise_equal=True, + canaries_intact=True, + rejected_non_q8_shapes_per_rank=rejected_shapes, + changing_input_cycles=args.cycles, + rank_skew_cycles=20000, + medians_ms=[statistics.median(x) for x in samples] + if not args.correctness_only + else None, + samples_ms=samples, + paired_saved_ms=[a - b for a, b in zip(*samples)], + library_sha256=hashlib.sha256(args.library.read_bytes()).hexdigest(), + communicator_sha256=hashlib.sha256( + Path(os.environ["VLLM_SM70_CUSTOM_AR_LIBRARY"]).read_bytes() + ).hexdigest(), + complete_round_performance=False, + auxiliary_library_sha256={ + str(p): hashlib.sha256(p.read_bytes()).hexdigest() + for p in auxiliary_libraries + }, + hypothesis=( + "Both arms use packed normalized column inputs and QPN2 " + "row publication. Candidate gate/up directly writes the " + "packed down-projection input; no separate transpose and " + "no arithmetic edits." + ), + coupled_gate_down_pairs=4, + shared_column_input_pack_timing_included=False, + ) + args.output.write_text(json.dumps(result, indent=2)) + print( + json.dumps( + {k: v for k, v in result.items() if k != "rank_weights"}, indent=2 + ), + flush=True, + ) + finally: + ca.close() + dist.destroy_process_group(group) + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_sm70_qpn2_publish.py b/benchmarks/kernels/benchmark_sm70_qpn2_publish.py new file mode 100644 index 0000000000..7728cce295 --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_qpn2_publish.py @@ -0,0 +1,219 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Validate QPN2 output publication using four ranks of real prepared weights. + +This screen uses synthetic activations and sixteen projections from consecutive +layers, with eight dependent all-reduces and an extra ordinary push call to +exercise mixed-protocol epoch transitions. It is not a complete model round. +Set VLLM_SM70_CUSTOM_AR_LIBRARY to a communicator built from the same header +as the candidate. Use Gloo process coordination for focused sanitizer runs. +""" + +import argparse +import hashlib +import json +import os +import statistics +from pathlib import Path + +import torch +import torch.distributed as dist + +from vllm import _sm70_ops as ops +from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--root", type=Path, required=True) + p.add_argument("--library", type=Path, required=True) + p.add_argument("--output", type=Path, required=True) + p.add_argument("--correctness-only", action="store_true") + p.add_argument("--cycles", type=int, default=5) + p.add_argument("--coordination-backend", choices=["nccl", "gloo"], default="nccl") + args = p.parse_args() + if args.cycles < 1: + p.error("--cycles must be positive") + args.output.parent.mkdir(parents=True, exist_ok=True) + rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(rank) + assert torch.cuda.get_device_capability() == (7, 0) + torch.ops.load_library(str(args.library)) + dist.init_process_group(args.coordination_backend) + assert dist.get_world_size() == 4 + group = dist.new_group(backend="gloo") + ca = CustomAllreduce(group, rank, max_size=128 * 1024) + assert not ca.disabled and ca.sm70_tp4_push_buffer_ptrs is not None + peers = ca.sm70_tp4_push_buffer_ptrs + generator = torch.Generator(device="cuda").manual_seed(20260908 + rank) + items, provenance, guards = [], [], [] + + def output(shape): + storage = torch.full( + (shape[0] * shape[1] + 16,), -37, device="cuda", dtype=torch.float16 + ) + guards.append(storage) + return storage[8:-8].view(shape) + + try: + for path in sorted((args.root / f"rank{rank}").glob("*.pt")): + item = torch.load(path, map_location="cpu", weights_only=True) + assert item["rank"] == rank + provenance.append( + dict( + file=path.name, sha256=hashlib.sha256(path.read_bytes()).hexdigest() + ) + ) + item["codes"] = item["codes"].cuda() + item["scales"] = item["scales"].cuda() + item["input"] = torch.randn( + (8, item["k"]), dtype=torch.float16, device="cuda", generator=generator + ) + width = item["n"] // 2 if item["gated"] else item["n"] + item["output"] = [output((8, width)) for _ in range(2)] + item["row_parallel"] = not item["gated"] and item["n"] == 5120 + item["reduced"] = ( + [output((8, width)) for _ in range(2)] if item["row_parallel"] else None + ) + items.append(item) + assert len(items) == 16 and sum(i["row_parallel"] for i in items) == 8 + tail_input = torch.randn( + (8, 5120), dtype=torch.float16, device="cuda", generator=generator + ) + tail_outputs = [output((8, 5120)) for _ in range(2)] + + def run(arm): + for item in items: + arguments = [ + item["output"][arm], + item["input"], + item["codes"], + item["scales"], + item["global_scale"], + item["split_k"], + item["nacc"], + ] + if arm and item["row_parallel"]: + torch.ops._qpn2_candidate.publish(*arguments, peers, rank) + torch.ops._qpn2_candidate.consume( + item["output"][arm], item["reduced"][arm], peers, rank + ) + else: + op = ( + ops.nvfp4_qpn2_gated_sm70_out + if item["gated"] + else ops.nvfp4_qpn2_gemm_sm70_out + ) + op(*arguments) + if item["row_parallel"]: + ca.all_reduce( + item["output"][arm], + out=item["reduced"][arm], + registered=True, + ) + # Odd collective count plus an ordinary push call exercises epoch + # transitions across graphs and the two publication mechanisms. + ca.all_reduce(tail_input, out=tail_outputs[arm], registered=True) + + graphs = [] + for arm in range(2): + torch.cuda.synchronize() + dist.barrier() + graph = torch.cuda.CUDAGraph() + with ca.capture(), torch.cuda.graph(graph): + run(arm) + graphs.append(graph) + + def check(): + pairs = ( + [i["output"] for i in items] + + [i["reduced"] for i in items if i["row_parallel"]] + + [tail_outputs] + ) + for index, (a, b) in enumerate(pairs): + assert torch.isfinite(a).all() and torch.isfinite(b).all(), ( + "nonfinite", + rank, + index, + ) + assert torch.equal(a.view(torch.uint8), b.view(torch.uint8)), ( + "mismatch", + rank, + index, + int(torch.count_nonzero(a != b)), + ) + for g in guards: + assert (g[:8] == -37).all() and (g[-8:] == -37).all() + + for cycle in range(args.cycles): + for item in items: + item["input"].normal_(generator=generator) + tail_input.normal_(generator=generator) + for g in guards: + g[8:-8].fill_(float("nan")) + dist.barrier() + for arm in (0, 1) if cycle % 2 == 0 else (1, 0): + if rank == cycle % 4: + torch.cuda._sleep(20000) + graphs[arm].replay() + torch.cuda.synchronize() + check() + samples = [[], []] + for trial in range(0 if args.correctness_only else 7): + for arm in (0, 1) if trial % 2 == 0 else (1, 0): + for _ in range(20): + graphs[arm].replay() + torch.cuda.synchronize() + dist.barrier() + start, end = [torch.cuda.Event(enable_timing=True) for _ in range(2)] + start.record() + for _ in range(50): + graphs[arm].replay() + end.record() + end.synchronize() + ranks = [None] * 4 + dist.all_gather_object(ranks, start.elapsed_time(end) / 50, group=group) + samples[arm].append(max(ranks)) + check() + ranks = [None] * 4 + dist.all_gather_object( + ranks, dict(rank=rank, snapshots=provenance), group=group + ) + if rank == 0: + result = dict( + rank_weights=ranks, + shape="TP4 B1/q8 four consecutive layers", + projections=16, + row_parallel=8, + extra_ordinary_push=1, + activation_kind="synthetic_fp16_normal", + all_outputs_bitwise_equal=True, + canaries_intact=True, + changing_input_cycles=args.cycles, + rank_skew_cycles=20000, + medians_ms=[statistics.median(x) for x in samples] + if not args.correctness_only + else None, + samples_ms=samples, + paired_saved_ms=[a - b for a, b in zip(*samples)], + library_sha256=hashlib.sha256(args.library.read_bytes()).hexdigest(), + communicator_sha256=hashlib.sha256( + Path(os.environ["VLLM_SM70_CUSTOM_AR_LIBRARY"]).read_bytes() + ).hexdigest(), + complete_round_performance=False, + ) + args.output.write_text(json.dumps(result, indent=2)) + print( + json.dumps( + {k: v for k, v in result.items() if k != "rank_weights"}, indent=2 + ), + flush=True, + ) + finally: + ca.close() + dist.destroy_process_group(group) + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_sm70_qpn2_working_set.py b/benchmarks/kernels/benchmark_sm70_qpn2_working_set.py new file mode 100644 index 0000000000..0a1d8a0a3c --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_qpn2_working_set.py @@ -0,0 +1,229 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Race QPN2 kernels over prepared weights from four consecutive TP4 layers. + +Snapshots contain the actual runtime codes/scales and dispatch parameters, +including fused projections and padding. Activations are frozen synthetic +FP16 inputs. This measures a sequential projection working set, not a complete +verification round. The optional candidate library registers _qpn2_candidate. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import statistics +import sys +from pathlib import Path + +import regex as re +import torch + + +def save_model_snapshot(model: torch.nn.Module, directory: Path, rank: int) -> int: + """Export prepared runtime weights once, before graph capture/profiling.""" + from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_w4a4_nvfp4 import ( # noqa: E501 + _SM70_NVFP4_QPN2_CONFIGS, + ) + + directory.mkdir(parents=True, exist_ok=True) + count = 0 + for name, module in model.named_modules(): + match = re.search(r"\.layers\.(\d+)\.", name) + if ( + match is None + or int(match[1]) >= 4 + or not getattr(module, "sm70_nvfp4_qpn2", False) + ): + continue + k = module.input_size_per_partition + n = module.sm70_nvfp4_qpn2_output_size + gated = module.sm70_nvfp4_qpn2_gated_silu + split_k, nacc = _SM70_NVFP4_QPN2_CONFIGS[k, n, gated] + torch.save( + { + "name": name, + "layer": int(match[1]), + "rank": rank, + "k": k, + "n": n, + "gated": gated, + "split_k": split_k, + "nacc": nacc, + "global_scale": module.sm70_nvfp4_qpn2_global_scale, + "codes": module.sm70_nvfp4_qpn2_codes.detach().cpu(), + "scales": module.sm70_nvfp4_qpn2_scales.detach().cpu(), + }, + directory / f"{count:02}-{name}.pt", + ) + count += 1 + if count != 16: + raise ValueError(f"Expected sixteen runtime QPN2 projections, got {count}") + return count + + +def main() -> None: + from vllm import _sm70_ops as ops + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("snapshots", type=Path) + parser.add_argument("--candidate-library", type=Path) + parser.add_argument("--candidate-namespace", default="_qpn2_candidate") + parser.add_argument("--output", required=True, type=Path) + parser.add_argument("--iterations", type=int, default=50) + parser.add_argument("--trials", type=int, default=7) + args = parser.parse_args() + if args.iterations < 1 or args.trials < 2: + parser.error("positive iterations and at least two alternating trials required") + if torch.cuda.get_device_capability() != (7, 0): + raise ValueError("SM70 is required") + if args.candidate_library: + torch.ops.load_library(str(args.candidate_library)) + candidate_ops = getattr(torch.ops, args.candidate_namespace) + torch.manual_seed(20260908) + items, provenance = [], [] + for path in sorted(args.snapshots.glob("*.pt")): + item = torch.load(path, map_location="cpu", weights_only=True) + provenance.append( + { + "file": path.name, + "sha256": hashlib.sha256(path.read_bytes()).hexdigest(), + **{k: v for k, v in item.items() if not isinstance(v, torch.Tensor)}, + } + ) + item["codes"] = item["codes"].cuda() + item["scales"] = item["scales"].cuda() + item["candidate_scales"] = item["scales"] + if args.candidate_library and hasattr(candidate_ops, "prepare_scales"): + item["candidate_scales"] = candidate_ops.prepare_scales( + item["scales"], item["global_scale"] + ) + item["input"] = torch.randn(8, item["k"], device="cuda", dtype=torch.float16) + width = item["n"] // 2 if item["gated"] else item["n"] + item["outputs"] = [ + torch.empty(8, width, device="cuda", dtype=torch.float16) for _ in range(2) + ] + items.append(item) + shapes = {(x["k"], x["n"]) for x in items} + if ( + len(items) != 16 + or {x["layer"] for x in items} != {0, 1, 2, 3} + or len({x["rank"] for x in items}) != 1 + or shapes + != {(1536, 5120), (5120, 3584), (5120, 4128), (4352, 5120), (5120, 8704)} + ): + raise ValueError("Expected all five shapes in sixteen consecutive projections") + + def run(arm: int) -> None: + for item in items: + if arm and args.candidate_library: + op = candidate_ops.gated if item["gated"] else candidate_ops.gemm + else: + op = ( + ops.nvfp4_qpn2_gated_sm70_out + if item["gated"] + else ops.nvfp4_qpn2_gemm_sm70_out + ) + op( + item["outputs"][arm], + item["input"], + item["codes"], + item["candidate_scales"] if arm else item["scales"], + item["global_scale"], + item["split_k"], + item["nacc"], + ) + + graphs = [] + for arm in range(2): + for _ in range(3): + run(arm) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run(arm) + graphs.append(graph) + + def check_outputs() -> None: + exact = [ + torch.equal( + x["outputs"][0].view(torch.uint8), x["outputs"][1].view(torch.uint8) + ) + for x in items + ] + finite = all( + torch.isfinite(x["outputs"][a]).all() for x in items for a in range(2) + ) + if not finite or not all(exact): + args.output.write_text( + json.dumps({"bitwise_equal": exact, "finite": bool(finite)}) + ) + raise RuntimeError("Copy/load candidate changed numerical output") + + for graph in graphs: + graph.replay() + check_outputs() + for _ in range(20): + for graph in graphs: + graph.replay() + samples = [[], []] + for trial in range(args.trials): + for arm in range(2) if trial % 2 == 0 else range(1, -1, -1): + start, end = (torch.cuda.Event(enable_timing=True) for _ in range(2)) + start.record() + for _ in range(args.iterations): + graphs[arm].replay() + end.record() + end.synchronize() + samples[arm].append(start.elapsed_time(end) / args.iterations) + check_outputs() + control_library = Path(sys.modules["vllm._C"].__file__).resolve() + result = { + "snapshots": provenance, + "activation_seed": 20260908, + "activation_kind": "synthetic_fp16_normal", + "shape": "B1/q8 TP4 four consecutive layers", + "weight_working_set_bytes": sum( + x[k].numel() * x[k].element_size() + for x in items + for k in ("codes", "scales") + ), + "candidate_weight_working_set_bytes": sum( + x[k].numel() * x[k].element_size() + for x in items + for k in ("codes", "candidate_scales") + ), + "candidate_scale_dtype": str(items[0]["candidate_scales"].dtype), + "candidate_namespace": args.candidate_namespace + if args.candidate_library + else None, + "candidate_library_sha256": hashlib.sha256( + args.candidate_library.read_bytes() + ).hexdigest() + if args.candidate_library + else None, + "all_outputs_bitwise_equal": True, + "parity_checked": "after first graph replay and after all timing trials", + "control_library": str(control_library), + "control_library_sha256": hashlib.sha256( + control_library.read_bytes() + ).hexdigest(), + "torch_version": torch.__version__, + "torch_cuda_version": torch.version.cuda, + "gpu": torch.cuda.get_device_name(), + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "iterations": args.iterations, + "trials": args.trials, + "working_set_samples_ms": samples, + "working_set_medians_ms": [statistics.median(s) for s in samples], + "paired_saved_ms": [a - b for a, b in zip(*samples)], + "complete_round_performance": False, + } + args.output.write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps({k: v for k, v in result.items() if k != "snapshots"}, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/benchmark_sm70_sparse_dense_topk.py b/benchmarks/kernels/benchmark_sm70_sparse_dense_topk.py new file mode 100644 index 0000000000..b5b49d6f49 --- /dev/null +++ b/benchmarks/kernels/benchmark_sm70_sparse_dense_topk.py @@ -0,0 +1,234 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Check an SM70 sparse top-k candidate against PyTorch's dense tie order. + +This private operator screen uses frozen PyTorch 2.10.0, a 62080-token local +vocabulary, 64 reranked FP32 candidates, seven/eight rows and top-k 16/20/21. +Candidate IDs must be unique and in the local vocabulary. Reproduce the native +sort wrapper with build_sm70_native_sort_candidate.py. No serving route is +installed; this primitive screen does not measure a complete DFlash2 round. + +The selection order follows PyTorch v2.10.0 TensorTopK.cu's multiblock gather: +all keys above the cutoff in vocabulary order, then cutoff ties in vocabulary +order. The native sort wrapper preserves the final unstable tie permutation. +""" + +import argparse +import hashlib +import json +import os +import statistics +from pathlib import Path + +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit +def compact_gather( + IDS, + LOGITS, + VALUES, + OUT_IDS, + K: tl.constexpr, + ROW_STRIDE: tl.constexpr, + VOCAB_START: tl.constexpr, +): + row = tl.program_id(0) + c = tl.arange(0, 64) + ids = tl.load(IDS + row * ROW_STRIDE + c) + values = tl.load(LOGITS + row * ROW_STRIDE + c) + bits = values.to(tl.uint32, bitcast=True) + keys = bits ^ tl.where((bits & 0x80000000) != 0, 0xFFFFFFFF, 0x80000000).to( + tl.uint32 + ) + keys = tl.where(values != values, 0xFFFFFFFF, keys).to(tl.uint32) + ordered = tl.sort(keys, descending=True) + kth = tl.sum(tl.where(c == K - 1, ordered, 0), 0) + above = keys > kth + equal = keys == kth + n_above = tl.sum(above.to(tl.int32), 0) + before = ids[:, None] > ids[None, :] + above_rank = tl.sum((before & above[None, :]).to(tl.int32), 1) + equal_rank = tl.sum((before & equal[None, :]).to(tl.int32), 1) + pos = tl.where(above, above_rank, n_above + equal_rank) + selected = above | (equal & (pos < K)) + # If the kth value is -Inf, implicit dense background entries also tie. + # All selected non-background values still precede those entries. + if kth == 0x007FFFFF: + tl.store(VALUES + row * K + pos, values, above) + tl.store(OUT_IDS + row * K + pos, ids + VOCAB_START, above) + low_id = c.to(tl.int64) + prior_above = tl.sum( + ((ids[None, :] < low_id[:, None]) & above[None, :]).to(tl.int32), 1 + ) + is_above = ( + tl.sum(((ids[None, :] == low_id[:, None]) & above[None, :]).to(tl.int32), 1) + != 0 + ) + low_pos = n_above + c - prior_above + keep = (c < K) & ~is_above & (low_pos < K) + tl.store(VALUES + row * K + low_pos, -float("inf"), keep) + tl.store(OUT_IDS + row * K + low_pos, low_id + VOCAB_START, keep) + else: + tl.store(VALUES + row * K + pos, values, selected) + tl.store(OUT_IDS + row * K + pos, ids + VOCAB_START, selected) + + +def select(candidate_ids, logits, values, ids, vocab_start=0): + compact_gather[(logits.shape[0],)]( + candidate_ids, + logits, + values, + ids, + K=values.shape[1], + ROW_STRIDE=logits.stride(0), + VOCAB_START=vocab_start, + num_warps=4, + ) + torch.ops.quasar_native_sort.sort_pairs(values, ids) + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--library", type=Path, required=True) + p.add_argument("--output", type=Path, required=True) + p.add_argument("--correctness-only", action="store_true") + a = p.parse_args() + assert torch.cuda.get_device_capability() == (7, 0) + torch.ops.load_library(str(a.library)) + torch.manual_seed(20260908) + cases = [] + width = 62080 + for rows in (7, 8): + for k in (16, 20, 21): + for scenario in ( + "random", + "ties", + "boundary_ties", + "zeros", + "signed_zero", + "nan", + "infinity", + "negative_infinity", + "few_finite", + ): + cid = torch.stack( + [torch.randperm(width, device="cuda")[:64] for _ in range(rows)] + ) + logits = torch.randn(rows, 64, device="cuda") + if scenario == "ties": + logits = logits.round() + if scenario == "boundary_ties": + logits[:, k - 5 : k + 10] = 3.0 + if scenario in ("zeros", "signed_zero"): + logits.zero_() + if scenario == "signed_zero": + logits[:, ::2] = -0.0 + if scenario == "nan": + logits[:, ::3] = float("nan") + if scenario == "infinity": + logits[:, ::3] = float("inf") + if scenario == "negative_infinity": + logits.fill_(-float("inf")) + if scenario == "few_finite": + logits[:, 3:] = -float("inf") + values = [torch.empty(rows, k, device="cuda") for _ in range(2)] + ids = [ + torch.empty(rows, k, device="cuda", dtype=torch.int64) + for _ in range(2) + ] + dense = torch.full((rows, width), -float("inf"), device="cuda") + dense.scatter_(1, cid, logits) + torch.topk(dense, k, sorted=True, out=(values[0], ids[0])) + select(cid, logits, values[1], ids[1]) + exact_values = torch.equal( + values[0].view(torch.uint8), values[1].view(torch.uint8) + ) + exact_ids = torch.equal(ids[0], ids[1]) + result = dict( + rows=rows, + k=k, + scenario=scenario, + values_equal=exact_values, + ids_equal=exact_ids, + ) + if not exact_values or not exact_ids: + result.update( + control_values=values[0].tolist(), + candidate_values=values[1].tolist(), + control_ids=ids[0].tolist(), + candidate_ids=ids[1].tolist(), + ) + cases.append(result) + a.output.write_text( + json.dumps(dict(cases=cases, passed=False), indent=2) + ) + raise RuntimeError(f"Dense order mismatch: {rows}, {k}, {scenario}") + cases.append(result) + report = dict( + cases=cases, + passed=True, + library_sha256=hashlib.sha256(a.library.read_bytes()).hexdigest(), + torch_version=torch.__version__, + local_vocab_size=width, + cuda_visible_devices=os.environ.get("CUDA_VISIBLE_DEVICES"), + source_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(), + complete_round_performance=False, + ) + if not a.correctness_only: + rows, k = 8, 20 + cid = torch.stack( + [torch.randperm(width, device="cuda")[:64] for _ in range(rows)] + ) + logits = torch.randn(rows, 64, device="cuda") + values = [torch.empty(rows, k, device="cuda") for _ in range(2)] + ids = [torch.empty(rows, k, device="cuda", dtype=torch.int64) for _ in range(2)] + dense = torch.empty(rows, width, device="cuda") + + def run(arm): + if arm: + select(cid, logits, values[1], ids[1]) + else: + dense.fill_(-float("inf")) + dense.scatter_(1, cid, logits) + torch.topk(dense, k, sorted=True, out=(values[0], ids[0])) + + graphs = [] + for arm in range(2): + run(arm) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run(arm) + graphs.append(graph) + for _ in range(3000): + for graph in graphs: + graph.replay() + samples = [[], []] + for trial in range(7): + local = [[], []] + for arm in (0, 1, 1, 0) if trial % 2 == 0 else (1, 0, 0, 1): + start, end = [torch.cuda.Event(enable_timing=True) for _ in range(2)] + start.record() + for _ in range(300): + graphs[arm].replay() + end.record() + end.synchronize() + local[arm].append(start.elapsed_time(end) / 300) + for arm in range(2): + samples[arm].append(statistics.mean(local[arm])) + assert torch.equal(values[0].view(torch.uint8), values[1].view(torch.uint8)) + assert torch.equal(ids[0], ids[1]) + report.update( + samples_ms=samples, + medians_ms=[statistics.median(s) for s in samples], + paired_saved_ms=[x - y for x, y in zip(*samples)], + ) + a.output.write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps({k: v for k, v in report.items() if k != "cases"}, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/build_sm70_grouped_attention_candidate.py b/benchmarks/kernels/build_sm70_grouped_attention_candidate.py new file mode 100644 index 0000000000..42dda02c61 --- /dev/null +++ b/benchmarks/kernels/build_sm70_grouped_attention_candidate.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""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. +""" + +import argparse +import hashlib +import json +import shutil +from pathlib import Path + + +def replace_once(source: str, old: str, new: str) -> str: + if source.count(old) != 1: + raise ValueError(f"Expected one source anchor: {old}") + return source.replace(old, new) + + +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("--build", action="store_true") + args = parser.parse_args() + root = Path(__file__).resolve().parents[2] / "flash-attention-v100" + original = root / "kernel/flash_decode_paged.cu" + source = original.read_text() + 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: + source = replace_once( + source, + "constexpr int kGroupedVerifyRows = 48;", + "constexpr int kGroupedVerifyRows = 16;", + ) + source = replace_once( + source, + "constexpr int kGroupedVerifyThreads = 512;", + "constexpr int kGroupedVerifyThreads = 256;", + ) + source = replace_once( + source, + "kernel<<>>", + "kernel<<>>", + ) + source = replace_once( + source, + "flash_attention_grouped_e4m3_fp32_paged(", + "private_grouped_e4m3_fp32_paged(", + ) + source += ( + "\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n" + ' m.def("run", &private_grouped_e4m3_fp32_paged);\n}\n' + ) + directory = args.output_dir.resolve() + sources = directory / "sources" + sources.mkdir(parents=True, exist_ok=True) + for name in ("include", "kernel"): + target = sources / name + target.mkdir(exist_ok=True) + for pattern in ("*.h", "*.cuh"): + for header in (root / name).glob(pattern): + shutil.copy2(header, target) + shutil.copy2(root / "LICENSE", sources) + path = sources / "kernel/grouped-attention.cu" + path.write_text(source) + # Retain Flash-V100's existing math flags; this is a scheduling candidate. + 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 = { + "input_source_sha256": hashlib.sha256(original.read_bytes()).hexdigest(), + "source_sha256": hashlib.sha256(path.read_bytes()).hexdigest(), + "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, + "extra_cuda_cflags": flags, + "scope": "Private operator candidate; full-model admission required", + } + if args.build: + from torch.utils.cpp_extension import load + + build = directory / "build" + build.mkdir(exist_ok=True) + library = Path( + load( + name="sm70_grouped_attention_candidate", + sources=[str(path)], + build_directory=str(build), + extra_cuda_cflags=flags, + 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_native_sort_candidate.py b/benchmarks/kernels/build_sm70_native_sort_candidate.py new file mode 100644 index 0000000000..164cbccd6d --- /dev/null +++ b/benchmarks/kernels/build_sm70_native_sort_candidate.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Build a private wrapper around the frozen PyTorch CUDA key/value sorter. + +This has no CUDA implementation of its own: the loaded PyTorch library owns +the final sort, including its unstable tie order. Build with an empty +CUDA_VISIBLE_DEVICES and an explicit CUDA_HOME; no GPU context is needed. +The wrapper is only for the paired SM70 sparse/dense top-k experiment. +""" + +import argparse +import hashlib +import json +from pathlib import Path + +SOURCE = r""" +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +#include +#include +#include +#include + +void sort_pairs(torch::Tensor values, torch::Tensor ids) { + TORCH_CHECK(values.is_cuda() && ids.device() == values.device()); + TORCH_CHECK(values.scalar_type() == at::kFloat && ids.scalar_type() == at::kLong); + TORCH_CHECK(values.dim() == 2 && ids.sizes() == values.sizes()); + TORCH_CHECK(values.is_contiguous() && ids.is_contiguous()); + const at::cuda::OptionalCUDAGuard guard(device_of(values)); + at::native::sortKeyValueInplace(values, ids, 1, true, false); +} +TORCH_LIBRARY_FRAGMENT(quasar_native_sort, m) { + m.def("sort_pairs(Tensor(a!) values, Tensor(b!) ids) -> ()"); + m.impl("sort_pairs", torch::kCUDA, &sort_pairs); +} +""" + + +def main(): + import torch + from torch.utils.cpp_extension import load + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-dir", required=True, type=Path) + parser.add_argument("--build", action="store_true") + args = parser.parse_args() + if torch.__version__.split("+")[0] != "2.10.0": + parser.error("This experiment requires the frozen PyTorch 2.10.0 sorter") + args.output_dir.mkdir(parents=True, exist_ok=True) + source = args.output_dir / "native_sort.cpp" + source.write_text(SOURCE.lstrip()) + torch_root = Path(torch.__file__).resolve().parent + header = torch_root / "include/ATen/native/cuda/Sort.h" + result = { + "torch_version": torch.__version__, + "torch_git_version": torch.version.git_version, + "cuda_version": torch.version.cuda, + "source_sha256": hashlib.sha256(source.read_bytes()).hexdigest(), + "native_header_sha256": hashlib.sha256(header.read_bytes()).hexdigest(), + "library_sha256": None, + "serving_route_installed": False, + } + if args.build: + build_dir = args.output_dir / "build" + build_dir.mkdir(exist_ok=True) + library = load( + name="quasar_native_sort", + sources=[str(source.resolve())], + build_directory=str(build_dir.resolve()), + extra_cflags=["-O3"], + with_cuda=True, + is_python_module=False, + verbose=True, + ) + result["library_sha256"] = hashlib.sha256( + Path(library).read_bytes() + ).hexdigest() + (args.output_dir / "manifest.json").write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(result, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/build_sm70_qpn2_chunked_candidate.py b/benchmarks/kernels/build_sm70_qpn2_chunked_candidate.py new file mode 100644 index 0000000000..a00e03d91b --- /dev/null +++ b/benchmarks/kernels/build_sm70_qpn2_chunked_candidate.py @@ -0,0 +1,206 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# CUDA source anchors retain their original spelling. +# ruff: noqa: E501 +"""Build a private two-chunk QPN2 publisher with independent packet channels. + +The existing generator supplies the unchanged dot-product arithmetic and +packet protocol. Only column indexing and the consumer's output addressing +change. A caller must complete each local publisher before its consumer and +join both consumers before the next dependent projection. This does not +install a model route or change the communicator's registered storage. +""" + +import argparse +import hashlib +import json +import subprocess +import sys +from pathlib import Path + + +def replace_once(source: str, old: str, new: str) -> str: + assert source.count(old) == 1, (old, source.count(old)) + return source.replace(old, new) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-dir", required=True, type=Path) + parser.add_argument("--build", action="store_true") + args = parser.parse_args() + root = Path(__file__).resolve().parents[2] + directory = args.output_dir.resolve() + template_dir = directory / "template" + subprocess.run( + [ + sys.executable, + str(Path(__file__).with_name("build_sm70_qpn2_publish_candidate.py")), + "--output-dir", + str(template_dir), + ], + check=True, + ) + template = (template_dir / "sources/qpn2-publish.cu").read_text() + producer_start = template.index( + "template \n" + "__global__ void nvfp4_qpn2_publish_sm70_kernel" + ) + producer_end = template.index("\n}\n\nnamespace vllm {", producer_start) + producer = template[producer_start:producer_end] + producer = producer.replace( + "nvfp4_qpn2_publish_sm70_kernel", "nvfp4_qpn2_chunk_publish_sm70_kernel" + ) + producer = replace_once( + producer, + "const uint32_t* epochs) {", + "const uint32_t* epochs, int column_start, int chunk_columns) {", + ) + producer = replace_once( + producer, + "const int tile = blockIdx.x;", + "const int tile = blockIdx.x + column_start / 32;", + ) + producer = replace_once( + producer, + "const int packed_index = output_index / P::size;", + "const int packed_index = (output_row * chunk_columns + tile * 32 + output_col - column_start) / P::size;", + ) + consumer_start = template.index("template \n__global__", producer_end) + consumer_end = template.index( + "\n}\n\nvllm::RankData qpn2_peer_pointers", consumer_start + ) + consumer = template[consumer_start:consumer_end].replace( + "qpn2_consume_published", "qpn2_consume_chunk" + ) + consumer = replace_once( + consumer, + "int packed_size) {", + "int packed_size, int column_start, int chunk_columns) {", + ) + consumer = replace_once( + consumer, + "reinterpret_cast(output)[offset] =", + "reinterpret_cast(output)[(offset / (chunk_columns / P::size)) * (5120 / P::size) + column_start / P::size + offset % (chunk_columns / P::size)] =", + ) + # Keep separately named controls available without colliding with the + # frozen production publisher loaded in the same benchmark process. + source = template.replace("_qpn2_candidate", "_qpn2_chunked_base") + source += "\nnamespace {\n" + producer + "\n}\n" + source += "\nnamespace vllm {\n" + consumer + "\n}\n" + source += r""" +void qpn2_chunk_initialize(torch::Tensor anchor, int64_t pointer) { + TORCH_CHECK(anchor.is_cuda() && pointer != 0 && pointer % 16 == 0, + "same-device anchor and aligned local channel required"); + const at::cuda::OptionalCUDAGuard guard(device_of(anchor)); + const auto stream = at::cuda::getCurrentCUDAStream(); + auto* local = reinterpret_cast(pointer); + C10_CUDA_CHECK(cudaMemsetAsync(local, 0, + vllm::kSm70Tp4PushAllreduceSignalBytes, stream)); + C10_CUDA_CHECK(cudaMemsetAsync(local + vllm::kSm70Tp4PushAllreduceSignalBytes, + vllm::kSm70Tp4PushAllreduceSentinelByte, + vllm::kSm70Tp4PushAllreduceGenericBufferBytes - + vllm::kSm70Tp4PushAllreduceSignalBytes, stream)); +} + +void qpn2_chunk_publish(torch::Tensor out, torch::Tensor input, + torch::Tensor codes, torch::Tensor scales, double global_scale, + int64_t split_k, int64_t nacc, std::vector pointers, + int64_t rank, int64_t column_start) { + check_qpn2_tensors(out, input, codes, scales, false); + const int k = input.size(1); + TORCH_CHECK(input.size(0) == 8 && out.size(1) == 5120 && nacc == 2 && + ((k == 1536 && split_k == 8) || (k == 4352 && split_k == 16)), + "exact TP4 q8 row projection required"); + TORCH_CHECK(column_start == 0 || column_start == 2560, + "two complete 2560-column chunks required"); + const at::cuda::OptionalCUDAGuard guard(device_of(input)); + const auto stream = at::cuda::getCurrentCUDAStream(); + const auto peers = qpn2_peer_pointers(pointers, rank); + const auto* in = reinterpret_cast(input.data_ptr()); + auto* dst = reinterpret_cast(out.data_ptr()); + const auto* epochs = reinterpret_cast(peers.ptrs[rank]); + if (split_k == 8) + nvfp4_qpn2_chunk_publish_sm70_kernel<8, 2, 1><<<80, 256, 0, stream>>>( + codes.data_ptr(), scales.data_ptr(), in, dst, + 5120, k, 8, global_scale, peers, rank, epochs, column_start, 2560); + else + nvfp4_qpn2_chunk_publish_sm70_kernel<16, 2, 1><<<80, 512, 0, stream>>>( + codes.data_ptr(), scales.data_ptr(), in, dst, + 5120, k, 8, global_scale, peers, rank, epochs, column_start, 2560); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void qpn2_chunk_consume(torch::Tensor out, std::vector pointers, + int64_t rank, int64_t column_start) { + TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kFloat16 && + out.is_contiguous() && out.dim() == 2 && out.size(0) == 8 && + out.size(1) == 5120, "exact q8 FP16 output required"); + TORCH_CHECK(column_start == 0 || column_start == 2560, + "two complete 2560-column chunks required"); + const at::cuda::OptionalCUDAGuard guard(device_of(out)); + const auto stream = at::cuda::getCurrentCUDAStream(); + const auto peers = qpn2_peer_pointers(pointers, rank); + vllm::qpn2_consume_chunk<4><<<20, 128, 0, stream>>>( + peers.ptrs[rank], nullptr, + reinterpret_cast(out.data_ptr()), rank, + 2560, column_start, 2560); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +TORCH_LIBRARY_FRAGMENT(_qpn2_chunked, ops) { + ops.def("initialize(Tensor anchor, int pointer) -> ()"); + ops.impl("initialize", torch::kCUDA, &qpn2_chunk_initialize); + ops.def("publish(Tensor(a!) out, Tensor input, Tensor codes, Tensor scales, float global_scale, int split_k, int nacc, int[] pointers, int rank, int column_start) -> ()"); + ops.impl("publish", torch::kCUDA, &qpn2_chunk_publish); + ops.def("consume(Tensor(a!) out, int[] pointers, int rank, int column_start) -> ()"); + ops.impl("consume", torch::kCUDA, &qpn2_chunk_consume); +} +""" + path = template_dir / "sources/qpn2-chunked.cu" + path.write_text(source) + flags = [ + "-O3", + "-lineinfo", + "-gencode=arch=compute_70,code=sm_70", + "-DVLLM_NVFP4_QPN2_STANDALONE", + "-DVLLM_NVFP4_QPN2_BENCHMARK_CANDIDATE", + "-Xptxas=-v", + ] + manifest = dict( + source=str(path), + source_sha256=hashlib.sha256(path.read_bytes()).hexdigest(), + template_manifest=json.loads((template_dir / "manifest.json").read_text()), + builder_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(), + repository=str(root), + chunks=2, + columns_per_chunk=2560, + separate_channels=True, + producer_waits=False, + local_producer_event_required=True, + flags=flags, + installed=False, + ) + if args.build: + from torch.utils.cpp_extension import load + + build = directory / "build" + build.mkdir(exist_ok=True) + library = load( + name="qpn2_chunked_candidate", + sources=[str(path)], + build_directory=str(build), + extra_cuda_cflags=flags, + is_python_module=False, + verbose=True, + ) + manifest.update( + library=str(library), + library_sha256=hashlib.sha256(Path(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_qpn2_cooperative_mlp.py b/benchmarks/kernels/build_sm70_qpn2_cooperative_mlp.py new file mode 100644 index 0000000000..c92d8ae6ae --- /dev/null +++ b/benchmarks/kernels/build_sm70_qpn2_cooperative_mlp.py @@ -0,0 +1,191 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Literal CUDA templates retain the validated generated-source spelling. +# ruff: noqa: E501 +"""Private cooperative gate/up -> row publication, with the original consumer.""" + +import argparse +import hashlib +import json +import subprocess +import sys +from pathlib import Path +from textwrap import dedent + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--output-dir", type=Path, required=True) + p.add_argument("--build", action="store_true") + args = p.parse_args() + W = Path(__file__).resolve().parents[2] + D = args.output_dir.resolve() + subprocess.run( + [ + sys.executable, + str(W / "benchmarks/kernels/build_sm70_qpn2_publish_candidate.py"), + "--output-dir", + str(D), + ], + check=True, + ) + path = D / "sources/qpn2-publish.cu" + s = path.read_text() + parent_sha = hashlib.sha256(path.read_bytes()).hexdigest() + + def function(name): + start = s.index( + "template \n__global__ void " + name + ) + opening = s.index("{", start) + level = 1 + i = opening + 1 + while level: + if s[i] == "{": + level += 1 + elif s[i] == "}": + level -= 1 + i += 1 + return s[start:i] + + gate = function("nvfp4_qpn2_gated_sm70_kernel") + publish = function("nvfp4_qpn2_publish_sm70_kernel") + gate = gate.replace( + "__global__ void nvfp4_qpn2_gated_sm70_kernel", + "__device__ __forceinline__ void coop_mlp_gate_device", + ) + publish = publish.replace( + "__global__ void nvfp4_qpn2_publish_sm70_kernel", + "__device__ __forceinline__ void coop_mlp_publish_device", + ) + # Fixed q8 only. All accumulators, K partitions, activation boundaries and + # packet payloads are taken from the exact existing source. + for old, new in [ + ( + "const int row_base = blockIdx.y * kQpn2RowsPerCta * RowTiles;", + "const int row_base = 0;", + ) + ]: + assert gate.count(old) == publish.count(old) == 1 + gate = gate.replace(old, new) + publish = publish.replace(old, new) + s = "#include \n" + s + s += ( + "\nnamespace {\n" + + gate + + "\n" + + publish + + dedent(r""" + __global__ __launch_bounds__(512,2) void qpn2_coop_mlp_kernel( + const uint8_t* gate_codes, const uint8_t* gate_scales, + const uint8_t* down_codes, const uint8_t* down_scales, + const half* input, half* gate_output, half* output, + float gate_global, float down_global, vllm::RankData peers, + int rank, const uint32_t* epochs) { + if (blockIdx.x < 136) { + coop_mlp_gate_device<8,2,1>(gate_codes,gate_scales,input,gate_output, + 4352,5120,8,gate_global); + } + cooperative_groups::this_grid().sync(); + // Every producer finishes before the separate, frozen consumer launches. + // There is no cross-rank wait or polling in this cooperative kernel. + coop_mlp_publish_device<16,2,1>(down_codes,down_scales,gate_output,output, + 5120,4352,8,down_global,peers,rank,epochs); + } + } + void qpn2_coop_mlp(torch::Tensor gate, torch::Tensor out, torch::Tensor input, + torch::Tensor gate_codes, torch::Tensor gate_scales, + torch::Tensor down_codes, torch::Tensor down_scales, + double gate_global, double down_global, std::vector pointers, int64_t rank) { + check_qpn2_tensors(gate,input,gate_codes,gate_scales,true); + check_qpn2_tensors(out,gate,down_codes,down_scales,false); + TORCH_CHECK(input.size(0)==8 && input.size(1)==5120 && gate.size(1)==4352 && + out.size(0)==8 && out.size(1)==5120, "exact TP4 q8 MLP required"); + const at::cuda::OptionalCUDAGuard guard(device_of(input)); + auto* props=at::cuda::getCurrentDeviceProperties(); + TORCH_CHECK(props->major==7 && props->minor==0 && props->cooperativeLaunch, + "SM70 cooperative launch required"); + int resident=0; + C10_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&resident,qpn2_coop_mlp_kernel,512,0)); + TORCH_CHECK(resident*props->multiProcessorCount>=160,"160 resident CTAs required"); + auto peers=qpn2_peer_pointers(pointers,rank); + auto* gc=gate_codes.data_ptr();auto* gs=gate_scales.data_ptr(); + auto* dc=down_codes.data_ptr();auto* ds=down_scales.data_ptr(); + auto* x=reinterpret_cast(input.data_ptr()); + auto* g=reinterpret_cast(gate.data_ptr()); + auto* y=reinterpret_cast(out.data_ptr()); + float gg=static_cast(gate_global),dg=static_cast(down_global); + int r=static_cast(rank); + auto* epochs=reinterpret_cast(peers.ptrs[rank]); + void* parameters[]={&gc,&gs,&dc,&ds,&x,&g,&y,&gg,&dg,&peers,&r,&epochs}; + C10_CUDA_CHECK(cudaLaunchCooperativeKernel(reinterpret_cast(qpn2_coop_mlp_kernel), + dim3(160),dim3(512),parameters,0,at::cuda::getCurrentCUDAStream())); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + std::vector qpn2_coop_resources() { + cudaFuncAttributes attr{};int resident=0; + C10_CUDA_CHECK(cudaFuncGetAttributes(&attr,qpn2_coop_mlp_kernel)); + C10_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&resident,qpn2_coop_mlp_kernel,512,0)); + return {attr.numRegs,static_cast(attr.sharedSizeBytes),static_cast(attr.localSizeBytes),resident}; + } + TORCH_LIBRARY_FRAGMENT(_qpn2_candidate,ops) { + ops.def("mlp(Tensor(a!) gate, Tensor(b!) out, Tensor input, Tensor gate_codes, Tensor gate_scales, Tensor down_codes, Tensor down_scales, float gate_global, float down_global, int[] pointers, int rank) -> ()"); + ops.impl("mlp",torch::kCUDA,&qpn2_coop_mlp); + ops.def("resources() -> int[]"); + ops.impl("resources",&qpn2_coop_resources); + } + """) + ) + # Avoid cross-DSO symbol preemption. The independently frozen publisher is + # loaded separately by the benchmark, and retains its own operator namespace. + s = s.replace("_qpn2_candidate", "_qpn2_coop_mlp") + s = ( + s.replace("nvfp4_qpn2_", "coopbase_nvfp4_qpn2_") + .replace("qpn2_publish", "coopbase_qpn2_publish") + .replace("qpn2_consume", "coopbase_qpn2_consume") + .replace("qpn2_peer_pointers", "coopbase_qpn2_peer_pointers") + ) + path = D / "sources/qpn2-coop-mlp.cu" + path.write_text(s) + flags = [ + "-O3", + "-lineinfo", + "-gencode=arch=compute_70,code=sm_70", + "-DVLLM_NVFP4_QPN2_STANDALONE", + "-DVLLM_NVFP4_QPN2_BENCHMARK_CANDIDATE", + "-Xptxas=-v", + "--maxrregcount=64", + ] + report = dict( + source=str(path), + source_sha256=hashlib.sha256(path.read_bytes()).hexdigest(), + parent_source_sha256=parent_sha, + flags=flags, + installed=False, + scope="Private TP4 q8 gate/up -> published down, fixed 160 resident CTAs and one grid barrier; original arithmetic and peer packet protocol", + ) + if args.build: + from torch.utils.cpp_extension import load + + build = D / "build" + build.mkdir(exist_ok=True) + lib = Path( + load( + name="qpn2_coop_mlp", + sources=[str(path)], + build_directory=str(build), + extra_cuda_cflags=flags, + is_python_module=False, + verbose=True, + ) + ) + report.update( + library=str(lib), + library_sha256=hashlib.sha256(lib.read_bytes()).hexdigest(), + ) + (D / "cooperative-manifest.json").write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/build_sm70_qpn2_dual_norm.py b/benchmarks/kernels/build_sm70_qpn2_dual_norm.py new file mode 100644 index 0000000000..34b14a248f --- /dev/null +++ b/benchmarks/kernels/build_sm70_qpn2_dual_norm.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Derive a private dual-store norm; preserve the original expressions.""" + +import argparse +import ast +import hashlib +import json +from pathlib import Path + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + W = Path(__file__).resolve().parents[2] + source_path = W / "vllm/model_executor/layers/layernorm.py" + source = source_path.read_text() + tree = ast.parse(source) + parts = ["from vllm.triton_utils import tl, triton\n"] + for name in ( + "_sm70_dflash2_fixed_gemma_rms_kernel", + "_sm70_dflash2_gemma_fused_add_rms_kernel", + ): + node = next( + n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == name + ) + segment = "@triton.jit\n" + ast.get_source_segment(source, node) + segment = segment.replace(name, name + "_dual_store", 1) + segment = segment.replace( + " residual_out,\n", " residual_out,\n packed_out,\n", 1 + ) + store = next(n for n in reversed(node.body) if isinstance(n, ast.Expr)) + tail = ast.get_source_segment(source, store) + old_index = ( + "normalized_out + row * 5120 + cols" + if "fixed" in name + else "normalized_out + row * hidden_size + cols" + ) + assert tail.count(old_index) == 1 + tail = tail.replace( + old_index, "packed_out + (cols // 16) * 128 + row * 16 + cols % 16" + ) + segment += "\n " + tail + "\n" + parts.append(segment) + output = args.output.resolve() + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text("\n".join(parts)) + manifest = dict( + parent_source_sha256=hashlib.sha256(source_path.read_bytes()).hexdigest(), + source_sha256=hashlib.sha256(output.read_bytes()).hexdigest(), + scope="private dual-store q8 node candidate, not installed", + ) + output.with_suffix(".manifest.json").write_text( + json.dumps(manifest, indent=2) + "\n" + ) + print(json.dumps(manifest)) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/build_sm70_qpn2_packed_input_candidate.py b/benchmarks/kernels/build_sm70_qpn2_packed_input_candidate.py new file mode 100644 index 0000000000..1d6163e98c --- /dev/null +++ b/benchmarks/kernels/build_sm70_qpn2_packed_input_candidate.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Literal CUDA anchors retain their production source spelling. +# ruff: noqa: E501 +"""Build private q8 QPN2 layout experiments without changing their arithmetic. + +The input's physical layout is [K/16, 8, 16]. With --pack-gated-output, +the fused gate/up projection also writes [hidden/16, 8, 16]. These tensors +must only be passed to consumers of the corresponding private layout. +This does not install a serving operator or enable a model route. +""" + +import argparse +import hashlib +import json +import shutil +from pathlib import Path + + +def replace_once(source: str, old: str, new: str, count: int = 1) -> str: + if source.count(old) != count: + raise ValueError(f"Expected {count} occurrences of source anchor: {old}") + return source.replace(old, new) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--pack-gated-output", action="store_true") + parser.add_argument("--build", action="store_true") + args = parser.parse_args() + root = Path(__file__).resolve().parents[2] + original = root / "csrc/sm70_turbomind/ops/nvfp4_qpn2_sm70.cu" + directory = args.output_dir.resolve() + sources = directory / "sources" + sources.mkdir(parents=True, exist_ok=True) + source = replace_once( + original.read_text(), + """const half* input_row = input + static_cast(row) * k; + input01 = *reinterpret_cast(input_row + group * 16); + input23 = *reinterpret_cast(input_row + group * 16 + 8);""", + """const half* input_row = input + static_cast(group) * 128 + row * 16; + input01 = *reinterpret_cast(input_row); + input23 = *reinterpret_cast(input_row + 8);""", + count=2, + ) + source = replace_once( + source, + "m >= 1 && m <= kQpn2MaxRows && out.size(0) == m", + "m == 8 && out.size(0) == m", + ) + source = replace_once( + source, + '"NVFP4 QPN2 requires M in [1, ", kQpn2MaxRows, "]"', + '"Private packed-input NVFP4 QPN2 requires M=8"', + ) + if args.pack_gated_output: + source = replace_once( + source, + """output[static_cast(output_row) * hidden + blockIdx.x * 32 + + output_col] = __hmul(silu, up_half);""", + """const int logical_col = blockIdx.x * 32 + output_col; + output[static_cast(logical_col / 16) * 128 + + output_row * 16 + logical_col % 16] = __hmul(silu, up_half);""", + ) + namespace = "_qpn2_packed_mlp" if args.pack_gated_output else "_qpn2_packed_input" + prefix = "packedmlp_" if args.pack_gated_output else "packedinput_" + source = source.replace("_qpn2_candidate", namespace) + source = source.replace("nvfp4_qpn2_", prefix + "nvfp4_qpn2_") + path = sources / "qpn2-packed-input.cu" + path.write_text(source) + shutil.copy2(original.parent / "LICENSE.v100-skinny", sources) + flags = [ + "-O3", + "-lineinfo", + "-gencode=arch=compute_70,code=sm_70", + "-DVLLM_NVFP4_QPN2_STANDALONE", + "-DVLLM_NVFP4_QPN2_BENCHMARK_CANDIDATE", + "-Xptxas=-v", + ] + manifest = { + "input_source_sha256": hashlib.sha256(original.read_bytes()).hexdigest(), + "source_sha256": hashlib.sha256(path.read_bytes()).hexdigest(), + "namespace": namespace, + "input_layout": "[K/16, 8, 16]", + "gated_output_layout": "[hidden/16, 8, 16]" + if args.pack_gated_output + else "[8, hidden]", + "extra_cuda_cflags": flags, + "math_mode": "default", + "scope": "Private q8 operator; producer/consumer pairing required", + } + if args.build: + from torch.utils.cpp_extension import load + + build = directory / "build" + build.mkdir(exist_ok=True) + library = Path( + load( + name=namespace.removeprefix("_") + "_candidate", + sources=[str(path)], + build_directory=str(build), + extra_cuda_cflags=flags, + is_python_module=False, + verbose=True, + ) + ) + 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_qpn2_publish_candidate.py b/benchmarks/kernels/build_sm70_qpn2_publish_candidate.py new file mode 100644 index 0000000000..ce3c04e918 --- /dev/null +++ b/benchmarks/kernels/build_sm70_qpn2_publish_candidate.py @@ -0,0 +1,281 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Literal CUDA templates and source anchors retain their original spelling. +# ruff: noqa: E501 +"""Generate a private SM70 QPN2 publisher and matching consumer for experiments. + +The generator extracts production arithmetic and the existing packet protocol, +then changes only publication placement and parameter passing. It deliberately +fails if its source anchors change. Its generated source/DSO is a benchmark +candidate, not an installed vLLM operator. Use one candidate library per process. +""" + +import argparse +import hashlib +import json +import shutil +from pathlib import Path + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--build", action="store_true") + parser.add_argument( + "--packed-input", + action="store_true", + help="Read private [K/16, 8, 16] input in the q8 publisher only", + ) + parser.add_argument( + "--use-fast-math", + action="store_true", + help="Reproduce historical experiments; changes gated SiLU math", + ) + args = parser.parse_args() + cuda_flags = [ + "-O3", + "-lineinfo", + "-gencode=arch=compute_70,code=sm_70", + "-DVLLM_NVFP4_QPN2_STANDALONE", + "-DVLLM_NVFP4_QPN2_BENCHMARK_CANDIDATE", + "-Xptxas=-v", + ] + if args.use_fast_math: + cuda_flags.append("--use_fast_math") + A = args.output_dir.resolve() + D = A / "sources" + D.mkdir(parents=True, exist_ok=True) + W = Path(__file__).resolve().parents[2] + for name in ( + "custom_all_reduce.cuh", + "cub_helpers.h", + "sm70_tile_runtime_signal.cuh", + ): + shutil.copy2(W / "csrc" / name, D / name) + source = (W / "csrc/sm70_turbomind/ops/nvfp4_qpn2_sm70.cu").read_text() + start = source.index( + "template \n__global__ void nvfp4_qpn2_sm70_kernel" + ) + end = source.index( + "template \n__global__ void nvfp4_qpn2_gated_sm70_kernel", + start, + ) + producer = source[start:end].replace( + "nvfp4_qpn2_sm70_kernel", "nvfp4_qpn2_publish_sm70_kernel" + ) + if args.packed_input: + old = """const half* input_row = input + static_cast(row) * k; + input01 = *reinterpret_cast(input_row + group * 16); + input23 = *reinterpret_cast(input_row + group * 16 + 8);""" + assert producer.count(old) == 1 + producer = producer.replace( + old, + """const half* input_row = input + static_cast(group) * 128 + row * 16; + input01 = *reinterpret_cast(input_row); + input23 = *reinterpret_cast(input_row + 8);""", + ) + old = "int m, float global_scale) {" + assert producer.count(old) == 1 + producer = producer.replace( + old, "int m, float global_scale, vllm::RankData peers, int rank) {" + ) + old = """ output[static_cast(output_row) * n + tile * 32 + output_col] = + __float2half(value);""" + assert producer.count(old) == 1 + producer = producer.replace( + old, + """ const int output_index = output_row * n + tile * 32 + output_col; + const half rounded = __float2half(value); + output[output_index] = rounded; + // Every active warp owns one complete 32-column row fragment. Pack + // eight adjacent FP16 results without modifying their payload bits. + const unsigned bits = __half_as_ushort(rounded); + const unsigned pair = bits | (__shfl_down_sync(0xffffffffu, bits, 1) << 16); + const uint4 packet = make_uint4( + pair, __shfl_down_sync(0xffffffffu, pair, 2), + __shfl_down_sync(0xffffffffu, pair, 4), + __shfl_down_sync(0xffffffffu, pair, 6)); + if ((lane & 7) == 0) { + using P = vllm::packed_t::P; + P payload = *reinterpret_cast(&packet); +#pragma unroll + for (int i = 0; i < P::size; ++i) + vllm::sm70_push_escape_sentinel(payload.data[i]); + const int packed_index = output_index / P::size; + const int consumer_block = packed_index / vllm::kSm70Tp4PushAllreduceThreads; + const auto* epochs = reinterpret_cast(peers.ptrs[rank]); + const uint32_t epoch = epochs[consumer_block]; + constexpr int stride = vllm::kSm70Tp4PushAllreduceMaxBytes / sizeof(P); + const int epoch_offset = (epoch * 4 + rank) * stride; +#pragma unroll + for (int peer = 0; peer < 4; ++peer) { + auto* base = const_cast(reinterpret_cast(peers.ptrs[peer])); + void* destination = base + vllm::kSm70Tp4PushAllreduceSignalBytes + epoch_offset * sizeof(P); + vllm::sm70_push_store_volatile_16b(payload, destination, packed_index); + } + }""", + ) + header = (D / "custom_all_reduce.cuh").read_text() + start_c = header.index( + "template \n__global__ void __launch_bounds__(1024, 1)\n sm70_cross_device_reduce_1stage_push(" + ) + end_c = header.index( + "template \n__global__ void __launch_bounds__(1024, 1)\n sm70_cross_device_reduce_sum2_1stage_push(", + start_c, + ) + consumer = header[start_c:end_c].replace( + "sm70_cross_device_reduce_1stage_push", "qpn2_consume_published" + ) + start_publish = consumer.index( + " P value = reinterpret_cast(input)[offset];" + ) + end_publish = consumer.index(" P peer_values[ngpus];", start_publish) + consumer = consumer[:start_publish] + consumer[end_publish:] + source = '#include "custom_all_reduce.cuh"\n' + source + # Original producer arithmetic and standalone namespace stay available as + # independent controls. The new code uses the identical header/constants as + # the isolated communicator that owns the IPC buffers. + source += "\nnamespace {\n" + producer + "\n}\n" + source += "\nnamespace vllm {\n" + consumer + "\n}\n" + source += r""" +vllm::RankData qpn2_peer_pointers(const std::vector& pointers, int64_t rank) { + TORCH_CHECK(pointers.size() == 4 && rank >= 0 && rank < 4, "TP4 pointers/rank required"); + vllm::RankData peers{}; + for (int i = 0; i < 4; ++i) { + TORCH_CHECK(pointers[i] != 0 && pointers[i] % 16 == 0, "invalid IPC pointer"); + peers.ptrs[i] = reinterpret_cast(pointers[i]); + } + return peers; +} + +void qpn2_publish(torch::Tensor out, torch::Tensor input, torch::Tensor codes, + torch::Tensor scales, double global_scale, int64_t split_k, + int64_t nacc, std::vector pointers, int64_t rank) { + check_qpn2_tensors(out, input, codes, scales, false); + const int k = input.size(1); + TORCH_CHECK(input.size(0) == 8 && out.size(1) == 5120 && nacc == 2 && + ((k == 1536 && split_k == 8) || (k == 4352 && split_k == 16)), "exact TP4 q8 row projection required"); + const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const auto peers = qpn2_peer_pointers(pointers, rank); + const auto* in = reinterpret_cast(input.data_ptr()); + auto* dst = reinterpret_cast(out.data_ptr()); + if (split_k == 8) + nvfp4_qpn2_publish_sm70_kernel<8, 2, 1><<<160, 256, 0, stream>>>(codes.data_ptr(), scales.data_ptr(), in, dst, 5120, k, 8, global_scale, peers, rank); + else + nvfp4_qpn2_publish_sm70_kernel<16, 2, 1><<<160, 512, 0, stream>>>(codes.data_ptr(), scales.data_ptr(), in, dst, 5120, k, 8, global_scale, peers, rank); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void qpn2_consume(torch::Tensor projected, torch::Tensor out, + std::vector pointers, int64_t rank) { + TORCH_CHECK(projected.is_cuda() && out.is_cuda() && projected.device() == out.device() && + projected.scalar_type() == torch::kFloat16 && out.scalar_type() == torch::kFloat16 && + projected.is_contiguous() && out.is_contiguous() && projected.sizes() == out.sizes() && + out.dim() == 2 && out.size(0) == 8 && out.size(1) == 5120, "exact q8 FP16 output required"); + const at::cuda::OptionalCUDAGuard device_guard(device_of(projected)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const auto peers = qpn2_peer_pointers(pointers, rank); + vllm::qpn2_consume_published<4><<<80, 128, 0, stream>>>(peers, nullptr, + reinterpret_cast(out.data_ptr()), rank, 5120); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +TORCH_LIBRARY_FRAGMENT(_qpn2_candidate, ops) { + ops.def("publish(Tensor(a!) out, Tensor input, Tensor codes, Tensor scales, float global_scale, int split_k, int nacc, int[] pointers, int rank) -> ()"); + ops.impl("publish", torch::kCUDA, &qpn2_publish); + ops.def("consume(Tensor projected, Tensor(a!) out, int[] pointers, int rank) -> ()"); + ops.impl("consume", torch::kCUDA, &qpn2_consume); +} +""" + # Keep the local pointer out of dynamically indexed by-value arrays. + # Otherwise ptxas materializes a 64-byte local stack in both kernels. + old = "int m, float global_scale, vllm::RankData peers, int rank) {" + assert source.count(old) == 1 + source = source.replace( + old, + "int m, float global_scale, vllm::RankData peers, int rank, const uint32_t* epochs) {", + ) + old = " const auto* epochs = reinterpret_cast(peers.ptrs[rank]);\n" + assert source.count(old) == 1 + source = source.replace(old, "") + old = "8, global_scale, peers, rank);" + assert source.count(old) == 2 + source = source.replace( + old, + "8, global_scale, peers, rank, reinterpret_cast(peers.ptrs[rank]));", + ) + source = source.replace( + "qpn2_consume_published(RankData push_buffers,", + "qpn2_consume_published(const void* local_pointer,", + ) + source = source.replace( + "reinterpret_cast(push_buffers.ptrs[rank])", + "reinterpret_cast(local_pointer)", + ) + source = source.replace( + "qpn2_consume_published<4><<<80, 128, 0, stream>>>(peers, nullptr,", + "qpn2_consume_published<4><<<80, 128, 0, stream>>>(peers.ptrs[rank], nullptr,", + ) + if args.packed_input: + source = source.replace("_qpn2_candidate", "_qpn2_packed_row") + source = source.replace("nvfp4_qpn2_", "packedrow_nvfp4_qpn2_") + source = source.replace("qpn2_publish", "qpn2_packedrow_publish") + source = source.replace("qpn2_consume", "qpn2_packedrow_consume") + source = source.replace("qpn2_peer_pointers", "qpn2_packedrow_peer_pointers") + shutil.copy2( + W / "csrc/sm70_turbomind/ops/LICENSE.v100-skinny", D / "LICENSE.v100-skinny" + ) + p = D / "qpn2-publish.cu" + p.write_text(source) + manifest = { + str(f.name): hashlib.sha256(f.read_bytes()).hexdigest() + for f in D.iterdir() + if f.is_file() + } + (A / "manifest.json").write_text( + json.dumps( + dict( + sources=manifest, + extra_cuda_cflags=cuda_flags, + math_mode="fast" if args.use_fast_math else "default", + publisher_input_layout="[K/16, 8, 16]" + if args.packed_input + else "[8, K]", + input_sources={ + name: hashlib.sha256((W / name).read_bytes()).hexdigest() + for name in ( + "csrc/sm70_turbomind/ops/nvfp4_qpn2_sm70.cu", + "csrc/custom_all_reduce.cuh", + ) + }, + hypothesis="Publish completed FP16 output fragments during QPN2 epilogue, without polling or waiting in producer CTAs. The separate consumer keeps the established FP32 rank order, two epochs and sentinel cleanup. This differs from the previously rejected producer-poll fusion.", + ), + indent=2, + ) + ) + print(p) + + if args.build: + from torch.utils.cpp_extension import load + + build_dir = A / "build" + build_dir.mkdir(exist_ok=True) + library = load( + name="qpn2_publish_candidate", + sources=[str(p)], + build_directory=str(build_dir), + extra_cuda_cflags=cuda_flags, + is_python_module=False, + verbose=True, + ) + result = json.loads((A / "manifest.json").read_text()) + result["library"] = str(library) + result["library_sha256"] = hashlib.sha256( + Path(library).read_bytes() + ).hexdigest() + (A / "manifest.json").write_text(json.dumps(result, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/kernels/build_sm70_qpn2_q8_candidate.py b/benchmarks/kernels/build_sm70_qpn2_q8_candidate.py new file mode 100644 index 0000000000..0f30b63a96 --- /dev/null +++ b/benchmarks/kernels/build_sm70_qpn2_q8_candidate.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Build a private fixed-q8 QPN2 candidate with a 64-register cap. + +Only row-bound handling and compiler register allocation differ from the +production source. Keep the reduction chains, weight decoding and ordinary +CUDA math. The host entry points reject every row count except eight. +This builder does not install a serving route or enable a default. +""" + +import argparse +import hashlib +import json +import shutil +from pathlib import Path + + +def replace(source: str, old: str, new: str, count: int) -> str: + if source.count(old) != count: + raise ValueError(f"Expected {count} occurrences of source anchor: {old}") + return source.replace(old, new) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--build", action="store_true") + args = parser.parse_args() + root = Path(__file__).resolve().parents[2] + original = root / "csrc/sm70_turbomind/ops/nvfp4_qpn2_sm70.cu" + directory = args.output_dir.resolve() + sources = directory / "sources" + sources.mkdir(parents=True, exist_ok=True) + source = replace( + original.read_text(), + "const int row_base = blockIdx.y * kQpn2RowsPerCta * RowTiles;", + "const int row_base = 0;", + 2, + ) + source = replace(source, "if (row < m) {", "{", 2) + source = replace(source, "if (output_row < m) {", "{", 2) + source = replace( + source, + "m >= 1 && m <= kQpn2MaxRows && out.size(0) == m", + "m == 8 && out.size(0) == m", + 1, + ) + source = replace( + source, + '"NVFP4 QPN2 requires M in [1, ", kQpn2MaxRows, "]"', + '"Private fixed-q8 QPN2 requires M=8"', + 1, + ) + for gated in ("false", "true"): + anchor = f" check_qpn2_tensors(out, input, codes, scales, {gated});" + source = replace( + source, + anchor, + ' TORCH_CHECK(input.size(0) == 8, "benchmark candidate is q8-only");' + "\n" + anchor, + 1, + ) + source = source.replace("_qpn2_candidate", "_qpn2_capped") + source = source.replace("nvfp4_qpn2_", "q8capped_nvfp4_qpn2_") + path = sources / "qpn2-fixed-cap64-sidecar.cu" + path.write_text(source) + shutil.copy2(original.parent / "LICENSE.v100-skinny", sources) + flags = [ + "-O3", + "-lineinfo", + "-gencode=arch=compute_70,code=sm_70", + "-DVLLM_NVFP4_QPN2_STANDALONE", + "-DVLLM_NVFP4_QPN2_BENCHMARK_CANDIDATE", + "-Xptxas=-v", + "--maxrregcount=64", + ] + manifest = { + "input_source_sha256": hashlib.sha256(original.read_bytes()).hexdigest(), + "source_sha256": hashlib.sha256(path.read_bytes()).hexdigest(), + "namespace": "_qpn2_capped", + "extra_cuda_cflags": flags, + "math_mode": "default", + "row_count": 8, + "scope": "Private operator candidate; full-model admission required", + } + if args.build: + from torch.utils.cpp_extension import load + + build = directory / "build" + build.mkdir(exist_ok=True) + library = Path( + load( + name="qpn2_fixed_cap64_sidecar", + sources=[str(path)], + build_directory=str(build), + extra_cuda_cflags=flags, + is_python_module=False, + verbose=True, + ) + ) + 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/sm70_context_probe_candidate_route.py b/benchmarks/kernels/sm70_context_probe_candidate_route.py new file mode 100644 index 0000000000..fee8360617 --- /dev/null +++ b/benchmarks/kernels/sm70_context_probe_candidate_route.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Explicit experimental installation of context work behind the target probe. + +The unchanged CPU cutoff guard consumes a pinned copy while the original +context graph runs on its original stream. Importing this file changes no +dispatch. KV storage, proposal, sampling arithmetic and RNG are untouched. +""" + +import functools + +import torch + + +def install_context_probe_candidate(*, shadow: bool = False) -> None: + from vllm.v1.worker.gpu.model_runner import GPUModelRunner + from vllm.v1.worker.gpu.spec_decode.dflash2 import sparse_rejection + from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import DFlash2Speculator + + capture_original = DFlash2Speculator.capture + prepare_original = DFlash2Speculator.prepare_target_context + sample_original = GPUModelRunner.sample + guard_original = sparse_rejection._compact_target_requires_reference + active = None + + class PendingContext: + def __init__(self, owner): + self.owner = owner + self.arguments = None + self.probe = torch.empty((8, 21), dtype=torch.float32, pin_memory=True) + self.ready = torch.cuda.Event() + self.calls = 0 + + def flush(self): + if self.arguments is not None: + arguments = self.arguments + self.arguments = None + prepare_original(self.owner, *arguments) + if shadow: + hidden = self.owner.hidden_states[:8].clone() + kv = [t.clone() for t in self.owner._context_projected_kv] + prepare_original(self.owner, *arguments) + assert torch.equal( + hidden.view(torch.uint8), + self.owner.hidden_states[:8].view(torch.uint8), + ) + for left, right in zip(kv, self.owner._context_projected_kv): + assert torch.equal( + left.view(torch.uint8), right.view(torch.uint8) + ) + + @functools.wraps(capture_original) + def capture(self): + result = capture_original(self) + if self._context_compute_graph is not None: + self._context_probe_candidate = PendingContext(self) + return result + + @functools.wraps(prepare_original) + def prepare(self, batch, hidden, aux): + context = getattr(self, "_context_probe_candidate", None) + if context is not None: + assert context.arguments is None, "unconsumed previous context" + if ( + context is None + or batch.num_reqs != 1 + or batch.num_tokens != 8 + or batch.num_draft_tokens != 7 + or batch.is_prefilling_np[0] + ): + return prepare_original(self, batch, hidden, aux) + self._prepared_context_batch = None + context.arguments = (batch, hidden, aux) + + @functools.wraps(guard_original) + def guard(probe_logits, temperature, top_p): + context = active + if context is None or context.arguments is None: + return guard_original(probe_logits, temperature, top_p) + if ( + probe_logits.shape != (8, 21) + or probe_logits.dtype != torch.float32 + or probe_logits.device != context.owner.device + or not probe_logits.is_contiguous() + ): + context.flush() + return guard_original(probe_logits, temperature, top_p) + context.probe.copy_(probe_logits.detach(), non_blocking=True) + context.ready.record(torch.cuda.current_stream(context.owner.device)) + context.flush() + # Wait for the probe copy, not for the context graph queued after it. + context.ready.synchronize() + result = guard_original(context.probe, temperature, top_p) + if shadow: + original_probe = probe_logits.detach().cpu() + assert torch.equal( + context.probe.view(torch.uint8), original_probe.view(torch.uint8) + ) + assert result == guard_original(original_probe, temperature, top_p) + context.calls += 1 + if context.calls == 1 or (shadow and context.calls % 64 == 0): + print( + "CONTEXT_PROBE_ROUTE " + f"rank={torch.distributed.get_rank()} calls={context.calls} " + f"shadow={shadow} guard={guard_original.__module__}", + flush=True, + ) + return result + + @functools.wraps(sample_original) + def sample(self, *args, **kwargs): + nonlocal active + context = getattr(self.speculator, "_context_probe_candidate", None) + if context is None: + return sample_original(self, *args, **kwargs) + assert active is None, "nested sampling would reuse the pinned probe" + active = context + try: + result = sample_original(self, *args, **kwargs) + # Structured output, full-vocabulary fallback and teacher forcing + # may skip the compact guard. Complete their deferred preparation + # before the caller mutates state or submits proposal consumers. + context.flush() + return result + except BaseException: + context.arguments = None + raise + finally: + active = None + + DFlash2Speculator.capture = capture + DFlash2Speculator.prepare_target_context = prepare + GPUModelRunner.sample = sample + sparse_rejection._compact_target_requires_reference = guard diff --git a/benchmarks/kernels/sm70_dflash2_common_candidate_route.py b/benchmarks/kernels/sm70_dflash2_common_candidate_route.py new file mode 100644 index 0000000000..3417734d0e --- /dev/null +++ b/benchmarks/kernels/sm70_dflash2_common_candidate_route.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Explicit DFlash2 schedules independent of target weight quantization. + +The manifest selects independently reversible GDN/context/attention/selection +routes. It never loads a QPN2 projection library or inspects quantization names. +Existing tensor, state and shape guards decide whether an operation qualifies. +Set VLLM_SM70_DFLASH2_COMMON_MANIFEST before worker startup and select +CommonDFlash2Extension to use this experimental entry point. Importing without +that setting changes no route. Model/precision admission remains separate. +""" + +import hashlib +import importlib.util +import json +import os +from pathlib import Path + +import torch + + +def _library(entry: dict) -> Path: + path = Path(entry["library"]).resolve() + digest = hashlib.sha256(path.read_bytes()).hexdigest() + if digest != entry["sha256"]: + raise ValueError(f"DFlash2 library digest mismatch: {path.name}") + return path + + +def install_sparse_dense_order(select) -> None: + from vllm.model_executor.layers import vocab_parallel_embedding as embedding + + original = embedding._sm70_dflash2_dense_order_topk + seen = set() + + def run( + sparse_logits, + candidate_ids, + candidate_logits, + values, + ids, + selector_k, + vocab_start_index, + ): + eligible = ( + candidate_logits.dtype == torch.float32 + and candidate_logits.shape[0] in (7, 8) + and candidate_logits.shape[1] == 64 + and sparse_logits.shape[1] == 62080 + and selector_k in (16, 20, 21) + and candidate_logits.is_contiguous() + and candidate_ids.is_contiguous() + and values.is_contiguous() + and ids.is_contiguous() + ) + if torch.compiler.is_compiling() or not eligible: + return original( + sparse_logits, + candidate_ids, + candidate_logits, + values, + ids, + selector_k, + vocab_start_index, + ) + key = (candidate_logits.shape[0], selector_k) + if key not in seen: + seen.add(key) + print( + f"COMMON_DFLASH2_SORT_ROUTE rank={torch.distributed.get_rank()} " + f"rows={key[0]} k={key[1]}", + flush=True, + ) + return select(candidate_ids, candidate_logits, values, ids, vocab_start_index) + + embedding._sm70_dflash2_dense_order_topk = run + + +def install_common_routes(manifest: dict) -> None: + if manifest.get("schema_version") != 1: + raise ValueError("Unsupported common-route manifest version") + allowed = {"schema_version", "gdn_value_tile", "context_probe", "attention", "sort"} + if set(manifest) - allowed: + raise ValueError("Unknown common-route manifest entries") + # Validate every native dependency before modifying Python dispatch. + libraries = { + key: _library(manifest[key]) + for key in ("attention", "sort") + if manifest.get(key) is not None + } + if "sort" in libraries and torch.__version__.split("+")[0] != "2.10.0": + raise ValueError("Exact sparse tie ordering requires frozen PyTorch 2.10.0") + if "sort" in libraries: + from benchmarks.kernels.benchmark_sm70_sparse_dense_topk import select + + torch.ops.load_library(str(libraries["sort"])) + install_sparse_dense_order(select) + if "attention" in libraries: + from benchmarks.kernels.sm70_grouped_attention_candidate_route import ( + install_grouped_attention_candidate, + ) + + path = libraries["attention"] + name = manifest["attention"]["module_name"] + spec = importlib.util.spec_from_file_location(name, path) + if spec is None or spec.loader is None: + raise ImportError(f"Cannot load grouped attention: {path.name}") + native = importlib.util.module_from_spec(spec) + spec.loader.exec_module(native) + logged = False + + def attention(q, *args, **kwargs): + nonlocal logged + if not logged: + logged = True + print( + "COMMON_DFLASH2_ATTENTION_CAPTURE " + f"rank={torch.distributed.get_rank()} shape={tuple(q.shape)}", + flush=True, + ) + return native.run(q, *args, **kwargs) + + install_grouped_attention_candidate(attention) + if manifest.get("context_probe", False): + from benchmarks.kernels.sm70_context_probe_candidate_route import ( + install_context_probe_candidate, + ) + + install_context_probe_candidate() + if manifest.get("gdn_value_tile", False): + from benchmarks.kernels.sm70_gdn_value_tile_candidate_route import ( + install_gdn_value_tile_candidate, + ) + + install_gdn_value_tile_candidate() + + +class CommonDFlash2Extension: + """Worker extension for the explicitly selected, quantization-free routes.""" + + +if manifest_path := os.getenv("VLLM_SM70_DFLASH2_COMMON_MANIFEST"): + install_common_routes(json.loads(Path(manifest_path).read_text())) diff --git a/benchmarks/kernels/sm70_dflash2_qpn2_candidate_route.py b/benchmarks/kernels/sm70_dflash2_qpn2_candidate_route.py new file mode 100644 index 0000000000..9ecc5fdc3a --- /dev/null +++ b/benchmarks/kernels/sm70_dflash2_qpn2_candidate_route.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Optional NVFP4 projections alongside independent common DFlash2 schedules. + +VLLM_SM70_DFLASH2_QPN2_MANIFEST selects hashed cap64/publisher libraries and +their matching all-reduce library. Layers qualify through their actual QPN2 +representation and dimensions. Other quantizations retain their projections +while the common worker extension can still install its independent routes. +This experimental entry point does not change a serving default. +""" + +import functools +import json +import os +from pathlib import Path + +import torch + +from benchmarks.kernels.sm70_dflash2_common_candidate_route import ( + CommonDFlash2Extension, + _library, +) + + +def install_qpn2_routes(manifest: dict) -> None: + from vllm.distributed import get_tp_group + from vllm.model_executor.layers.linear import RowParallelLinear + from vllm.model_executor.layers.quantization.compressed_tensors.schemes import ( + compressed_tensors_w4a4_nvfp4 as quant, + ) + from vllm.v1.worker.gpu.model_runner import GPUModelRunner + + if set(manifest) != {"schema_version", "capped", "publisher", "all_reduce"}: + raise ValueError("Incomplete or unknown QPN2 manifest entries") + if manifest["schema_version"] != 1: + raise ValueError("Unsupported QPN2 manifest version") + libraries = { + key: _library(manifest[key]) for key in ("capped", "publisher", "all_reduce") + } + configured_ar = os.getenv("VLLM_SM70_CUSTOM_AR_LIBRARY") + if not configured_ar or Path(configured_ar).resolve() != libraries["all_reduce"]: + raise ValueError("QPN2 publication requires the matching all-reduce library") + torch.ops.load_library(str(libraries["publisher"])) + torch.ops.load_library(str(libraries["capped"])) + columns = {} + rows = {} + seen = set() + original_apply = quant.CompressedTensorsW4A4Fp4._apply_qpn2 + original_forward = RowParallelLinear.forward + original_load = GPUModelRunner.load_model + + @torch.library.custom_op("quasar_capped::column", mutates_args=()) + def column(x: torch.Tensor, prefix: str, gated: bool) -> torch.Tensor: + layer = columns[prefix] + if ( + x.ndim == 2 + and x.shape[0] == 8 + and x.dtype == torch.float16 + and torch.cuda.is_current_stream_capturing() + ): + assert x.shape[1] == 5120 and x.is_contiguous() + divisor = 2 if gated else 1 + logical = layer.output_size_per_partition // divisor + physical = int(layer.sm70_nvfp4_qpn2_output_size) // divisor + output = torch.empty((8, physical), device=x.device, dtype=x.dtype) + split, nacc = quant._SM70_NVFP4_QPN2_CONFIGS[ + (5120, physical * divisor, gated) + ] + op = torch.ops._qpn2_capped.gated if gated else torch.ops._qpn2_capped.gemm + op( + output, + x, + layer.sm70_nvfp4_qpn2_codes, + layer.sm70_nvfp4_qpn2_scales, + float(layer.sm70_nvfp4_qpn2_global_scale), + split, + nacc, + ) + if prefix not in seen: + seen.add(prefix) + print(f"QPN2_CAP64_CAPTURE prefix={prefix}", flush=True) + return output[:, :logical] + return original_apply(layer, x, None, gated_silu=gated) + + @column.register_fake + def column_fake(x, prefix, gated): + layer = columns[prefix] + divisor = 2 if gated else 1 + logical = layer.output_size_per_partition // divisor + physical = int(layer.sm70_nvfp4_qpn2_output_size) // divisor + return x.new_empty((x.shape[0], physical))[:, :logical] + + def apply(layer, x, bias, *, gated_silu): + if layer.prefix in columns: + assert bias is None + return column(x, layer.prefix, gated_silu) + return original_apply(layer, x, bias, gated_silu=gated_silu) + + @torch.library.custom_op("quasar_qpn2::row_parallel", mutates_args=()) + def row_parallel(input_: torch.Tensor, prefix: str) -> torch.Tensor: + layer, communicator = rows[prefix] + if ( + input_.ndim == 2 + and input_.shape[0] == 8 + and input_.dtype == torch.float16 + and input_.is_contiguous() + and torch.cuda.is_current_stream_capturing() + ): + peers = communicator.sm70_tp4_push_buffer_ptrs + assert peers is not None and communicator.world_size == 4 + assert communicator.fully_connected + projected = torch.empty( + (8, 5120), device=input_.device, dtype=torch.float16 + ) + output = torch.empty_like(projected) + torch.ops._qpn2_candidate.publish( + projected, + input_, + layer.sm70_nvfp4_qpn2_codes, + layer.sm70_nvfp4_qpn2_scales, + float(layer.sm70_nvfp4_qpn2_global_scale), + int(layer.sm70_nvfp4_qpn2_split_k), + int(layer.sm70_nvfp4_qpn2_nacc), + peers, + layer.tp_rank, + ) + torch.ops._qpn2_candidate.consume(projected, output, peers, layer.tp_rank) + if prefix not in seen: + seen.add(prefix) + print(f"QPN2_PUBLISH_CAPTURE prefix={prefix}", flush=True) + return output + result = original_forward(layer, input_) + return result[0] if isinstance(result, tuple) else result + + @row_parallel.register_fake + def row_fake(input_, prefix): + return input_.new_empty((*input_.shape[:-1], 5120)) + + @functools.wraps(original_forward) + def forward(self, input_): + if self.prefix in rows: + result = row_parallel(input_, self.prefix) + return (result, None) if self.return_bias else result + return original_forward(self, input_) + + @functools.wraps(original_load) + def load(self, *args, **kwargs): + result = original_load(self, *args, **kwargs) + communicator = getattr(get_tp_group().device_communicator, "ca_comm", None) + for _name, layer in self.model.named_modules(): + if not getattr(layer, "sm70_nvfp4_qpn2", False): + continue + if ( + getattr(layer, "input_size_per_partition", None) == 5120 + and layer.prefix.rsplit(".", 1)[-1] + in ("in_proj_qkvz", "qkv_proj", "gate_up_proj") + and layer.bias is None + and layer.tp_size == 4 + ): + physical = int(layer.sm70_nvfp4_qpn2_output_size) + if physical in (4128, 3584, 8704): + columns[layer.prefix] = layer + if ( + isinstance(layer, RowParallelLinear) + and layer.input_is_parallel + and layer.reduce_results + and layer.tp_size == 4 + and layer.bias is None + and layer.output_size_per_partition == 5120 + and layer.input_size_per_partition in (1536, 4352) + and int(layer.sm70_nvfp4_qpn2_nacc) == 2 + ): + assert communicator is not None and not communicator.disabled + assert communicator.sm70_tp4_push_buffer_ptrs is not None + rows[layer.prefix] = (layer, communicator) + print( + f"QPN2_ROUTES_READY rank={torch.distributed.get_rank()} " + f"columns={len(columns)} rows={len(rows)}", + flush=True, + ) + return result + + quant.CompressedTensorsW4A4Fp4._apply_qpn2 = staticmethod(apply) + RowParallelLinear.forward = forward + GPUModelRunner.load_model = load + + +class DFlash2KernelExtension(CommonDFlash2Extension): + """Optional QPN2 projections plus separately configured common schedules.""" + + +if manifest_path := os.getenv("VLLM_SM70_DFLASH2_QPN2_MANIFEST"): + install_qpn2_routes(json.loads(Path(manifest_path).read_text())) diff --git a/benchmarks/kernels/sm70_draft_column_candidate_route.py b/benchmarks/kernels/sm70_draft_column_candidate_route.py new file mode 100644 index 0000000000..c25cc08ea1 --- /dev/null +++ b/benchmarks/kernels/sm70_draft_column_candidate_route.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Explicit, experimental q8 draft GEMMs with column-major FP16 weights. + +Importing this module does not install a route. The candidate changes GEMM +arithmetic and is not admitted by the operator reference-error screen alone. +Only the draft's twenty captured query projections are eligible; context, +prefill and unsupported shapes retain their original implementation. +""" + +import functools + +import torch + + +def install_draft_column_candidate() -> None: + from vllm.model_executor.layers.linear import UnquantizedLinearMethod + from vllm.v1.worker.gpu.model_runner import GPUModelRunner + + original_apply = UnquantizedLinearMethod.apply + original_load = GPUModelRunner.load_model + registry = {} + captured = set() + expected_shapes = { + "qkv_proj": (1536, 5120), + "o_proj": (5120, 1024), + "gate_up_proj": (8704, 5120), + "down_proj": (5120, 4352), + } + + @torch.library.custom_op("quasar_draft_column::linear", mutates_args=()) + def linear(x: torch.Tensor, prefix: str) -> torch.Tensor: + layer, column_weight = registry[prefix] + if ( + x.ndim != 2 + or x.shape != (8, layer.weight.shape[1]) + or not x.is_contiguous() + or x.dtype != torch.float16 + or not x.is_cuda + or not torch.cuda.is_current_stream_capturing() + ): + return original_apply(layer.quant_method, layer, x, None) + previous = torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction + try: + # This selects the strict kernel only during this captured call. + # Restore the process setting before any other projection runs. + torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False + output = torch.mm(x, column_weight.T) + finally: + torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = previous + if prefix not in captured: + captured.add(prefix) + print( + f"DRAFT_COLUMN_CAPTURE rank={torch.distributed.get_rank()} " + f"prefix={prefix} input={tuple(x.shape)} " + f"weight={tuple(column_weight.shape)} " + f"stride={column_weight.stride()} strict_reduction=1", + flush=True, + ) + return output + + @linear.register_fake + def fake(x: torch.Tensor, prefix: str) -> torch.Tensor: + return x.new_empty((*x.shape[:-1], registry[prefix][0].weight.shape[0])) + + @functools.wraps(original_apply) + def apply(self, layer, x, bias=None): + if layer.prefix in registry: + assert bias is None + return linear(x, layer.prefix) + return original_apply(self, layer, x, bias) + + @functools.wraps(original_load) + def load(self, *args, **kwargs): + result = original_load(self, *args, **kwargs) + assert torch.distributed.get_world_size() == 4 + assert torch.cuda.get_device_capability(self.device) == (7, 0) + for name, layer in self.speculator.model.named_modules(): + leaf = name.rsplit(".", 1)[-1] + if ".layers." not in "." + name or leaf not in expected_shapes: + continue + assert isinstance(layer.quant_method, UnquantizedLinearMethod) + assert layer.bias is None and layer.weight.dtype == torch.float16 + assert tuple(layer.weight.shape) == expected_shapes[leaf] + column_weight = layer.weight.T.contiguous().T + assert torch.equal(layer.weight, column_weight) + assert layer.prefix not in registry + registry[layer.prefix] = (layer, column_weight) + assert len(registry) == 20 + print( + f"DRAFT_COLUMN_READY rank={torch.distributed.get_rank()} " + f"layers={len(registry)} strict_reduction=1", + flush=True, + ) + return result + + UnquantizedLinearMethod.apply = apply + GPUModelRunner.load_model = load diff --git a/benchmarks/kernels/sm70_gdn_value_tile_candidate_route.py b/benchmarks/kernels/sm70_gdn_value_tile_candidate_route.py new file mode 100644 index 0000000000..cd551a697d --- /dev/null +++ b/benchmarks/kernels/sm70_gdn_value_tile_candidate_route.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Explicit TP4 q8 value-tile experiment for the existing packed GDN kernel. + +The query/key reduction shape, one-warp schedule, recurrence arithmetic and +FP32 state remain unchanged. This module does not install a default route. +""" + +import functools +import importlib + +import torch + + +def install_gdn_value_tile_candidate() -> None: + from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( + QwenGatedDeltaNetAttention, + ) + + module = importlib.import_module( + "vllm.model_executor.layers.fla.ops.fused_sigmoid_gating" + ) + original = module._select_fused_sigmoid_launch + original_forward = QwenGatedDeltaNetAttention._forward_dflash2_packed_gdn_verify + logged = False + eligible = False + + @functools.wraps(original_forward) + def forward(self, **kwargs): + nonlocal eligible + previous = eligible + eligible = ( + self.tp_size == 4 + and kwargs["num_spec_decodes"] == 1 + and kwargs["mixed_qkv"].shape == (8, 2560) + and kwargs["mixed_qkv"].dtype == torch.float16 + and kwargs["ssm_state"].dtype == torch.float32 + ) + try: + return original_forward(self, **kwargs) + finally: + eligible = previous + + def select(V, N, HV, T, device, *, match_recurrent_schedule): + nonlocal logged + baseline = original( + V, N, HV, T, device, match_recurrent_schedule=match_recurrent_schedule + ) + if ( + eligible + and (V, N, HV, T) == (128, 1, 12, 8) + and match_recurrent_schedule + and torch.cuda.is_current_stream_capturing() + and torch.cuda.get_device_capability(device) == (7, 0) + ): + assert baseline[:2] == (8, 1), baseline + if not logged: + logged = True + print( + f"GDN_VALUE_TILE_CAPTURE rank={torch.distributed.get_rank()} " + "TP4/B1/q8 HV=12 V=128 BV=2 warps=1 original_BV=8", + flush=True, + ) + return 2, 1, baseline[2] + return baseline + + module._select_fused_sigmoid_launch = select + QwenGatedDeltaNetAttention._forward_dflash2_packed_gdn_verify = forward diff --git a/benchmarks/kernels/sm70_grouped_attention_candidate_route.py b/benchmarks/kernels/sm70_grouped_attention_candidate_route.py new file mode 100644 index 0000000000..cd1fd757f5 --- /dev/null +++ b/benchmarks/kernels/sm70_grouped_attention_candidate_route.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Bind a private q8 experiment to the Flash-V100 interface's native module.""" + +from collections.abc import Callable +from functools import wraps +from typing import Any + +import torch + + +def install_grouped_attention_candidate(candidate: Callable[..., Any]) -> None: + """Install explicitly; importing this module does not change serving. + + The package-relative and top-level extension imports can load the same + DSO into distinct Python module objects. Patch the object actually used + by the public interface, rather than independently importing its name. + The caller owns the candidate's build/quality gates and experiment flag. + """ + from flash_attn_v100 import flash_attn_interface + + native = flash_attn_interface.flash_attn_v100_cuda + original = native.grouped_e4m3_fp32_paged_fwd + + @wraps(original) + def run(q: torch.Tensor, *args: Any, **kwargs: Any) -> Any: + if q.shape == (8, 6, 256) and torch.cuda.is_current_stream_capturing(): + return candidate(q, *args, **kwargs) + return original(q, *args, **kwargs) + + native.grouped_e4m3_fp32_paged_fwd = run diff --git a/benchmarks/kernels/sm70_qpn2_graph_layout.py b/benchmarks/kernels/sm70_qpn2_graph_layout.py new file mode 100644 index 0000000000..680f158318 --- /dev/null +++ b/benchmarks/kernels/sm70_qpn2_graph_layout.py @@ -0,0 +1,241 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Pair dual-store q8 norms with packed QPN2 consumers after graph capture.""" + +import torch + +from benchmarks.kernels.sm70_qpn2_graph_nodes import ( + clone_params, + nodes, + only_kernel, + set_params, +) +from vllm.model_executor.layers import layernorm as norm + +OriginalGraph = torch.cuda.CUDAGraph + + +def capture_one(fn, parameter_count): + fn() + torch.accelerator.synchronize() + graph = OriginalGraph(keep_graph=True) + with torch.cuda.graph(graph): + fn() + graph.instantiate() + return graph, only_kernel(graph, parameter_count) + + +class LayoutTemplates: + def __init__(self, projection_items, dual_module): + projection_items = list(projection_items) + configs = {(4128, False): 16, (3584, False): 16, (8704, True): 8} + for item in projection_items: + assert item["k"] == 5120 + assert item["nacc"] == 2 + assert item["split_k"] == configs[(item["n"], item["gated"])] + dual = dual_module + self.keep = [] + self.norms = {} + self.projections = {} + x = torch.empty((8, 5120), device="cuda", dtype=torch.float16) + weight = torch.empty((5120,), device="cuda", dtype=torch.float16) + out = torch.empty_like(x) + packed = torch.empty_like(x) + # Finite warmups avoid incidental NaNs; values are not used by the gate. + x.zero_() + weight.zero_() + for kind in ("none", "half", "float"): + residual = ( + None + if kind == "none" + else torch.zeros_like( + x, dtype=torch.float16 if kind == "half" else torch.float32 + ) + ) + residual_out = ( + None if residual is None else torch.empty_like(x, dtype=torch.float32) + ) + kwargs = dict(num_stages=1) + if kind == "float": + old = norm._sm70_dflash2_gemma_fused_add_rms_kernel + new = dual._sm70_dflash2_gemma_fused_add_rms_kernel_dual_store + kwargs.update( + hidden_size=5120, BLOCK_SIZE=8192, epsilon=1e-6, num_warps=8 + ) + else: + old = norm._sm70_dflash2_fixed_gemma_rms_kernel + new = dual._sm70_dflash2_fixed_gemma_rms_kernel_dual_store + kwargs.update( + HAS_RESIDUAL=residual is not None, + epsilon=1e-6, + num_warps=16, + enable_fp_fusion=True, + ) + operands = (x, residual, weight, out, residual_out) + old_graph, old_node = capture_one( + lambda old=old, operands=operands, kwargs=kwargs: old[(8,)]( + *operands, **kwargs + ), + dict(none=5, half=7, float=8)[kind], + ) + new_graph, new_node = capture_one( + lambda new=new, operands=operands, kwargs=kwargs: new[(8,)]( + *operands, packed, **kwargs + ), + dict(none=6, half=8, float=9)[kind], + ) + self.keep.extend([old_graph, new_graph, operands, packed]) + mapping = [] + for value in new_node.args: + if ( + len(value) == 8 + and int.from_bytes(value, "little") == packed.data_ptr() + ): + mapping.append(("packed", None)) + elif value in old_node.args: + indices = [i for i, arg in enumerate(old_node.args) if arg == value] + # Only the trailing zero scratch pointers may be duplicated. + assert len(indices) == 1 or value == bytes(len(value)), ( + kind, + indices, + ) + mapping.append(("old", indices[0])) + else: + raise AssertionError(("Unmapped parameter", kind, value.hex())) + assert sum(kind == "packed" for kind, _ in mapping) == 1 + out_index = old_node.args.index(out.data_ptr().to_bytes(8, "little")) + key = (old_node.name, tuple(map(len, old_node.args))) + assert key not in self.norms + self.norms[key] = (old_node, new_node, mapping, out_index) + for item in projection_items: + if item["k"] != 5120: + continue + gated = item["gated"] + width = item["n"] // (2 if gated else 1) + result = torch.empty((8, width), device="cuda", dtype=torch.float16) + op = ( + torch.ops._qpn2_packed_input.gated + if gated + else torch.ops._qpn2_packed_input.gemm + ) + args = ( + result, + packed, + item["codes"], + item["scales"], + item["global_scale"], + item["split_k"], + item["nacc"], + ) + graph, node = capture_one(lambda op=op, args=args: op(*args), 8) + key = (gated, width, item["split_k"], item["nacc"]) + self.projections[key] = node + self.keep.extend([graph, result, args]) + + def describe(self): + return dict( + norms=[ + dict( + name=k[0], + sizes=k[1], + new_name=v[1].name, + new_sizes=list(map(len, v[1].args)), + mapping=v[2], + out_index=v[3], + ) + for k, v in self.norms.items() + ], + projections=[ + dict(key=k, name=v.name, sizes=list(map(len, v.args))) + for k, v in self.projections.items() + ], + ) + + def patch(self, graph): + kernel_nodes, parents = nodes(graph) + edits, norms_seen, storage = [], {}, [] + for handle, projection in kernel_nodes.items(): + if ( + "qpn2" not in projection.name + or "sm70_kernel" not in projection.name + or len(projection.args) != 8 + or projection.integer(5) != 5120 + or projection.integer(6) != 8 + ): + continue + gated = "gated" in projection.name + split = projection.params.blockDim.x // (64 if gated else 32) + key = (gated, projection.integer(4), split, 2) + assert key in self.projections, (projection.name, key) + template = self.projections[key] + input_pointer = projection.args[2] + frontier, visited, found = list(parents[handle]), set(), None + while frontier and found is None: + matches, next_frontier = [], [] + for parent in frontier: + if parent in visited: + continue + visited.add(parent) + candidate = kernel_nodes.get(parent) + if candidate is not None and candidate.grid == (8, 1, 1): + norm_key = (candidate.name, tuple(map(len, candidate.args))) + if norm_key in self.norms: + out_index = self.norms[norm_key][3] + if candidate.args[out_index] == input_pointer: + matches.append((candidate, self.norms[norm_key])) + next_frontier.extend(parents[parent]) + assert len(matches) <= 1, "Ambiguous norm producer" + found = matches[0] if matches else None + frontier = next_frontier + assert found is not None, ("Missing q8 norm producer", projection.name) + producer, (_, new_norm, mapping, _) = found + if producer.handle not in norms_seen: + # Preserve the original normalized output for every other reader. + canary = torch.full( + (8 * 5120 + 16,), -37, device="cuda", dtype=torch.float16 + ) + packed = canary[8:-8].view(8, 5120) + arguments = [ + packed.data_ptr().to_bytes(8, "little") + if kind == "packed" + else producer.args[index] + for kind, index in mapping + ] + changed, owned = clone_params( + producer.params, arguments, template=new_norm.params + ) + edits.append((producer, changed)) + storage.extend([canary, packed, owned]) + norms_seen[producer.handle] = (packed, canary) + packed, _ = norms_seen[producer.handle] + arguments = list(projection.args) + arguments[2] = packed.data_ptr().to_bytes(8, "little") + changed, owned = clone_params( + projection.params, arguments, template=template.params + ) + edits.append((projection, changed)) + storage.append(owned) + return LayoutEdit(graph, edits, storage, norms_seen) + + +class LayoutEdit: + def __init__(self, graph, edits, storage, norms): + self.graph, self.edits, self.storage, self.norms = graph, edits, storage, norms + self.mode = "control" + + def switch(self, mode): + assert mode in ("control", "candidate") + # Call only between requests / isolated replays; no concurrent launches. + torch.accelerator.synchronize() + for original, changed in self.edits: + set_params( + self.graph, + original.handle, + changed if mode == "candidate" else original.params, + ) + self.mode = mode + return len(self.edits) + + def check_canaries(self): + for _, canary in self.norms.values(): + assert bool((canary[:8] == -37).all() and (canary[-8:] == -37).all()) diff --git a/benchmarks/kernels/sm70_qpn2_graph_nodes.py b/benchmarks/kernels/sm70_qpn2_graph_nodes.py new file mode 100644 index 0000000000..44d8a210e9 --- /dev/null +++ b/benchmarks/kernels/sm70_qpn2_graph_nodes.py @@ -0,0 +1,206 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Private executable-node substitutions. Raw graph and its edges stay frozen.""" + +import ctypes as C +from dataclasses import dataclass + +rt = C.CDLL("libcuda.so.1") + + +class Dim3(C.Structure): + _fields_ = [("x", C.c_uint), ("y", C.c_uint), ("z", C.c_uint)] + + +class Params(C.Structure): + _fields_ = [ + ("func", C.c_void_p), + ("gridDim", Dim3), + ("blockDim", Dim3), + ("sharedMemBytes", C.c_uint), + ("kernelParams", C.POINTER(C.c_void_p)), + ("extra", C.POINTER(C.c_void_p)), + ("kern", C.c_void_p), + ("ctx", C.c_void_p), + ] + + +class EdgeData(C.Structure): + _fields_ = [ + ("from_port", C.c_ubyte), + ("to_port", C.c_ubyte), + ("type", C.c_ubyte), + ("reserved", C.c_ubyte * 5), + ] + + +for name, types in { + "cuGraphGetNodes": [C.c_void_p, C.POINTER(C.c_void_p), C.POINTER(C.c_size_t)], + "cuGraphGetEdges_v2": [ + C.c_void_p, + C.POINTER(C.c_void_p), + C.POINTER(C.c_void_p), + C.POINTER(EdgeData), + C.POINTER(C.c_size_t), + ], + "cuGraphNodeGetType": [C.c_void_p, C.POINTER(C.c_int)], + "cuGraphKernelNodeGetParams_v2": [C.c_void_p, C.POINTER(Params)], + "cuGraphExecKernelNodeSetParams_v2": [C.c_void_p, C.c_void_p, C.POINTER(Params)], + "cuFuncGetName": [C.POINTER(C.c_char_p), C.c_void_p], + "cuFuncGetParamInfo": [ + C.c_void_p, + C.c_size_t, + C.POINTER(C.c_size_t), + C.POINTER(C.c_size_t), + ], +}.items(): + getattr(rt, name).argtypes = types + getattr(rt, name).restype = C.c_int +rt.cuGetErrorString.argtypes = [C.c_int, C.POINTER(C.c_char_p)] +rt.cuGetErrorString.restype = C.c_int + + +def check(status): + if status: + message = C.c_char_p() + rt.cuGetErrorString(status, C.byref(message)) + raise RuntimeError((status, message.value)) + + +def param_sizes(func, count): + assert 0 < count <= 16 + sizes = [] + for index in range(count): + offset, size = C.c_size_t(), C.c_size_t() + check(rt.cuFuncGetParamInfo(func, index, C.byref(offset), C.byref(size))) + assert 0 < size.value <= 16, (index, size.value) + sizes.append(size.value) + return sizes + + +def known_count(name, params): + # Frozen public signatures, not error-driven queries past the last argument. + if "nvfp4_qpn2_sm70_kernel" in name or "nvfp4_qpn2_gated_sm70_kernel" in name: + return 8 + if name == "_sm70_dflash2_gemma_fused_add_rms_kernel": + return 8 # five data pointers, epsilon, two Triton scratch pointers + if name == "_sm70_dflash2_fixed_gemma_rms_kernel": + # Both valid signatures have at least five pointers. The no-residual + # specialization has three data pointers then two null scratch pointers; + # the residual specialization has five nonnull data pointers then scratch. + assert param_sizes(params.func, 5) == [8] * 5 + values = [ + int.from_bytes(C.string_at(params.kernelParams[i], 8), "little") + for i in range(5) + ] + assert all(values[:3]) + if values[3:] == [0, 0]: + return 5 + assert values[3] and values[4], "Unsupported fixed-norm signature" + return 7 + return None + + +def clone_params(params, arguments, *, template=None): + value = Params.from_buffer_copy(params) + if template is not None: + value.func, value.kern, value.ctx = template.func, template.kern, template.ctx + value.blockDim = template.blockDim + value.sharedMemBytes = template.sharedMemBytes + expected = param_sizes(value.func, len(arguments)) + assert list(map(len, arguments)) == expected, (list(map(len, arguments)), expected) + buffers = [C.create_string_buffer(arg, len(arg)) for arg in arguments] + pointers = (C.c_void_p * len(buffers))(*(C.addressof(b) for b in buffers)) + value.kernelParams, value.extra = pointers, None + return value, (buffers, pointers) + + +@dataclass +class Node: + handle: int + name: str + params: Params + args: list[bytes] + storage: object + + def integer(self, index): + return int.from_bytes(self.args[index], "little") + + @property + def grid(self): + return (self.params.gridDim.x, self.params.gridDim.y, self.params.gridDim.z) + + +def nodes(graph, *, single_parameter_count=None): + ptr = graph.raw_cuda_graph() + count = C.c_size_t() + check(rt.cuGraphGetNodes(ptr, None, C.byref(count))) + handles = (C.c_void_p * count.value)() + check(rt.cuGraphGetNodes(ptr, handles, C.byref(count))) + result = {} + for handle in handles: + kind = C.c_int() + check(rt.cuGraphNodeGetType(handle, C.byref(kind))) + if kind.value != 0: + continue + params = Params() + check(rt.cuGraphKernelNodeGetParams_v2(handle, C.byref(params))) + name = C.c_char_p() + check(rt.cuFuncGetName(C.byref(name), params.func)) + text_name = name.value.decode() + recognized = ( + "nvfp4_qpn2_sm70_kernel" in text_name + or "nvfp4_qpn2_gated_sm70_kernel" in text_name + or text_name + in ( + "_sm70_dflash2_fixed_gemma_rms_kernel", + "_sm70_dflash2_gemma_fused_add_rms_kernel", + ) + ) + if single_parameter_count is None and not recognized: + continue + # cuBLAS may use the packed launch-parameter buffer convention. It is + # outside this rewrite and must be skipped before assuming pointer args. + assert params.kernelParams and not params.extra + count_params = ( + single_parameter_count + if single_parameter_count is not None + else known_count(name.value.decode(), params) + ) + if count_params is None: + continue + sizes = param_sizes(params.func, count_params) + arguments = [ + C.string_at(params.kernelParams[i], n) for i, n in enumerate(sizes) + ] + owned, storage = clone_params(params, arguments) + result[handle] = Node(handle, name.value.decode(), owned, arguments, storage) + edge_count = C.c_size_t() + check(rt.cuGraphGetEdges_v2(ptr, None, None, None, C.byref(edge_count))) + if edge_count.value == 0: + return result, {handle: [] for handle in handles} + sources = (C.c_void_p * edge_count.value)() + targets = (C.c_void_p * edge_count.value)() + data = (EdgeData * edge_count.value)() + check(rt.cuGraphGetEdges_v2(ptr, sources, targets, data, C.byref(edge_count))) + assert all((d.from_port, d.to_port, d.type) == (0, 0, 0) for d in data), ( + "Unsupported graph dependency type" + ) + parents = {handle: [] for handle in handles} + for a, b in zip(sources, targets): + parents[b].append(a) + return result, parents + + +def set_params(graph, node, params): + check( + rt.cuGraphExecKernelNodeSetParams_v2( + graph.raw_cuda_graph_exec(), node, C.byref(params) + ) + ) + + +def only_kernel(graph, parameter_count): + kernels, _ = nodes(graph, single_parameter_count=parameter_count) + assert len(kernels) == 1, [n.name for n in kernels.values()] + return next(iter(kernels.values())) diff --git a/benchmarks/replay_sm70_dflash2_packed_state.py b/benchmarks/replay_sm70_dflash2_packed_state.py new file mode 100644 index 0000000000..224bfd3572 --- /dev/null +++ b/benchmarks/replay_sm70_dflash2_packed_state.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Replay real q8 GDN inputs through split and packed verification kernels. + +Uses StateAuditExtension captures. Exercises every accepted-slot selector, +non-monotonic slot IDs (including zero), an empty padded request and a strided +state pool. This is an operator parity gate, never a complete-round speed or +acceptance-length benchmark. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import torch + + +def replay_capture(path: Path, *, row_strided: bool = False) -> list[dict]: + from vllm.model_executor.layers.fla.ops.fused_recurrent import ( + fused_recurrent_gated_delta_rule, + ) + from vllm.model_executor.layers.fla.ops.fused_sigmoid_gating import ( + fused_sigmoid_gating_delta_rule_update_mixed_qkv_out, + ) + + record = torch.load(path, weights_only=True, mmap=True, map_location="cpu") + if record["phase"] != "verify" or record["num_draft_tokens"] != 7: + raise ValueError(f"{path}: requires a q8 verification capture") + results = [] + for layer in record["expected_layers"]: + prefix = f"verify/layer{layer}/recurrent/" + states = { + k.removeprefix(prefix).split(":")[0]: v + for k, v in record["states"].items() + if k.startswith(prefix) + } + q, k, v, g, beta = ( + states[name].cuda() for name in ("q", "k", "v", "g", "beta") + ) + if q.shape[:2] != (1, 8): + raise ValueError(f"{path}: expected B1/q8 tensors, got {q.shape}") + if not bool(states["input_state/valid"][0]): + raise ValueError(f"{path}: missing live incoming state") + incoming = states["input_state/values"][0].cuda() + mixed = torch.cat([tensor.reshape(8, -1) for tensor in (q, k, v)], dim=1) + projection = mixed + if row_strided: + # QPN2 pads the QKVZBA projection width to a multiple of eight. + row_width = ( + (mixed.shape[1] + v.shape[2] * v.shape[3] + 2 * v.shape[2] + 7) // 8 + ) * 8 + projection = torch.full( + (8, row_width), -42.0, dtype=mixed.dtype, device=mixed.device + ) + projection[:, : mixed.shape[1]].copy_(mixed) + mixed = projection[:, : mixed.shape[1]] + projection_before = projection.clone() + heads, width, depth = incoming.shape + # The extra columns mimic Mamba pool padding and must stay untouched. + backing = torch.full( + (16, incoming.numel() + 64), -42.0, device="cuda", dtype=incoming.dtype + ) + padded_table = torch.tensor( + [[7, 2, 11, 0, 9, 5, 13, 3], [-1] * 8], + device="cuda", + dtype=torch.int32, + ) + dummy = torch.zeros(heads, device="cuda", dtype=torch.float32) + + def view(storage, size=heads * width * depth, shape=(16, heads, width, depth)): + return storage[:, :size].view(shape) + + for padded, selector in [ + *((False, s) for s in range(1, 9)), + (True, 1), + (True, 8), + ]: + table = padded_table if padded else padded_table[:1] + cu = torch.tensor( + [0, 8, 8] if padded else [0, 8], device="cuda", dtype=torch.int32 + ) + seed = backing.clone() + view(seed)[int(table[0, selector - 1])].copy_(incoming) + left, right = seed.clone(), seed.clone() + selectors = torch.tensor( + [selector, 1] if padded else [selector], + device="cuda", + dtype=torch.int32, + ) + reference, _ = fused_recurrent_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=view(left), + inplace_final_state=True, + cu_seqlens=cu, + ssm_state_indices=table, + num_accepted_tokens=selectors, + use_qk_l2norm_in_kernel=True, + ) + actual = torch.empty((8, 1, heads, width), device="cuda", dtype=v.dtype) + fused_sigmoid_gating_delta_rule_update_mixed_qkv_out( + A_log=dummy, + a=g, + b=beta, + dt_bias=dummy, + mixed_qkv=mixed, + num_q_heads=q.shape[2], + num_v_heads=heads, + head_k_dim=depth, + head_v_dim=width, + out=actual, + initial_state=view(right), + cu_seqlens=cu, + ssm_state_indices=table, + num_accepted_tokens=selectors, + use_qk_l2norm_in_kernel=True, + precomputed_g=g, + precomputed_beta=beta, + quantize_state_each_step=False, + match_recurrent_schedule=True, + match_recurrent_numerics=True, + ) + # Compare backing storage too: an equal live result cannot excuse + # a write into padding or a retired slot in either implementation. + live = torch.zeros_like(seed, dtype=torch.bool) + live[table[0].long(), : incoming.numel()] = True + output_exact = torch.equal( + actual.transpose(0, 1).view(torch.uint8), reference.view(torch.uint8) + ) + states_exact = torch.equal(left.view(torch.uint8), right.view(torch.uint8)) + untouched = all(torch.equal(t[~live], seed[~live]) for t in (left, right)) + result = { + "capture": path.name, + "layer": layer, + "selector": selector, + "empty_padded_request": padded, + "qkv_row_stride": mixed.stride(0), + "input_projection_untouched": torch.equal( + projection.view(torch.uint8), projection_before.view(torch.uint8) + ), + "reference_matches_capture": None + if padded + else torch.equal( + reference.view(torch.uint8), + states["output"].cuda().view(torch.uint8), + ), + "output_bitwise_equal": output_exact, + "state_bitwise_equal": states_exact, + "padding_and_retired_slots_untouched": untouched, + } + results.append(result) + return results + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("captures", nargs="+", type=Path) + parser.add_argument("--output", required=True, type=Path) + parser.add_argument("--row-strided", action="store_true") + args = parser.parse_args() + if torch.cuda.get_device_capability() != (7, 0): + raise ValueError("This audit targets SM70") + torch.set_num_threads(4) + results = [] + for path in args.captures: + results.extend(replay_capture(path, row_strided=args.row_strided)) + args.output.write_text(json.dumps(results, indent=2) + "\n") + print(f"Replayed {path.name}", flush=True) + passed = all( + row[key] + for row in results + for key in ( + "output_bitwise_equal", + "state_bitwise_equal", + "padding_and_retired_slots_untouched", + "input_projection_untouched", + ) + ) and all(row["reference_matches_capture"] is not False for row in results) + print(json.dumps({"cases": len(results), "all_exact": passed})) + if not passed: + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/sm70_dflash2_state_audit.py b/benchmarks/sm70_dflash2_state_audit.py new file mode 100644 index 0000000000..d796483c19 --- /dev/null +++ b/benchmarks/sm70_dflash2_state_audit.py @@ -0,0 +1,461 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Opt-in state/proposal diagnostic, never a performance measurement. + +Set VLLM_SM70_DFLASH2_AUDIT_ROOT and use StateAuditExtension as the worker +extension. The root's active-case.json contains name, prompt_ids and token_ids +(the complete forced tape). Remove that file for ordinary requests. Each worker +records prefill and verifier inputs, state selection, outputs and native logits. +Set force_tokens=false to observe real sampling instead of forcing token_ids. +Natural mode also records auxiliary states, proposal scores and acceptance; +its extra snapshots still require a separate diagnostic-perturbation check. +""" + +from __future__ import annotations + +import functools +import json +import os +import sys +from pathlib import Path + +import torch + + +def gather_state(pool: torch.Tensor, indices: torch.Tensor) -> dict[str, torch.Tensor]: + """Snapshot indexed slots without disguising padding as live slot zero.""" + indices = indices.reshape(-1).to(device=pool.device, dtype=torch.int64) + valid = (indices >= 0) & (indices < pool.shape[0]) + values = pool.index_select(0, indices.clamp(0, pool.shape[0] - 1)) + mask = valid.reshape((-1,) + (1,) * (values.ndim - 1)) + return { + "indices": indices.clone(), + "valid": valid, + "values": torch.where(mask, values, 0), + } + + +def selected_ssm_slots(indices: torch.Tensor, selectors: torch.Tensor) -> torch.Tensor: + """The recurrent verifier reads the preceding accepted slot, not column 0.""" + if indices.ndim != 2 or selectors.numel() != indices.shape[0]: + raise ValueError("Expected [requests, slots] and one selector per request") + columns = selectors.reshape(-1, 1).to(torch.int64) - 1 + valid = (columns >= 0) & (columns < indices.shape[1]) + slots = indices.gather(1, columns.clamp(0, indices.shape[1] - 1)) + return torch.where(valid, slots, -1).reshape(-1) + + +def cpu_request_slots(values: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: + return values.index_select(0, indices.to(torch.int64)).detach().cpu().clone() + + +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 + from vllm.v1.worker.gpu.model_runner import GPUModelRunner + from vllm.v1.worker.gpu.sample.output import SamplerOutput + from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import DFlash2Speculator + + root = Path(os.environ["VLLM_SM70_DFLASH2_AUDIT_ROOT"]) + mode = os.environ.get("VLLM_SM70_DFLASH2_AUDIT_MODE", "control") + directory = root / "captures" / mode + directory.mkdir(parents=True, exist_ok=True) + layers = { + int(x) + for x in os.environ.get("VLLM_SM70_DFLASH2_AUDIT_LAYERS", "0,1").split(",") + } + buffers: dict[str, torch.Tensor] = {} + markers: dict[str, torch.Tensor] = {} + epoch: torch.Tensor | None = None + runner = None + + def caller_layer() -> int | None: + if epoch is None: + return None + # Functionalization may clone cache arguments or reuse their addresses. + # Bind observations to the actual GDN call, never a storage pointer. + frame = sys._getframe(2) + for _ in range(8): + module = frame.f_locals.get("self") + if isinstance(module, gd.QwenGatedDeltaNetAttention): + layer = gd._sm70_gdn_layer_idx(module.prefix) + return layer if layer in layers else None + frame = frame.f_back + if frame is None: + break + return None + + def record(key: str, tensor: torch.Tensor | None) -> None: + if tensor is None or epoch is None: + return + key = f"{key}:{tuple(tensor.shape)}" + if key not in buffers: + buffers[key] = torch.empty_like(tensor) + markers[key] = torch.empty_like(epoch) + buffers[key].copy_(tensor) + markers[key].copy_(epoch) + + def record_slots(key: str, pool: torch.Tensor, indices: torch.Tensor) -> None: + for label, tensor in gather_state(pool, indices).items(): + record(f"{key}/{label}", tensor) + + initialize = GPUModelRunner.initialize_kv_cache + + @functools.wraps(initialize) + def initialize_kv_cache(self, *args, **kwargs): + nonlocal epoch, runner + runner = self + result = initialize(self, *args, **kwargs) + epoch = torch.full((1,), -1, device=self.device, dtype=torch.int64) + for name, module in self.model.named_modules(): + if not isinstance(module, gd.QwenGatedDeltaNetAttention): + continue + layer = gd._sm70_gdn_layer_idx(name) + if layer not in layers: + continue + original = module.chunk_gated_delta_rule.forward + + def chunk(*args, _original=original, _layer=layer, **kwargs): + key = f"prefill/layer{_layer}/recurrent" + for label in ( + "q", + "k", + "v", + "g", + "beta", + "cu_seqlens", + "has_initial_state", + ): + record(f"{key}/{label}", kwargs.get(label)) + state = kwargs.get("initial_state") + indices = kwargs.get("state_indices") + if state is not None: + if indices is None: + record(f"{key}/input_state", state) + else: + record_slots(f"{key}/input_state", state, indices) + out = _original(*args, **kwargs) + record(f"{key}/output", out[0]) + if indices is not None: + record_slots(f"{key}/output_state", state, indices) + else: + record(f"{key}/output_state", out[1]) + return out + + module.chunk_gated_delta_rule.forward = chunk + return result + + GPUModelRunner.initialize_kv_cache = initialize_kv_cache + recurrent = gd.fused_recurrent_gated_delta_rule + + @functools.wraps(recurrent) + def recurrent_wrapper(*args, **kwargs): + state = kwargs.get("initial_state") + indices = kwargs.get("ssm_state_indices") + selectors = kwargs.get("num_accepted_tokens") + layer = caller_layer() + if layer is None or indices is None or selectors is None: + return recurrent(*args, **kwargs) + key = f"verify/layer{layer}/recurrent" + record(f"route/verify/layer{layer}/split", epoch) + for label in ("q", "k", "v", "g", "beta", "cu_seqlens"): + record(f"{key}/{label}", kwargs.get(label)) + record(f"{key}/slot_table", indices) + record(f"{key}/selectors", selectors) + record_slots( + f"{key}/input_state", state, selected_ssm_slots(indices, selectors) + ) + out = recurrent(*args, **kwargs) + record(f"{key}/output", out[0]) + record_slots(f"{key}/output_states", state, indices) + return out + + gd.fused_recurrent_gated_delta_rule = recurrent_wrapper + packed = gd.fused_sigmoid_gating_delta_rule_update_mixed_qkv_out + + @functools.wraps(packed) + def packed_wrapper(*args, **kwargs): + layer = caller_layer() + if layer is None or kwargs.get("precomputed_g") is None: + return packed(*args, **kwargs) + # Keep the same raw metadata coverage as the split verifier. The + # packed operator receives only the live slice of these parent args. + frame = sys._getframe(1) + for _ in range(8): + if frame.f_code.co_name == "_forward_dflash2_packed_gdn_verify": + indices = frame.f_locals["spec_state_indices_tensor"] + selectors = frame.f_locals["spec_state_slot_selectors"] + break + frame = frame.f_back + if frame is None: + raise RuntimeError("Missing packed-verifier metadata provenance") + else: + raise RuntimeError("Missing packed-verifier caller") + key = f"verify/layer{layer}/recurrent" + record(f"route/verify/layer{layer}/packed", epoch) + mixed = kwargs["mixed_qkv"] + tokens = mixed.shape[0] + q_heads, v_heads = kwargs["num_q_heads"], kwargs["num_v_heads"] + dk, dv = kwargs["head_k_dim"], kwargs["head_v_dim"] + q, k, v = mixed.split([q_heads * dk, q_heads * dk, v_heads * dv], dim=1) + for label, tensor in ( + ("q", q.reshape(1, tokens, q_heads, dk)), + ("k", k.reshape(1, tokens, q_heads, dk)), + ("v", v.reshape(1, tokens, v_heads, dv)), + ("g", kwargs["precomputed_g"].reshape(1, tokens, v_heads)), + ("beta", kwargs["precomputed_beta"].reshape(1, tokens, v_heads)), + ("cu_seqlens", kwargs["cu_seqlens"]), + ("slot_table", indices), + ("selectors", selectors), + ): + record(f"{key}/{label}", tensor) + state = kwargs["initial_state"] + record_slots( + f"{key}/input_state", state, selected_ssm_slots(indices, selectors) + ) + out = packed(*args, **kwargs) + record(f"{key}/output", out[0].transpose(0, 1)) + record_slots(f"{key}/output_states", state, indices) + return out + + gd.fused_sigmoid_gating_delta_rule_update_mixed_qkv_out = packed_wrapper + for function_name, phase in ( + ("causal_conv1d_fn", "prefill"), + ("causal_conv1d_update", "verify"), + ): + original = getattr(gd, function_name) + + @functools.wraps(original) + def convolution(*args, _original=original, _phase=phase, **kwargs): + state = kwargs.get("conv_states") if _phase == "prefill" else args[1] + indices = kwargs.get( + "cache_indices" if _phase == "prefill" else "conv_state_indices" + ) + layer = caller_layer() + if layer is None or indices is None: + return _original(*args, **kwargs) + key = f"{_phase}/layer{layer}/conv" + record(f"{key}/input", args[0]) + for label in ( + "has_initial_state", + "num_accepted_tokens", + "query_start_loc", + ): + record(f"{key}/{label}", kwargs.get(label)) + record_slots(f"{key}/input_state", state, indices) + out = _original(*args, **kwargs) + record(f"{key}/output", out) + record_slots(f"{key}/output_state", state, indices) + return out + + setattr(gd, function_name, convolution) + + def active(self, batch): + if not hasattr(self, "_state_audit_requests"): + self._state_audit_requests = {} + if not batch.req_ids: + return None + request_id = batch.req_ids[0] + if request_id not in self._state_audit_requests: + path = root / "active-case.json" + entry = None + if path.exists(): + if batch.num_reqs != 1: + raise ValueError("State audit requires exactly one request") + case = json.loads(path.read_text()) + entry = { + "case": case, + "tape": torch.tensor( + case["token_ids"], device=self.device, dtype=torch.int64 + ) + if case.get("force_tokens", True) + else None, + "step": 0, + } + self._state_audit_requests[request_id] = entry + return self._state_audit_requests[request_id] + + prepare = GPUModelRunner.prepare_inputs + + @functools.wraps(prepare) + def prepare_inputs(self, *args, **kwargs): + batch = prepare(self, *args, **kwargs) + if epoch is not None: + epoch.add_(1) + entry = active(self, batch) + if entry is not None and entry["tape"] is not None: + batch.input_ids[: batch.num_tokens].copy_( + entry["tape"][batch.positions[: batch.num_tokens]] + ) + return batch + + GPUModelRunner.prepare_inputs = prepare_inputs + sample = GPUModelRunner.sample + + @functools.wraps(sample) + def sample_fixed_prefix(self, hidden, batch, grammar): + entry = active(self, batch) + if entry is None: + return sample(self, hidden, batch, grammar) + rank = torch.distributed.get_rank() + step = entry["step"] + phase = "prefill" if step == 0 else "verify" + hs = hidden[batch.logits_indices] + positions = batch.positions[batch.logits_indices] + native = self.model.compute_logits(hs) + + def cpu(tensor): + return tensor.detach().cpu().clone() + + assert epoch is not None + current_epoch = int(epoch.item()) + keys = list(markers) + epochs = torch.cat([markers[key] for key in keys]).cpu().tolist() + fresh = { + key for key, observed in zip(keys, epochs) if observed == current_epoch + } + tensors = { + key: cpu(value) + for key, value in buffers.items() + if key.startswith(phase + "/") and key in fresh + } + layer_tensors = { + key: {**qn._SM70_QWEN_LAYER_GRAPH_META[key], "tensor": cpu(value)} + for key, value in qn._SM70_QWEN_LAYER_GRAPH_BUFFERS.items() + if value.ndim and value.shape[0] in (batch.num_tokens, hidden.shape[0]) + } + result = { + "rank": rank, + "step": step, + "case": entry["case"]["name"], + "phase": phase, + "positions": cpu(positions), + "input_ids": cpu(batch.input_ids[: batch.num_tokens]), + "hidden": cpu(hs), + "native_logits": cpu(native) if rank == 0 else None, + "states": tensors, + "verifier_routes": sorted( + key.split(":")[0] for key in fresh if key.startswith("route/") + ), + "tensors": layer_tensors, + "cuda_rng": torch.cuda.get_rng_state(self.device), + "cpu_rng": torch.get_rng_state(), + "num_draft_tokens": batch.num_draft_tokens, + "capture_epoch": current_epoch, + "expected_layers": sorted(layers), + "sampling": { + name: cpu( + getattr(self.sampler.sampling_states, name).gpu.index_select( + 0, batch.idx_mapping.to(torch.int64) + ) + ) + for name in ("seeds", "temperature", "top_k", "top_p", "min_p") + }, + } + natural_output = None + if entry["tape"] is None: + # Observe the actual proposal/rejection path without replacing its + # 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") + result["aux_hidden_states"] = [cpu(t) for t in aux] + if batch.num_draft_tokens: + result["draft_logits"] = cpu( + self.speculator.draft_logits.index_select( + 0, batch.idx_mapping.to(torch.int64) + ) + ) + natural_output = sample(self, hidden, batch, grammar) + result["sampled_token_ids"] = cpu(natural_output[0].sampled_token_ids) + result["num_sampled"] = cpu(natural_output[1]) + result["num_rejected"] = cpu(natural_output[2]) + torch.save( + result, directory / f"{entry['case']['name']}-rank{rank}-step{step}.pt" + ) + entry["step"] += 1 + if natural_output is not None: + return natural_output + next_ids = entry["tape"][positions + 1].view(1, -1).to(torch.int32) + count = torch.full( + (1,), next_ids.shape[1], device=self.device, dtype=torch.int32 + ) + # Teacher forcing is a diagnostic control, never an acceptance metric. + return ( + SamplerOutput(next_ids, None, None, count), + count, + torch.zeros_like(count), + ) + + GPUModelRunner.sample = sample_fixed_prefix + propose = DFlash2Speculator.propose + + @functools.wraps(propose) + def propose_observed(self, input_batch, *args, **kwargs): + batch = input_batch + out = propose(self, batch, *args, **kwargs) + if runner is None or not batch.req_ids: + return out + entry = active(runner, batch) + if entry is None or entry["tape"] is not None: + return out + step = entry["step"] - 1 + if step < 0: + return out + rank = torch.distributed.get_rank() + result = { + "rank": rank, + "step": step, + "case": entry["case"]["name"], + "draft_tokens": out.detach().cpu().clone(), + "idx_mapping": batch.idx_mapping.detach().cpu().clone(), + } + for name in ("_cached_candidate_ids", "_cached_candidate_scores"): + tensor = getattr(self, name) + if tensor is not None: + result[name] = ( + tensor.index_select(0, batch.idx_mapping.to(torch.int64)) + .detach() + .cpu() + .clone() + ) + # Reuse the existing proposal shadow buffers when explicitly enabled. + # They are refreshed by the captured draft graph, including replays. + for name in ( + "_debug_backbone_hidden_states", + "_debug_candidate_ids", + "_debug_unary_logits", + "_debug_lattice_scores", + ): + tensor = getattr(self, name, None) + if tensor is not None: + result[name] = tensor[: batch.num_reqs].detach().cpu().clone() + result["projected_context"] = ( + self.hidden_states[: batch.num_tokens].detach().cpu().clone() + ) + result["sample_pos"] = ( + self.sample_pos[: batch.num_reqs * self.draft_block].detach().cpu().clone() + ) + # Positions are packed per draft row, but seeds/temperature are indexed + # by request slot. Prefix slicing those arrays reads inactive warmup rows. + result["sampling_layout"] = "request_gathered_v1" + for name in ("temperature", "seeds"): + result[name] = cpu_request_slots(getattr(self, name), batch.idx_mapping) + torch.save( + result, + directory / f"proposal-{entry['case']['name']}-tp{rank}-forward{step}.pt", + ) + return out + + DFlash2Speculator.propose = propose_observed + + +class StateAuditExtension: + pass + + +if os.getenv("VLLM_SM70_DFLASH2_AUDIT_ROOT"): + install() diff --git a/docs/design/sm70_dflash2_acceptance_20260909.md b/docs/design/sm70_dflash2_acceptance_20260909.md new file mode 100644 index 0000000000..ffb0ae4e42 --- /dev/null +++ b/docs/design/sm70_dflash2_acceptance_20260909.md @@ -0,0 +1,209 @@ +# DFlash2 acceptance and quantization-independent schedules, 2026-09-09 + +The current approximately 16.2/15.8-ms combination exposes independently +selectable schedules and records dataset-level decode speed, acceptance and +quality. Quantization-independent optimizations are available to other weight +formats through the common entry point. The user has explicitly requested +main integration of the current PR. Experimental routes remain disabled by +default; source integration does not certify the pending runtime gates below. +The original sub-15-ms performance objective is not claimed achieved. + +## Current capacity contract: 256K + +The user subsequently required 256K context without the evaluation's artificial +16K generation cutoff. The server already used `--max-model-len 262144`; the +cutoff came from the client request's `max_tokens=16384`. The old seed-one job +was stopped deliberately and its partial records retained. Its exit code 143 +is an authorized protocol transition, not a numerical failure. + +The new `acceptance-256k` campaign retains the same prompts, seeds, weights, +sampling and four GPUs. Before natural generation, the client uses the server's +`/tokenize` renderer, verifies `max_model_len=262144`, and explicitly sets +`max_tokens=262144-prompt_tokens`. The input is never truncated and EOS remains +natural. A 135-token prompt therefore has a 262009-token output budget. The +actual response's prompt count must match the tokenizer result. There is no +separate 16K/32K generation cap. Reaching the model's total context limit is +still reported as a length stop, never presented as natural completion. + +All three paired dataset launches are restarted with this policy; previous +truncated cases run first. FP8, public-route and whole-stack controls use the +same remaining-capacity policy. Speed requests also use the full remaining +capacity while retaining the canonical natural-EOS fixtures. One-token prefix +warmups and bounded teacher-forced operator diagnostics remain explicitly +excluded from natural-generation quality and performance results. + +The corpus, capacity-policy checks, launch-time source archives and new results +are separate from `acceptance-16ms`. The historical 16K-cap results later in +this document do not certify the new 256K-capacity campaign. + +## Current 256K-capacity observations and integration scope + +The first independent startup measures five requests after five warmups per +fixture and arm, with no profiler or tensor dump: + +| Fixture | BV8 / BV2 complete-round median | BV2 pure decode median | Accepted drafts / emitted tokens per round | +| --- | ---: | ---: | ---: | +| release1k | 16.571456 / 16.339483 ms | 182.259 token/s | 1.989011 / 2.989011 | +| MBPP28 | 16.144814 / 15.907431 ms | 306.098 token/s | 3.876923 / 4.876923 | + +These are medians of request-average complete-round costs. All tokens, natural +EOS and acceptance match. This A/B isolates BV8/BV2 with the other performance +candidates shared; it is not the full-stack all-off comparison. + +The first five completed long-output pairs also match token IDs, acceptance, +finish reasons and semantic tool calls exactly. Every response stops naturally. +Their candidate measurements are: + +| Case | Output tokens | Mean complete round | Pure decode | +| --- | ---: | ---: | ---: | +| HumanEval/10 | 21162 | 19.486480 ms | 186.554 token/s | +| LiveCodeBench/21 | 76955 | 26.916502 ms | 122.378 token/s | +| LiveCodeBench/64 | 34520 | 21.619425 ms | 148.984 token/s | +| LiveCodeBench/93 | 70725 | 26.381223 ms | 132.853 token/s | +| LiveCodeBench/131 | 52704 | 23.985040 ms | 131.885 token/s | + +Within LiveCodeBench/21, median client inter-chunk intervals rise from +17.161768 ms over the first 1024 intervals to 37.595053 ms over the last 1024. +The interval count matches the draft-round count, but these transport timings +are not GPU instrumentation or a per-window token/s measurement. They establish +a same-response latency trend without assigning its cost to an individual +operator. Across the whole request, acceptance is identical between BV8 and +BV2. Long-context target/draft attention needs a separate context sweep and +trace before attributing the slowdown or selecting another optimization. + +Evidence: `results/v4-accept256-datasets-seed0-switch.json`, +`acceptance-256k/v4-accept256-datasets-seed0-pairs.json` and the retained +`acceptance-256k/long-generation-cost-progress-20260909.json` snapshot under +the campaign artifact root recorded in the companion worklogs. Full natural +generation scoring, independent startups, whole-stack comparisons, all-layer +repeatability, the public installer, FP8 and long-prefix gates remain pending. +Five matching long generations are not proof of universally unchanged quality +or a 256K-context speed claim. + +Integration retains the audited source and opt-in manifests, the grouped +attention synchronization repair and dependency #563's singleton-prefill fix. +It does not change serving defaults, weights, sampling semantics or context +capacity. No pending or rejected arithmetic candidate is promoted by this +integration. The frozen evaluation checkout and native libraries remain intact. + +## Frozen evaluation + +The running reference checkout remains detached at +`a7cc5ae305149d7a9ffdf42fb224dff34e5606aa` in +`/home/ymzx/桌面/1cat-vllm/worktrees/v100-quasar-dflash2-15ms-20260907-161715`. +Implementation continues on the same owned PR #556 branch in +`/home/ymzx/桌面/1cat-vllm/worktrees/v100-quasar-dflash2-acceptance-20260909`. +Main `b6d91d61ff` was merged into that branch without conflicts; this is not a +merge of the candidate into main. New code does not alter the running reference +checkout or its native libraries. + +Artifacts are under +`/data/minimax-h3/task-cache/v100-quasar-dflash2-15ms-20260908/acceptance-16ms`. +`freeze.json` records runtime Python hashes, source, sampling and GPU ownership. +Each startup retains its own four-worker runtime-library inventory after +measurement. Physical GPUs 4–7, TP4/B1/q8, E4M3 target KV, FP32 logits/state, +FP16 draft transport, weights, T1/k20/p.95/xhigh and natural EOS remain fixed. + +The first fresh independent startup, with five warmups and five measurements +per fixture/arm, records: + +| Fixture | BV8 / BV2 complete-round median | BV2 pure decode | BV2 TTFT | +| --- | ---: | ---: | ---: | +| release1k | 16.454723 / 16.219526 ms | 183.607 token/s | 350.235 ms | +| MBPP28 | 16.314348 / 16.096967 ms | 302.494 token/s | 127.903 ms | + +Both arms use the already-frozen attention/context/QPN2/sparse-selection stack; +this comparison isolates GDN BV2. It does not substitute for a whole-stack +quality comparison. All measured token IDs, natural EOS and acceptance match. +MBPP28 is slower than the previous 15.872776-ms observation; retain the new +samples rather than selecting only the earlier minimum. + +## Dataset protocol and open findings + +The immutable corpus hash is +`6756091e4061b0b092ceeac71e691a79b2015ef2030548f74f7cc7cc2d1cb5ed`. +It contains 32 prompts each from GSM8K, MATH500, HumanEval and MBPP, 16 from +the existing stratified LiveCodeBench v6 subset, and four JSON/tool fixtures. +Seeds are 0, 1 and 2. These are subset results, not full benchmark scores. +Each paired case has a separate one-token prefix warmup per arm, excluded from +quality/performance scoring. Measured generations have a 16384-token cap and +do not ignore EOS. Pair order alternates. All responses, including failures, +remain retained. + +Report actual accepted/proposed draft tokens, accepted draft tokens per round, +emitted tokens per round, position-specific acceptance, request-average complete +rounds, TTFT and pure decode separately. Aggregate decode is +`sum(output_tokens - 1) / sum(engine_decode_seconds)`; stream chunk intervals +are transport observations rather than instrumented GPU-round percentiles. + +Seed zero completes all 148 pairs with identical token IDs, acceptance counts +(including per-position counts), finish reasons and semantic tool calls. Its +unprofiled aggregate decode measurements are: + +| Subset | Cases | BV8 / BV2 pure decode (token/s) | Accepted / proposed | Accepted drafts / emitted tokens per round | +| --- | ---: | ---: | ---: | ---: | +| GSM8K | 32 | 309.791 / 316.749 | 57.9715% | 4.058008 / 5.058848 | +| MATH500 | 32 | 243.960 / 246.660 | 48.2514% | 3.377595 / 4.377883 | +| HumanEval | 32 | 225.650 / 229.195 | 42.8062% | 2.996437 / 3.996309 | +| MBPP | 32 | 231.882 / 234.678 | 43.3363% | 3.033541 / 4.033895 | +| LiveCodeBench v6 | 16 | 175.055 / 177.259 | 32.9424% | 2.305965 / 3.305909 | +| JSON/tool fixtures | 4 | 214.813 / 217.170 | 42.0974% | 2.824121 / 3.824121 | + +These measurements isolate the GDN change inside the shared candidate stack. +They are not complete-stack acceptance evidence. Seed-zero mathematics is +provisionally scored at GSM8K 30/32 and MATH500 31/32 in both arms. Six cases +reach the 16K cap without final content: HumanEval/10 and LiveCodeBench subset +indices 21, 64, 93, 131 and 162. They remain failures at that cap, with no credit +from reasoning-only code. Disabling all performance candidates reproduces +identical tokens for the three mathematics errors and HumanEval/10. A separate +32K-cap diagnostic makes HumanEval/10 end naturally at 21162 tokens in both +arms, again with identical tokens. It does not replace the original truncated +sample. The five LiveCodeBench failures have a separate whole-stack check. + +The retained LiveCodeBench wrapper incorrectly treated negative error codes as +truthy passes. The campaign's private corrected wrapper follows official +`lcb_runner/evaluation/pass_k_utils.py` at commit +`28fef95ea8c9f7a547c8329f2cd3d32b92c1fa24`: every test result must be greater +than zero. Eight synthetic sentinel checks pass. The original wrapper, both +source hashes and the correction are recorded; no affected score is credited. +Original HumanEval/MBPP assertion scores and EvalPlus scores are separate, with +EvalPlus's eligible-subset denominator reported explicitly. Other seeds, +executable scores and the final acceptance verdict remain pending. + +A separate same-startup control/candidate/control diagnostic is queued for +fixed prefixes, all 48 GDN layers and all 64 target layer observations. It +retains full vocabulary logits and lossless SHA256 fingerprints of the other +tensor bytes to bound disk consumption. It is not a timing service and is not +yet quality evidence. Long-context admission remains open. + +## Common schedules + +`sm70_dflash2_common_candidate_route.py` provides an explicit manifest-based +entry point for GDN value tiling, context/probe overlap, grouped E4M3 attention +and exact sparse candidate gathering. It loads no QPN2 projection library and +does not require a target quantization name. Existing guards continue to require +the audited shapes, activation/state types and applicable graph path. +Native dependencies are hashed before installation. Unsupported calls retain +the existing operator. + +`sm70_dflash2_qpn2_candidate_route.py` separately packages optional cap64 column +projections and TP4 row publication. Its representation and dimension guards +match the audited QPN2 calls. A model with other weight formats can install the +common routes without importing or loading any QPN2 projection implementation. +The combined public entry point is queued for an NVFP4 comparison with the +frozen private installer before admission. + +The integration branch includes dependency PR #563 at +`b4334fc028593942d658854e94461546c40ee21b`. It preserves prefill classification +for initial one-token requests in speculative GDN, preventing reads from a +previous request's recycled state. Its 32 focused CPU metadata tests pass; +GPU singleton/history-reuse diagnostics are queued on the fixed integration +source. This PR carries the dependency into main; the unsafe original singleton +case is not rerun on the unpatched frozen evaluation checkout. + +The FP8 model snapshot has all 66 indexed shards present. Independent control +and common-route candidate model jobs are queued, including the two speed +fixtures and 20 real quality cases. Other-quantization performance/quality is +not yet established. QPN2 compressed weight decoding remains NVFP4-specific. +No new route is enabled by default. The user has requested source integration +of PR #556 while the remaining runtime gates continue on frozen artifacts. diff --git a/docs/design/sm70_quasar_dflash2_15ms.md b/docs/design/sm70_quasar_dflash2_15ms.md new file mode 100644 index 0000000000..1b4f4cf946 --- /dev/null +++ b/docs/design/sm70_quasar_dflash2_15ms.md @@ -0,0 +1,1136 @@ +# QUASAR DFlash2 TP4: quality-preserving 15 ms campaign + +## Frozen contract + +The target is a complete, unprofiled B1 verification round below 15 ms on +release1k and MBPP28. A round includes target execution, sampling, state +handling and the DFlash2 proposal. Target-only graph time is not this metric. + +- Integration base: `56f534e672657a6c7599afd6c0dcb2e2c211b2e3`, `onecat/main`. +- Four V100-SXM2-32GB GPUs, TP4; no TP8 substitution. +- QUASAR NVFP4 target revision `d8e6fbfa3e3a78899b440222b827430045a05b44`; + DFlash2 revision `dedf8df68adfb1afeaf7b7480c0a0243108177b4`. +- FP16 activation, E4M3 target KV, FP16 draft KV, FP32 logits, seven draft + tokens and eight verification rows. Preserve the recurrent-state dtype. +- CUDA 12.8, Torch 2.10.0+cu128, Python 3.12.13; V2 runner and Flash-V100 + target/draft graphs. Model limit 262144, token budget 4096, capacity four, + one live request, memory utilization 0.8, prefix caching and Mamba align. +- Temperature 1, top-k 20, top-p .95, natural EOS, thinking `xhigh`. + release1k retains seed 20260925 and MBPP28 seed 0; speed output cap 1024. + +The retained precision-preserving results, 18.892/18.435 ms, predate this +integration base. They establish the optimization gap, not the new baseline. +New measurements freeze source overlays and every loaded native library. + +## First change: observe the actual recurrent input + +`benchmarks.sm70_dflash2_state_audit.StateAuditExtension` is an explicit +diagnostic worker extension. It is enabled only when +`VLLM_SM70_DFLASH2_AUDIT_ROOT` is set. Its active-case file supplies a complete +forced token tape. It records native logits but forces continuation and +acceptance; these requests must never be used for speed, acceptance, or task +quality claims. + +The earlier fixed-prefix audit did not preserve incoming conv/SSM state and +excluded prefill layer tensors with its eight-token dump limit. This extension +records convolution inputs and states, recurrent q/k/v/g/beta and states, +slot tables/selectors, positions and RNG states. In particular, the recurrent +input comes from `slot_table[request, accepted_selector - 1]`; column zero is +not a valid replacement. Padding remains marked invalid, distinct from live +slot zero. Snapshots own their storage and graph replays refresh them. + +State tensors are copied to persistent device buffers inside existing opaque +GDN calls and exported at the sampler boundary. Set the ordinary Qwen layer +dump token limit high enough to include prefill, and collect all four ranks. +The extension changes diagnostic work and allocation; a diagnostic result is +not evidence that an uninstrumented serving path has identical timing or +arithmetic selection. Require same-configuration repeats before attribution. + +## Ordered promotion gates + +1. Close fixed-prefix A/A repeatability, including the first prefill difference. +2. Test the existing packed verifier and eliminate confirmed layout/state copies. +3. Profile and optimize the five actual TP4 QPN2 projection shapes. +4. Extend communication/Gemma normalization fusion to q8, then assess overlap. +5. Optimize draft FP16 computation and complete-round graph scheduling. + +Copy/layout/scheduling changes require exact affected tensors, states, logits, +probabilities and acceptance. Arithmetic changes require an independent +FP32/FP64 oracle, distribution/EOS analysis, long-output checks and paired +acceptance noninferiority with no preallocated loss margin. The previous +4.33% same-configuration TV is an unresolved defect in reproducibility, never +a tolerance. Unexplained token or EOS changes block promotion. + +Require three independent paired startups and five measured requests after +warmup per speed fixture. Report mean-round-cost medians, round tails, TTFT, +pure decode throughput, accepted drafts per round and emitted tokens per round +separately. Profiler service sums and overlapping phases are not additive +end-to-end savings. Long-context quality and performance remain separate gates. + +## Provenance and current status + +This scope differs from the open FP8 target campaign (#405) and independent +batch FlashInfer ports (#515/#523): it targets QUASAR NVFP4 B1/q8 and first +repairs the missing state evidence. FlashInfer mechanisms are studied at +`91bda04c66f7cb851e1ab3b78b9fecea644b9844`; no upstream SM75+ binary is used +as an SM70 replacement. + +Artifacts for this campaign are retained under +`/data/minimax-h3/task-cache/v100-quasar-dflash2-15ms-20260908`. +`baseline-manifest.json` records the source overlay, native-library SHA256s +and loaded libraries on all four workers. The base vLLM DSO is an archived +compatible build, not a full rebuild of main. Flash-V100 and FlashQLA were +rebuilt from the frozen tree with CUDA 12.8/GCC 12. GPU clocks remain dynamic. +The initial campaign uses devices 4--7 with independent telemetry. Following +the September 8 host reboot, devices 4--7 host another service; new diagnostic +pairs use a fixed lease on devices 0--3. Results from the two GPU groups are +kept separate, and final speed pairs require a fresh baseline on the same group. +The user subsequently reserved devices 0--3 for other work: all subsequent +campaign GPU execution is restricted to devices 4--7. The temporary 0--3 lease +and telemetry are released; those diagnostics remain historical evidence only. + +Fresh uninstrumented baseline, one independent startup and five measured +requests per fixture after warmup: + +| Fixture | Median request-average complete round | Output tokens | Rounds | +| --- | ---: | ---: | ---: | +| release1k | 19.017 ms | 248 | 82 | +| MBPP28 | 18.567 ms | 270 | 60 | + +Within this startup, each fixture's five output hashes match. The JSON smoke +passes. Three existing long-code cases stop naturally at 5235, 1084 and 1312 +tokens; EvalPlus reports base 3/3 and plus 1/3. This small subset is a baseline, +not evidence of a quality improvement. The three-startup performance and +acceptance promotion gates remain outstanding. + +### First reproducibility defect: autotuned prefill reduction + +The corrected `audit-a1r3` and `audit-a2` captures each contain 144 records: +two fixed tapes, full prefill and 17 subsequent forwards, and four ranks. +The sampler may observe an extra pipelined forward beyond the API output cap; +these forced-accept diagnostics are not acceptance measurements. + +`results/audit-aa.json` reports maximum post-sampling TV 0.043373242, +five changed top-p support rows and no top-1 flips. The first causal difference +is layer 0's prefill input RMSNorm on rank 2, before the GDN projection. +The input hidden states are bitwise equal, but 9 FP16 outputs differ for +MBPP28 and 17 for MBPP3, by at most 0.0009765625. Differences then enter the +prefill conv/SSM states and the first verifier's incoming SSM state. + +Retained per-rank Inductor `.best_config` files establish that `audit-a1r3` +rank 2 selected R0_BLOCK=2048/16 warps, while the other three ranks and all +four `audit-a2` ranks selected R0_BLOCK=8192/16 warps. Replaying the actual +generated kernel with checkpoint weights and the captured MBPP28 input +exactly reproduces each arm, respectively (zero output-element mismatches). +Thus this observed drift comes from changing FP32 reduction order, not a +first difference in verifier state addressing. It does not establish that +every historical quality issue has the same cause. + +Both reductions differ from a rounded FP64 oracle (46 and 51 elements in +this captured tensor). Selecting the more common configuration alone is not +a precision argument. An environment-only attempt with +`TORCHINDUCTOR_DETERMINISTIC=1` did not reach the current AOT compile path: +the generated kernel metadata still says `deterministic=False`. Its 2.517% +A/A TV therefore does not evaluate the actual deterministic mechanism. It is +recorded as a failed route hit, not a rejected numerical implementation. + +### Opt-in fixed Gemma reduction + +`VLLM_SM70_DFLASH2_FIXED_GEMMA_RMS=1` selects a fixed 8192-element, 16-warp +reduction for contiguous FP16 `[M, 5120]` inputs and FP16 weights, with either +no residual or an FP16 residual. The latter retains FP32 residual output. +The established FP32-residual fused path and unsupported shapes keep their +existing dispatch. The new flag defaults to zero. + +The initial fixed-kernel A/A (`audit-fixed-norm-1r2` versus +`audit-fixed-norm-2`) has 144 records per arm: zero differing intermediates, +bitwise-equal logits and zero full/sampling TV. Comparing that candidate to +`audit-a2` exposed another arithmetic detail at MBPP3 step 8: three values +in layer 0 post-attention norm differ, eventually producing maximum sampling +TV 0.003018199. Top-p support and top-1 stay unchanged, which is insufficient +for acceptance. Its masked square and residual materialization boundary had +been removed, changing FMA contraction even with the same tile and warp count. + +The corrected kernel preserves those boundaries. With the same checkpoint +weights and captured inputs, both norms and the residual now exactly match +`audit-a2` for all 36 case/step combinations (two full prefills plus all 17 +verification steps per tape). The final corrected model comparison, +`audit-a2` versus `audit-fixed-norm-3`, also passes: all 144 records have +bitwise-equal intermediates and native logits, zero full/sampling TV and no +support or top-1 changes (`results/a2-versus-fixed-norm-3.json`). +No serving default is promoted. +The first AOT attempt also exposed an unresolved imported `tldevice` alias in +generated code. Using `tl.rsqrt` fixes code generation, and the test now runs +the actual Inductor backend rather than only Dynamo's eager backend. + +Current focused norm/state tests: **28 pass on V100**, including graph replay, +irregular prefill versus q1/q8 row invariance, residual storage/precision and +FP64-reference checks. Recorded operator replay, source hashes and A/A results +are under `results/fixed-norm-*.json`; the actual model runs use isolated +compiler caches and the frozen native libraries. + +Focused tests: seven pass on V100, including accepted-slot indexing, invalid +padding, owned snapshot storage, CUDA graph replay with changing selectors, +and incomplete/nonfinite capture rejection. Invalid early attempts are +retained separately: `audit-a1` failed to wrap a module; `audit-a1r2` had +stale warmup buffers and unreliable address-based layer identity. Neither +is used for attribution. Current records carry per-forward epochs and use +the caller's layer identity. Unused prefill conv-history bytes are not +automatically treated as live-state corruption. + +Hardware NCU counters are currently unavailable: the driver sets +`RmProfilingAdminOnly=1`, and the available root helper only manages GPU +clocks. This does not block state, numerical, CUDA-event or Nsight Systems +work, but no counter-based bottleneck claim is made without those counters. + +### Natural-output gate and packed verifier integration + +The first uninstrumented fixed-norm startup has median complete-round costs +19.035 ms (release1k) and 18.719 ms (MBPP28), with five measured requests after +warmup. Its three long-code cases score base 3/3 and plus 1/3, matching the +initial baseline; JSON and all nine seed/structured-fixture pairs pass, +including the existing parallel-tool premature-EOS fixture. Token sequences +change versus the original unpinned startup: first flips are output positions +187 (release1k) and 8 (MBPP28), zero-based. Accepted drafts/round are 1.917 and +3.934, respectively. These are observations, not an acceptance noninferiority +pass; the small score set cannot clear the changed-output gate. + +The initial packed on/off model run (`audit-packed-1`) is excluded as packed +parity evidence: it did not log an actual route hit and produced no additional +packed-verifier kernel specialization. Inspection finds two integration bugs: +the Qwen3.5 projection and in-place convolution retain a wider QKVZBA row +stride, rejected by the contiguous-only gate; and the packed bridge retains +the old FP16 beta default while the standard speculative path uses FP32 beta. + +The candidate reads contiguous-feature, row-strided QKV directly using its +actual row stride, and explicitly materializes FP32 beta. The existing +default-off verifier flag still controls dispatch. Sixteen GPU parity cases +cover the old FP16-beta component contract and the actual runtime bridge with +FP32 beta, wider projection rows, q4/q8, B1/B2 and FP16/FP32 states. Projection +storage remains unchanged. The diagnostic now records per-forward kernel +route markers; the comparator can require a packed hit on every observed +layer/rank/verification step. Eight state-audit tests pass, including rejection +of an equal-output comparison with no candidate hit. The complete model +comparison (`audit-fixed-norm-3` versus `audit-packed-2`) now passes with +required per-forward packed hits: all 144 records, intermediate tensors, +states and logits are bitwise equal, with zero TV/support/top-1 changes. +See `results/packed-stride-ab.json` and the pinned source overlay in +`results/audit-packed-2-source.json`. + +Twenty additional real-state cases exercise the new 4128-element QKV row +stride; all outputs, states, padding and projection storage match exactly. +Task-local NVIDIA Compute Sanitizer 2025.1.0 (CUDA package 12.8.93-1) memcheck +reports zero errors on these cases; racecheck reports zero errors or warnings. +This is an operator memory gate, not +a long-context or acceptance noninferiority gate. + +Real-state component replay also covers 480 combinations of two layers, two +tapes, four ranks, three verifier steps, all eight accepted-slot selectors, +non-monotonic state IDs including zero, empty padded requests, strided state +pools and untouched retired slots. Those original runs supplied captured FP32 +gates and contiguous QKV; they do not validate the previously incorrect runtime +bridge. All evidence remains independent of performance and natural acceptance. + +The first uninstrumented packed startup measures 18.278 ms / 17.762 ms on +release1k / MBPP28 (five requests after warmup). It is **not promoted**: +relative to `fixed-speed-1`, first output flips occur at positions 123 / 8, +and accepted drafts/round change from 1.917 / 3.934 to 1.989 / 3.500. +The three-code subset still scores base 3/3 and plus 1/3, and nine structured +seed/fixture pairs pass, but these cannot clear the changed acceptance gate. +The original unpinned baseline also selected different first-layer reduction +blocks across ranks: 2048 on ranks 0/1/3 and 8192 on rank 2. The retained +compiler configurations are in `results/natural-baseline-norm-configs.json`. + +The diagnostic supports `force_tokens=false` cases. These keep +the actual sampler outputs and record target auxiliary hidden states, incoming +draft logits, proposal candidates/scores and sampled/rejected counts. Requests +are bounded probes with synchronization and full-vocabulary dumps; neither +their latency nor their forced length cap is a performance/text-health result. +The first natural-mode startup failed because its proposal wrapper did not +preserve the `input_batch` keyword used by warmup; the signature is corrected. +The next failed before model loading because another task occupied GPU4--7. +Neither failed run contains a usable natural-sampling comparison. Per-launch +GPU availability is now rechecked in addition to the existing advisory locks. + +The first successful natural control on GPU0--3 preserves the previous +uninstrumented fixed-norm output prefixes: all 32 MBPP28 and 144 release1k +tokens match. This bounds the observed diagnostic perturbation; it is not a +complete-output, acceptance noninferiority or performance result. The natural +comparator checks complete four-rank target/proposal coverage before finding +the first observed difference. Eleven CPU tests pass (one CUDA graph test +skipped), including missing-proposal rejection, proposal-before-target ordering +and exclusion of unwritten sampled-output padding. + +The completed natural pair has 232 target records per arm (nine MBPP28 and +49 release1k forwards, each on four ranks). Its first observed difference is +already at the prefill target boundary, before the first packed q8 verifier: +layer 0/1 GDN observations remain equal, while all five auxiliary hidden-state +tensors and final logits differ. The first sampled row has zero post-top-k/p +TV in both cases, despite nonzero full-vocabulary TV. Later token/acceptance +changes therefore cannot be dismissed based on that first sampled row. +The investigation has moved to a bounded eight-token probe with layer 2/3 +observations, including the first full-attention layer. Q/K reduction autotune +choices are being checked; no Q/K normalization cause is established yet. + +The extended eight-token pair has 24 target records per arm. Its logical +target, auxiliary, proposal and acceptance tensors match exactly after aligning +physical request slots. The early proposal observer incorrectly sliced the +per-request seed/temperature arrays as if they were packed per draft row; +it now gathers by request slot. The comparator also aligns those older retained +captures and reports physical slot mappings separately. Twelve CPU tests pass +(one CUDA graph test skipped). The passing short pair does not clear the +earlier prefill drift or the complete acceptance gate. + +From the user's subsequent scope clarification, this branch concentrates on +DFlash2 complete-round cost. Independent quality-root-cause investigation is +left for the other agents, with retained artifacts in `results/quality-handoff.md`. +Numerical parity and acceptance remain mandatory candidate promotion gates. + +At 2026-09-08 01:24:11 UTC, main merged PR #560 as +`e5d63c51f0fcc1ddf75d229e3df06bf52df206f5`. It routes DFlash2 E4M3 q8 to FP32 +attention intermediates and changes the scalar/q1 precision path. The frozen +campaign results above predate that change. The cost branch integrates that +main as `631780fb4229e3cc4f384571135f6fd86996ce3f`, with the old overlay and +libraries retained for the separate numerical investigation. An independent +build from the exact merge tree is active under `flash-v4-source` / +`flash-v4-build`. Its import reports precision revision 4; 138 attention-policy +tests pass on CPU, with one GPU test skipped. Establish a new unprofiled +baseline before interpreting complete-round gains. The separate FP8-target +model gate for #560 does not validate QUASAR. +The revision-4 Flash-V100 library SHA256 is +`a751fed902279b0de23537c4aad2dc4fee360146d7fce7ef0c4f255a77f48b02`; +the matching paged-KV utility SHA256 is +`571fe2a96b70d76737375eaed9fb8ad1cac3bc7eefadf139ea3d2437e0cfdb7d`. + +### QPN2 cost measurement + +`benchmarks/kernels/benchmark_sm70_qpn2_working_set.py` exports the prepared +runtime codes, scales and actual dispatch parameters from four consecutive +TP4 layers. Its benchmark replays all sixteen projections in model order, +covering the five production shapes. Activations are explicitly frozen +synthetic FP16 inputs; weights must come from the real loaded model. Both +arms are checked for finite, bitwise-equal outputs after the first graph +replay and after all alternating timing trials. Source-library and snapshot +hashes accompany the result. This is a working-set measurement, never a +complete-round result or a replacement for the candidate's full model gates. + +The task-local trace parser now discovers the captured steady rounds and +kernel counts. A regression against the retained September 6 trace exactly +reproduces the recorded phase and complete-round timings. + +### FP32-attention baseline and confirmed layout saving, September 8 + +The user authorized stopping the service on GPU4--7. That service is stopped; +the campaign holds the rear-four lease and does not allocate GPU0--3. +The following independent startups use main `e5d63c51f0`, the pinned revision-4 +attention DSOs above, fixed Gemma reduction, and no profiler or tensor dumps. +Each fixture has one warmup and five measured requests. These are initial +screens, not the required three paired startups or final promotion evidence. + +| Fixture | Packed off | Packed on | Packed + combined split | +| --- | ---: | ---: | ---: | +| release1k complete round | 19.336 ms | 18.506 ms | 17.366 ms | +| MBPP28 complete round | 18.923 ms | 18.024 ms | 16.953 ms | + +Every output hash matches within and across these three arms. release1k emits +272 tokens in 91 rounds (1.989 accepted drafts/round; 2.989 emitted/round). +MBPP28 emits 634 tokens in 130 rounds (3.877 accepted drafts/round; 4.877 +emitted/round). Both stop naturally. The combined-split startup has median +TTFT 354.2/115.7 ms and pure decode throughput 171.5/287.2 tokens/s. +See `results/v4-layout-initial-ab.json` and +`results/v4-combined-initial-ab.json`. Loaded DSO inventories for the packed +and combined startups were hashed after measured requests. The first baseline +retains the pinned library manifest but predates that extra process-map capture. + +The opt-in `VLLM_SM70_DFLASH2_FUSED_GDN_COMBINED_SPLIT=1` reuses the existing +bitwise split kernel for the all-NVFP4 QKVZBA allocation. The old split flag +only covered checkpoints with a separate b/a projection. In the real combined +layout, each of 48 GDN layers instead launched three index creations and three +separate tail copies. The new path copies z/b/a in one launch before convolution +mutates QKV. It is limited to SM70 DFlash2 TP4/hidden5120, FP16 q8 and +QKV/z/ba widths 2560/1536/12; other shapes retain their existing route. The +new flag defaults to zero and is not automatically enabled by DFlash2 setup. +Nine GPU tests pass, including the real 4120-column view with row stride 4128, +aliased input arguments, compiled and ordinary CUDA graph replay, changed input +values, preserved tails after QKV mutation, and untouched padding. + +The fresh packed trace is `profile/v4-packed-tp4.nsys-rep` and its SQLite; +`results/v4-packed-trace.json` retains the analysis. Ten steady rounds on all +four ranks show QPN2 service 7.033 ms, target communication 1.652 ms, target +copies 1.339 ms and normalization 0.831 ms. Draft service is 3.774 ms, including +1.831 ms of dense GEMMs. The 96 tiny b/a index-select kernels each launch one +eight-thread block. These measurements motivate the combined split above. +Profiler critical-rank round time is 21.618 ms; it is diagnostic, not an +endpoint performance result. The profiler stop/export request is also excluded. + +Four QPN2 candidates were screened on rank-0 runtime weights from four +consecutive layers (214,087,680 bytes), all five production shapes, seven +alternating A/B trials and 50 graph replays per trial. All remain bitwise exact; +none is faster, so none is integrated into serving: + +| Candidate | Control/candidate median working-set ms | Decision | +| --- | ---: | --- | +| Static K specialization | 0.389 / 0.405 | Reject; every paired trial slower | +| Precombined FP16 scales | 0.379 / 0.417 | Reject; scale traffic grows | +| 16-column CTA remapping | 0.406 / 0.429 | Reject; every paired trial slower | +| One-group codes/scale prefetch | 0.390 / 0.408 | Reject; every paired trial slower | + +These are synthetic activations with real runtime weights, not full-layer or +model-quality results. The precombined-scale working set is 237,875,200 bytes; +the benchmark now records candidate scale dtype and footprint. The unchanged +source-library control also matches the archived production QPN2 outputs. +Source and DSO SHA256s accompany each `results/qpn2-*-real.json` result. + +The existing small-message push route also passed its native gate after a +task-local build of the current main communicator: all four ranks, 13 message +sizes, 32 changing-input cycles for each of random/zero/special patterns, +mixed graph order, interleaved sum2, delayed ranks and canaries. For 128 q8 +collectives, 80 versus 40 blocks measures 0.884 versus 0.882 ms in the same +communicator lifetime. This gain is too small to justify a full model candidate; +it is not promoted. See `results/custom-ar-v4-mixed-size-gate.json` and +`results/q8-push-grid-race.json`. No normalization arithmetic was changed. + +### QPN2 publication and communication arrival audit + +The private publication candidate moves the established 16-byte packet writes +into the unchanged QPN2 arithmetic epilogue. Producer CTAs never poll or wait. +A separate consumer preserves rank-ordered FP32 addition, FP16 materialization, +sentinel escaping/cleanup and both epochs of the existing push pool. The normal +projection output remains materialized. This differs from producer-poll fusion; +the first implementation spilled its dynamically indexed peer-pointer array +and was slower. Passing the local pointer directly eliminates those spills. +The exact q8 row-projection kernels use 48 registers with no stack or spills; +the consumer uses 40 registers with no stack or spills. + +`benchmarks/kernels/build_sm70_qpn2_publish_candidate.py` generates the private +candidate from production source anchors and records source/DSO hashes. It does +not replace a serving operator. `--build` now keeps default CUDA math; +`--use-fast-math` explicitly reproduces historical experiments and is an +arithmetic change for gated SiLU. All extra CUDA flags are recorded in the +manifest. The generated CUDA source SHA256 is +`7360d578080c96350970b9ceb42fa5e470627949c28219013d1daa0861acc42a`. +Use a communicator built from the same header and set +`VLLM_SM70_CUSTOM_AR_LIBRARY` to that sidecar; do not mix opaque communicator +objects between libraries. The measured sidecar SHA256 is +`e32f156f606c47dc5863a7785065f1fef9ff3668788d9c44e5f85c2264e8a7d1`. + +`benchmarks/kernels/benchmark_sm70_qpn2_publish.py` exercises sixteen real +prepared projections on each of four ranks, with eight dependent all-reduces +and an extra ordinary push call. Five changing-input cycles, alternating graph +order, delayed ranks and output canaries pass byte-for-byte checks on projected +and reduced tensors. The working-set screen measures 0.491/0.465 ms for +control/candidate; it is not a complete model round. A focused memcheck and +racecheck pass with zero errors using Gloo process coordination. The first +NCCL-coordinated sanitizer run exited on CUDA API 209 during NCCL's kernel +capability probing; it was not counted as a pass. Racecheck is a shared-memory +check, not proof of all inter-GPU global-memory ordering. + +The model-screen DSO SHA256 is +`3e5afdbdb176cf4ed7460e125f48f0bb1f36e85fae7107d33f48deb157919a37`. +The reusable builder produces identical CUDA source; its independently rebuilt +DSO `007298ccc086ac3181a29b9d745f32c76809a9d2b90158af474b2b94f3efc457` +also passes the four-rank five-cycle correctness gate. Original sanitizer +evidence is tied to the model-screen DSO, not relabeled as a rebuilt-DSO run. + +The task-local model integration preserves one opaque row-projection boundary +in both arms, reuses the same compiler/autotuner cache, and enables publication +only during q8 CUDA capture. All 128 target row projections hit on every rank; +draft projections retain their existing path. One independent startup per arm, +one warmup and five measured requests per fixture, without instrumentation: + +| Fixture | Matching publication control | Publication enabled | Saved | +| --- | ---: | ---: | ---: | +| release1k complete round | 17.232 ms | 17.073 ms | 0.159 ms | +| MBPP28 complete round | 16.785 ms | 16.660 ms | 0.126 ms | + +All output hashes, natural EOS and both acceptance-length definitions match the +earlier arms above. The matching control includes the new row-op boundary and +rebuilt communicator: do not attribute its difference from the earlier +17.366/16.953 ms result to publication. See +`results/v4-publish-initial-ab.json`, per-startup source manifests and mapped +library inventories. The model integration remains task-local and disabled +by default; full state/distribution and repeated-startup gates remain open. + +Ordinal-matched communication in the four-rank packed trace shows that its +first target push has mean kernel duration 0.388 ms, rank arrival skew 0.684 ms, +and last-arrival-to-last-finish time 0.010 ms. The first draft push similarly +measures 0.236/0.473/0.009 ms. Much of these particular kernel durations is rank +waiting. Graph-node profiling can itself inflate arrival skew; these numbers +do not establish unprofiled host overhead. The CPU sparse-target probe span +includes waiting for queued target GPU work and must not be added again as +independent CPU cost. See `results/v4-collective-arrival-audit.json`. + +Further exact QPN2 screens (shared-partial bank swizzle, cache policy and a +33-percent shared-memory carveout hint) show either noise-level savings or +regressions and remain unpromoted. A vectorized peer-read/Gemma prototype is +exact but slower (eight joins: 0.100/0.264 ms); a local-push consumer design +requires its own evidence. The existing draft TurboMind FP16 GEMM screen saves +about 0.200 ms across twenty real-weight projections but is not bitwise equal. +Its independent FP64-reference errors do not worsen in that primitive screen; +model-distribution and acceptance gates are still required, so it is not enabled. + +The subsequent forced-tape publication pair (`v4-publish-audit-control` / +`v4-publish-audit-speed`) has 140 records per arm across all four ranks and +two 128-token tapes. Requested layer 0/1 intermediates and conv/SSM state, +target boundary tensors, and native logits are bitwise equal. +Sampling TV is zero with no changed top-p support or top-1 rows. This gate +retains the complete prefill records; it is not natural acceptance evidence. +See `results/v4-publish-audit-comparison.json` and its separate manifest. + +The lower-overhead whole-graph trace in `profile/v4-publish-graph/tp4.sqlite` +does not collect individual graph nodes. Ten steady rounds have diagnostic +critical-rank mean interval 18.465 ms, GPU union 17.214 ms and uncovered time +1.252 ms. Target graph mean duration is 12.308 ms. Its host launch skew is +0.685 ms, but GPU start skew is only 0.005 ms because launches are queued. +The main draft graph has host/GPU start skew 0.292/0.297 ms and mean duration +3.862 ms. Do not treat the earlier target-node arrival skew as an established +unprofiled saving. The request containing profiler stop/export is excluded +from endpoint performance claims. Analysis: +`results/v4-publish-graph-trace.json` and +`results/v4-publish-graph-arrival-audit.json`. + +Additional independent screens remain unpromoted: + +| Candidate | Control/candidate working-set ms | Result | +| --- | ---: | --- | +| Local published-packet consumer + Gemma | 0.508 / 0.515 | Exact; slower in every pair | +| Global QPN2 partials, four warps per CTA | 0.389 / 0.539 | Exact; added traffic/launches do not pay back | +| Fixed-q8 input/output bounds | 0.380 / 0.477 | Exact; compiled register use rises to 70--72 | +| Fixed-q8 bounds, unroll two | 0.393 / 0.408 | Exact; 64 registers, still slower | +| K-group-major weight codes/scales | 0.379 / 0.392 | Exact; same footprint, still slower | + +The local consumer preserves the existing packet protocol and Gemma reduction +topology and returns the materialized reduced tensor. Its successful gate +covers changing inputs, graph order, delayed ranks and canaries. It follows +two retained harness failures: incorrect packed inline-assembly return +constraints and an unregistered warmup-only ordinary collective buffer. +Neither failed run is counted as correctness or speed evidence. + +An independent-stream context experiment gives each arm an identical context +capture stream, private graph pool and cuBLAS workspace. Only replay placement +differs. Context scratch reads wait for target output; accepted-slot KV writes +stay on the main stream and wait for context completion. The unprofiled model +pair is slower: release1k 17.088/17.547 ms and MBPP28 16.659/17.130 ms. All +output hashes and acceptance lengths match. See +`results/v4-context-overlap-ab.json`; the experiment remains disabled and no +further quality promotion work is justified by this negative speed result. + +Rear-four telemetry during decode reports 1530-MHz SM clocks, 877-MHz memory, +roughly 171--183 W draw under the unchanged 300-W limit, and no active clock +event reason in the checked samples. Clock headroom is not credited as a +remaining optimization. + +### Direct attention output: speed candidate held at the numerical gate + +`VLLM_SM70_DFLASH2_DIRECT_ATTENTION_OUTPUT` defaults to zero. It is armed only +for SM70, DFlash2, TP4, FP16 dense Qwen3.5 with hidden size 5120. The decoder +can consume the projection tensor already returned by attention. The new GDN +opaque entry returns that allocation while retaining the full-forward operation +order and explicit conv/SSM mutation arguments. It does not enable the existing +long-prefill collective/norm switch. Other models keep the existing path. + +The initial artifact prototype failed during compilation because the existing +GDN full-forward boundary required an output buffer even though the decoder +could accept a direct return. The new return-valued opaque entry resolves that +interface issue. Its schema marks both caches mutated and its return unaliased; +the fake implementation passes shape/dtype checks at 1, 8, 135 and 4097 rows. +Scoped Python lint, format and bytecode checks pass. + +One unprofiled artifact A/B startup per arm, five measured requests per fixture, +gives 17.091/16.824 ms on release1k and 16.561/16.402 ms on MBPP28. All output +hashes, natural EOS and acceptance lengths match. The source-integrated version +also reaches this range, but **this candidate is not numerically cleared**. +Its 140-record fixed-prefix comparison has identical captured layer 0/1 +intermediates and conv/SSM state, while all target-boundary hidden records and +native logits differ. Maximum sampling TV is 0.0104069 with no changed top-p +support or top-1 rows. These observations do not yet distinguish later-layer +arithmetic/compiler effects from diagnostic perturbation. See +`results/v4-direct-source-audit-comparison.json` and +`results/direct-output-quality-hold.json`. A bounded eight-token probe captures +GDN layer 2 and the first full-attention layer 3. Until the difference is +localized and resolved, exclude this candidate from promoted combinations. + +The QPN2 input-layout screen is separate: arranging the same FP16 q8 input as +`[K/16, 8, 16]` gives 0.387/0.357 ms across the real four-layer weight working +set, bitwise equal in every projection. That screen excludes packing time and +does not establish model speed. Gemma producers that write this layout directly +pass 45 changing-input graph cases across no-residual, FP16-residual and +FP32-residual contracts and three magnitude ranges. The subsequent model +experiment is based on the cleared publication combination, with the direct +attention-output switch disabled. See `results/qpn2-input-packed-real.json` +and `results/packed-gemma-gate.json`. + +The first complete-model packed-input pair fails admission: release1k changes +from 272 to 210 tokens and 1.989 to 1.800 accepted drafts/round; MBPP28 changes +from 634 to 754 tokens and 3.877 to 3.303 accepted drafts/round. Its apparent +16.823/16.443 ms timing is not a promoted gain. A diagnostic shadow then uses +the actual input and weight of every selected operator: per rank, 383 fresh +comparisons cover 128 norms, 127 residuals and 128 column projections. Norms, +residuals and ungated projections all match. Gated projections in later layers +show a few differing bytes, often one FP16 ULP, which the first-four-layer +synthetic-input screen missed. Retained evidence is +`results/packed-input-operator-shadow-summary.json` and each rank's raw report. + +The experimental packed-input DSO used `--use_fast_math`. Disassembly of its +gated kernel has no FFMA correction instructions; the archived production +gated kernel and the default-math rebuild each contain 22. The current CMake +QPN2 path obtains Torch's common CUDA flags without adding fast math. A +same-input layer-22 shadow separates raw gate/up GEMM from activation: both +raw GEMMs are bitwise equal on all four ranks, and the control fused activation +matches native `silu_and_mul`. Only the experimental activation differs. The +builder's former implicit fast-math default has therefore been removed and +its math mode made explicit. Historical SHA256s and measured results are not +relabeled as default-math results. Publication's serving path used only its +nongated producer and already passed its complete 140-record comparison. + +The default-math packed-input rebuild has source SHA256 +`13618190405372caed28f15391c4783c884ffca02a10af25a210da28668e216a` +and DSO SHA256 +`257af8ceb4230428f874ac7429387ffff1697925cdb92bf1227477d5a1eed564`. +All 383 fresh actual-input comparisons on each of four ranks now match +bitwise, including all 128 column projections. One unprofiled startup per +arm gives release1k 17.011715/16.934972 ms and MBPP28 +16.549580/16.415542 ms, each the median of five measured requests after +warmup. All output hashes, natural EOS and acceptance counts match the +control: 272 tokens / 91 rounds / 181 accepted drafts and +634 tokens / 130 rounds / 504 accepted drafts. The isolated saving is +0.076743/0.134038 ms, not a sub-15-ms result or a repeated-startup admission. +See `results/packed-input-strict-shadow-summary.json` and +`results/v4-packed-input-strict-ab.json`. The direct-attention switch stays +disabled in both arms. + +The longer fixed-prefix pair does not clear admission. After interpreting the +544 packed norm snapshots in their logical layout, all captured layer-0/1 +states and layer outputs match. However, all 144 target-boundary records and +native logits differ, with maximum sampling TV 0.0441784 and three changed +top-p support rows. The first difference already occurs at the prefill target +boundary, before the q8-only layout is active. This experiment cannot attribute +that difference to packing or supersede the separate prefill/compiler +repeatability investigation. Keep the layout candidate held. See +`results/v4-packed-input-strict-audit-canonical-comparison.json` and +`results/packed-input-quality-hold.json`. The first candidate startup was +terminated by another task before producing a result; its identical recovery +run supplies the candidate captures. The interrupted run is not a gate result. + +An exhaustive activation check covers all 63,488 finite FP16 gate values, +with the up input fixed at one. Default CUDA math exactly matches native +SiLU. Explicit fast division/exponential differs at precisely two inputs: +`-2.724609375` and `-4.921875`. A separate experimental helper retains native +math for those inputs and every nonfinite input. It matches native output +bits for all 65,536 FP16 bit patterns, including signed zeros and NaN payloads. +This is a bounded SM70/CUDA-12.8 activation contract check, not a claim about +other compilers or the performance of a complete gated projection. The +helper remains a private candidate until actual projection and model gates +pass. See `results/silu-math-contract-gate.json` and +`results/silu-corrected-contract-gate.json`. + +The corrected activation subsequently matches every projection in the real +four-layer working set, including three changed input magnitudes, but does +not improve timing: 0.367063/0.367616 ms. It is rejected for speed and is not +advanced to a model combination (`results/qpn2-packed-silu-exact-real.json`). + +### Packed MLP boundary and graph scheduling screens + +The next layout candidate lets gate/up write its FP16 output directly as +`[hidden/16, 8, 16]`, then lets the down-projection publisher read that layout. +Both arms use the preceding packed normalized input and existing publication +protocol. No extra transpose, SiLU change or accumulation change is introduced. +Four ranks pass five changing-input cycles with delayed-rank replays, canaries, +all eight reductions and an ordinary ninth push. All seven timing pairs favor +the candidate in the four-layer working set. See +`results/qpn2-packed-mlp-chain-real.json`. + +One unprofiled startup per arm, warmup plus five requests per fixture, gives +release1k 16.849092/16.802114 ms and MBPP28 16.566333/16.280088 ms. Every output +hash and acceptance count matches the control. These are isolated screen +results, not a completed repeated-startup gate. They inherit the preceding +packed-input quality hold. See `results/v4-packed-mlp-ab.json`. + +The private packed-input and gated-output DSOs are reproducible with +`benchmarks/kernels/build_sm70_qpn2_packed_input_candidate.py`, using +`--pack-gated-output` for the latter. The builder retains the production +arithmetic, restricts these entry points to M=8, records source/DSO SHA256s, +and uses ordinary CUDA math. The existing publication builder accepts +`--packed-input` for its publisher only; its standalone control GEMMs retain +ordinary inputs. `benchmarks/kernels/benchmark_sm70_qpn2_packed_mlp.py` accepts +all four DSO paths explicitly and checks the coupled gate/down boundary. +These tools do not install a serving route or enable a default. + +The source-rebuilt versions pass the four-rank coupled gate with five changed +input cycles, skewed rank launches and intact canaries. Each rank also rejects +twelve non-q8 operator calls before kernel launch. Rebuilt DSO SHA256s are +`6f718757d5918ae413e3f6977605a451649dfc31e3b3e2c709f53b59a14e77ec` +(packed input), +`d9a3990222b7a2d98243ab8707b582fdbe3de22a759096a0015dec131c2389ba` +(packed gated output) and +`5adec45ecceecbd0c8c7f7f37eac2d0f507e0c3ed09d2867eb817dac5b64e8fd` +(packed publisher). This rebuild check has no timing claim. See +`results/qpn2-packed-mlp-versioned-gate.json`. Without `--packed-input`, the +publication builder still generates the previously recorded source SHA256 +`7360d578080c96350970b9ceb42fa5e470627949c28219013d1daa0861acc42a`. + +CUDA 12.8 conditional IF/ELSE graphs were exercised on the rear V100 with ten +changing-condition replays. NVIDIA documents the conditional-body node +restrictions in its [CUDA 12.8 runtime interface](https://docs.nvidia.com/cuda/archive/12.8.1/cuda-runtime-api/structcudaConditionalNodeParams.html). +This capability differs from PDL and was checked on SM70 directly. A follow-up +prototype retains the original NumPy boundary decision via pinned transfers +and a graph host callback. All 24 changing-input decisions match, but its tiny +round-trip screen is slower: 0.091187/0.173937 ms. It is not integrated into +the model (`results/host-conditional-graph-gate.json`). An independent draft +tail experiment composes existing KV-store, metadata and query graphs in +their original order. PyTorch rejected the first nested-replay capture before +measurement. The follow-up retains the original graph handles and composes +them with native child-graph APIs. Sixteen changing-input primitive replays +match. Its unprofiled complete-round pair is 16.986402/16.991194 ms for +release1k and 16.540575/16.560563 ms for MBPP28. All five requests per fixture +retain the canonical output hashes and acceptance counts, but neither +fixture improves. Reject this scheduling candidate for speed; see +`results/v4-draft-tail-native-v2-ab.json`. + +The bounded direct-output probe has eight four-rank records: GDN layer 2, +full-attention layer 3, target hidden and native logits all match, with zero TV. +This does not clear the earlier 140-record drift; the direct-output switch +remains disabled. See `results/v4-direct-tiny-comparison.json`. + +### Publication node trace and bounded follow-up screens + +The next node-level trace uses the publication combination with packed GDN, +combined split and fixed Gemma RMSNorm. Direct attention output and the held +packed-input experiment remain disabled. The raw SQLite SHA256 is +`13a16f21fba7a0fd46e41bccb741694f3d395048101f6e3e3edacf31305e8b37`. +Ten complete steady rounds, four ranks, give the following CUDA service +attribution in `results/v4-publish-nodes-trace.json`: + +| Work | Mean service ms / rank / round | Calls / rank / round | +|---|---:|---:| +| Target QPN2 gate/up and SiLU | 2.968 | 64 | +| Target QPN2 published row projections | 2.887 | 128 | +| Target QPN2 other projections | 1.562 | 64 | +| Draft dense GEMMs and reductions | 1.832 | See phase attribution | +| GDN recurrent update | 0.954 | 48 | +| Target grouped attention | 0.952 | 32 | +| Target normalization and residual | 0.837 | See phase attribution | +| Gather/scatter/copy across phases | 0.351 | 63 | + +This confirms projection work as the largest remaining target. The trace +also records 128 published-packet consumers with 0.656 ms of service and +12 ordinary push collectives with 1.717 ms. Those durations include waiting +for other ranks; they are not independently removable work. Node tracing +perturbs scheduling: the diagnostic critical interval is 22.175 ms, whereas +the previous whole-graph trace gave 18.465 ms and the unprofiled paired +endpoint measurements were lower. Neither trace is a new performance +baseline. The profiler-stop/export request is excluded from speed evidence. + +A private two-stream experiment starts the existing packet consumer before +the producer and joins it before the next dependent projection. All four +ranks finish capturing both graphs, and the control graph replays. The +candidate hangs in its first replay and reaches the bounded timeout. No +numerical or timing result exists for it. The cause is not yet localized; +do not label it a measured overlap benefit or a proven occupancy failure. +See `logs/qpn2-publish-overlap-diagnostic.log`. + +A NUMA scheduling pair binds the control to both CPU nodes and the candidate +to the rear GPUs' local CPU node. The isolated medians are +17.059970/17.033062 ms and 16.582857/16.485381 ms. Both arms agree with each +other but share a changed trajectory relative to the earlier canonical +baseline: release1k 283 tokens / 94 rounds / 189 accepted drafts, and MBPP28 +270 / 60 / 210. Common mapped dynamic libraries have identical hashes. +Do not compare that shortened MBPP28 request with the earlier 634-token +performance or attribute the common trajectory change to local CPU binding. +This is another unresolved baseline-repeatability observation, not an +admitted scheduling change (`results/v4-numa-ab.json`). + +An independent FP64 arithmetic screen decodes the actual rank-0 QPN2 +weights into the unchanged FP16 operands and evaluates all five matrix +shapes, four consecutive layers and three activation magnitudes. Every +single-accumulator-chain configuration expands at least one registered +reference error and is rejected. Some row-projection configurations retain +the checked max, p99 and relative-L2 bounds with small working-set gains, +but have no model/acceptance evidence and remain disabled. See +`results/qpn2-calibration-fp64.json`; these are not full-round gains. + +The earlier fixed-q8 specialization crossed a register-use boundary and was +slower. A new build retains its accumulation order and ordinary CUDA math +while capping registers at 64. Its four-layer working-set pair is +0.405852/0.392275 ms, with all seven paired differences positive. An actual +MBPP28 q8 step checks all 128 affected target column projections on each of +four ranks: every output byte matches the existing operator. This is a +bounded operator check, not full-prefix admission. The model sidecar source +SHA256 is +`a3b480efe2f671cf05ef39bd775f18c0f4cfc6b67bcf92fc96aec57960b06514` +and its DSO SHA256 is +`a62fa06fecb0f67a9011e011f2112f8006e691018e95c218c0bc092818303a4d`. +See `results/qpn2-fixed-q8-cap64-real.json` and +`results/qpn2-cap64-shadow-summary.json`. Its unprofiled complete-round +pair improves release1k from 17.051370 to 16.888148 ms and MBPP28 from +16.652563 to 16.456873 ms. Each arm has one independent startup and five +measured requests after warmup. All ten candidate requests retain the +canonical token hashes, natural EOS and acceptance counts. The subsequent +144-record fixed-prefix pair also matches: all captured layer-0/1 outputs, +conv/recurrent states, target hidden states and native logits are byte-equal; +sampling TV, changed support rows and top-1 changes are all zero. This is not +final repeated-startup admission or a sub-15-ms result. See +`results/v4-qpn2-cap64-ab.json` and +`results/v4-qpn2-cap64-audit-comparison.json`. The source is reproducible +with `benchmarks/kernels/build_sm70_qpn2_q8_candidate.py`; the working-set +benchmark accepts its private namespace through `--candidate-namespace`. +The source-rebuilt DSO has SHA256 +`6a7e3e3f06f1f7cd2cadec4d4205381b0e6b73f8904ea65782306a54afd2c9ee`. +It reproduces the recorded source hash, matches the sixteen-projection +working-set outputs and rejects twelve non-q8 calls before launching a +kernel (`results/qpn2-cap64-versioned-gate.json`). This rebuild has no +additional performance claim. + +An independent L2-prefetch candidate avoids keeping future decoded weights +in registers. It preserves all checked output bits, but the four-layer +working set slows from 0.378491 to 0.433336 ms, with every paired trial +slower. Reject it; see `results/qpn2-l2-prefetch-real.json`. + +### Follow-up draft and attention candidates + +A different GDN output-copy experiment preserves the original opaque GDN +operator and its explicit state/output mutations. Its row-projection consumer +writes directly into the existing output buffer. All 48 GDN layers hit on +each rank, and both fixtures preserve canonical tokens, natural EOS and +acceptance counts. However, release1k changes from 16.962519 to 17.066488 ms +and MBPP28 from 16.517022 to 16.913165 ms. Reject this candidate for complete- +round speed, without extending its quality tests +(`results/v4-gdn-sink-ab.json`). This does not clear or reuse the earlier +direct-attention-return candidate. + +The draft's live B1 query is eight rows. Its five query attention layers use +fused QKV projections with TP4 shape N=1536, K=5120, rather than a standalone +N=1024 query projection. The first private loader correctly stopped when +only fifteen of twenty expected projections matched its module selection; +it produced no model measurement. The corrected twenty-projection screen +includes all three Q/K/V shards, unchanged FP16 weights, and the existing +separate BF16-emulation/activation boundaries. Its raw FP16 HMMA implementation +uses two FP32 accumulator chains. With four K partitions, the working-set +pair is 1.458278/1.151078 ms, with non-increased FP64 reference error across +three synthetic activation magnitudes. The earlier q-only projection screen +is not a complete runtime draft-path measurement. + +Actual model inputs reveal why that synthetic gate is insufficient. The +control-fed shadow captures five q8 query steps, twenty projections per step, +on all four ranks. Among 400 comparisons with independent FP64 products, +the four-partition candidate expands a reference metric three times: one +gate/up maximum error and two QKV relative-L2 errors. It remains disabled +despite improved aggregate errors. Re-evaluating exactly those saved inputs +with eight or sixteen K partitions gives no expanded max, p99 or relative-L2 +metric in all 400 comparisons. The uniform eight-partition candidate then reaches a complete-model A/B, +but it is not admitted: release1k changes from 16.890287 to 16.549846 ms +while emitted tokens / rounds / accepted drafts change from 272 / 91 / 181 +to 270 / 91 / 179. The accepted-draft count falls and the token hash changes. +MBPP28 changes from 16.458697 to 16.117234 ms, but its control trajectory +is 457 / 96 / 361 instead of the canonical 634 / 130 / 504, which the +candidate happens to reproduce. Do not treat this as a matched-output +speedup or use its lower cost as the numerically cleared best result. See +`results/v4-draft-q8-split8-ab.json`; distribution/acceptance gates stay open. +See `results/draft-f16-q8-fused-real.json`, +`results/draft-q8-actual-reference.json` and +`results/draft-q8-actual-calibration.json`. + +An attention scheduling candidate divides six query heads into three groups +of two, preserving each head's QK/PV arithmetic, probability compensation, +K partitioning and FP32 numerator/max/sum storage. Forty-five changing-row- +length graph states match output, the entire partial/max/sum workspace and +canaries byte-for-byte, including a 262144-token case. A sixteen-layer screen +with cache-eviction work between layers saves approximately 0.067 ms; this +is not a complete-round result. Native memcheck reports zero errors. +Racecheck reports shared-memory access warnings: isolated baseline-only and +candidate-only probes each reproduce two warning sites. The unsynchronized +candidate stays held. Adding an explicit warp barrier before lane 0 updates the online-softmax +row state clears both isolated warning sites. Both the original CTA layout +plus the barrier and the regrouped layout plus the barrier pass all 45 +output/partial/max/sum/guard comparisons, including zero rows and the +262144-token operator case. Isolated racecheck reports zero hazards, errors +and warnings for both variants. Neither this bounded check nor the original +warnings establish a text-quality root cause. See `results/attention-headsplit-gate.json`, +`results/attention-headsplit-memcheck.json`, and the corresponding isolated +`logs/attention-race-{baseline,candidate}.log` files. + +The source fix in `flash_decode_paged.cu` orders each warp's shared-state +reads before lane 0 overwrites the row maximum. NVIDIA documents +[`__syncwarp` memory ordering](https://docs.nvidia.com/cuda/archive/12.8.0/cuda-c-programming-guide/index.html); +a shuffle's synchronization does not provide that shared-memory ordering. +The source SHA256 is +`4e8a2ea7fe5315f30cdc66e4c60a90a9460b28e4ec723a6549892d03a2d5d9bd`. +The standalone original-layout DSO is +`21520f8d573bd7b8ec74943cac15385671c3f226bde9640d92b34ff3825ed472`; +the regrouped-layout DSO is +`690300fffc5f265642d8effbebfe89992fcaaed5a5a1764464ed4615eb8bc6ab`. +See `results/attention-baseline-sync-gate.json`, +`results/attention-headsplit-sync-gate.json`, and +`logs/attention-{baseline,headsplit}-sync-racecheck.log`. Both variants also pass native memcheck over thirty short changing-row +states and isolated synccheck, each with zero errors. The complete native FA rebuild has SHA256 +`bc8410cf09e87e6ca31679886fbe85e0a37bd010d7d6a317208e8b2a56ce9e01`. +It passes 89 targeted grouped E4M3-FP32, E4M3 and legacy grouped-verifier +tests, including the added short q8 multi-tile racecheck fixture. The first +test launcher failed before importing the extension because Torch had not +yet loaded `libc10`; the corrected entry explicitly imports Torch first. +See `results/flash-sync-build.json` and `logs/flash-sync-native-tests-v2.log`. + +A QPN2 experiment assigns the two original FP32 accumulation chains to +separate warps, preserving the final pairwise sum and K-partition order. +All sixteen real-weight projection outputs match, but the four-layer +working set slows from 0.387072 to 0.442921 ms; every paired trial is slower +(`results/qpn2-chain-warps32-real.json`). Chain-specific code packing also +passes the sixteen-output numerical screen on physical GPU 5. That screen +ran with another task on GPUs 4/7 and only one timing iteration; its timing +is excluded from performance evidence. Two isolated seven-pair screens reject the packed version as well: +default carveout changes 0.406610 to 0.471409 ms, and 100% shared-memory +preference changes 0.381563 to 0.514908 ms. All outputs remain byte-equal. +CUDA's occupancy API gives the same theoretical block limits (2 / 4 / 2 +for gated-S8 / GEMM-S8 / GEMM-S16) before and after the carveout preference; +these are resource estimates, not measured achieved occupancy. This path is +closed (`results/qpn2-chain-packed-{default,carveout100}.json`). + +The earlier head-regrouping model pair preserves both canonical token hashes, +natural EOS and acceptance counts across all five measured requests per +fixture. Its raw medians are 17.064797/16.912649 ms for release1k and +16.479594/16.451247 ms for MBPP28. **Withdraw attribution of these differences +to head regrouping:** the later node trace and CPU module-identity probe +show that the hook patched a different Python extension module from the one +used by the model. Both names resolve to the same frozen DSO, but their +module objects and function bindings differ. The subsequent 144-record +comparison remains a valid equality observation of the executed paths, +not evidence that the regrouped model route was active. See +`results/v4-attention-headsplit-sync-ab.json` and +`results/v4-attention-headsplit-sync-audit-comparison.json` and +`results/attention-headsplit-binding-identity.json`. The isolated candidate +operator gates remain separate evidence. Corrected route, repeated-startup +acceptance and full-model long-context validation remain open. + +An adjacent-K16 weight/scale packing experiment retains the current q8 +accumulation chains and uses vector loads for pairs of groups. Sixteen +projection outputs match the original operator; the four-layer working set +changes from 0.379699 to 0.374784 ms with all seven paired differences +positive. Because the candidate includes fixed-q8/cap64 as well, a direct cap64 pair is inconclusive: the arm medians are +0.393216/0.400855 ms while six of seven paired differences favor the +candidate. The samples have substantial time variation. The sustained-warmup ABBA screen resolves the paired direction: all seven +trials favor the candidate, with medians 0.376730/0.375122 ms. Each post-trial +sample reports 1530/877-MHz SM/memory clocks. The approximately 0.0016-ms +four-layer benefit is too small to justify another model route now; no model +promotion follows. This does not establish what caused the earlier time +variation (`results/qpn2-pair-load64-steady.json`). This is not an additional +admitted gain over cap64 (`results/qpn2-pair-load64-real.json` and +`results/qpn2-pair-load64-vs-cap64.json`). + +Decoding half a K16 group's weights immediately before its corresponding +MMA pair retains all sixteen projection outputs but is slower in every +paired trial; reject it (`results/qpn2-late-decode64-vs-cap64.json`). A +fixed-q8 publication producer increases compiled register use to 72 before +capping. Its 64-register build preserves the serial protocol and all checked +four-rank outputs/canaries, but gives only 0.453878/0.452792 ms across the +four-layer collective screen with mixed-sign differences. No robust speed +claim or model promotion follows (`results/qpn2-publish-fixed64-real.json`). + +A bounded device-polling probe completes one K4352 row projection under both +serial and overlapping graph schedules on all four ranks. All 5120 packets +per rank arrive, epochs are uniform, and projected/reduced output bytes +match. The maximum observed candidate poll interval is 76757 device cycles. +Because the probe changes the polling kernel, this does not clear the +original multi-projection hang or measure a speedup. Restoring the consecutive projections and ordinary-push transition exposes +all 5120 first-collective packets timing out on every rank. Once that +perturbed poll exits, all seven later collectives receive their packets. +The expected nonfinite-output assertion rejects the run; it is diagnostic +evidence, not quality admission. It narrows the missing condition to the +first collective in the larger graph, without proving a scheduling or +memory-ordering root cause. See `results/qpn2-publish-overlap-probe.json`, +`logs/qpn2-publish-overlap-multi-probe.log` and the per-rank saved probes. + +`benchmarks/kernels/build_sm70_grouped_attention_candidate.py` reproduces +the private attention scheduling experiment from the repaired source. It +supports the original one-group layout and the three-group candidate, +preserves the native validation/math contract, copies the required headers +and license, and records source/header/library hashes. The three-group +source rebuild has DSO SHA256 +`e40120cb4788a7b1443fd48cd2e348df95c6782997d9040efd7aa778031acfeb`. +All 45 output/workspace/guard comparisons pass for this rebuilt DSO +(`results/attention-headsplit-versioned-correctness.json`). It installs no +serving route; its operator gate is not an additional model speed claim. + +Rotating physical warp assignments to logical K partitions across N tiles +preserves the matrix arithmetic but is slower in all seven steady ABBA +pairs: 0.374835/0.375613 ms. Reject that schedule +(`results/qpn2-staggered-k64-steady.json`). + +A private fused producer/consumer instead uses cooperative kernel launch and +checks the complete 160-CTA grid against the device's admission capacity. +[NVIDIA documents cooperative launch in CUDA Graphs](https://developer.nvidia.com/blog/cuda-11-features-revealed/). +It keeps the original K accumulation, FP16 projection rounding, FP32 rank +reduction, packet cleanup and eighty epoch counters. Five changing-input +cycles across four ranks, delayed-rank cases and the ordinary ninth-push +transition pass bytewise output and canary checks. However, the complete +four-layer collective screen regresses from 0.445501 to 0.500654 ms in all +seven pairs. A second version replaces the grid-wide barrier with per-epoch +last-arrival counters. It also preserves outputs and resets its counters, +but remains slower in every pair (approximately 0.4445/0.4980 ms). Both are +rejected for speed; these bounded results do not diagnose all causes of the +separate split-stream hang. See `results/qpn2-cooperative64-real.json` and +`results/qpn2-cooperative-counts64-real.json`. No fused consumer is enabled. + +The minimal node trace of the bounded split-stream probe records the first +candidate producer starting approximately 77.07 ms after its consumer, +only near the latter's timeout. CUPTI's driver-selected shared-memory size +is 64 KiB for the first S8 producer and 0 KiB for its consumer; S16 +producers use 32 KiB. See `results/qpn2-overlap-first-collective-trace.json`. +[NVIDIA describes possible synchronization when cache preferences change](https://docs.nvidia.com/cuda/archive/12.8.0/cuda-driver-api/group__CUDA__EXEC.html). +An intervention sets a matching 25% shared-memory preference on the bounded +probe's two producer kernels and consumer. The previously failing complete +four-layer graph then finishes: all four ranks match all outputs/canaries, +with no packet timeouts in its eight collectives. This supports the cache- +configuration explanation. It is a one-cycle diagnostic with modified +polling, not a production liveness guarantee or a speed result +(`results/qpn2-compatible-carveout-probe.json`). A separate bounded early consumer leaves epoch updates to a main-stream +completion kernel after joining the producer and early consumer. Five changing- +input cycles match all four-rank outputs and canaries, both with normal overlap +and with early polling forced to finish before the producer (all packets then +use completion). The three-arm medians are 0.457175 ms for original serial, +0.458035 ms for aligned serial and 0.508580 ms for bounded overlap. Every +overlap pair is slower. Reject this route; no liveness assumption is added to +serving (`results/qpn2-bounded-overlap-{forced,real}.json`). + +An exact FP16 operand-lookup candidate precomputes the original decode for +every scale/code-pair combination and retains the original S16 column/S8 +gated splits. The first S8-only host gate correctly rejects the frozen +column shapes before measurement. After correcting that dispatch, all +sixteen projection outputs match, but the steady four-layer screen slows +from 0.377795 to 0.787712 ms in every pair. Reject the lookup path +(`results/qpn2-lookup64v2-steady.json`); changing decode representation +alone does not imply lower cost. + +A warp-specialized QPN2 double buffer adds eight loader warps for sixteen +compute warps, retaining both original accumulation chains inside each compute +warp. All sixteen real-weight outputs match, but the sustained seven-pair +working set regresses from 0.374917 to 0.480748 ms, slower in every pair. +Reject it (`results/qpn2-staged40-steady.json`). A capture-time equal-cache +preference screen uses identical serial publication kernels in both arms and +restores the prior context preference after capture. Its five-cycle four-rank +output gate passes, but medians 0.454779/0.456745 ms and mixed-sign differences +show no gain. Reject the screen; it does not prove the preference survives +graph instantiation (`results/qpn2-context-cache-real.json`). + +The generated FP32-residual Gemma kernel exchanges blocked layouts solely for +its residual store. Scalar and vector inline stores eliminate that exchange: +static PTX barriers fall from ten to two and dynamic shared scratch from +8192 to 32 bytes, while the fifteen FMA contractions remain. All 128 varied- +scale normalized/residual outputs match. Nevertheless, the 128-call working +set regresses from 0.526305 to 0.981023 ms with scalar stores and from +0.525332 to 0.607007 ms with vector stores. Both are rejected. An initial +inline-assembly pointer/type compilation failure is retained separately; +these timings come from the corrected kernels (`results/gemma-direct-residual- +{v2,vector}.json`). + +The completed cap64 fixed-prefix tapes are retained in lossless archives to +make space for the next paired audit. Every file was hashed, then read back +and verified through decompression before removing its raw duplicate: 144 +files per arm, approximately 16.19 GB reduced to 2.42 GB per archive. +`archives/v4-qpn2-cap64-audit-control.tar.zst` has SHA256 +`4bcd3f5f764e3914b2f0245fa76059802220a00287f964c9ef01741d4c6dcab3`; +the candidate archive has SHA256 +`cda7db98115c2fb5f6a4f8a4c46fdc6f3926cce04808350298e1dcd69dd0409a`. +Their adjacent manifests retain per-file checksums and original sizes. + +### Sparse rerank selection screen + +The FP32 rerank currently scatters 64 candidates into a 62080-token local +vocabulary before calling dense top-k. The private compact gather reproduces +[PyTorch v2.10.0 multiblock collection](https://github.com/pytorch/pytorch/blob/v2.10.0/aten/src/ATen/native/cuda/TensorTopK.cu): +keys above the cutoff are collected in vocabulary order, followed by cutoff +ties in that order. The unchanged native PyTorch key/value sorter then +preserves its final tie permutation. Implicit -Inf background entries and +canonical NaN radix keys are handled explicitly. No float dot, rounding, +probability, sampling, or acceptance arithmetic is changed. + +`benchmarks/kernels/benchmark_sm70_sparse_dense_topk.py` covers seven/eight +rows, top-k 16/20/21 and nine value families, including ties, signed zeros, +NaN/Inf and fewer finite values than k. All 54 cases match values and IDs +bytewise at the actual 62080-token width; native memcheck reports zero +errors. The seven-pair graph screen reduces this primitive from approximately +0.073213 to 0.010491 ms. The earlier 62464-width synthetic screen is labelled as +such. Its initially incorrect model eligibility guard did not hit the route +and provides no model evidence. + +`benchmarks/kernels/build_sm70_native_sort_candidate.py` reproduces the +private wrapper against the frozen PyTorch 2.10.0 native sorter and records +Torch/header/source/DSO provenance. Its rebuilt DSO SHA256 is +`1ada9a86172d9a524cb2828b32fbbd35a917130abf8a1740b50ac3be4a30c2fc`; +all 54 value/ID cases pass again (`results/sparse-dense-order-versioned-gate.json`). +Neither utility installs a serving route. + +The first effective four-rank shadow checks 120 actual target calls plus +eight draft warmups bytewise. Draft replay diagnostic copies subsequently +contain out-of-range IDs and invalid data, so those failed diagnostic runs +are excluded rather than attributed to the operator. Reading the model's +persistent candidate buffers after actual replay avoids the invalid diagnostic +copies. The completed v6 check compares 128 actual eager calls and 64 actual +draft replay inputs across all four ranks, with identical FP32 values and IDs +for both independently recomputed selectors. This does not prove which graph +allocation behavior invalidated the earlier extra buffers. The first uninstrumented +complete-round pair and a separate five-warmup pair are complete; see `results/sparse-dense-order-62080-gate.json`, +`results/sparse-dense-order-shadow-v2.json`, +`results/sparse-dense-order-shadow-v6.json`, and the excluded shadow v3--v5 logs. + +A separate filtered top-k port uses the installed FlashInfer source matching +[official revision 064d9aa](https://github.com/flashinfer-ai/flashinfer/blob/064d9aa268fe8d2f4d7c9f3c5ca83ecb02fb2c9c/include/flashinfer/topk.cuh). +Shrinking its index buffers from 128 to 64 KiB allows SM70 compilation; its +15760-byte static shared scratch also fits. A negative-NaN mismatch is fixed +by canonicalizing half NaN radix keys. All twenty checked cases then match +the frozen PyTorch selection after restoring collection order. The isolated +screen measures 0.041341/0.040759 ms, too little gain to justify adding this +native path. It remains disabled (`results/flashinfer-filtered-sm70-nan-gate.json`). +The separate GPU 5 checks are numerical only and have no timing claims. + +The complete-round target below 15 ms, full distribution/state comparison for +the final combination, repeated-startup acceptance gates, and long-context +validation remain open. No new serving default or merge is claimed. + +### Whole-round resource audit and route correction, 2026-09-09 + +The five-warmup sparse-selector pair measures 17.017221/16.761147 ms for +release1k and 16.624768/16.373416 ms for MBPP28. All five measured token IDs, +EOS and acceptance counts match in both fixtures. MBPP28 retains a +16.938744-ms outlier. The earlier one-warmup release regression is retained, +not replaced or trimmed. The full fixed-prefix pair now passes: 144 records +per arm, all captured intermediates/state/native logits byte-equal, TV zero, +no support or top-1 changes (`results/v4-sparse-dense-order-audit-comparison.json`). + +The new trace confirms the compact collector in both actual target and draft +execution. QPN2 remains 7.342 ms of service. Draft attention has only eight +CTAs on an eighty-SM GPU, with 97920 bytes of shared memory per CTA. The +head-regrouping hook was attached to the wrong native Python module; withdraw +its previous model speed attribution. The new explicit benchmark installer +patches the interface's actual module and preserves other shapes/eager calls. +See [the complete resource audit](sm70_quasar_dflash2_resource_audit_20260909.md) +for endpoint statistics, critical-rank closure, launch resources, the unchanged +sampling quality guard and trace provenance limits. No new default is enabled. + +The corrected attention module binding now has actual graph proof: 640 partial +launches over forty rank-rounds use a 3 × 80 grid, 256 threads, 234 registers +and 30464-byte shared memory. Grouped-attention service is 0.909762 ms versus +0.951279 ms in the preceding diagnostic trace; no whole-round gain follows +from that comparison. The profiled release output remains canonical. A +separate unprofiled pair subsequently completes as recorded below. +The owned trace client was recovered from +a job-name mismatch without reloading the model, and the final library +manifest uses process ancestry to include Nsight's separate child group. +The obsolete waiting client was then stopped; the wrapper exit 143 remains +recorded rather than relabelled successful. + +### Actual attention admission and chunked-publication screen + +The actual grouped-attention pair measures release1k 16.797233/16.637915 ms +and MBPP28 16.416132/16.248439 ms after five warmups, with five measured +requests per arm. All measured natural tokens and acceptance remain canonical. +The actual-route four-rank fixed-prefix pair now passes: 144 records per arm, +zero captured intermediate differences, byte-equal logits, TV zero and no +support or top-1 changes. This closes the inactive-route evidence gap, not the +final multi-seed, repeated-startup or long-context gate. + +The new two-chunk QPN2 publication implementation passes nine changing-input +four-rank cycles, including rank skew and allocation canaries. Its real-weight +working-set median regresses from 0.456499 ms to 0.511037 ms serialized and +0.557527 ms overlapped. Keep it disabled and do not extend to four chunks. +See the [resource audit](sm70_quasar_dflash2_resource_audit_20260909.md) for +the independent channel protocol, source/DSO provenance and evidence limits. diff --git a/docs/design/sm70_quasar_dflash2_resource_audit_20260909.md b/docs/design/sm70_quasar_dflash2_resource_audit_20260909.md new file mode 100644 index 0000000000..230daa6139 --- /dev/null +++ b/docs/design/sm70_quasar_dflash2_resource_audit_20260909.md @@ -0,0 +1,671 @@ +# QUASAR + DFlash2 complete-round resource audit, 2026-09-09 + +The 15-ms goal is not met. The latest same-startup unprofiled GDN BV2 isolation +measures 16.280 ms for release1k and 15.873 ms for MBPP28. Both use +rear GPUs 4–7, +TP4/B1/q8, E4M3 target KV, FP32 logits/state, the frozen model and natural +EOS. One startup pair with five warmups and five measured requests per fixture +does not complete the final performance or quality gates. + +## Unprofiled endpoint evidence + +All values below summarize the five measured requests; no profiler or tensor +dump is active. Complete-round cost is engine decode time divided by draft +round count, and includes target, sampling, state and draft. + +| Fixture / arm | Complete-round mean | Median | p90 | p99 | TTFT median | Pure decode median | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | +| release1k / control | 17.023 ms | 17.017 ms | 17.034 ms | 17.039 ms | 353.351 ms | 175.00 token/s | +| release1k / candidate | 16.757 ms | 16.761 ms | 16.774 ms | 16.780 ms | 353.295 ms | 177.67 token/s | +| mbpp28 / control | 16.627 ms | 16.625 ms | 16.634 ms | 16.638 ms | 126.470 ms | 292.89 token/s | +| mbpp28 / candidate | 16.491 ms | 16.373 ms | 16.741 ms | 16.919 ms | 116.408 ms | 297.39 token/s | + +The p90/p99 columns above describe request-average round cost, not individual +GPU rounds. Five requests are insufficient to establish tail reliability. +Candidate MBPP28 retains a 16.939-ms outlier. All post-request telemetry samples +show 1530/877-MHz SM/memory clocks; those samples do not exclude transient +events during requests. The earlier one-warmup pair is retained separately: +release1k 17.002/17.081 ms, MBPP28 16.560/16.381 ms. No requests were discarded. + +Measured tokens, natural EOS and acceptance match bytewise across arms: + +| Fixture | Output tokens | Rounds | Accepted drafts | Accepted drafts / round | Emitted tokens / round | +| --- | ---: | ---: | ---: | ---: | ---: | +| release1k | 272 | 91 | 181 | 1.989011 | 2.989011 | +| MBPP28 | 634 | 130 | 504 | 3.876923 | 4.876923 | + +See `results/v4-sparse-dense-order-warm5-ab.json` and its four hashed input +reports. The primitive passes 54 boundary cases, native memcheck and 192 real +four-rank input comparisons. The complete four-rank fixed-prefix pair now passes: 144 records per arm, +no captured intermediate differences, all native logits byte-equal, TV zero +and no top-p support or top-1 changes. This includes the captured layer 0/1 +conv/SSM state and metadata; it is not an all-layer operator oracle. See +`results/v4-sparse-dense-order-audit-comparison.json`. + +## Whole-round trace closure + +The node trace contains twelve complete four-rank rounds; discard the edge +rounds and analyze rounds 9–18. Select the longest worker interval in each +round, then close that same rank with GPU event union plus uncovered time. + +| Same critical rank / round | Mean | p50 | p90 | p99 | +| --- | ---: | ---: | ---: | ---: | +| Worker round interval | 18.651 ms | 18.471 ms | 18.809 ms | 19.922 ms | +| GPU event union | 16.774 ms | 16.639 ms | 16.934 ms | 17.949 ms | +| Time without GPU events | 1.876 ms | 1.859 ms | 2.038 ms | 2.114 ms | + +GPU activity covers 89.94% of this instrumented interval. This measures the +presence of GPU work, not achieved SM occupancy, issue rate, Tensor Core use +or HBM efficiency. NCU counters are unavailable. Profiled gaps and collective +waiting are not directly recoverable latency. These values do not replace the +16.761-ms unprofiled endpoint median. + +| Phase | Mean GPU service per rank | Mean GPU envelope per rank | Kernel calls / rank / round | +| --- | ---: | ---: | ---: | +| target_graph | 12.295 ms | 12.762 ms | 952 | +| target_head_sampling | 0.543 ms | 0.908 ms | 27 | +| request_state | 0.013 ms | 0.115 ms | 3 | +| draft_propose | 3.556 ms | 3.960 ms | 193 | +| input_metadata | 0.057 ms | 0.107 ms | 14 | +| context_and_output | 0.230 ms | 5.664 ms | 13 | + +Context/output work is interleaved with sampling and draft; its envelope +spans those phases. Do not sum phase envelopes or compare independent rank +maxima as a single critical path. Native memcpy/memset events are included in +service, while the call count column counts kernels. + +## Large costs and weak launch parallelism + +These are observed launch resources. Grid counts constrain work distribution +but do not establish achieved occupancy or a particular stall reason. + +| Work | GPU service / rank / round | Observed launch | Implication / next bounded step | +| --- | ---: | --- | --- | +| QPN2 gate/up | 2.963 ms | 136 CTAs, 512 threads, 64 registers, 16 KiB shared | Largest individual family; retain HMMA chains and test only loading/layout ideas supported by real-weight working sets. | +| Published QPN2 row projections | 2.889 ms | 160 CTAs; 256/512 threads; 48 registers | Preserve rank reduction order and epoch lifetime. Earlier bounded overlap and cooperative consumers were slower. | +| Other QPN2 columns | 1.489 ms | 112/129 CTAs, 512 threads, 64 registers | Limited grid alongside finite register residency; cap64 is active. Tile/chain changes need separate error gates. | +| Draft dense projections/reductions | 1.834 ms | Main WMMA kernel uses 32-thread CTAs; common grids have 320 CTAs | Small-row GEMM work is spread over few warps per SM. Earlier arithmetic candidate changed acceptance and remains off. | +| Target grouped attention | 0.951 ms | Original partial: 80 CTAs, 512 threads, 128 registers, 56832-byte shared | Correct the inactive experiment binding, then verify 240-CTA/256-thread candidate in the actual replay. | +| Target normalization/residual | 0.835 ms | Dominant fused Gemma kernel has 8 CTAs, 256 threads | Small grid and many dependent launches. Prior direct residual stores were slower; no new fusion benefit assumed. | +| Draft attention | 0.490 ms | 8 CTAs, 512 threads, 97920-byte shared | At most 8 of 80 SMs receive a CTA per invocation. Investigate output-work partitioning without changing QK/softmax order; changing KV splits is arithmetic. | +| GDN convolution | 0.283 ms | 10 CTAs, 128 threads | Small work per invocation; fusion must retain each token state and rollback boundaries. | +| Target KV write | 0.231 ms | 8 CTAs, 32 threads | Compare direct producer layout only with exact cache/slot checks. | + +QPN2 totals 7.342 ms of service. This remains the main performance target; +resource-thin attention and small kernels are complementary opportunities, +not a claim that their service time can all be removed. + +## Host gaps and quality-sensitive decisions + +The same critical-rank gap closure assigns 0.468 ms per round to gaps between +target-graph nodes (the largest individual such gap is only 0.001344 ms), +0.422 ms inside draft, 0.388 ms inside target sampling, and 0.208 ms between +state handling and draft. Numerous short node gaps cannot be treated as one +large idle segment. The largest sampling gap lies between the probe memcpy +and sparse rejection: usually about 0.27–0.31 ms, with a 0.461-ms sample. + +The source copies the 21-candidate probe to CPU and checks top-20 cutoff ties, +ties crossing the nucleus and FP32 CDF proximity before selecting compact or +full-vocabulary rejection. Keep this guard and fallback. Eliminating its wait +requires preserving the decision and dependent RNG/acceptance state; simply +removing the CPU branch is not an admissible optimization. + +## Route correction and evidence limits + +The compact FP32 collector appears once per target and once per draft round +on every rank (80 calls across the forty analyzed rank-rounds). Its final +native sorter is also present. That proves active target/draft dispatch. + +The head-regrouping hook instead patched top-level `flash_attn_v100_cuda`. +The model interface calls `flash_attn_v100.flash_attn_v100_cuda`. Both resolve +to DSO SHA256 `a751fed902279b0de23537c4aad2dc4fee360146d7fce7ef0c4f255a77f48b02`, +but CPU identity checks prove separate module objects and function bindings. +No regrouped capture marker or 240-CTA launch is present. Withdraw the earlier +head-regrouping speed attribution and its model-level candidate quality claim; +keep the raw measurements and isolated operator gates. + +`benchmarks/kernels/sm70_grouped_attention_candidate_route.py` now resolves the +actual native object through the interface. It installs only when explicitly +called by an experiment and delegates non-q8/eager calls to the original. +The corrected route is now proven in ten steady rounds across four ranks: +640 grouped partial kernels use 240 CTAs and 256 threads. Their measured +launch footprint is 234 registers/thread and 30464-byte shared memory per CTA. +Thus more CTAs do not by themselves establish better achieved occupancy; +register pressure remains a constraint. Grouped attention service changes +from 0.951279 to 0.909762 ms in the two diagnostic traces. This is not an +unprofiled full-round improvement. The canonical release token IDs and +acceptance remain unchanged in the profiled request. The separate five-warmup +unprofiled pair and four-rank fixed-prefix comparison are now complete. + +Raw evidence: `profile/v4-sparse-dense-order-nodes/tp4.{nsys-rep,sqlite}`, +`results/v4-sparse-dense-order-nodes-trace.json`, +`results/v4-sparse-dense-order-resource-trace.json`, and +`results/attention-headsplit-binding-identity.json`. The trace capture and export +completed, but the wrapper then failed its runtime-map ownership-name assertion. +Therefore this trace lacks its own final map manifest; separate unprofiled +four-worker DSO manifests are retained. The corrected-route trace has a separate four-worker/360-library manifest. +Its original client waited on a mismatched ownership-name suffix, so a +corrected client completed the capture against the existing owned service. +Map collection was expanded to verified descendants because Nsight gives +the application a separate process group. After saving the manifest, the +obsolete waiting client was stopped and the wrapper cleaned its service, +exiting 143. The client/capture completed; the wrapper did not exit cleanly. +See `results/attention-bound-profile-harness-recovery.json`, +`results/attention-bound-route-hit.json` and +`results/nsys-v4-attention-bound-nodes-runtime-libraries.json`. + +Three independent startup pairs, acceptance non-inferiority and model +long-context gates remain open. No +15-ms result, default promotion, merge or 256K performance claim follows. + +## Actual attention quality and unprofiled follow-up + +The completed pair keeps five warmups and five measured requests per fixture. +Request-average complete-round medians change 16.797233 -> 16.637915 ms for +release1k and 16.416132 -> 16.248439 ms for MBPP28. Every measured token hash, +natural EOS, accepted-draft count and emitted-token count remains canonical. +This is one startup pair, not the final three-pair gate. The input reports and +their hashes are in `results/v4-attention-bound-warm5-ab.json`. + +Both actual-route fixed-prefix jobs exit zero and collect 144 records each. +The comparison finds no captured intermediate differences, all native logits +byte-equal, TV zero, and no support or top-1 changes. The recorded conv/SSM +states and metadata cover layers 0/1, not every layer's operator internals. +Both arms retain mapped-library manifests; candidate capture logs prove the +actual module binding. See `results/v4-attention-bound-audit-comparison.json`. +Completed raw tapes are retired only after lossless archive reconstruction +verifies each of the 144 per-file SHA256 values. + +## Two-chunk QPN2 publication screen: rejected + +The new private builder partitions the 5120 output columns into two 2560-column +chunks, preserving each output's original dot product and rank reduction. +Each chunk has separate two-epoch storage. Its consumer waits on the local +producer's completion event; the main stream joins both consumers before +dependent work. This does not reuse the rejected pre-producer polling scheme. + +Four ranks, sixteen real consecutive-layer projection weights, nine changing +synthetic-input cycles, rank start delays and mixed ordinary-push calls pass +bytewise output comparisons with intact allocation canaries. Seven alternating +working-set measurements give 0.456499 ms for frozen publication, 0.511037 ms +for serial chunks and 0.557527 ms for overlapped chunks. All paired differences +are regressions. Therefore neither candidate gets a model run or a four-chunk +extension; there is no end-to-end speed claim and no default change. + +`benchmarks/kernels/build_sm70_qpn2_chunked_candidate.py` and +`benchmarks/kernels/benchmark_sm70_qpn2_chunked.py` reproduce the screen. +The native DSO SHA256 is +`745a2bf88bef7c5bd5284f1f45ebc36575f2cb1a320a5a3a04e6db817e224688`. +The original publisher and communicator remain independently frozen and are +hashed in `results/qpn2-two-chunks-real.json`. Kernel-level race and memory +sanitizer admission is not claimed for this rejected route. + +NCU 2022.4.1 exists at `/usr/bin/ncu`, but the driver reports +`RmProfilingAdminOnly: 1` and this task's noninteractive sudo attempt requires +a password. Other campaigns' counters do not establish access for this task; +its occupancy and memory-throughput counter gap remains explicit. + +## Context computation behind the target probe: model screen + +The explicit benchmark installer defers eligible q8 context preparation until +after the target's 21-candidate probe is copied to preallocated pinned memory. +It records a copy event, submits the original context graph on the original +stream, waits only for the copy, and calls the unchanged CPU cutoff predicate. +Full-vocabulary/structured-output paths flush any pending preparation before +the caller updates request state or proposes drafts. KV stores retain their +acceptance-dependent ordering. There is no additional CUDA compute stream. + +The CPU dependency/fallback gate passes 256 predicate inputs, including ties, +and checks missing-guard fallback, unsupported probe layout, prefill and error +cleanup. The actual natural-sampling shadow then checks at least 1280 calls +on each rank: probe bytes, cutoff decisions, staged hidden states and projected +context K/V match. All release/MBPP measured output hashes and acceptance +counts remain canonical. Shadow executes additional reference work and is +not performance evidence. See `results/context-probe-cpu-dispatch.json` and +`results/context-probe-actual-shadow.json`. + +The first uninstrumented five-warmup pair measures release1k +17.038700 -> 16.554804 ms and MBPP28 16.767940 -> 16.217768 ms. Its control is +slower than the preceding actual-attention pair; do not attribute that entire +difference to the pipeline. The reversed candidate completed, but its control was interrupted by a host +reboot and produced no endpoint result. It is not a paired comparison. A fresh +post-reboot pair measures release1k 16.779740 -> 16.591267 ms and MBPP28 +16.466058 -> 16.106911 ms, with five warmups and five measured requests per +fixture. Every measured token hash and acceptance count remains canonical. The candidate remains experimental; +this does not clear distribution/state, final performance or long-context gates. + +## Draft cuBLAS layout screen: numerical rejection of broad changes + +All 400 retained four-rank raw projection controls reproduce bytewise with the +original layout. A column-major weight view changes 300/400 outputs and expands +FP64 reference error in 298 cases. Padding queries to sixteen rows changes +200/400 outputs and expands reference error in 102 cases. Combining both changes +has 300 differences and 298 expanded-error cases. The aggregate working-set +medians 1.458115/1.283830/1.409249/1.355162 ms do not admit these broad routes. + +Only `o_proj` with column-major weights and `down_proj` with padded row-major +weights retain byte parity in their respective 100-case subsets. The separate +screen includes the required input copy and leaves QKV and gate/up unchanged. +Its complete twenty-projection working set regresses from 1.451684 ms to +1.485681 ms when combining the byte-equal subsets. Either subset alone also +regresses. These exact-layout routes are rejected before model testing. `benchmarks/kernels/benchmark_sm70_draft_f16_layout.py` reproduces the +full numerical screen; `results/draft-f16-layout-real.json` retains every case, +FP64 metric, original snapshot hash and aggregate timing. + +## Post-reboot context trace and closure + +The machine rebooted at 2026-09-09 02:34:57 UTC. The old lease and unfinished +reverse-control processes were gone. The user-authorized rear-GPU default +service was stopped, and the task lease was restarted with explicit physical +GPU order 4,5,6,7. Frozen candidate/control DSOs were rehashed. The interrupted +job is retained as `zz240-v4-context-probe-warm5-reverse-control.interrupted.json`; +`host-recovery-20260909.json` records ownership recovery. + +The new node trace exits cleanly and retains its own four-worker library +manifest. Its ten steady rounds have critical-rank interval mean 19.553882 ms, +p50 18.468254 ms, GPU union mean 17.426705 ms and uncovered mean 2.127177 ms. +Two roughly 24-ms rounds remain in those aggregates: one has 7.417044 ms of +uncovered time and the next has 22.366826 ms of GPU activity, including waits. +Their cause is not assigned to a source change. The trace is diagnostic and +cannot replace the separate unprofiled results above. + +GPU correlation identifies exactly one six-kernel context graph after the +672-byte target probe on each of forty rank-rounds. Its mean service/envelope +is 0.076562/0.082549 ms; mean overlap with the remaining host sampling span is +0.080438 ms. That span includes the unchanged CPU guard and rejection launch +handling, so it is not a pure predicate timer. The roughly 11-ms event wait +includes queued target work, not just probe transfer. QPN2 still totals +7.351498 ms per rank-round. No achieved-occupancy or HBM counter claim is made. + +Evidence: `results/v4-context-probe-nodes-trace.json`, +`results/v4-context-probe-nodes-resource-trace.json`, +`results/v4-context-probe-overlap-proof.json`, and +`results/nsys-v4-context-probe-nodes-runtime-libraries.json`. + +## Strict draft column GEMM: arithmetic gate passes, model trajectory held + +The bounded follow-up disables reduced-precision FP16 GEMM reduction only +while selecting each candidate kernel, then restores the process setting. +The original layout still reproduces all 400 retained controls. Strict column +weights change 300 outputs but expand none of the registered FP64 max, p99 or +relative-L2 errors. Strict padded-row weights still expand three cases and +are rejected. Twenty-projection medians are 1.469460 ms for the original and +1.307750 ms for strict column weights. This is a local arithmetic screen. +The switch follows the documented PyTorch 2.10 reduction control; its effect +here is measured, not a diagnosis inferred solely from the documentation. + +The explicit `sm70_draft_column_candidate_route.py` installer selects only the +twenty captured q8 query projections and retains original prefill/context +calls. In the first natural model run, release1k changes 272 -> 248 emitted +tokens, with first token difference at zero-based offset 123; MBPP28 changes +634 -> 357, first differing at offset 194. Accepted drafts per round are +1.989011 -> 2.024390 and 3.876923 -> 4.100000, respectively. These changed +trajectories do not establish acceptance non-inferiority or preserved quality. +The apparent 16.369764/15.902543-ms medians are not admitted performance gains. +The arithmetic route remains closed pending causal distribution/acceptance +and broader quality evidence. No score is used to excuse these differences. + +See `results/draft-f16-layout-strict-real.json`, +`results/draft-f16-layout-gated-real.json`, and +`results/draft-column-model-screen.json`. The operator benchmark restores the +original reduction property in a `finally` block; no global serving default +is changed. Official reference: +[PyTorch 2.10 numerical accuracy](https://docs.pytorch.org/docs/2.10/notes/numerical_accuracy.html). + +## Native FlashInfer fragment draft prototype: rejected for speed + +A private B1/H8/q8/D128 FP16 paged prototype reuses the project's native +FlashInfer Volta WMMA fragments with the frozen Flash-V100 K176 and online +softmax schedule. It reduces the query tile to sixteen rows and uses one or +four independent output-column partitions, retaining every output's QK, FP16 +probability and FP32 PV order. Each version passes 28 changing-length, +permuted-page, tail, graph-replay and output-canary comparisons bytewise +against frozen native attention. This screen uses 16-token pages; it is not +an actual 1648-token model-page or long-context admission. + +The first variant has 128 registers and a four-byte spill. A second keeps PV +accumulators in registers across K tiles and uses the actual one-resident-CTA +launch bound. Its four-part kernel has 147 registers, 73088 bytes shared and +zero spill stores/loads. It also passes all 28 comparisons, but still has no +stable speed gain. At 4096 keys, baseline/one-part/four-part medians are +0.382853/0.463544/0.420690 ms. Neither version is installed in a model, and no +sanitizer or model-quality admission is claimed for these rejected routes. + +Retained source, flags and library hashes are in +`candidates/draft-fi-q8-p{1,4}{,-register}/manifest.json`; direct gate reports +are `results/draft-fi-q8-gate.json` and +`results/draft-fi-q8-register-gate.json`. Register four-part DSO SHA256: +`76766685df4a15c1ecc60f8dff6d890dd2578b58076d7b88b0070c2b8f9cdd6e`. +This reuse does not claim that an unmodified upstream FlashInfer kernel was +run. The native project route remains a valid porting base; GPU support-list +membership is not used to reject further implementations. + +## QPN2 gate/up CTA redistribution: rejected + +Another bounded screen replaces each 136-CTA/512-thread fused gate/up with +272 256-thread GEMM CTAs followed by the original native SiLU. Split-K eight, +two accumulator chains, reduction order and FP16 activation boundaries remain +unchanged. Four real consecutive-layer weights and nine changing-input cycles +pass all 36 bytewise comparisons, but seven alternating working-set medians +regress 0.159058 -> 0.171489 ms. No model run or default change follows. +See `results/qpn2-unfused-gated-real.json`. + +## GDN value-tile candidate + +The current TP4 trace launches 192 one-warp GDN CTAs with 80 registers/thread, +covering twelve value heads with BV=8. A separate exact-shape screen adapts the +value-tiling mechanism to TP4; it does not transfer another TP size's timings. +BV 8/4/2/1 each preserve all output and FP32 state bits for eight acceptance +selectors and two changing graph replays per selector, including strided QKV, +strided state pools, padding canaries and untouched retired slots. +Sixteen distinct state working sets measure 0.381416/0.312884/0.305196/0.308504 ms. +BV2 retains the original K reduction shape and one-warp schedule while exposing +768 CTAs. Focused BV8/BV2 memcheck and racecheck both exit zero. + +The explicit `sm70_gdn_value_tile_candidate_route.py` installer is limited to +captured TP4/B1/q8 with twelve value heads and FP32 state. Other shapes/dtypes +retain the original schedule. The live same-input/state shadow now passes all 48 GDN +layers on every rank, with at least 2230 calls per layer. Output/state bits, +finite-value checks and active-slot validity all match. Original outputs drive +shadow generation, and both natural trajectories remain canonical. This is +diagnostic evidence, not a timing result. Evidence is `results/tp4-gdn-bv-screen.json` and +`results/tp4-gdn-bv2-{memcheck,racecheck}.json`. Reusing FP32 Q/K normalization +across value tiles is a separate unadmitted screen and is not combined yet. + +The first separate-startup unprofiled GDN pair retains five warmups and five +measurements per fixture. Release1k changes 17.085427 -> 16.333210 ms and MBPP28 +16.551539 -> 15.969318 ms, with identical canonical token hashes, natural EOS, +accepted-draft counts and emitted-token counts. Its control is slower than the +prior context pair; the entire gap cannot yet be credited to BV2. The reversed +pair measures control/candidate 16.585019/16.552420 ms for release1k and +16.326226/16.324162 ms for MBPP28. All trajectories remain canonical, but +these 0.032598/0.002064-ms differences do not establish a stable whole-round +gain. Six CPU mocked checks also confirm q8 dispatch, other dtype, +TP size, query width, eager and head-count fallbacks, and restored scope. +See `results/gdn-value-tile-live-shadow-admission.json`, +`results/gdn-value-tile-first-pair.json`, and +`results/gdn-value-tile-cpu-dispatch.json`. + +Actual CUDA graph node tracing independently confirms the candidate route: +all recurrent launches use grid `(1,64,12)`, one warp and 55 registers, versus +the frozen `(1,16,12)` and 80 registers. The forty steady rank-rounds each +contain all 48 recurrent calls. Their mean summed service falls from +0.955449 to 0.682117 ms between the retained context and BV2 traces; QPN2 +service remains about 7.35 ms. These are profiled observations, not endpoint +speed evidence or measured occupancy. The BV2 trace retains host/rank-wait +outliers: critical-round p50 is 17.876601 ms and mean 18.513127 ms, with +2.048493 ms mean uncovered GPU time. See +`results/v4-gdn-value-tile-nodes-{trace,resource-trace}.json` and the four-worker +runtime-map manifest. Separate-startup variability still needs isolation +before the route is promoted. + +The first same-startup diagnostic captures adjacent BV2/BV8 kernels on the +same buffers, disabling one state mutation before the first replay. Dependency +edges identify each pair; both node-enable states are read back after changes. +Twenty-four changing-input/selector switches first pass complete-state, output, +padding and retired-slot checks. The owned model client then switches only +between requests, with five warmups per arm and five interleaved measurements. +Release1k changes 16.493886 -> 16.279546 ms and MBPP28 16.088556 -> 15.872776 ms, +with canonical token IDs, accepted drafts and natural EOS in every request. +This isolates an approximately 0.21-ms whole-round gain while retaining startup, +prefill, allocations and GEMM choices. One such startup is not final admission. +See `results/gdn-value-tile-within-start-1-summary.json` and +`results/gdn-pair-graph-gate.json`. The diagnostic uses CUDA's documented +[individual-node enable behavior](https://docs.nvidia.com/cuda/cuda-programming-guide/04-special-topics/cuda-graphs.html#individual-node-enable); +disabled nodes retain dependencies and behave as empty nodes. + +## Q/K reuse exposes a recursive FP32 rounding boundary + +The first private normalization-reuse screen leaves the immediate FP16 output +unchanged but changes 62767 FP32 state elements on its first candidate case. +Timing is skipped. A diagnostic tap is first checked against the frozen +recurrence: both its output and complete state pool remain byte-equal. +On the exact failing input, the standalone and in-recurrence normalized Q/K +also match bytewise. Thus this candidate's first state difference is downstream +of Q/K normalization, not in those normalized operands. + +PTX identifies a changed contraction boundary in the local four-element +`h dot k` reduction. The frozen kernel first rounds the product at local K1, +then contracts K0, K2 and K3 through FMA. Loading materialized normalized K +allows the compiler to choose K0 as the initially rounded product. The +subsequent warp reduction has the same shape, but these programs need not +produce identical FP32 state. The output's FP16 rounding initially conceals it. + +The corrected private candidate makes the K1 product's rounding explicit with +`mul.rn.f32` and retains the remaining reduction. All 48 checked cases across +BV8, BV2 and corrected reuse now have zero output and complete-state bit +differences. Sixteen-state working-set medians are 0.379136/0.303040/0.291456 ms; +the extra gain over BV2 is only 0.011584 ms and is not a model gain. Further +kernel safety, real-model shadow and whole-round validation remain open. +This diagnosis concerns the new reuse experiment; it does not resolve the +previous unrelated 4.33% repeat-start distribution discrepancy. + +Evidence: `results/tp4-gdn-precomputed-qk-screen.json`, +`results/gdn-norm-tap-failure-comparison.json`, +`results/gdn-norm-tap-failure-operands.pt`, and +`results/tp4-gdn-precomputed-qk-fmafix-screen.json`. The corrected derived-source +SHA256 is `ce200148aa03ab9da6692d69e07f5d48f8e58b1f08763aaa88972ab3b0e7acee`. +The artifact root retains original/corrected Triton source, TTGIR and PTX. + +## Cooperative MLP publication: native gates pass + +A new private candidate executes gate/up and the dependent down projection +inside a 160-CTA cooperative kernel, with one grid barrier between them. It +retains the original dot-product chains, FP16 SiLU boundaries and packet +publisher. The original consumer remains separate and starts after the local +producer finishes; this does not revive the rejected pre-producer polling +scheme. The launcher checks that all 160 CTAs can be resident before launch. +Compilation reports 64 registers, 32768-byte shared storage and no local stack +or spills; runtime confirms two resident CTAs per SM. Four ranks, four real +consecutive-layer weight sets, nine changing-input cycles, skewed ranks and an +additional ordinary push all preserve gate, down and reduced output bits and +buffer canaries. Seven paired working-set trials measure 0.457871 -> 0.451072 +ms, about 1.5% locally; both arms drift during the trials, so raw samples are +retained. Four-rank memcheck and racecheck exit zero, with zero reported errors +or hazards. The first private model shadow executes no candidate calls: its +outer Python shape guard is specialized away during dynamic model compilation. +Those results are explicitly excluded. Moving eligibility into the opaque +runtime custom op and asserting all 64 captured prefixes fixes the route. +All four ranks then pass at least 1338 live comparisons in each of 64 layers, +with zero gate/final-output bit differences or nonfinite values. Original +outputs drive generation; both natural fixtures stay canonical. + +The separate-startup five-warmup/five-measurement pair measures release1k +16.445409 -> 16.453666 ms and MBPP28 16.129142 -> 15.966142 ms. All token IDs, +acceptance counts and natural EOS match, but the gain is workload-dependent +and only one pair is available. This is not a promoted combination or a +sub-15-ms result. `results/coop-mlp-live-shadow-admission.json` and the +`v4-coop-mlp-warm5-{control,candidate}-speed-*` reports retain the evidence. + +Reports are `results/qpn2-coop-mlp-{real,memcheck,racecheck}.json` with sanitizer +logs and manifests alongside them. The private library SHA256 is +`22a91bd9f9e8aa0cc1324b0482c0fc4d6fc695ef7b553935801a935e47194c31`; +`candidates/qpn2-coop-mlp/cooperative-manifest.json` retains the source and flags. + +The versioned `build_sm70_qpn2_cooperative_mlp.py` accepts an explicit private +output directory and reproduces the exact validated CUDA source SHA256 +`21f5a7448cb71ec3b847f21e41068b64f7eae44db0e5e2ec365b96a2fbd99b65`. +Its companion benchmark keeps real consecutive-layer weights, changing inputs, +rank skew and mixed-protocol epochs, and adds six rejected non-q8 row counts. +Its four-rank rerun passes both changing-input cycles, all six rejected shapes +on every rank and all output/canary comparisons; the generated source matches +the previously built and sanitized DSO. The rerun is recorded in +`results/qpn2-coop-mlp-versioned-gate.json`. +The benchmark is an operator/communication gate, not a model quality score. + +## Independent draft query-row partitions: rejected + +A further native FlashInfer-fragment experiment partitions the eight query +rows among four or eight CTAs per head, retaining K176, per-row reduction and +FP16 probability boundaries. This differs from the earlier output-column +partitions. Both versions match frozen FP16 output in 40 graph/canary cases at +the actual 1648-token page size, including 1647/1648/1649 and 3295/3296/3297 key +lengths. The candidates' FP32 LSE also matches each other; that check is not an +independent FP32-score reference. +Yet both regress: at 1024 keys, control/four/eight-part medians are +0.074189/0.077496/0.080691 ms; at 4096 keys they are +0.322202/0.369172/0.380150 ms. Increasing the grid from eight to 32/64 CTAs does +not itself improve latency. No serving hook or model admission follows. +`results/draft-fi-query-rows-gate.json` and the two candidate manifests retain +the frozen native DSO hash, generated source, raw samples and compile resources. + +The upstream QPN2 source was rechecked at +[`v100-skinny` 5b589c0](https://github.com/dnv2003/v100-skinny/blob/5b589c0dc81223e0ba65bcb3e755874723f8b515/kernels/skinny_kernels.cu) +and the independent +[`ninfer-v100` 8fd0e2e implementation](https://github.com/geoffwatts/ninfer-v100/blob/8fd0e2efdea77bab944991f2394309c07b8baffe/src/ops/linear/nvfp4/nvfp4_volta_qpn_gemm.cuh). +Their prepacked quadpair-on-N layout and independent accumulator mechanism +are already represented in this campaign; their weight/KV contracts and +published timings are not imported as this model's performance evidence. + +## Direct native m8 draft QK/PV: no whole-round admission + +The next native FlashInfer-fragment screen uses Volta's +`mma.sync.aligned.m8n8k4` for the eight actual query rows. It preserves the +original K4 accumulation order, K176 schedule, FP32 softmax and FP16 +probability boundary. A separate QK oracle compares original WMMA and native +FP32 scores before softmax in 45 cases. Extending the same mechanism to PV +adds 45 comparisons with nonzero FP32 initial accumulators. All 90 FP32 +comparisons are byte-equal; independent FP64 references are also retained. +Twenty complete attention cases cover the actual 1648-token pages and changing +graph inputs, with matching frozen FP16 outputs and intact canaries. Candidate +FP32 LSE matches the parent fragment implementation, not an independent +original-native LSE oracle. + +The QK-only implementation reports 180 registers and the combined QK/PV +implementation 138, both with 73088-byte shared memory and zero spills. +At 512/1024/4096 keys, frozen/combined medians are respectively +0.043653/0.042240, 0.074793/0.074117 and 0.322703/0.360151 ms. +The small short-context difference and longer-context regression do not +justify a serving route. No sanitizer or model admission follows. These +results do not establish the memory/issue bottleneck without counters. + +Evidence is `results/draft-{qk,qkpv}-m8-gate.json`. Generated source SHA256s +are `b98ee65791dcafe46e0ccb7529feff7a555b9d5ab02798bb70fc2dd54a8bb6f0` +and `1f6d113deb1144838de6f49aefa4328f9f5a45aeacdcde3793d9a94689d575a8`; +their native DSO SHA256s are +`a15460c1078136d84bd9630055fce7095b45b65f7c0b09ca08dea9e9d178d084` +and `be9844d71c36b1ed6e9309fd8faf8388e07fa4ebea32a5d23c3169d9a9c86f9c`. +Operand mapping follows NVIDIA's +[PTX m8n8k4 fragment documentation](https://docs.nvidia.com/cuda/archive/11.0/parallel-thread-execution/index.html). + +## Positive QPN2 scale decoding: no stable gain + +A separate screen checks every actual scale byte in the four-rank, +four-consecutive-layer working set before removing a redundant sign-bit +construction. All scale bytes are below 128, and exhaustive bit mapping of +those 128 codes matches the original. The two FP16 multiplies, HMMA chains, +SiLU and rank reduction remain unchanged. Nine changing-input cycles, skewed +ranks, dependent gate/down projections, mixed epochs and canaries all pass +bytewise comparison. + +Seven paired working-set trials nevertheless measure 0.458691/0.459366 ms +for control/candidate, with mixed signs in the paired differences. Keep this +specialization off and do not advance it to a model test without new evidence. +The result is `results/qpn2-positive-scales-real.json`. Column/row source +SHA256s are `ef64e021cd88403acd2dfa676653fa293244aa280330338760e91c8b344198ee` +and `ab2adf4c76298186eed97c684c461ed792c8d5c46c945f4be4e225ba465fe5d6`; +DSO SHA256s are `ed02ddf8baac4d537caaf328d181ff9dca6fedaac15a7453410a37beb58aef9d` +and `be15316b1f8471063803a3f87d4a5aee1c63f9955c5a57c17bb1c7dab019b968`. + +## QPN2 layout substitution within the original graph + +The preceding packed-input quality hold first diverges during prefill, before +its q8 layout is active. A new diagnostic therefore retains the original +Python/FX path and changes only executable q8 CUDA graph nodes after capture. +Dependency ancestry and buffer addresses pair a norm with its QPN2 consumer. +The replacement norm keeps its original row-major output and FP32 residual, +and additionally writes private `[320,8,16]` storage. Only the paired projection +receives the packed pointer. Original raw graph nodes and edges remain intact; +the caller synchronizes and switches executable parameters between requests. +This uses CUDA's documented +[kernel-node parameter update interface](https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__GRAPH.html), +not a change to the compilation boundary or prefill implementation. + +Four ranks each pass 216 real-weight projection cases across three residual +contracts, three magnitudes and changing inputs. Eighty-one control/candidate/ +control replays per rank preserve output/residual bits, logical packed values, +allocation canaries and the raw graph fingerprint. M1/M7 norm-plus-QPN2 graphs +remain unmodified. M9/M32 checks cover norms only: the independently frozen raw +QPN2 entry correctly rejects M greater than eight. The first script mistakenly +used that entry for larger rows and is not recorded as a complete gate pass. +The eight-column/norm working set has only a small local difference: rank 0 +medians are 0.294416/0.291616 ms. It excludes row projections, communication and +the model round; no end-to-end saving is inferred. + +An initial graph reader mishandles the zero-edge single-node case. A later +parameter-count probe intentionally reaches an invalid API index, producing +4912 memcheck API errors despite passing data comparisons. Neither is a clean +admission. Reading the known kernel signatures removes that probe; the next +memcheck and racecheck each exit zero, with zero errors/hazards. Reports retain +these separate attempts rather than filtering the earlier errors. + +The first model startup then stops during draft graph capture: the reader +assumes pointer-array arguments for an unrelated cuBLAS kernel using the packed +launch-parameter convention. No endpoint timing or model quality result is +produced. The reader is narrowed to registered kernels before accessing their +arguments; a separate unmodified-cuBLAS graph check is added. A subsequent +model attempt remains necessary. The scoped installer also checks the target +norm weights/epsilon and verifies packed producers were written on candidate +requests and left untouched on control requests. These checks occur between +requests, not inside the timed round. + +The explicit builder `build_sm70_qpn2_dual_norm.py` reproduces generated source +SHA256 `346e063dfaf185650c279394588eac2663e12b886e179a7a6a71a899a6b9f096` +from the pinned norm expressions. `benchmark_sm70_qpn2_graph_layout.py` accepts +the two frozen projection libraries, generated norm module and real-weight +root explicitly; its `sm70_qpn2_graph_{nodes,layout}.py` helpers install no +serving default. The first versioned four-rank rerun passes 72 cases per rank. +The final helper revision adds the cuBLAS fallback check and uses the standard +accelerator synchronization API. Its four-rank rerun passes another 72 cases +per rank, including the untouched cuBLAS graph. The matching private helper's +memcheck/racecheck reruns both exit zero with zero errors/hazards and no invalid +API queries. A second model attempt finds the 128-pair full graph plus temporary +compiler/piecewise graphs, so its broad count assertion fails before serving. +The installer is then scoped directly to the model manager's owned +`FULL / num_tokens=8 / num_reqs=1 / uniform_token_count=8` descriptor. + +The third attempt completes the same-startup five-warmup/five-measurement pair. +All four ranks match 128 norm/projection pairs on that exact descriptor. The +between-request sentinel check confirms all candidate producers were written +and control requests leave their packed buffers untouched. Both fixtures keep +canonical tokens, acceptance and natural EOS. Release1k medians are +16.348511/16.328542 ms and MBPP28 16.001308/15.983139 ms. The paired MBPP savings +include two regressions; approximately 0.02 ms is not a substantial or admitted +full-round gain. Keep this candidate off. The original fixed-prefix quality +hold is not cleared by these natural trajectories alone. + +Evidence is retained in +`results/qpn2-layout-graph-*`, the corresponding queue records, and +`candidates/qpn2-layout-model-v{1,2}`. The full pair and raw samples are in +`results/qpn2-layout-within-start-summary.json` and +`results/v4-qpn2-layout-within-start-3-switch.json`, with their own four-worker +runtime-library manifest. Final fixed-prefix, acceptance and complete-round +admission remain open; this does not resolve the old repeat-start TV discrepancy +by itself. + +## Draft WMMA output-tile grouping: exact reconstruction, no speed gain + +The real cuBLAS trace uses different split-K rounding contracts for draft +projections. QKV uses three FP32 partials and a separate reduction; gate/up +uses two serial partitions with an FP16 intermediate output. O/down write +their FP16 output directly. A parallel FP16-partial gate/up reconstruction +does not match this contract. A separate one-partition reduction also erases +one negative-zero output in the actual O-projection corpus. Neither mismatch +is waived by a numerical tolerance. + +Using the corresponding serial/direct/FP32-parallel contracts, both one-warp +and four-output-warp CUTLASS prototypes pass all 400 real projection cases: +four ranks, five layers, four projections and five input snapshots. Outputs +match the captured frozen cuBLAS bytes, QKV partials match across warp grouping, +and workspace/output canaries remain intact, including graph replay. + +Seven paired timings over the twenty distinct rank-zero consecutive-layer +weights give medians of 1.450368 ms for frozen cuBLAS, 1.895968 ms for one warp +and 1.505536 ms for four warps. The four-warp prototype remains slower than the +frozen library, so it receives no serving route or end-to-end admission. +Reference CUTLASS commit is `b2dd65dc864e09688245b316ac46c4a6cd07e15c`. +Serial/parallel DSO hashes are +`47dc2f9f1978777428247bbc1970eb497d25dcbd2335e261b8e85645497b9b8a` and +`5c832dde87a11e51338b0851eff90cf43c76c156996cec4012c5c8371d1b7b33`. +Retained evidence includes `results/draft-wmma-serial-oracle.json`, the failed +first working-set gate and `results/draft-wmma-working-set-v2.json`. diff --git a/flash-attention-v100/kernel/flash_decode_paged.cu b/flash-attention-v100/kernel/flash_decode_paged.cu index 3d9df0ea23..a98e2974e3 100644 --- a/flash-attention-v100/kernel/flash_decode_paged.cu +++ b/flash-attention-v100/kernel/flash_decode_paged.cu @@ -2391,6 +2391,9 @@ __launch_bounds__(kGroupedVerifyThreads, 1) void flash_attention_grouped_verify_ shared_prob_residual[row * kResidualStride + lane_id] = __float2half_rn((probability - rounded) * 2048.0f); } + // Finish every lane's shared-state reads before lane 0 overwrites the + // 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; diff --git a/tests/kernels/attention/test_sm70_grouped_e4m3_fp32.py b/tests/kernels/attention/test_sm70_grouped_e4m3_fp32.py index 338b8c79e7..bd37fe599f 100644 --- a/tests/kernels/attention/test_sm70_grouped_e4m3_fp32.py +++ b/tests/kernels/attention/test_sm70_grouped_e4m3_fp32.py @@ -40,6 +40,8 @@ def test_precision_capability_rejects_stale_binary(monkeypatch, has_entry, versi (8, 1616, 65536), (5, 1648, 131072), (5, 3296, 262144), + # Small multi-tile q8 case for the online-softmax warp-state racecheck. + (8, 3296, 512), # DFlash2 q8 uses the same repaired arithmetic at each context boundary. (8, 3296, 8192), (8, 3296, 65536), diff --git a/tests/kernels/core/test_sm70_dflash2_gemma_rms.py b/tests/kernels/core/test_sm70_dflash2_gemma_rms.py index 87c0eda3da..56a4258950 100644 --- a/tests/kernels/core/test_sm70_dflash2_gemma_rms.py +++ b/tests/kernels/core/test_sm70_dflash2_gemma_rms.py @@ -6,11 +6,85 @@ from vllm.model_executor.layers.layernorm import ( GemmaRMSNorm, + _sm70_dflash2_fixed_gemma_rms_norm, _sm70_dflash2_gemma_fused_add_rms_norm, + _use_sm70_dflash2_fixed_gemma_rms, _use_sm70_dflash2_gemma_fused_add_rms, ) +@pytest.mark.parametrize("has_residual", [False, True]) +def test_fixed_gemma_norm_is_row_invariant_and_preserves_residual(has_residual): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0): + pytest.skip("SM70 CUDA device required") + torch.manual_seed(20260908) + x = torch.randn((153, 5120), dtype=torch.float16, device="cuda") + weight = torch.randn(5120, dtype=torch.float16, device="cuda").mul_(0.05) + residual = torch.randn_like(x) if has_residual else None + x_before = x.clone() + residual_before = residual.clone() if residual is not None else None + actual = _sm70_dflash2_fixed_gemma_rms_norm(x, residual, weight, 1e-6) + if residual is not None: + actual, residual_out = actual + torch.testing.assert_close( + residual_out, x.float() + residual.float(), atol=0, rtol=0 + ) + assert residual_out.dtype == torch.float32 + torch.testing.assert_close(residual, residual_before, atol=0, rtol=0) + torch.testing.assert_close(x, x_before, atol=0, rtol=0) + + # The same rows must not acquire another reduction order at q1/q8 or at + # an irregular prefill boundary. Include the final masked tile's values. + for first, last in ((0, 1), (0, 8), (8, 143), (143, 153)): + r = residual[first:last] if residual is not None else None + part = _sm70_dflash2_fixed_gemma_rms_norm(x[first:last], r, weight, 1e-6) + if has_residual: + part = part[0] + assert torch.equal(part, actual[first:last]) + + values = x.double() + if residual is not None: + # Residual storage retains the existing FP32 addition contract. + values = (x.float() + residual.float()).double() + reference = values * torch.rsqrt(values.square().mean(-1, keepdim=True) + 1e-6) + reference = (reference * (weight.double() + 1)).half() + lower = torch.nextafter(reference, torch.full_like(reference, -float("inf"))) + upper = torch.nextafter(reference, torch.full_like(reference, float("inf"))) + assert torch.all((actual >= lower) & (actual <= upper)) + + +def test_fixed_gemma_norm_fullgraph_replay_and_dispatch(monkeypatch): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0): + pytest.skip("SM70 CUDA device required") + monkeypatch.setenv("VLLM_SM70_DFLASH2_FIXED_GEMMA_RMS", "1") + monkeypatch.setenv("VLLM_SM70_FLASH_V100_0DOT3_COMPILE_GRAPH", "1") + x = torch.randn((8, 5120), device="cuda", dtype=torch.float16) + weight = torch.zeros(5120, device="cuda", dtype=torch.float16) + residual = torch.randn_like(x) + + @torch.compile(backend="inductor", fullgraph=True) + def apply(values, skip, norm_weight): + if _use_sm70_dflash2_fixed_gemma_rms(values, skip, norm_weight): + return _sm70_dflash2_fixed_gemma_rms_norm(values, skip, norm_weight, 1e-6) + return values, skip + + apply(x, residual, weight) + assert not _use_sm70_dflash2_fixed_gemma_rms(x, residual.float(), weight) + assert not _use_sm70_dflash2_fixed_gemma_rms(x[:, :128], None, weight[:128]) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual, skip = apply(x, residual, weight) + for _ in range(3): + x.copy_(torch.randn_like(x)) + residual.copy_(torch.randn_like(residual)) + graph.replay() + expected, expected_skip = _sm70_dflash2_fixed_gemma_rms_norm( + x, residual, weight, 1e-6 + ) + assert torch.equal(actual, expected) + assert torch.equal(skip, expected_skip) + + @pytest.mark.parametrize("num_tokens", [1, 8, 32, 137, 512]) @pytest.mark.parametrize("weight_dtype", [torch.float16, torch.bfloat16, torch.float32]) def test_sm70_dflash2_gemma_fused_add_rms_is_within_one_fp16_ulp( diff --git a/tests/kernels/core/test_sm70_dflash2_state_audit.py b/tests/kernels/core/test_sm70_dflash2_state_audit.py new file mode 100644 index 0000000000..f6698e0db6 --- /dev/null +++ b/tests/kernels/core/test_sm70_dflash2_state_audit.py @@ -0,0 +1,222 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from benchmarks.compare_sm70_dflash2_natural_audit import compare_natural +from benchmarks.compare_sm70_dflash2_state_audit import compare, tensor_difference +from benchmarks.sm70_dflash2_state_audit import ( + cpu_request_slots, + gather_state, + selected_ssm_slots, +) + + +@pytest.fixture +def captures(tmp_path): + directories = (tmp_path / "left", tmp_path / "right") + for directory in directories: + directory.mkdir() + for rank in range(4): + for step, phase in enumerate(("prefill", "verify")): + states = { + f"{phase}/layer0/{name}:(1, 2)": torch.ones(1, 2) + for name in ( + "conv/input", + "conv/output", + "recurrent/q", + "recurrent/input_state", + "recurrent/output", + ) + } + data = { + "case": "test", + "rank": rank, + "step": step, + "phase": phase, + "num_draft_tokens": 7 if step else 0, + "positions": torch.tensor([step]), + "input_ids": torch.tensor([1]), + "hidden": torch.ones(1, 2), + "states": states, + "tensors": {}, + "sampling": {"seeds": torch.tensor([0])}, + "native_logits": torch.arange(32).float().reshape(1, -1) + if rank == 0 + else None, + "expected_layers": [0], + "capture_epoch": step, + } + torch.save(data, directory / f"test-rank{rank}-step{step}.pt") + return directories + + +def test_audit_comparator_requires_all_ranks(captures): + left, right = captures + assert compare(left, right)["summary"]["all_logits_bitwise_equal"] + # Even equal, incomplete arms cannot be reported as a successful A/A. + for directory in captures: + (directory / "test-rank3-step1.pt").unlink() + with pytest.raises(ValueError, match="four TP ranks"): + compare(left, right) + + +def test_audit_comparator_requires_actual_candidate_route(captures): + left, right = captures + with pytest.raises(ValueError, match="missing packed route hit"): + compare(left, right, right_verifier_route="packed") + for path in right.glob("*-step1.pt"): + data = torch.load(path, weights_only=True) + data["verifier_routes"] = ["route/verify/layer0/packed"] + torch.save(data, path) + assert compare(left, right, right_verifier_route="packed")["summary"][ + "all_logits_bitwise_equal" + ] + + +def test_audit_comparator_rejects_missing_layer_observations(captures): + left, right = captures + for directory in captures: + path = directory / "test-rank0-step1.pt" + data = torch.load(path, weights_only=True) + data["expected_layers"] = [0, 1] + torch.save(data, path) + with pytest.raises(ValueError, match="missing verify/layer1"): + compare(left, right) + + +def test_audit_comparator_rejects_nonfinite_logits(captures): + left, right = captures + path = right / "test-rank0-step1.pt" + data = torch.load(path, weights_only=True) + data["native_logits"][0, 0] = float("nan") + torch.save(data, path) + with pytest.raises(ValueError, match="nonfinite"): + compare(left, right) + + +@pytest.fixture +def natural_captures(captures): + for directory in captures: + for path in directory.glob("*-rank*-step*.pt"): + row = torch.load(path, weights_only=True) + row.update( + control="natural_sampling", + aux_hidden_states=[torch.ones(1, 2)], + num_sampled=torch.tensor([1]), + num_rejected=torch.tensor([7 if row["step"] else 0]), + sampled_token_ids=torch.tensor([[3, -1]]), + ) + torch.save(row, path) + torch.save( + { + **{k: row[k] for k in ("case", "rank", "step")}, + "draft_tokens": torch.tensor([[4, 5, 6]]), + "projected_context": torch.ones(1, 2), + }, + directory / f"proposal-test-tp{row['rank']}-forward{row['step']}.pt", + ) + return captures + + +def test_natural_audit_locates_proposal_before_next_target(natural_captures): + left, right = natural_captures + assert compare_natural(left, right)["cases"][0]["all_logical_tensors_equal"] + path = right / "proposal-test-tp2-forward0.pt" + row = torch.load(path, weights_only=True) + row["draft_tokens"][0, 0] += 1 + torch.save(row, path) + path = right / "test-rank2-step1.pt" + row = torch.load(path, weights_only=True) + row["input_ids"][0] += 1 + torch.save(row, path) + result = compare_natural(left, right)["cases"][0] + assert not result["all_logical_tensors_equal"] + first = result["first_observed_difference"] + assert (first["step"], first["phase"]) == (0, "proposal") + assert first["differences"][0]["name"] == "draft_tokens" + + +@pytest.mark.parametrize( + "missing", ["test-rank3-step1.pt", "proposal-test-tp3-forward1.pt"] +) +def test_natural_audit_rejects_incomplete_equal_arms(natural_captures, missing): + for directory in natural_captures: + (directory / missing).unlink() + with pytest.raises(ValueError, match="four TP ranks|missing proposal"): + compare_natural(*natural_captures) + + +def test_natural_audit_ignores_unwritten_output_padding(natural_captures): + left, right = natural_captures + path = right / "test-rank0-step1.pt" + row = torch.load(path, weights_only=True) + row["sampled_token_ids"][0, 1] = 100 + torch.save(row, path) + assert compare_natural(left, right)["cases"][0]["all_logical_tensors_equal"] + + +def test_natural_audit_aligns_request_slots(natural_captures): + for index, directory in enumerate(natural_captures): + for path in directory.glob("proposal-*.pt"): + row = torch.load(path, weights_only=True) + row["idx_mapping"] = torch.tensor([index + 1]) + row["seeds"] = torch.tensor([999, 888, 777, 666]) + row["seeds"][index + 1] = 0 + torch.save(row, path) + result = compare_natural(*natural_captures)["cases"][0] + assert result["all_logical_tensors_equal"] + assert result["different_request_slot_mappings"] + values = torch.tensor([999, 0, 777]) + actual = cpu_request_slots(values, torch.tensor([1])) + values[1] = 100 + assert actual.tolist() == [0] + + +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] + assert selected_ssm_slots(table, torch.tensor([0, 5])).tolist() == [-1, -1] + + +def test_audit_compares_singleton_strided_metadata(): + wide = torch.empty_strided((1,), (8,), dtype=torch.int64).fill_(7) + assert wide.is_contiguous() and wide.stride() == (8,) + assert tensor_difference(wide, torch.tensor([7]))["bitwise_equal"] + + +def test_state_audit_distinguishes_padding_from_live_slot_zero(): + pool = torch.arange(24, dtype=torch.float32).reshape(4, 2, 3) + snapshot = gather_state(pool, torch.tensor([0, -1, 3, 4])) + assert snapshot["indices"].tolist() == [0, -1, 3, 4] + assert snapshot["valid"].tolist() == [True, False, True, False] + torch.testing.assert_close(snapshot["values"][0], pool[0], atol=0, rtol=0) + torch.testing.assert_close(snapshot["values"][2], pool[3], atol=0, rtol=0) + assert not snapshot["values"][[1, 3]].count_nonzero() + pool.fill_(100) + assert snapshot["values"][0, 0, 0] == 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_state_audit_graph_replays_changed_slots_without_mutating_pool(): + pool = torch.arange(48, device="cuda", dtype=torch.float32).reshape(8, 2, 3) + table = torch.tensor([[7, 0, 6, 2]], device="cuda", dtype=torch.int32) + selector = torch.tensor([1], device="cuda", dtype=torch.int32) + # Warm CUDA allocation and selection before graph capture. + gather_state(pool, selected_ssm_slots(table, selector)) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + snapshot = gather_state(pool, selected_ssm_slots(table, selector)) + for selected, expected in ((1, 7), (2, 0), (4, 2)): + selector.fill_(selected) + before = pool.clone() + graph.replay() + torch.testing.assert_close( + snapshot["values"][0], pool[expected], atol=0, rtol=0 + ) + torch.testing.assert_close(pool, before, atol=0, rtol=0) + selector.fill_(0) + graph.replay() + assert not snapshot["valid"].any() + assert not snapshot["values"].count_nonzero() diff --git a/tests/kernels/test_fused_sigmoid_gating_delta_rule.py b/tests/kernels/test_fused_sigmoid_gating_delta_rule.py index 2ac073b410..6b3dd8b80c 100644 --- a/tests/kernels/test_fused_sigmoid_gating_delta_rule.py +++ b/tests/kernels/test_fused_sigmoid_gating_delta_rule.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import pytest import torch import torch.nn.functional as F @@ -12,6 +14,7 @@ fused_sigmoid_gating_delta_rule_update_mixed_qkv_out, ) from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( + QwenGatedDeltaNetAttention, fused_gdn_gating, ) from vllm.platforms import current_platform @@ -347,10 +350,12 @@ def test_fused_sigmoid_gating_delta_rule_update_spec( @pytest.mark.parametrize("num_reqs", [1, 2]) @pytest.mark.parametrize("num_speculative_tokens", [3, 7]) @pytest.mark.parametrize("state_dtype", [torch.float16, torch.float32]) -def test_dflash2_packed_verify_matches_split_fp16_contract( +@pytest.mark.parametrize("runtime_bridge", [False, True]) +def test_dflash2_packed_verify_matches_split_contract( num_reqs: int, num_speculative_tokens: int, state_dtype: torch.dtype, + runtime_bridge: bool, ) -> None: if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0): pytest.skip("the DFlash2 verifier fast path is SM70-only") @@ -365,7 +370,11 @@ def test_dflash2_packed_verify_matches_split_fp16_contract( num_tokens = num_reqs * tokens_per_req qkv_width = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim - mixed_qkv = torch.randn(num_tokens, qkv_width, dtype=dtype) + # Runtime Qwen3.5 retains the QKVZBA projection row stride after conv. + padding = num_v_heads * head_v_dim + 2 * num_v_heads if runtime_bridge else 0 + projection = torch.randn(num_tokens, qkv_width + padding, dtype=dtype) + mixed_qkv = projection[:, :qkv_width] + projection_before = projection.clone() query, key, value = torch.split( mixed_qkv, [ @@ -397,7 +406,9 @@ def test_dflash2_packed_verify_matches_split_fp16_contract( dtype=state_dtype, ) - g, beta = fused_gdn_gating(A_log, a, b, dt_bias) + g, beta = fused_gdn_gating( + A_log, a, b, dt_bias, beta_dtype=torch.float32 if runtime_bridge else dtype + ) reference_state = initial_state.clone() reference_out, _ = fused_recurrent_gated_delta_rule( q=query, @@ -415,29 +426,55 @@ def test_dflash2_packed_verify_matches_split_fp16_contract( fused_state = initial_state.clone() fused_out = torch.empty(num_tokens, 1, num_v_heads, head_v_dim, dtype=dtype) - fused_sigmoid_gating_delta_rule_update_mixed_qkv_out( - A_log=A_log, - a=a, - b=b, - dt_bias=dt_bias, - mixed_qkv=mixed_qkv, - num_q_heads=num_k_heads, - num_v_heads=num_v_heads, - head_k_dim=head_k_dim, - head_v_dim=head_v_dim, - out=fused_out, - initial_state=fused_state, - cu_seqlens=cu_seqlens, - ssm_state_indices=state_indices, - num_accepted_tokens=num_accepted_tokens, - use_qk_l2norm_in_kernel=True, - precomputed_g=g, - precomputed_beta=beta, - match_recurrent_schedule=True, - match_recurrent_numerics=True, - ) + if runtime_bridge: + module = SimpleNamespace( + A_log=A_log, + dt_bias=dt_bias, + num_k_heads=num_k_heads * 4, + num_v_heads=num_v_heads * 4, + tp_size=4, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + ) + QwenGatedDeltaNetAttention._forward_dflash2_packed_gdn_verify( + module, + mixed_qkv=mixed_qkv, + a=a, + b=b, + core_attn_out=fused_out.squeeze(1), + ssm_state=fused_state, + spec_query_start_loc=cu_seqlens, + spec_state_indices_tensor=state_indices, + spec_state_slot_selectors=num_accepted_tokens, + num_spec_decodes=num_reqs, + ) + else: + fused_sigmoid_gating_delta_rule_update_mixed_qkv_out( + A_log=A_log, + a=a, + b=b, + dt_bias=dt_bias, + mixed_qkv=mixed_qkv, + num_q_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + out=fused_out, + initial_state=fused_state, + cu_seqlens=cu_seqlens, + ssm_state_indices=state_indices, + num_accepted_tokens=num_accepted_tokens, + use_qk_l2norm_in_kernel=True, + precomputed_g=g, + precomputed_beta=beta, + match_recurrent_schedule=True, + match_recurrent_numerics=True, + ) torch.accelerator.synchronize() + assert torch.equal( + projection.view(torch.uint8), projection_before.view(torch.uint8) + ) torch.testing.assert_close(fused_out.transpose(0, 1), reference_out, rtol=0, atol=0) torch.testing.assert_close(fused_state, reference_state, rtol=0, atol=0) diff --git a/tests/kernels/test_sm70_qwen35_gdn_split.py b/tests/kernels/test_sm70_qwen35_gdn_split.py index e8a65c0002..a4e86a77bc 100644 --- a/tests/kernels/test_sm70_qwen35_gdn_split.py +++ b/tests/kernels/test_sm70_qwen35_gdn_split.py @@ -81,6 +81,50 @@ def test_qwen35_gdn_split_graph_replay_reads_current_projection_values(): assert torch.equal(actual_a, mixed_ba[:, ba_size:].contiguous()) +@pytest.mark.parametrize("compiled", [False, True]) +def test_qwen35_combined_split_preserves_tails_across_graph_replay(compiled): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for the Qwen3.5 GDN split kernel") + + # Real TP4 NVFP4 geometry: 4120 logical columns in a 4128-column row. + # Both input arguments alias this allocation, and b/a start at an offset. + storage = torch.empty((8, 4128), dtype=torch.float16, device="cuda") + projection = storage[:, :4120] + + def split(values): + return _sm70_materialize_qwen35_gdn_splits( + values, values[:, 4096:4120], 2560, 1536, 12 + ) + + if compiled: + split = torch.compile(split, fullgraph=True) + storage.normal_() + split(projection) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = split(projection) + # The consumer may overwrite QKV in the same graph. Tail outputs must + # be independent snapshots, including the two small gating vectors. + projection[:, :2560].zero_() + + for seed in (41, 42, 43): + torch.manual_seed(seed) + storage.normal_() + expected = [ + projection[:, start:end].clone() + for start, end in ((2560, 4096), (4096, 4108), (4108, 4120)) + ] + padding = storage[:, 4120:].clone() + graph.replay() + torch.cuda.synchronize() + for result, reference in zip(actual, expected): + assert result.is_contiguous() + assert torch.equal(result.view(torch.uint8), reference.view(torch.uint8)) + assert torch.count_nonzero(projection[:, :2560]) == 0 + assert torch.equal(storage[:, 4120:], padding) + + @pytest.mark.parametrize("num_rows", [1, 8]) def test_qwen35_gdn_qkv_pack_is_bitwise_exact(num_rows: int): if not torch.cuda.is_available(): diff --git a/tests/v1/attention/test_gdn_metadata_builder.py b/tests/v1/attention/test_gdn_metadata_builder.py index c7207e55aa..e0d925ef78 100644 --- a/tests/v1/attention/test_gdn_metadata_builder.py +++ b/tests/v1/attention/test_gdn_metadata_builder.py @@ -311,9 +311,14 @@ def _build( batch_spec: BatchSpec, num_decode_draft_tokens: list[int] | None = None, use_common_metadata: bool = False, + is_prefilling: list[bool] | None = None, ) -> GDNAttentionMetadata: """Build GDN attention metadata, optionally with spec-decode kwargs.""" common = create_common_attn_metadata(batch_spec, BLOCK_SIZE, DEVICE) + if is_prefilling is not None: + common = common.replace( + is_prefilling=torch.tensor(is_prefilling, dtype=torch.bool, device=DEVICE) + ) kwargs: dict = {} if num_decode_draft_tokens is not None: num_decode_draft_tokens_cpu = torch.tensor( @@ -411,6 +416,97 @@ def test_gdn_build_classification( assert meta.num_spec_decodes == test_case.expected_num_spec_decodes +@pytest.mark.parametrize("use_full_cuda_graph", [False, True]) +@pytest.mark.parametrize( + ( + "num_speculative_tokens", + "query_len", + "seq_len", + "is_prefilling", + "expected_prefills", + "expected_initial_state", + ), + [ + pytest.param(7, 1, 1, True, 1, [False], id="fresh-singleton"), + pytest.param(7, 1, 17, True, 1, [True], id="cached-singleton-extension"), + pytest.param(7, 1, 17, False, 0, None, id="regular-decode"), + pytest.param(7, 2, 2, True, 1, [False], id="two-token-prefill"), + pytest.param(7, 17, 17, True, 1, [False], id="longer-prefill"), + pytest.param(7, 1, 1, None, 0, None, id="legacy-metadata"), + pytest.param(0, 1, 1, True, 0, None, id="non-speculative-unchanged"), + ], +) +def test_singleton_prefill_state_initialization( + local_gdn_model: str, + use_full_cuda_graph: bool, + num_speculative_tokens: int, + query_len: int, + seq_len: int, + is_prefilling: bool | None, + expected_prefills: int, + expected_initial_state: list[bool] | None, +): + builder = _create_gdn_builder( + local_gdn_model, + num_speculative_tokens=num_speculative_tokens, + use_full_cuda_graph=use_full_cuda_graph, + max_cudagraph_capture_size=8, + ) + meta = _build( + builder, + BatchSpec(seq_lens=[seq_len], query_lens=[query_len]), + num_decode_draft_tokens=[-1] if num_speculative_tokens else None, + is_prefilling=None if is_prefilling is None else [is_prefilling], + ) + + assert meta.num_spec_decodes == 0 + assert meta.num_prefills == expected_prefills + assert meta.num_decodes == 1 - expected_prefills + assert meta.num_prefill_tokens == query_len * expected_prefills + if expected_initial_state is None: + assert meta.has_initial_state is None + else: + assert meta.has_initial_state is not None + assert meta.has_initial_state.tolist() == expected_initial_state + + +@pytest.mark.parametrize("poison", [10.0, float("nan")], ids=["recycled", "nan"]) +def test_singleton_prefill_masks_recycled_conv_state(local_gdn_model: str, poison): + from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( + causal_conv1d_torch, + ) + + builder = _create_gdn_builder(local_gdn_model, num_speculative_tokens=7) + meta = _build( + builder, + BatchSpec(seq_lens=[1], query_lens=[1]), + num_decode_draft_tokens=[-1], + is_prefilling=[True], + ) + assert meta.num_prefills == 1 + assert meta.has_initial_state is not None + assert meta.non_spec_query_start_loc is not None + assert meta.non_spec_state_indices_tensor is not None + states = torch.full((1, 1, 3), poison, dtype=torch.float32, device=DEVICE) + output = causal_conv1d_torch( + x=torch.ones((1, 1), dtype=torch.float32, device=DEVICE), + weight=torch.ones((1, 4), dtype=torch.float32, device=DEVICE), + bias=None, + conv_states=states, + query_start_loc=meta.non_spec_query_start_loc, + cache_indices=torch.zeros_like(meta.non_spec_state_indices_tensor), + has_initial_state=meta.has_initial_state, + activation=None, + ) + torch.testing.assert_close(output, torch.ones_like(output), rtol=0, atol=0) + torch.testing.assert_close( + states, + torch.tensor([[[0.0, 0.0, 1.0]]], device=DEVICE), + rtol=0, + atol=0, + ) + + @pytest.mark.parametrize( "test_case", [ diff --git a/vllm/envs.py b/vllm/envs.py index 3bcb9e35b4..84327507f0 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -243,10 +243,13 @@ VLLM_SM70_DFLASH2_FUSED_GDN_VERIFY: bool = False VLLM_SM70_DFLASH2_FUSED_GDN_NORM: bool = False VLLM_SM70_DFLASH2_FUSED_GDN_SPLIT: bool = False + VLLM_SM70_DFLASH2_FUSED_GDN_COMBINED_SPLIT: bool = False + VLLM_SM70_DFLASH2_DIRECT_ATTENTION_OUTPUT: bool = False VLLM_SM70_DFLASH2_FUSED_SMALLQ_METADATA: bool = False VLLM_SM70_DFLASH2_GROUPED_SMALLQ_METADATA: bool = False VLLM_SM70_DFLASH2_FUSED_QKV_PACK: bool = False VLLM_SM70_DFLASH2_FUSED_GEMMA_RMS: bool = False + VLLM_SM70_DFLASH2_FIXED_GEMMA_RMS: bool = False VLLM_SM70_DFLASH2_SPARSE_TARGET_REJECTION: bool = False VLLM_SM70_DFLASH2_SHARDED_CONTEXT_FC: bool = False VLLM_SM70_DFLASH2_CONTEXT_KV_GRAPH: bool = False @@ -2234,6 +2237,15 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_DFLASH2_FUSED_GDN_SPLIT": lambda: bool( int(os.getenv("VLLM_SM70_DFLASH2_FUSED_GDN_SPLIT", "0")) ), + # Independently gate the TP4 q8 all-NVFP4 QKVZBA projection layout. + "VLLM_SM70_DFLASH2_FUSED_GDN_COMBINED_SPLIT": lambda: bool( + int(os.getenv("VLLM_SM70_DFLASH2_FUSED_GDN_COMBINED_SPLIT", "0")) + ), + # Return the existing projection tensor across the GDN opaque boundary. + # This does not enable collective/norm fusion or change state arithmetic. + "VLLM_SM70_DFLASH2_DIRECT_ATTENTION_OUTPUT": lambda: bool( + int(os.getenv("VLLM_SM70_DFLASH2_DIRECT_ATTENTION_OUTPUT", "0")) + ), # Build Flash-V100 small-query verifier rows directly in their persistent # graph buffers. This replaces four repeat_interleave scans per KV group. # The matched TP4 trace is token/acceptance exact and cuts the synchronized @@ -2258,6 +2270,13 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_SM70_DFLASH2_FUSED_GEMMA_RMS": lambda: bool( int(os.getenv("VLLM_SM70_DFLASH2_FUSED_GEMMA_RMS", "0")) ), + # Experimental fixed 8192/16-warp reduction for the FP16 no-residual and + # FP16-residual Gemma norms not covered by the existing FP32-residual path. + # Prevents per-rank/startup autotune from changing reduction order. Keep + # disabled until fixed-prefix, natural-output and full-round gates pass. + "VLLM_SM70_DFLASH2_FIXED_GEMMA_RMS": lambda: bool( + int(os.getenv("VLLM_SM70_DFLASH2_FIXED_GEMMA_RMS", "0")) + ), # Avoid materializing/gathering full-vocabulary target logits when the # DFlash2 proposal and target sampling distributions both have compact # top-k support. Default-off until paired output/acceptance and end-to-end diff --git a/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py b/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py index 49f1e3aad1..0316e03e5f 100644 --- a/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py +++ b/vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py @@ -649,10 +649,11 @@ def fused_sigmoid_gating_delta_rule_update_mixed_qkv_out( ``precomputed_g`` and ``precomputed_beta`` retain the split verifier's exact gating materialization while still removing the packed-QKV rearrange. Omitting both computes gating inside the recurrent kernel. + Row-strided QKV views with contiguous features are consumed without a copy. """ if mixed_qkv.ndim != 2: raise ValueError("mixed_qkv must have shape [T, qkv_hidden].") - if not mixed_qkv.is_contiguous(): + if mixed_qkv.stride(1) != 1: mixed_qkv = mixed_qkv.contiguous() if cu_seqlens is None: raise ValueError("cu_seqlens is required for mixed_qkv_out.") @@ -677,6 +678,7 @@ def fused_sigmoid_gating_delta_rule_update_mixed_qkv_out( raise ValueError( f"mixed_qkv width {mixed_qkv.shape[1]} != expected {qkv_stride}." ) + qkv_stride = mixed_qkv.stride(0) if out.shape != (T, 1, HV, V): raise ValueError(f"out must have shape {(T, 1, HV, V)}, got {out.shape}.") if scale is None: diff --git a/vllm/model_executor/layers/layernorm.py b/vllm/model_executor/layers/layernorm.py index 624a207613..feaca4def4 100644 --- a/vllm/model_executor/layers/layernorm.py +++ b/vllm/model_executor/layers/layernorm.py @@ -21,6 +21,100 @@ logger = init_logger(__name__) +@triton.jit +def _sm70_dflash2_fixed_gemma_rms_kernel( + x, + residual, + weight, + normalized_out, + residual_out, + HAS_RESIDUAL: tl.constexpr, + epsilon: tl.constexpr, +): + # Pin both the reduction extent and warp count. Inductor's 2048/8192 + # autotune changes FP32 reduction order, including between TP ranks. + row = tl.program_id(0) + cols = tl.arange(0, 8192) + mask = cols < 5120 + values = tl.load(x + row * 5120 + cols, mask=mask, other=0.0).to(tl.float32) + if HAS_RESIDUAL: + values += tl.load(residual + row * 5120 + cols, mask=mask, other=0.0).to( + tl.float32 + ) + tl.store(residual_out + row * 5120 + cols, values, mask=mask) + # Preserve the masked square and residual materialization of the pinned + # Inductor reduction. Removing this boundary changes FMA contraction for + # sums of two FP16 inputs even with an identical reduction tile. + variance = tl.sum(tl.where(mask, values * values, 0.0), axis=0) / 5120.0 + inverse_rms = tl.rsqrt(variance + epsilon) + if HAS_RESIDUAL: + values = tl.load(residual_out + row * 5120 + cols, mask=mask, other=0.0) + gemma_weight = tl.load(weight + cols, mask=mask, other=0.0).to(tl.float32) + 1.0 + tl.store( + normalized_out + row * 5120 + cols, + values * inverse_rms * gemma_weight, + mask=mask, + ) + + +def _sm70_dflash2_fixed_gemma_rms_norm( + x: torch.Tensor, + residual: torch.Tensor | None, + weight: torch.Tensor, + variance_epsilon: float, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + normalized_out = torch.empty_like(x) + residual_out = ( + torch.empty_like(x, dtype=torch.float32) if residual is not None else None + ) + _sm70_dflash2_fixed_gemma_rms_kernel[(x.shape[0],)]( + x, + residual, + weight, + normalized_out, + residual_out, + HAS_RESIDUAL=residual is not None, + epsilon=variance_epsilon, + num_warps=16, + num_stages=1, + enable_fp_fusion=True, + ) + if residual_out is None: + return normalized_out + return normalized_out, residual_out + + +def _use_sm70_dflash2_fixed_gemma_rms( + x: torch.Tensor, + residual: torch.Tensor | None, + weight: torch.Tensor, +) -> bool: + return bool( + envs.VLLM_SM70_DFLASH2_FIXED_GEMMA_RMS + and envs.VLLM_SM70_FLASH_V100_0DOT3_COMPILE_GRAPH + and _sm70_gemma_long_prefill_available() + and x.is_cuda + and x.dtype == torch.float16 + and x.ndim == 2 + and x.shape[0] > 0 + and x.shape[1] == 5120 + and x.is_contiguous() + and weight.device == x.device + and weight.dtype == torch.float16 + and weight.shape == (5120,) + and weight.is_contiguous() + and ( + residual is None + or ( + residual.dtype == torch.float16 + and residual.device == x.device + and residual.shape == x.shape + and residual.is_contiguous() + ) + ) + ) + + @triton.jit def _sm70_dflash2_gemma_fused_add_rms_kernel( x, @@ -429,6 +523,10 @@ def forward_native( residual: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """PyTorch-native implementation equivalent to forward().""" + if _use_sm70_dflash2_fixed_gemma_rms(x, residual, self.weight): + return _sm70_dflash2_fixed_gemma_rms_norm( + x, residual, self.weight, self.variance_epsilon + ) if _use_sm70_dflash2_gemma_fused_add_rms(x, residual, self.weight): assert residual is not None return _sm70_dflash2_gemma_fused_add_rms_norm( @@ -483,6 +581,10 @@ def forward_cuda( x: torch.Tensor, residual: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if _use_sm70_dflash2_fixed_gemma_rms(x, residual, self.weight): + return _sm70_dflash2_fixed_gemma_rms_norm( + x, residual, self.weight, self.variance_epsilon + ) if _use_sm70_dflash2_gemma_fused_add_rms(x, residual, self.weight): assert residual is not None return _sm70_dflash2_gemma_fused_add_rms_norm( diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index fc99c3a880..a288bc6eb9 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -2443,6 +2443,13 @@ def __init__( and current_platform.is_device_capability(70) and _is_dflash2_spec_config(vllm_config) ) + self.enable_sm70_dflash2_fused_gdn_combined_split = bool( + envs.VLLM_SM70_DFLASH2_FUSED_GDN_COMBINED_SPLIT + and current_platform.is_device_capability(70) + and _is_dflash2_spec_config(vllm_config) + and self.tp_size == 4 + and self.hidden_size == 5120 + ) self.enable_sm70_dflash2_fused_qkv_pack = bool( envs.VLLM_SM70_DFLASH2_FUSED_QKV_PACK and current_platform.is_device_capability(70) @@ -3858,8 +3865,15 @@ def forward( ): conv_state_cache, ssm_state_cache = _resolve_qwen_gdn_kv_cache_args( layer_name, - output, + hidden_states if output is None else output, ) + if output is None: + return torch.ops.vllm.qwen_gdn_full_forward_direct( + hidden_states, + conv_state_cache, + ssm_state_cache, + layer_name, + ) torch.ops.vllm.qwen_gdn_full_forward( hidden_states, output, @@ -3873,7 +3887,7 @@ def forward( def _full_forward( self, hidden_states: torch.Tensor, - output: torch.Tensor, + output: torch.Tensor | None, ): return self._forward_method(hidden_states, output) @@ -5181,7 +5195,10 @@ def _can_use_dflash2_packed_gdn_verify( # The supported verifier contract keeps recurrent state in FP32; # an explicit FP16 cache override is also supported. and ssm_state.dtype in (torch.float16, torch.float32) - and mixed_qkv.is_contiguous() + # Qwen3.5's fused projection and in-place convolution retain the + # wider QKVZBA row stride. The packed consumer can read it directly. + and mixed_qkv.stride(1) == 1 + and mixed_qkv.stride(0) >= mixed_qkv.shape[1] and a.is_contiguous() and b.is_contiguous() and core_attn_out.is_contiguous() @@ -5202,7 +5219,9 @@ def _forward_dflash2_packed_gdn_verify( ) -> torch.Tensor: num_tokens = mixed_qkv.shape[0] out = core_attn_out[:num_tokens].unsqueeze(1) - g, beta = fused_gdn_gating(self.A_log, a, b, self.dt_bias) + g, beta = fused_gdn_gating( + self.A_log, a, b, self.dt_bias, beta_dtype=torch.float32 + ) fused_sigmoid_gating_delta_rule_update_mixed_qkv_out( A_log=self.A_log, a=a, @@ -7145,6 +7164,34 @@ def qwen_gdn_full_forward_fake( """Fake implementation for torch.compile.""" +def qwen_gdn_full_forward_direct( + hidden_states: torch.Tensor, + conv_state_cache: torch.Tensor, + ssm_state_cache: torch.Tensor, + layer_name: LayerNameType, +) -> torch.Tensor: + """Keep the full-forward order while returning its projection allocation.""" + layer_name = _resolve_layer_name(layer_name) + layer = get_forward_context().no_compile_layers[layer_name] + # Match the original opaque op's explicit recurrent-state dependencies. + # Its eager body accesses these same caches through the layer object. + _ = conv_state_cache, ssm_state_cache + output = layer._full_forward(hidden_states, None) + if output is None: + raise RuntimeError("Direct GDN full-forward did not return a projection") + _log_runtime_route_once("SM70 Qwen GDN direct full-forward output route hit.") + return output + + +def qwen_gdn_full_forward_direct_fake( + hidden_states: torch.Tensor, + conv_state_cache: torch.Tensor, + ssm_state_cache: torch.Tensor, + layer_name: LayerNameType, +) -> torch.Tensor: + return torch.empty_like(hidden_states) + + def qwen_gdn_output_projection( core_attn_out: torch.Tensor, z: torch.Tensor, @@ -7451,6 +7498,14 @@ def qwen_gdn_input_projection_fake( ) +direct_register_custom_op( + op_name="qwen_gdn_full_forward_direct", + op_func=qwen_gdn_full_forward_direct, + mutates_args=["conv_state_cache", "ssm_state_cache"], + fake_impl=qwen_gdn_full_forward_direct_fake, +) + + direct_register_custom_op( op_name="qwen_gdn_output_projection", op_func=qwen_gdn_output_projection, diff --git a/vllm/model_executor/models/qwen3_5.py b/vllm/model_executor/models/qwen3_5.py index dbc891fee5..c44e572a31 100644 --- a/vllm/model_executor/models/qwen3_5.py +++ b/vllm/model_executor/models/qwen3_5.py @@ -44,6 +44,7 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( QwenGatedDeltaNetAttention, + _is_dflash2_spec_config, _qwen_gdn_run_recurrent_core, _resolve_qwen_gdn_kv_cache_args, _sm70_compile_graph_slice_dim, @@ -66,6 +67,7 @@ maybe_remap_kv_scale_name, ) from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.transformers_utils.configs.qwen3_5 import ( Qwen3_5Config, @@ -477,9 +479,29 @@ def forward_cuda( ba_start = z_start + z_size a_start = ba_start + ba_size mixed_qkv = mixed_qkvzba[..., :qkv_size] - z = _sm70_compile_graph_slice_dim(mixed_qkvzba, -1, z_start, z_size) - b = _sm70_compile_graph_slice_dim(mixed_qkvzba, -1, ba_start, ba_size) - a = _sm70_compile_graph_slice_dim(mixed_qkvzba, -1, a_start, ba_size) + if ( + self.enable_sm70_dflash2_fused_gdn_combined_split + and num_tokens == 8 + and mixed_qkvzba.dtype == torch.float16 + and (qkv_size, z_size, ba_size) == (2560, 1536, 12) + ): + # Any combined projection may have a padded QKVZBA allocation. + # Copy its three tails together before convolution mutates QKV. + z, b, a = _sm70_materialize_qwen35_gdn_splits( + mixed_qkvzba, + mixed_qkvzba[:, ba_start : a_start + ba_size], + qkv_size, + z_size, + ba_size, + ) + logger.info_once( + "SM70 DFlash2 combined QKVZBA q8 split route hit.", + scope="local", + ) + else: + z = _sm70_compile_graph_slice_dim(mixed_qkvzba, -1, z_start, z_size) + b = _sm70_compile_graph_slice_dim(mixed_qkvzba, -1, ba_start, ba_size) + a = _sm70_compile_graph_slice_dim(mixed_qkvzba, -1, a_start, ba_size) mixed_qkv = _sm70_dump_gdn_projection_tensor( "split_mixed_qkv", layer_name, mixed_qkv @@ -529,6 +551,15 @@ def __init__( self.layer_type = layer_type self.layer_idx = extract_layer_index(prefix) + self.sm70_dflash2_direct_attention_output = bool( + envs.VLLM_SM70_DFLASH2_DIRECT_ATTENTION_OUTPUT + and current_platform.is_device_capability(70) + and _is_dflash2_spec_config(vllm_config) + and vllm_config.parallel_config.tensor_parallel_size == 4 + and model_config.dtype == torch.float16 + and config.hidden_size == 5120 + and config.model_type == "qwen3_5_text" + ) if self.layer_type == "linear_attention": self.linear_attn = Qwen3_5GatedDeltaNet( diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index b32171fcad..589ab002fa 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -751,8 +751,9 @@ def forward( ) use_direct_attention_output = ( - envs.VLLM_SM70_TP4_LONG_PREFILL_FUSED_NORM and torch.compiler.is_compiling() - ) + envs.VLLM_SM70_TP4_LONG_PREFILL_FUSED_NORM + or getattr(self, "sm70_dflash2_direct_attention_output", False) + ) and torch.compiler.is_compiling() self_attention_output = ( None if use_direct_attention_output else torch.empty_like(hidden_states) ) @@ -772,7 +773,7 @@ def forward( if use_direct_attention_output: if projected_attention_output is None: raise RuntimeError( - "SM70 TP4 fused prefill requires a direct attention output" + "SM70 TP4 direct attention route requires a projection tensor" ) hidden_states = projected_attention_output else: diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index 33f496826c..091ee1770a 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/vllm/v1/attention/backends/gdn_attn.py @@ -1452,7 +1452,15 @@ def build( # type: ignore[override] if spec_sequence_masks is None: num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( - split_decodes_and_prefills(m, decode_threshold=1) + split_decodes_and_prefills( + m, + decode_threshold=1, + # A singleton prefill must initialize speculative GDN state, + # rather than consume a recycled recurrent-state slot. + treat_short_extends_as_decodes=not ( + self.use_spec_decode and m.is_prefilling is not None + ), + ) ) num_spec_decode_tokens = 0 spec_token_indx = None