feat(kda): add bitwise-aligned SM100 CuTeDSL forward

benchmarks/bench_kda_sm100_recompute_wu.py

@@ -0,0 +1,118 @@
+# Copyright 2025-2026 Ant Group Co., Ltd.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Compare SM100 C++ and CuTe DSL recompute-WU kernels on the same inputs."""
+
+import argparse
+import pathlib
+import sys
+
+import torch
+
+sys.path.insert(0, str(pathlib.Path(__file__).parent.parent))
+
+import cula.cudac as cula_cuda
+from benchmarks.bench_recompute_wu import prepare_recompute_wu_inputs
+from benchmarks.utils import relative_rms_error_rel_max_mean_abs_rhs, triton_bench_fn
+from cula.ops.kda.sm100 import recompute_wu as recompute_wu_module
+from cula.ops.kda.sm100.recompute_wu import recompute_w_u_fwd
+
+
+def _run_cpp(k, v, beta, A, gk, cu_seqlens, chunk_indices):
+    w = torch.empty_like(v)
+    u = torch.empty_like(v)
+    kg = torch.empty_like(v)
+    cula_cuda.recompute_w_u_cuda(
+        k,
+        v,
+        beta,
+        A,
+        gk,
+        cu_seqlens,
+        chunk_indices,
+        w,
+        u,
+        kg,
+        A.shape[-1],
+        None,
+        None,
+    )
+    return w, u, None, kg
+
+
+def _max_error(ref, out):
+    stats = [relative_rms_error_rel_max_mean_abs_rhs(a, b) for a, b in zip(ref, out) if a is not None]
+    return tuple(max(values) for values in zip(*stats))
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--heads", type=int, default=64)
+    parser.add_argument("--beta-bf16", action="store_true")
+    parser.add_argument(
+        "--force-varlen",
+        action="store_true",
+        help="Keep packed uniform inputs on the varlen kernel path (diagnostic only)",
+    )
+    parser.add_argument("--lengths", type=int, nargs="+", default=[512, 1024, 4096, 8192, 16384, 32768])
+    parser.add_argument(
+        "--profile",
+        choices=("cpp", "ws"),
+        help="Warm up, then launch exactly one selected kernel between CUDA profiler markers",
+    )
+    args = parser.parse_args()
+
+    if args.force_varlen:
+        recompute_wu_module._uniform_problem = lambda _cu_seqlens: None
+
+    import benchmarks.bench_recompute_wu as common
+
+    common.H = args.heads
+    common.HV = args.heads
+    device = torch.device("cuda")
+    if args.profile:
+        if len(args.lengths) != 1:
+            parser.error("--profile requires exactly one value in --lengths")
+        T = args.lengths[0]
+        cu_seqlens = torch.tensor([0, T, 2 * T], dtype=torch.int32, device=device)
+        _q, k, v, cu_gk, beta, A, cu_seqlens, chunk_indices = prepare_recompute_wu_inputs(
+            2,
+            T,
+            device,
+            cu_seqlens=cu_seqlens,
+        )
+        if args.beta_bf16:
+            beta = beta.bfloat16()
+        runners = {
+            "cpp": lambda: _run_cpp(k, v, beta, A, cu_gk, cu_seqlens, chunk_indices),
+            "ws": lambda: recompute_w_u_fwd(k, v, beta, A, cu_gk, cu_seqlens, chunk_indices),
+        }
+        runner = runners[args.profile]
+        runner()
+        torch.cuda.synchronize()
+        torch.cuda.cudart().cudaProfilerStart()
+        runner()
+        torch.cuda.synchronize()
+        torch.cuda.cudart().cudaProfilerStop()
+        print(f"profiled {args.profile} at T={T}, H={args.heads}")
+        return
+
+    print(f"{'T':>8} {'C++ (ms)':>12} {'CuTeDSL (ms)':>14} {'C++/CuTeDSL':>14} {'rel_rmse':>12}")
+    for T in args.lengths:
+        cu_seqlens = torch.tensor([0, T, 2 * T], dtype=torch.int32, device=device)
+        _q, k, v, cu_gk, beta, A, cu_seqlens, chunk_indices = prepare_recompute_wu_inputs(2, T, device, cu_seqlens=cu_seqlens)
+        if args.beta_bf16:
+            beta = beta.bfloat16()
+        cpp = _run_cpp(k, v, beta, A, cu_gk, cu_seqlens, chunk_indices)
+        ws = recompute_w_u_fwd(k, v, beta, A, cu_gk, cu_seqlens, chunk_indices)
+        ws_err = _max_error(cpp, ws)[0]
+        if not torch.isfinite(torch.tensor(ws_err)):
+            raise AssertionError(f"non-finite error at T={T}: ws={ws_err}")
+
+        cpp_ms = triton_bench_fn(lambda: _run_cpp(k, v, beta, A, cu_gk, cu_seqlens, chunk_indices))
+        ws_ms = triton_bench_fn(lambda: recompute_w_u_fwd(k, v, beta, A, cu_gk, cu_seqlens, chunk_indices))
+        print(f"{T:8d} {cpp_ms:12.4f} {ws_ms:14.4f} {cpp_ms / ws_ms:14.3f} {ws_err:12.6g}")
+
+
+if __name__ == "__main__":
+    main()

