[Perf] Reduce Flash-Next MTP4 verifier batch inflation on V100

benchmarks/flashnext_prefill/analyze_mtp_verifier.py

@@ -0,0 +1,101 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Diagnostic MTP4 verifier phases, aligned by step ordinal across all TP4 ranks."""
+
+import argparse
+import json
+import statistics
+from pathlib import Path
+
+import regex as re
+
+p = argparse.ArgumentParser()
+p.add_argument("log", type=Path)
+p.add_argument("requests", type=Path)
+p.add_argument("output", type=Path)
+a = p.parse_args()
+
+
+def stats(v):
+    x = sorted(v)
+    return dict(
+        n=len(x),
+        mean=statistics.mean(x),
+        p50=statistics.median(x),
+        p90=x[round(0.9 * (len(x) - 1))],
+        p99=x[round(0.99 * (len(x) - 1))],
+        max=x[-1],
+    )
+
+
+raw = a.log.read_bytes()
+result = []
+for request in json.loads(a.requests.read_text()):
+    c = request["concurrency"]
+    groups = {}
+    for line in (
+        raw[request["profile_byte_start"] : request["profile_byte_end"]]
+        .decode(errors="replace")
+        .splitlines()
+    ):
+        match = re.search(
+            r"Worker_TP(\d+).*SM70 V2 MTP profile interval_avg_ms "
+            r"calls=(\d+) interval_calls=1 tokens=(\d+) drafts=(\d+) (.*)",
+            line,
+        )
+        if not match:
+            continue
+        rank, ordinal, tokens, drafts = map(int, match.groups()[:4])
+        if (tokens, drafts) != (5 * c, 4 * c):
+            continue
+        values = {k: float(v) for k, v in re.findall(r"([a-z_]+)=([0-9.]+)", match[5])}
+        if rank in groups.setdefault(ordinal, {}):
+            raise ValueError("Duplicate rank ordinal")
+        groups[ordinal][rank] = values
+    complete = [(n, g) for n, g in sorted(groups.items()) if set(g) == {0, 1, 2, 3}]
+    if request["repetition"] < 0:
+        continue
+    if len(complete) < 16:
+        raise ValueError(("Too few full-concurrency verifier rounds", c, len(complete)))
+    records = []
+    for ordinal, ranks in complete:
+        critical = max(ranks, key=lambda r: ranks[r]["target_verifier_gpu"])
+        rec = dict(
+            ordinal=ordinal, critical_rank=critical, ranks=ranks, **ranks[critical]
+        )
+        rec["rank_max_verifier_wall_ms"] = max(
+            r["target_verifier_wall_cpu"] for r in ranks.values()
+        )
+        records.append(rec)
+    fields = (
+        "target_verifier_gpu",
+        "rank_max_verifier_wall_ms",
+        "target_forward",
+        "target_sample",
+        "target_state_update",
+        "draft_total",
+        "total_gpu",
+    )
+    row = dict(
+        concurrency=c,
+        repetition=request["repetition"],
+        complete_rounds=len(complete),
+        incomplete_rounds=len(groups) - len(complete),
+        phases={k: stats([r[k] for r in records]) for k in fields},
+        per_round=records,
+    )
+    result.append(row)
+    print(c, request["repetition"], json.dumps(row["phases"]), flush=True)
+a.output.write_text(
+    json.dumps(
+        dict(
+            definition=(
+                "Per-ordinal TP4 maximum GPU event sum: target forward + sampling "
+                "+ state commit. Diagnostic events/fences are enabled. Draft is "
+                "separate; stage statistics are taken on the critical verifier rank."
+            ),
+            cases=result,
+        ),
+        indent=2,
+    )
+)

benchmarks/flashnext_prefill/cold_dma.py

@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Screen copying only occupied cold slots, retaining their staging offsets.
+
+Uses three real TP4 layer-zero shard sizes and a synthetic fragmented slot
+map. Timing includes Python/native submission of every contiguous range.
+This is a DMA screen; model transcode overlap must be measured separately.
+"""
+
+import argparse
+import json
+from functools import partial
+from pathlib import Path
+
+import torch
+from moe_tiles import artifact_identity, measure
+
+from vllm.model_executor.layers.quantization.gguf_expert_cache import _native
+from vllm.transformers_utils.gguf_tensor_reader import quant_size
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--output", type=Path, required=True)
+    args = parser.parse_args()
+    gate_block, gate_size = quant_size(18)
+    down_block, down_size = quant_size(20)
+    slot_bytes = [160 * (2560 // gate_block) * gate_size] * 2 + [
+        2560 * (160 // down_block) * down_size
+    ]
+    slots = 136
+    host = [
+        torch.empty(slots, size, dtype=torch.uint8).pin_memory() for size in slot_bytes
+    ]
+    for tensor in host:
+        tensor.random_(0, 256)
+    stage = [torch.empty_like(tensor, device="cuda") for tensor in host]
+    copy = _native().copy_h2d
+    stream = torch.cuda.current_stream().cuda_stream
+    report = {"artifact": artifact_identity(), "slot_bytes": slot_bytes, "cases": []}
+    for cold in (110, 104):
+        selected = torch.randperm(slots, generator=torch.Generator().manual_seed(9683))[
+            :cold
+        ]
+        occupied = set(selected.tolist())
+        extent = max(occupied) + 1
+        ranges = []
+        for slot in sorted(occupied):
+            if ranges and ranges[-1][1] == slot:
+                ranges[-1][1] += 1
+            else:
+                ranges.append([slot, slot + 1])
+
+        def submit(spans):
+            for src, dst, size in zip(host, stage, slot_bytes):
+                for begin, end in spans:
+                    copy(
+                        dst.data_ptr() + begin * size,
+                        src.data_ptr() + begin * size,
+                        (end - begin) * size,
+                        stream,
+                    )
+
+        bulk = partial(submit, [(0, extent)])
+        sparse = partial(submit, ranges)
+        for dst in stage:
+            dst.zero_()
+        sparse()
+        torch.accelerator.synchronize()
+        for src, dst in zip(host, stage):
+            actual = dst.cpu()
+            torch.testing.assert_close(actual[selected], src[selected], rtol=0, atol=0)
+            holes = torch.tensor(sorted(set(range(slots)) - occupied))
+            assert (actual[holes] == 0).all()
+        row = {
+            "cold_slots": cold,
+            "extent": extent,
+            "ranges": len(ranges),
+            "bulk_bytes": extent * sum(slot_bytes),
+            "selected_bytes": cold * sum(slot_bytes),
+            "bulk": measure(bulk, 10),
+            "selected": measure(sparse, 10),
+            "bulk_repeat": measure(bulk, 10),
+            "cold_exact_holes_untouched": True,
+        }
+        report["cases"].append(row)
+        print(json.dumps(row), flush=True)
+        args.output.write_text(json.dumps(report, indent=2) + "\n")
+
+
+if __name__ == "__main__":
+    main()

benchmarks/flashnext_prefill/collectives.py

@@ -0,0 +1,79 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""TP4 prefill-sized collective screen; launch with torchrun on four idle GPUs.
+
+Uses the same PyNccl and custom-AR implementations as serving. Custom AR's larger
+scratch allocation is experimental here, not a change to runtime admission.
+"""
+
+import json
+import os
+import statistics
+from functools import partial
+
+import torch
+import torch.distributed as dist
+
+from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce
+from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
+
+
+def measure(fn):
+    samples = []
+    for _ in range(5):
+        fn()
+    torch.accelerator.synchronize()
+    dist.barrier()
+    for _ in range(5):
+        begin, end = (torch.cuda.Event(enable_timing=True) for _ in range(2))
+        begin.record()
+        for _ in range(20):
+            fn()
+        end.record()
+        end.synchronize()
+        times = [None] * dist.get_world_size()
+        dist.all_gather_object(times, begin.elapsed_time(end) / 20)
+        samples.append(max(times))
+    return {"median_ms": statistics.median(samples), "samples_ms": samples}
+
+
+def main():
+    rank = int(os.environ["LOCAL_RANK"])
+    torch.accelerator.set_device_index(rank)
+    dist.init_process_group("gloo")
+    assert dist.get_world_size() == 4
+    nccl = PyNcclCommunicator(dist.group.WORLD, rank)
+    ca = CustomAllreduce(dist.group.WORLD, rank, max_size=48 * 1024**2)
+    assert not ca.disabled and not nccl.disabled
+    try:
+        for rows in (7344, 5347):
+            x = torch.randn(rows, 2560, device="cuda", dtype=torch.half) * 0.03
+            control, candidate = torch.empty_like(x), torch.empty_like(x)
+            run_nccl = partial(nccl.all_reduce, x, control)
+            run_custom = partial(ca.all_reduce, x, out=candidate)
+            result = {
+                "tokens": rows,
+                "reduce_bytes": x.numel() * x.element_size(),
+                "nccl_all_reduce": measure(run_nccl),
+                "custom_all_reduce_eager": measure(run_custom),
+            }
+            difference = (candidate.float() - control.float()).norm()
+            result["relative_l2"] = (difference / control.float().norm()).item()
+            result["finite"] = bool(torch.isfinite(candidate).all())
+            # Two HC packets per module: N=88 and N=640 per rank.
+            for n in (88, 640):
+                packet = torch.randn(rows, n, device="cuda", dtype=torch.half)
+                gathered = torch.empty(4 * rows, n, device="cuda", dtype=torch.half)
+                result[f"nccl_all_gather_n{n}"] = measure(
+                    partial(nccl.all_gather, gathered, packet)
+                )
+            if rank == 0:
+                print(json.dumps(result), flush=True)
+    finally:
+        torch.accelerator.synchronize()
+        ca.close()
+        dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

benchmarks/flashnext_prefill/dense_blas.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Compare existing fused and dequantize+cuBLAS dense providers on real TP4 weights.
+
+BLAS timing includes dequantization into its scratch on every call. This is a
+provider screen, not a model qualification or a persistent FP16 weight cache.
+"""
+
+import argparse
+import json
+from functools import partial
+from pathlib import Path
+
+import numpy as np
+import torch
+from moe_tiles import artifact_identity, measure, relative_l2
+
+from vllm.model_executor.layers.quantization.gguf_transcode import transcode_affine
+from vllm.transformers_utils.gguf_tensor_reader import GGUFReader, dequantize
+
+CASES = [
+    ("blk.0.attn_qkv.weight", 0),
+    ("blk.1.attn_qkv.weight", 0),
+    ("blk.0.ssm_out.weight", 1),
+    ("blk.1.attn_gate.weight", 0),
+    ("blk.5.ffn_up_shexp.weight", 0),
+]
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("model")
+    parser.add_argument("--tokens", type=int, nargs="+", default=[7344, 5347])
+    parser.add_argument("--iterations", type=int, default=5)
+    parser.add_argument("--output", type=Path, required=True)
+    args = parser.parse_args()
+    reader = GGUFReader(args.model)
+    report = {"artifact": artifact_identity(), "cases": []}
+    for name, axis in CASES:
+        tensor = next(t for t in reader.tensors if t.name == name)
+        raw = np.ascontiguousarray(tensor.data)
+        kind = int(tensor.tensor_type)
+        canonical = transcode_affine(raw, kind).tp_slice(0, 4, axis=axis)
+        n, k = canonical.codes.shape
+        dense = dequantize(raw, kind)
+        dense = dense[:n] if axis == 0 else dense[:, :k]
+        reference_weight = torch.from_numpy(np.ascontiguousarray(dense)).cuda()
+        w, s, meta = torch.ops._C.gguf_affine_sm70_prepare(
+            torch.from_numpy(canonical.codes).cuda(),
+            torch.from_numpy(canonical.scales).cuda(),
+            torch.from_numpy(canonical.mins).cuda(),
+            canonical.bits,
+            canonical.group_size,
+        )
+        k_ld, s_ld = meta.tolist()
+        scratch = torch.empty(k, n, dtype=torch.half, device="cuda")
+        for m in args.tokens:
+            torch.manual_seed(9683 + m)
+            x = (torch.randn(m, k, device="cuda") * 0.125).half()
+            fused_out = torch.empty(m, n, dtype=torch.half, device="cuda")
+            blas_out = torch.empty_like(fused_out)
+
+            fused = partial(
+                torch.ops._C.gguf_affine_gemm_sm70_out,
+                fused_out,
+                x,
+                w,
+                s,
+                canonical.bits,
+                k_ld,
+                s_ld,
+                canonical.group_size,
+            )
+            blas = partial(
+                torch.ops._C.gguf_affine_blas_sm70_out,
+                blas_out,
+                x,
+                w,
+                s,
+                canonical.bits,
+                scratch,
+                canonical.group_size,
+            )
+
+            fused()
+            blas()
+            expected = x[:16].float() @ reference_weight.float().T
+            errors = {
+                "fused_fp32_max_rel_l2": relative_l2(fused_out[:16], expected),
+                "blas_fp32_max_rel_l2": relative_l2(blas_out[:16], expected),
+                "paired_max_rel_l2": relative_l2(blas_out, fused_out),
+            }
+            passed = (
+                bool(torch.isfinite(fused_out).all() and torch.isfinite(blas_out).all())
+                and max(errors.values()) <= 0.003
+            )
+            row = {
+                "tensor": name,
+                "source_type": kind,
+                "tp_axis": axis,
+                "shape_mnk": [m, n, k],
+                "scratch_bytes": scratch.numel() * scratch.element_size(),
+                "fused": measure(fused, args.iterations),
+                "dequant_blas": measure(blas, args.iterations),
+                "fused_repeat": measure(fused, args.iterations),
+                **errors,
+                "passed": passed,
+            }
+            report["cases"].append(row)
+            print(json.dumps(row), flush=True)
+            args.output.write_text(json.dumps(report, indent=2) + "\n")
+    raise SystemExit(0 if all(row["passed"] for row in report["cases"]) else 1)
+
+
+if __name__ == "__main__":
+    main()