[KernelSwift 算子创新大赛][T1, T2, T3, T4]tuolajilatuo

benchmarks/_benchmark.py

@@ -0,0 +1,47 @@
+"""Shared Triton benchmark and autotune reporting helpers."""
+
+from collections.abc import Mapping
+
+import torch
+from triton.testing import do_bench
+
+
+def benchmark_mean(function):
+    """Return mean latency in microseconds using Triton's benchmark policy."""
+    with torch.inference_mode():
+        return (
+            do_bench(
+                function,
+                warmup=25,
+                rep=100,
+                return_mode="mean",
+            )
+            * 1000.0
+        )
+
+
+def selected_config(kernel, fallback=None):
+    """Return the most recently cached ``(num_warps, num_stages)`` pair."""
+    # The SSA-first NineToothed compiler records the selected candidate on
+    # the public handle instead of exposing Triton's Autotuner object.
+    candidate = getattr(kernel, "_selected_tuning_candidate", None)
+    if candidate is not None:
+        if isinstance(candidate, Mapping):
+            return int(candidate["num_warps"]), int(candidate["num_stages"])
+        return int(candidate.num_warps), int(candidate.num_stages)
+
+    # Keep compatibility with the legacy frontend, which emits a decorated
+    # Triton kernel and stores the winner in the decorator's cache.
+    globals_ = getattr(getattr(kernel, "_kernel", None), "__globals__", {})
+    autotuners = (
+        value
+        for name, value in globals_.items()
+        if name.endswith("_with_auto_tuning") and hasattr(value, "cache")
+    )
+    autotuner = next(autotuners, None)
+    if autotuner is None or not autotuner.cache:
+        if fallback is not None:
+            return tuple(int(value) for value in fallback)
+        raise RuntimeError("the kernel has not completed autotuning")
+    config = next(reversed(autotuner.cache.values()))
+    return int(config.num_warps), int(config.num_stages)

benchmarks/bench_block_scaled_fp8_mm.py