benchmarks/stress_kda_sm100_csrc_boundary_determinism.py

@@ -0,0 +1,174 @@
+#!/usr/bin/env python3
+# Copyright 2025-2026 Ant Group Co., Ltd.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Stress the bitwise-aligned SM100 CuTeDSL KDA forward path.
+
+The script first requires complete tensor equality with the csrc boundary for
+Aqk, Akk, KG, W, and U. It then captures the CuTeDSL intra, Akk inverse,
+recompute-WU, and an exact comparison against those csrc outputs in one CUDA
+graph. Every replay therefore validates every output element.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import pathlib
+import sys
+import time
+
+sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent))
+
+import torch
+import torch.nn.functional as F
+
+from cula.kda.chunk_intra import chunk_kda_fwd_intra as csrc_chunk_kda_fwd_intra
+from cula.ops.kda.sm100.intra_fused import BT, K_DIM, chunk_kda_fwd_intra_sm100_from_gk
+from cula.ops.kda.sm100.recompute_wu import recompute_w_u_fwd
+
+OUTPUT_NAMES = ("W", "U", "KG", "Aqk", "Akk")
+
+
+def _bitwise_stats(
+    outputs: tuple[torch.Tensor, ...], references: tuple[torch.Tensor, ...]
+) -> dict[str, dict[str, float | int | bool]]:
+    stats = {}
+    for name, output, reference in zip(OUTPUT_NAMES, outputs, references, strict=True):
+        mismatch_count = torch.count_nonzero(output != reference).item()
+        max_abs = (output.float() - reference.float()).abs().max().item()
+        stats[name] = {
+            "equal": mismatch_count == 0,
+            "mismatches": mismatch_count,
+            "max_abs": max_abs,
+        }
+    return stats
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--iterations", type=int, default=10_000_000)
+    parser.add_argument("--checkpoint", type=int, default=1_000_000)
+    parser.add_argument("--batch", type=int, default=1)
+    parser.add_argument("--seqlen", type=int, default=256)
+    parser.add_argument("--heads", type=int, default=4)
+    parser.add_argument("--beta-dtype", choices=("bfloat16", "float32"), default="bfloat16")
+    parser.add_argument("--seed", type=int, default=20260825)
+    parser.add_argument("--report-json")
+    args = parser.parse_args()
+
+    if args.iterations <= 0 or args.checkpoint <= 0:
+        raise ValueError("--iterations and --checkpoint must be positive")
+    if args.seqlen <= 0 or args.seqlen % (4 * BT) != 0:
+        raise ValueError(f"--seqlen must be a positive multiple of {4 * BT}")
+    if args.batch <= 0 or args.heads <= 0:
+        raise ValueError("--batch and --heads must be positive")
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0):
+        raise RuntimeError("an SM100 CUDA device is required")
+
+    torch.manual_seed(args.seed)
+    device = torch.device("cuda")
+    shape = (args.batch, args.seqlen, args.heads, K_DIM)
+    q = F.normalize(torch.randn(*shape, device=device).float(), dim=-1).bfloat16()
+    k = F.normalize(torch.randn(*shape, device=device).float(), dim=-1).bfloat16()
+    gk = torch.randn(*shape, device=device, dtype=torch.float32) * 0.02
+    beta_dtype = torch.bfloat16 if args.beta_dtype == "bfloat16" else torch.float32
+    beta = torch.randn(*shape[:-1], device=device).sigmoid().to(beta_dtype)
+    scale = K_DIM**-0.5
+
+    def run_cutedsl() -> tuple[torch.Tensor, ...]:
+        aqk, akk = chunk_kda_fwd_intra_sm100_from_gk(
+            q=q,
+            k=k,
+            gk=gk,
+            beta=beta,
+            scale=scale,
+            fp32_akk_inv=True,
+        )
+        w, u, _, kg = recompute_w_u_fwd(k, k, beta, akk, gk)
+        return w, u, kg, aqk, akk
+
+    print(
+        f"device={torch.cuda.get_device_name(0)} shape={shape} beta_dtype={args.beta_dtype} iterations={args.iterations}",
+        flush=True,
+    )
+
+    w_ref, u_ref, _, kg_ref, aqk_ref, akk_ref = csrc_chunk_kda_fwd_intra(
+        q=q,
+        k=k,
+        v=k,
+        gk=gk,
+        beta=beta,
+        scale=scale,
+        chunk_size=BT,
+        safe_gate=True,
+    )
+    references = (w_ref, u_ref, kg_ref, aqk_ref, akk_ref)
+
+    outputs = run_cutedsl()
+    torch.cuda.synchronize()
+    bitwise = _bitwise_stats(outputs, references)
+    print("BITWISE_JSON=" + json.dumps(bitwise, sort_keys=True), flush=True)
+    if not all(values["equal"] for values in bitwise.values()):
+        raise AssertionError("CuTeDSL output is not bitwise equal to csrc")
+    if not all(torch.isfinite(output).all().item() for output in outputs):
+        raise AssertionError("CuTeDSL output contains NaN or Inf")
+
+    # Warm compilation, allocator, and concatenation paths before capture.
+    for _ in range(3):
+        warm = run_cutedsl()
+        torch.cat([tensor.reshape(-1) for tensor in warm])
+    torch.cuda.synchronize()
+
+    reference_flat = torch.cat([tensor.reshape(-1) for tensor in references])
+    mismatch_count = torch.zeros((), dtype=torch.int64, device=device)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = run_cutedsl()
+        captured_flat = torch.cat([tensor.reshape(-1) for tensor in captured])
+        mismatch_count.add_(torch.count_nonzero(captured_flat != reference_flat))
+
+    torch.cuda.synchronize()
+    if mismatch_count.item() != 0:
+        raise AssertionError(f"graph capture differed from csrc: mismatches={mismatch_count.item()}")
+
+    started = time.perf_counter()
+    completed = 0
+    while completed < args.iterations:
+        stop = min(completed + args.checkpoint, args.iterations)
+        for _ in range(completed, stop):
+            graph.replay()
+        torch.cuda.synchronize()
+        completed = stop
+        mismatches = mismatch_count.item()
+        elapsed = time.perf_counter() - started
+        print(
+            f"progress={completed}/{args.iterations} mismatches={mismatches} "
+            f"elapsed_s={elapsed:.3f} iterations_per_s={completed / elapsed:.1f}",
+            flush=True,
+        )
+        if mismatches != 0:
+            raise AssertionError(f"non-deterministic output after {completed} iterations: mismatches={mismatches}")
+
+    elapsed = time.perf_counter() - started
+    report = {
+        "status": "passed",
+        "device": torch.cuda.get_device_name(0),
+        "device_index_visible": torch.cuda.current_device(),
+        "shape": shape,
+        "beta_dtype": args.beta_dtype,
+        "iterations": args.iterations,
+        "mismatches": mismatch_count.item(),
+        "elapsed_seconds": elapsed,
+        "iterations_per_second": args.iterations / elapsed,
+        "bitwise_csrc": bitwise,
+    }
+    print("RESULT_JSON=" + json.dumps(report, sort_keys=True), flush=True)
+    if args.report_json:
+        with open(args.report_json, "w", encoding="utf-8") as handle:
+            json.dump(report, handle, indent=2, sort_keys=True)
+            handle.write("\n")
+
+
+if __name__ == "__main__":
+    main()