Add Qwen3 component harness: torch reference, equivalence helper, microbench helper

examples/qwen3/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.

examples/qwen3/equiv.py

@@ -0,0 +1,68 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+"""Equivalence-test helper: compare ARK output against torch reference.
+
+Provides richer mismatch diagnostics (bad-element count, per-tensor stats)
+than torch.testing.assert_close, useful for debugging ARK-vs-reference
+numerical differences.
+"""
+
+import torch
+
+
+def assert_close(
+    ark_out: torch.Tensor,
+    ref_out: torch.Tensor,
+    atol: float = 1e-2,
+    rtol: float = 1e-2,
+    msg: str = "",
+) -> None:
+    """Assert that two tensors are element-wise close.
+
+    On mismatch, reports shape, max absolute error, relative error,
+    and basic statistics for both tensors.
+
+    Args:
+        ark_out: Tensor produced by the ARK implementation.
+        ref_out: Tensor produced by the torch reference.
+        atol: Absolute tolerance.
+        rtol: Relative tolerance.
+        msg: Optional context message for the assertion.
+    """
+    if ark_out.shape != ref_out.shape:
+        raise AssertionError(
+            f"Shape mismatch: ark {ark_out.shape} vs ref {ref_out.shape}. {msg}"
+        )
+
+    ark_f = ark_out.float()
+    ref_f = ref_out.float()
+
+    abs_diff = (ark_f - ref_f).abs()
+    max_abs = abs_diff.max().item()
+    ref_abs = ref_f.abs().clamp(min=1e-12)
+    max_rel = (abs_diff / ref_abs).max().item()
+
+    close = abs_diff <= (atol + rtol * ref_abs)
+    if close.all():
+        return
+
+    n_bad = (~close).sum().item()
+    n_total = close.numel()
+
+    detail = (
+        f"Tensors not close. {n_bad}/{n_total} elements exceed tolerance "
+        f"(atol={atol}, rtol={rtol}).\n"
+        f"  max |diff|      = {max_abs:.6e}\n"
+        f"  max |diff|/|ref|= {max_rel:.6e}\n"
+        f"  ark  stats: mean={ark_f.mean().item():.4e}, "
+        f"std={ark_f.std().item():.4e}, "
+        f"min={ark_f.min().item():.4e}, max={ark_f.max().item():.4e}\n"
+        f"  ref  stats: mean={ref_f.mean().item():.4e}, "
+        f"std={ref_f.std().item():.4e}, "
+        f"min={ref_f.min().item():.4e}, max={ref_f.max().item():.4e}"
+    )
+    if msg:
+        detail = f"{msg}\n{detail}"
+
+    raise AssertionError(detail)

examples/qwen3/microbench.py