@@ -0,0 +1,139 @@
+"""Benchmark BlockWise1x128 x BlockWise128x128 FP8 projection shapes."""
+
+import importlib
+
+import torch
+import torch.nn.functional as F
+from _benchmark import benchmark_mean, selected_config
+
+import ntops
+
+_block_scaled_fp8_mm_module = importlib.import_module("ntops.torch.block_scaled_fp8_mm")
+
+
+def _column_major(value):
+    return value.t().contiguous().t()
+
+
+def _make_inputs(m, n, k):
+    mat_a = torch.randn((m, k), device="cuda").clamp(-3, 3).to(torch.float8_e4m3fn)
+    weight = torch.randn((n, k), device="cuda").clamp(-3, 3).to(torch.float8_e4m3fn)
+    k_blocks = k // 128
+    padded_k_blocks = ((k_blocks + 3) // 4) * 4
+    scale_a = _column_major(
+        torch.ones((m, k_blocks), device="cuda", dtype=torch.float32)
+    )
+    scale_b = _column_major(
+        torch.ones(
+            (padded_k_blocks, n // 128),
+            device="cuda",
+            dtype=torch.float32,
+        )
+    )
+    return mat_a, weight.t(), scale_a, scale_b
+
+
+def _torch_dtype_reference(mat_a, mat_b, scale_a, scale_b):
+    """Dequantize through PyTorch's FP8 dtype conversion, then run FP32 MM."""
+    m, k = mat_a.shape
+    n = mat_b.shape[1]
+    k_blocks = k // 128
+
+    dequant_a = (
+        mat_a.float().reshape(m, k_blocks, 128) * scale_a.float().unsqueeze(-1)
+    ).reshape(m, k)
+    expanded_scale_b = scale_b[:k_blocks].float().repeat_interleave(128, dim=1)
+    dequant_b = (
+        mat_b.float().reshape(k_blocks, 128, n) * expanded_scale_b.unsqueeze(1)
+    ).reshape(k, n)
+    return (dequant_a @ dequant_b).to(torch.bfloat16)
+
+
+def _select_reference(mat_a, mat_b, scale_a, scale_b):
+    native = getattr(F, "scaled_mm", None)
+    if callable(native):
+
+        def run_native():
+            return native(
+                mat_a,
+                mat_b,
+                scale_a,
+                ntops.torch.ScalingType.BlockWise1x128,
+                scale_b,
+                ntops.torch.ScalingType.BlockWise128x128,
+            )
+
+        try:
+            return "torch.nn.functional.scaled_mm", run_native, run_native()
+        except (NotImplementedError, RuntimeError, TypeError, ValueError):
+            pass
+
+    def run_dtype_reference():
+        return _torch_dtype_reference(mat_a, mat_b, scale_a, scale_b)
+
+    return "torch_dtype_dequant_mm", run_dtype_reference, run_dtype_reference()
+
+
+def main():
+    if not torch.cuda.is_available():
+        raise RuntimeError("a CUDA-compatible accelerator is required")
+
+    print(f"torch={torch.__version__}")
+    print(f"hip={torch.version.hip}")
+    print(f"device={torch.cuda.get_device_name()}")
+    print(
+        "scenario,M,N,K,reference_provider,reference_mean_us,ntops_mean_us,"
+        "best_num_warps,best_num_stages,speedup_vs_reference,tflops,"
+        "max_abs_error"
+    )
+    scenarios = (
+        ("attention_decode", 1, 4096, 4096),
+        ("moe_expert", 32, 14336, 4096),
+        ("linear_prefill", 128, 4096, 4096),
+    )
+    for name, m, n, k in scenarios:
+        mat_a, mat_b, scale_a, scale_b = _make_inputs(m, n, k)
+
+        def run_ntops():
+            return ntops.torch.block_scaled_fp8_mm(
+                mat_a,
+                mat_b,
+                scale_a,
+                ntops.torch.ScalingType.BlockWise1x128,
+                scale_b,
+                ntops.torch.ScalingType.BlockWise128x128,
+            )
+
+        with torch.inference_mode():
+            output = run_ntops()
+            reference_provider, run_reference, reference = _select_reference(
+                mat_a, mat_b, scale_a, scale_b
+            )
+            torch.testing.assert_close(output, reference, rtol=0.03, atol=0.03)
+            max_abs_error = (output.float() - reference.float()).abs().max().item()
+            torch.cuda.synchronize()
+
+        kernel = _block_scaled_fp8_mm_module._make_kernel(
+            mat_a.dtype,
+            torch.bfloat16,
+            None,
+            mat_a.device,
+        )
+        _, fixed_warps, fixed_stages, _ = (
+            _block_scaled_fp8_mm_module._kernel_tuning_config(mat_a.device)
+        )
+        fallback_warps = fixed_warps if isinstance(fixed_warps, int) else fixed_warps[0]
+        warps, stages = selected_config(kernel, fallback=(fallback_warps, fixed_stages))
+
+        reference_us = benchmark_mean(run_reference)
+        ntops_us = benchmark_mean(run_ntops)
+        tflops = 2.0 * m * n * k / (ntops_us * 1.0e6)
+        print(
+            f"{name},{m},{n},{k},{reference_provider},{reference_us:.3f},"
+            f"{ntops_us:.3f},{warps},{stages},{reference_us / ntops_us:.3f},"
+            f"{tflops:.3f},{max_abs_error:.6f}"
+        )
+
+
+if __name__ == "__main__":
+    main()

benchmarks/bench_fused_mla_rope_cache_write.py

@@ -0,0 +1,151 @@
+"""Benchmark MLA RoPE + compressed KV-cache write fusion."""
+
+import torch
+from _benchmark import benchmark_mean, selected_config
+
+import ntops
+from ntops.torch.utils import _cached_make
+
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+NUM_HEADS = 128  # documented MLA geometry; cache writer itself is head-shared
+KV_LORA_RANK = 512
+ROPE_DIM = 64
+CACHE_BLOCK_SIZE = 16
+TILE_SIZE = 128
+AUTOTUNE_WARPS = (1, 2, 4, 8)
+AUTOTUNE_STAGES = (1, 2)
+MAX_NUM_CONFIGS = 8
+
+
+def _make_inputs(tokens):
+    kv_c = torch.randn(tokens, KV_LORA_RANK, device=DEVICE, dtype=DTYPE)
+    k_pe = torch.randn(tokens, ROPE_DIM, device=DEVICE, dtype=DTYPE)
+    num_blocks = (tokens + CACHE_BLOCK_SIZE - 1) // CACHE_BLOCK_SIZE + 4
+    kv_cache = torch.empty(
+        num_blocks,
+        CACHE_BLOCK_SIZE,
+        KV_LORA_RANK + ROPE_DIM,
+        device=DEVICE,
+        dtype=DTYPE,
+    )
+    slot_mapping = torch.arange(tokens, device=DEVICE, dtype=torch.int64)
+    positions = torch.arange(tokens, device=DEVICE, dtype=torch.int64)
+    half = ROPE_DIM // 2
+    theta = 10000 ** (-2 * torch.arange(half, device=DEVICE) / ROPE_DIM)
+    phase = positions.to(torch.float32)[:, None] * theta[None, :]
+    cos_sin_cache = torch.cat((phase.cos(), phase.sin()), dim=-1)
+    return kv_c, k_pe, kv_cache, slot_mapping, positions, cos_sin_cache
+
+
+def _torch_rope(k_pe, positions, cos_sin_cache):
+    half = ROPE_DIM // 2
+    table = cos_sin_cache.index_select(0, positions)
+    cos = table[:, :half].to(torch.float32)
+    sin = table[:, half:].to(torch.float32)
+    x0 = k_pe[:, 0::2].to(torch.float32)
+    x1 = k_pe[:, 1::2].to(torch.float32)
+    output = torch.empty_like(k_pe, dtype=torch.float32)
+    output[:, 0::2] = x0 * cos - x1 * sin
+    output[:, 1::2] = x0 * sin + x1 * cos
+    return output.to(k_pe.dtype)
+
+
+def _torch_unfused(inputs):
+    kv_c, k_pe, cache, slots, positions, table = inputs
+    k_rot = _torch_rope(k_pe, positions, table)
+    block = slots // CACHE_BLOCK_SIZE
+    offset = slots % CACHE_BLOCK_SIZE
+    cache[block, offset, :KV_LORA_RANK] = kv_c
+    cache[block, offset, KV_LORA_RANK:] = k_rot
+
+
+def _torch_rope_only(inputs):
+    _, k_pe, _, slots, positions, table = inputs
+    del slots
+    return _torch_rope(k_pe, positions, table)
+
+
+def _torch_cache_only(inputs, k_rot):
+    kv_c, _, cache, slots, _, _ = inputs
+    block = slots // CACHE_BLOCK_SIZE
+    offset = slots % CACHE_BLOCK_SIZE
+    cache[block, offset, :KV_LORA_RANK] = kv_c
+    cache[block, offset, KV_LORA_RANK:] = k_rot
+
+
+def _kernel_handle(inputs):
+    entry_dim = KV_LORA_RANK + ROPE_DIM
+    tile_size = max(
+        1 << (TILE_SIZE - 1).bit_length(),
+        1 << (entry_dim - 1).bit_length(),
+    )
+    return _cached_make(
+        ntops.kernels.fused_mla_rope_cache_write.premake,
+        KV_LORA_RANK,
+        ROPE_DIM,
+        dtype=DTYPE,
+        block_size=tile_size,
+        cache_block_size=CACHE_BLOCK_SIZE,
+        cos_dtype=inputs[-1].dtype,
+        num_warps=AUTOTUNE_WARPS,
+        num_stages=AUTOTUNE_STAGES,
+        max_num_configs=MAX_NUM_CONFIGS,
+    )
+
+
+def main():
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is required for this benchmark")
+    print(f"torch={torch.__version__}")
+    print(f"device={torch.cuda.get_device_name()}")
+    print(
+        "scenario,tokens,heads,latent,rope,torch_rope_mean_us,"
+        "torch_cache_mean_us,torch_unfused_mean_us,best_num_warps,"
+        "best_num_stages,autotuned_mean_us,"
+        "speedup_vs_pytorch"
+    )
+
+    for scenario, tokens in (
+        ("decode", 1),
+        ("concurrent_decode", 10),
+        ("long_context", 2048),
+    ):
+        inputs = _make_inputs(tokens)
+
+        def rope_only():
+            return _torch_rope_only(inputs)
+
+        rotated = _torch_rope_only(inputs)
+
+        def cache_only():
+            return _torch_cache_only(inputs, rotated)
+
+        def unfused():
+            return _torch_unfused(inputs)
+
+        def autotuned():
+            return ntops.torch.fused_mla_rope_cache_write(
+                *inputs,
+                block_size=TILE_SIZE,
+                num_warps=AUTOTUNE_WARPS,
+                num_stages=AUTOTUNE_STAGES,
+                max_num_configs=MAX_NUM_CONFIGS,
+            )
+
+        autotuned()
+        torch.cuda.synchronize()
+        warps, stages = selected_config(_kernel_handle(inputs))
+        rope_us = benchmark_mean(rope_only)
+        cache_us = benchmark_mean(cache_only)
+        unfused_us = benchmark_mean(unfused)
+        auto_us = benchmark_mean(autotuned)
+        print(
+            f"{scenario},{tokens},{NUM_HEADS},{KV_LORA_RANK},{ROPE_DIM},"
+            f"{rope_us:.3f},{cache_us:.3f},{unfused_us:.3f},{warps},{stages},"
+            f"{auto_us:.3f},{unfused_us / auto_us:.3f}"
+        )
+
+
+if __name__ == "__main__":
+    main()