[Kernel] Finalize the audited TP2 DFlash2 31/29 ms endpoint

benchmarks/kernels/benchmark_sm70_tp2_e4m3_scalar.py

@@ -0,0 +1,144 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Compare TP2 scalar attention over distinct KV layer working sets.
+
+This measures attention operators, not model verification rounds. Run the
+bitwise kernel tests separately before considering a model experiment.
+"""
+
+import argparse
+import hashlib
+import json
+import os
+import statistics
+from pathlib import Path
+
+import torch
+
+FLAG = "VLLM_FLASH_V100_TP2_E4M3_SCALAR_FAST"
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--context-lengths", type=int, nargs="+", default=[270, 1100])
+    parser.add_argument("--layers", type=int, default=16)
+    parser.add_argument("--json-out", type=Path, required=True)
+    args = parser.parse_args()
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0):
+        raise RuntimeError("requires an owned SM70 GPU")
+    from flash_attn_v100 import flash_attn_interface as interface
+
+    native = interface.flash_attn_v100_cuda
+    if getattr(native, "tp2_e4m3_scalar_fast_version", lambda: 0)() < 2:
+        raise RuntimeError("rebuild Flash-V100 with TP2 scalar fast revision 2")
+    if args.layers < 1 or any(not 8 <= n <= 262144 for n in args.context_lengths):
+        raise ValueError("positive layer count and context lengths 8..262144 required")
+    torch.manual_seed(20260908)
+    report = {
+        "measurement": "distinct-KV attention working set, not complete model rounds",
+        "native_sha256": hashlib.sha256(Path(native.__file__).read_bytes()).hexdigest(),
+        "gpu": torch.cuda.get_device_name(),
+        "torch": torch.__version__,
+        "cuda": torch.version.cuda,
+        "layers": args.layers,
+        "rows": [],
+    }
+    previous = os.environ.get(FLAG)
+    try:
+        for length in args.context_lengths:
+            report["rows"].append(measure(native, length, args.layers))
+    finally:
+        if previous is None:
+            os.environ.pop(FLAG, None)
+        else:
+            os.environ[FLAG] = previous
+    args.json_out.parent.mkdir(parents=True, exist_ok=True)
+    args.json_out.write_text(json.dumps(report, indent=2) + "\n")
+    print(json.dumps(report, indent=2))
+
+
+def measure(native, length, layers):
+    page, parts = 3296, 256
+    pages = (length + page - 1) // page
+    operands = []
+    for _ in range(layers):
+        kv = torch.randn((pages, 2, page, 2, 256), device="cuda", dtype=torch.float16)
+        k, v = kv.to(torch.float8_e4m3fn).view(torch.uint8).unbind(1)
+        q = torch.randn((8, 12, 256), device="cuda", dtype=torch.float16)
+        table = torch.randperm(pages, device="cuda").int()[None].repeat(8, 1)
+        seq = torch.arange(length - 7, length + 1, device="cuda").int()
+        operands.append((q, k, v, table, seq))
+    out = torch.empty_like(operands[0][0])
+    tmp = torch.empty((8, 12, parts, 256), device="cuda")
+    maxima = torch.empty((8, 12, parts), device="cuda")
+    sums = torch.empty_like(maxima)
+    active = torch.full((1,), parts, device="cuda", dtype=torch.int32)
+
+    def call(operand):
+        q, k, v, table, seq = operand
+        native.decode_paged_fwd(
+            q,
+            k,
+            v,
+            out,
+            table,
+            seq,
+            tmp,
+            maxima,
+            sums,
+            active,
+            0.0625,
+            1024,
+            parts,
+            "fp8_e4m3",
+            0.5,
+            1.25,
+            -1,
+            -1,
+            None,
+            0,
+        )
+
+    # Compare every layer before capturing a shared-workspace timing graph.
+    for operand in operands:
+        os.environ[FLAG] = "0"
+        call(operand)
+        expected = out.clone()
+        os.environ[FLAG] = "1"
+        call(operand)
+        if not torch.equal(out.view(torch.int16), expected.view(torch.int16)):
+            raise AssertionError("candidate output differs from control")
+    graphs = {}
+    for enabled in ("0", "1"):
+        os.environ[FLAG] = enabled
+        graph = torch.cuda.CUDAGraph()
+        before = native.tp2_e4m3_scalar_fast_launch_count()
+        with torch.cuda.graph(graph):
+            for operand in operands:
+                call(operand)
+        count = native.tp2_e4m3_scalar_fast_launch_count() - before
+        assert count == (layers if enabled == "1" else 0)
+        graphs[enabled] = graph
+    samples = {enabled: [] for enabled in graphs}
+    for trial in range(7):
+        for enabled in ("0", "1") if trial % 2 == 0 else ("1", "0"):
+            graph = graphs[enabled]
+            graph.replay()
+            start = torch.cuda.Event(enable_timing=True)
+            end = torch.cuda.Event(enable_timing=True)
+            start.record()
+            for _ in range(4):
+                graph.replay()
+            end.record()
+            end.synchronize()
+            samples[enabled].append(start.elapsed_time(end) / 4)
+    return {
+        "context_length": length,
+        "all_layer_outputs_bitwise_equal": True,
+        "samples_ms": samples,
+        "median_ms": {k: statistics.median(v) for k, v in samples.items()},
+    }
+
+
+if __name__ == "__main__":
+    main()

benchmarks/kernels/build_sm70_tp2_matched_qpn2.py

@@ -0,0 +1,155 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Build an isolated TP2 QPN2 candidate with TurboMind's rounding/split order.
+
+This reproduces the retained TP2 projection experiment. It does not install a
+library or enable a serving route. The caller must use the split count observed
+for the matching TurboMind projection; this is not an autotuning interface.
+"""
+
+import argparse
+import hashlib
+import json
+import shutil
+from pathlib import Path
+
+
+def replace_exact(source: str, old: str, new: str, count: int) -> str:
+    if source.count(old) != count:
+        raise ValueError(f"QPN2 source anchor changed: {old!r}")
+    return source.replace(old, new)
+
+
+def generate(source: str) -> str:
+    source = source.replace("nvfp4_qpn2_", "tp2_qpn2_matched_").replace(
+        "TORCH_LIBRARY_FRAGMENT(_qpn2_candidate,",
+        "TORCH_LIBRARY_FRAGMENT(_tp2_qpn2_matched,",
+    )
+    source = replace_exact(
+        source,
+        "  const half2 global_scale2 = __float2half2_rn(global_scale * 16384.0f);\n",
+        "",
+        2,
+    )
+    source = replace_exact(
+        source,
+        """    const half2 scale = __hmul2(
+        fp8e4m3_to_half2(__ldg(scale_ptr + static_cast<size_t>(group) * 32)),
+        global_scale2);""",
+        """    const half2 raw_scale = fp8e4m3_to_half2(
+        __ldg(scale_ptr + static_cast<size_t>(group) * 32));
+    // Match TurboMind's effective FP16 scale before E2M1 multiplication.
+    // Do not round the global factor to FP16 before multiplying the group.
+    const half effective = __float2half_rn(__low2float(raw_scale) * global_scale);
+    const half2 scale = __hmul2(__halves2half2(effective, effective),
+                                __float2half2_rn(16384.0f));""",
+        2,
+    )
+    source = replace_exact(
+        source,
+        "  const int groups_per_warp = groups_k16 / SplitK;\n"
+        "  const int group_begin = warp * groups_per_warp;",
+        """  const int chunks = k / 64;
+  const int chunks_per_warp = chunks / SplitK;
+  const int extra_begin = SplitK - chunks % SplitK;
+  const int group_begin = (warp * chunks_per_warp + max(warp - extra_begin, 0)) * 4;
+  const int groups_per_warp = (chunks_per_warp + (warp >= extra_begin)) * 4;""",
+        2,
+    )
+    source = replace_exact(
+        source,
+        "split_k == 8 || split_k == 16 || split_k == 32",
+        "(split_k >= 1 && split_k <= 16) || split_k == 32",
+        1,
+    )
+    source = replace_exact(
+        source,
+        "(input.size(1) / 16) % split_k == 0",
+        "input.size(1) / 64 >= split_k",
+        2,
+    )
+    pieces = []
+    for split in range(1, 17):
+        if split in (5, 7, 8, 9, 16):
+            continue
+        condition = "if" if not pieces else "else if"
+        pieces.append(
+            f"  {condition} (split_k == {split}) {{\n"
+            f"    VLLM_LAUNCH_QPN2(1, {split}, 1);\n  }}"
+        )
+    dispatch = (
+        "\n".join(pieces)
+        + """ else if (split_k == 5) {
+    VLLM_LAUNCH_QPN2(1, 5, 1);
+  } else if (split_k == 7) {
+    VLLM_LAUNCH_QPN2(1, 7, 1);
+  } else if (split_k == 9) {
+    VLLM_LAUNCH_QPN2(1, 9, 1);
+  } else if (native_two_tile && split_k == 8 && accumulator_chains == 1) {
+    VLLM_LAUNCH_QPN2(2, 8, 1);"""
+    )
+    source = replace_exact(
+        source,
+        "  if (native_two_tile && split_k == 8 && accumulator_chains == 1) {\n"
+        "    VLLM_LAUNCH_QPN2(2, 8, 1);",
+        dispatch,
+        1,
+    )
+    return source.replace("qpn2_matched", "qpn2_matched_all")
+
+
+def main():
+    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]
+    parent = root / "csrc/sm70_turbomind/ops/nvfp4_qpn2_sm70.cu"
+    output = args.output_dir.resolve()
+    sources = output / "sources"
+    sources.mkdir(parents=True, exist_ok=True)
+    path = sources / "qpn2-matched-all.cu"
+    path.write_text(generate(parent.read_text()))
+    shutil.copy2(parent.parent / "LICENSE.v100-skinny", sources)
+    flags = [
+        "-O3",
+        "-std=c++17",
+        "-DVLLM_NVFP4_QPN2_STANDALONE",
+        "-DVLLM_NVFP4_QPN2_BENCHMARK_CANDIDATE",
+        "-U__CUDA_NO_HALF_OPERATORS__",
+        "-U__CUDA_NO_HALF_CONVERSIONS__",
+        "-U__CUDA_NO_HALF2_OPERATORS__",
+    ]
+    report = {
+        "parent_source_sha256": hashlib.sha256(parent.read_bytes()).hexdigest(),
+        "source_sha256": hashlib.sha256(path.read_bytes()).hexdigest(),
+        "cuda_flags": flags,
+        "supported_experiment": "TP2 q8, one accumulator, observed TM split",
+        "serving_route_enabled": False,
+    }
+    if args.build:
+        from torch.utils.cpp_extension import load
+
+        build = output / "build"
+        build.mkdir(exist_ok=True)
+        library = Path(
+            load(
+                name="tp2_qpn2_matched_all",
+                sources=[str(path)],
+                build_directory=str(build),
+                extra_cflags=["-O3"],
+                extra_cuda_cflags=flags,
+                is_python_module=False,
+                verbose=True,
+            )
+        )
+        report.update(
+            library=str(library),
+            library_sha256=hashlib.sha256(library.read_bytes()).hexdigest(),
+        )
+    (output / "manifest.json").write_text(json.dumps(report, indent=2) + "\n")
+    print(json.dumps(report, indent=2), flush=True)
+
+
+if __name__ == "__main__":
+    main()