@@ -0,0 +1,157 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+"""Microbenchmark helper: CUDA-graph capture, L2 flush, steady-state timing.
+
+Follows the gpu-kernel-perf-bench methodology:
+- L2 cache pollution buffer sized to 2x L2 cache.
+- CUDA-graph capture for launch-overhead elimination.
+- Pilot iteration tuning targeting 0.1-0.3 s total.
+- cuda.Event timing for all measurements (pilot, calibration, and measured runs).
+- Returns structured dict: mean_us, std_us, n_iters.
+"""
+
+from typing import Callable, Dict
+
+import torch
+
+
+def _l2_flush_buffer(device: torch.device) -> torch.Tensor:
+    """Allocate a buffer exceeding 2x typical L2 cache (128 MB covers A100's 40 MB)."""
+    nbytes = 128 * 1024 * 1024  # 128 MB
+    return torch.empty(nbytes // 4, dtype=torch.float32, device=device)
+
+
+def _flush_l2(buf: torch.Tensor) -> None:
+    """Touch the L2-flush buffer to evict cached data."""
+    buf.zero_()
+
+
+def _determine_iters(
+    fn: Callable[[], None],
+    target_secs: float = 0.2,
+    device: torch.device = None,
+) -> int:
+    """Pilot run: find iteration count for ~target_secs total execution time."""
+    if device is None:
+        device = torch.device("cuda")
+
+    # Warm up
+    fn()
+    torch.cuda.synchronize(device)
+
+    # Time a single call
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    fn()
+    end.record()
+    torch.cuda.synchronize(device)
+    single_ms = start.elapsed_time(end)
+
+    if single_ms <= 0:
+        return 100
+
+    n = max(1, int(target_secs * 1000 / single_ms))
+    return n
+
+
+def microbench(
+    fn: Callable[[], None],
+    device: torch.device = None,
+    n_iters: int = None,
+    use_cuda_graph: bool = True,
+    flush_l2: bool = True,
+) -> Dict[str, float]:
+    """Benchmark a CUDA callable and return timing statistics.
+
+    Args:
+        fn: Zero-argument callable that performs the GPU work.
+        device: CUDA device. Defaults to cuda:0.
+        n_iters: Override iteration count (else auto-tuned via pilot).
+        use_cuda_graph: Capture fn into a CUDA graph for replay.
+        flush_l2: Flush L2 cache between graph replays.
+
+    Returns:
+        Dict with keys: mean_us, std_us, n_iters.
+    """
+    if device is None:
+        device = torch.device("cuda")
+
+    # Pilot: determine iteration count
+    if n_iters is None:
+        n_iters = _determine_iters(fn, device=device)
+    n_iters = max(n_iters, 1)
+
+    flush_buf = _l2_flush_buffer(device) if flush_l2 else None
+
+    if use_cuda_graph:
+        # Warm-up run for CUDA graph capture
+        torch.cuda.synchronize(device)
+        fn()
+        torch.cuda.synchronize(device)
+
+        # Capture
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            fn()
+
+        # Determine per-graph batch to keep each replay > 1 ms
+        graph.replay()
+        torch.cuda.synchronize(device)
+        start_ev = torch.cuda.Event(enable_timing=True)
+        end_ev = torch.cuda.Event(enable_timing=True)
+        start_ev.record()
+        graph.replay()
+        end_ev.record()
+        torch.cuda.synchronize(device)
+        replay_ms = start_ev.elapsed_time(end_ev)
+
+        per_graph = max(1, int(1.0 / max(replay_ms, 1e-6)))
+        # With L2 flush, each replay must start cold.
+        if flush_l2:
+            per_graph = 1
+            n_replays = n_iters
+        else:
+            n_replays = max(1, n_iters // per_graph)
+
+        replay_fn = graph.replay
+    else:
+        per_graph = 1
+        n_replays = n_iters
+        replay_fn = fn
+
+    # Warm-up replay (not measured)
+    for _ in range(per_graph):
+        replay_fn()
+    torch.cuda.synchronize(device)
+
+    # Measured runs with cuda.Event timing
+    start_ev = torch.cuda.Event(enable_timing=True)
+    end_ev = torch.cuda.Event(enable_timing=True)
+    times_us: list[float] = []
+    for _ in range(n_replays):
+        if flush_l2 and flush_buf is not None:
+            _flush_l2(flush_buf)
+            torch.cuda.synchronize(device)
+        start_ev.record()
+        for _ in range(per_graph):
+            replay_fn()
+        end_ev.record()
+        torch.cuda.synchronize(device)
+        times_us.append(start_ev.elapsed_time(end_ev) * 1000.0)  # ms → us
+
+    mean_us = sum(times_us) / len(times_us) / per_graph
+    if len(times_us) > 1:
+        variance = sum((t / per_graph - mean_us) ** 2 for t in times_us) / (
+            len(times_us) - 1
+        )
+        std_us = variance**0.5
+    else:
+        std_us = 0.0
+
+    return {
+        "mean_us": mean_us,
+        "std_us": std_us,
+        "n_iters": total_invocations,
+    }

examples/qwen3/qwen3_config.py

@@ -0,0 +1,28 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+"""Qwen3-8B model configuration as a parameterized dataclass."""
+
+from dataclasses import dataclass
+
+
+@dataclass
+class Qwen3Config:
+    """Qwen3 model configuration with 8B defaults.
+
+    All fields are overridable. For example, a 32B variant is a one-liner:
+        Qwen3Config(n_layers=64, hidden_dim=5120, n_q_heads=40, n_kv_heads=8,
+                     intermediate_dim=15360)
+    """
+
+    n_layers: int = 36
+    hidden_dim: int = 4096
+    n_q_heads: int = 32
+    n_kv_heads: int = 8
+    head_dim: int = 128
+    intermediate_dim: int = 12288
+    vocab_size: int = 151936
+    rms_norm_eps: float = 1e-6
+    rope_theta: float = 1e6
+    max_seq_len: int = 4096
+    dtype: str = "float16"