diff --git a/examples/qwen3/__init__.py b/examples/qwen3/__init__.py new file mode 100644 index 00000000..9a045456 --- /dev/null +++ b/examples/qwen3/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. diff --git a/examples/qwen3/equiv.py b/examples/qwen3/equiv.py new file mode 100644 index 00000000..863e4ef5 --- /dev/null +++ b/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) diff --git a/examples/qwen3/microbench.py b/examples/qwen3/microbench.py new file mode 100644 index 00000000..76bbfaf5 --- /dev/null +++ b/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, + } diff --git a/examples/qwen3/qwen3_config.py b/examples/qwen3/qwen3_config.py new file mode 100644 index 00000000..de601b61 --- /dev/null +++ b/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" diff --git a/examples/qwen3/qwen3_ref.py b/examples/qwen3/qwen3_ref.py new file mode 100644 index 00000000..4a0aa31d --- /dev/null +++ b/examples/qwen3/qwen3_ref.py @@ -0,0 +1,255 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Pure-torch Qwen3-8B reference model (random weights, fixed seed, fp16). + +Implements: RMSNorm, RoPE, QK-norm, GQA attention, SwiGLU MLP, +TransformerBlock, and Qwen3Model. No ARK dependency. +""" + +import math +from typing import Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .qwen3_config import Qwen3Config + +_DEFAULT_SEED = 42 + + +def _get_dtype(cfg: Qwen3Config) -> torch.dtype: + dt = getattr(torch, cfg.dtype, None) + if not isinstance(dt, torch.dtype): + raise ValueError(f"Invalid dtype string in config: {cfg.dtype!r}") + return dt + + +class RMSNorm(nn.Module): + """Root-mean-square layer normalization.""" + + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + rms = torch.sqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + self.eps) + x_normed = x.float() / rms + return (x_normed * self.weight.float()).to(x.dtype) + + +def precompute_rope_freqs( + head_dim: int, max_seq_len: int, theta: float = 1e6 +) -> torch.Tensor: + """Precompute complex RoPE frequencies of shape (max_seq_len, head_dim//2).""" + freqs = 1.0 / ( + theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim) + ) + t = torch.arange(max_seq_len, dtype=torch.float32) + angles = torch.outer(t, freqs) # (seq, head_dim//2) + return torch.polar(torch.ones_like(angles), angles) # complex64 + + +def apply_rope(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + """Apply rotary position embeddings. + + Args: + x: (batch, n_heads, seq, head_dim) — real tensor. + freqs: (seq, head_dim//2) — complex tensor. + """ + # Reshape to pairs and view as complex + batch, n_heads, seq, hd = x.shape + x_complex = torch.view_as_complex( + x.float().reshape(batch, n_heads, seq, hd // 2, 2) + ) + freqs = freqs[:seq].unsqueeze(0).unsqueeze(0) # (1, 1, seq, hd//2) + x_rotated = torch.view_as_real(x_complex * freqs) + return x_rotated.reshape(batch, n_heads, seq, hd).to(x.dtype) + + +class QKNorm(nn.Module): + """Per-head RMS normalization applied to Q and K projections.""" + + def __init__(self, head_dim: int, eps: float = 1e-6): + super().__init__() + self.q_norm = RMSNorm(head_dim, eps=eps) + self.k_norm = RMSNorm(head_dim, eps=eps) + + def forward( + self, q: torch.Tensor, k: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Normalize Q and K. Inputs: (batch, n_heads, seq, head_dim).""" + orig_shape = q.shape + q = self.q_norm(q.reshape(-1, orig_shape[-1])).reshape(orig_shape) + k = self.k_norm(k.reshape(-1, orig_shape[-1])).reshape(k.shape) + return q, k + + +class GQAAttention(nn.Module): + """Grouped-query attention with QK-norm and RoPE.""" + + def __init__(self, cfg: Qwen3Config): + super().__init__() + self.n_q_heads = cfg.n_q_heads + self.n_kv_heads = cfg.n_kv_heads + self.head_dim = cfg.head_dim + self.n_rep = self.n_q_heads // self.n_kv_heads + + dtype = _get_dtype(cfg) + self.q_proj = nn.Linear( + cfg.hidden_dim, cfg.n_q_heads * cfg.head_dim, bias=False + ).to(dtype) + self.k_proj = nn.Linear( + cfg.hidden_dim, cfg.n_kv_heads * cfg.head_dim, bias=False + ).to(dtype) + self.v_proj = nn.Linear( + cfg.hidden_dim, cfg.n_kv_heads * cfg.head_dim, bias=False + ).to(dtype) + self.o_proj = nn.Linear( + cfg.n_q_heads * cfg.head_dim, cfg.hidden_dim, bias=False + ).to(dtype) + self.qk_norm = QKNorm(cfg.head_dim, eps=cfg.rms_norm_eps) + + def forward( + self, + x: torch.Tensor, + rope_freqs: torch.Tensor, + mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + batch, seq, _ = x.shape + + q = self.q_proj(x).reshape(batch, seq, self.n_q_heads, self.head_dim) + k = self.k_proj(x).reshape(batch, seq, self.n_kv_heads, self.head_dim) + v = self.v_proj(x).reshape(batch, seq, self.n_kv_heads, self.head_dim) + + # (batch, heads, seq, head_dim) + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + + # QK-norm before RoPE + q, k = self.qk_norm(q, k) + + # RoPE + q = apply_rope(q, rope_freqs) + k = apply_rope(k, rope_freqs) + + # Expand KV heads for GQA + if self.n_rep > 1: + k = k.repeat_interleave(self.n_rep, dim=1) + v = v.repeat_interleave(self.n_rep, dim=1) + + # Scaled dot-product attention + scale = 1.0 / math.sqrt(self.head_dim) + attn_weights = torch.matmul(q, k.transpose(-2, -1)) * scale + + if mask is not None: + attn_weights = attn_weights + mask + + attn_weights = F.softmax(attn_weights.float(), dim=-1).to(x.dtype) + out = torch.matmul(attn_weights, v) + + out = out.transpose(1, 2).reshape(batch, seq, -1) + return self.o_proj(out) + + +class SwiGLUMLP(nn.Module): + """SwiGLU feed-forward network.""" + + def __init__(self, cfg: Qwen3Config): + super().__init__() + dtype = _get_dtype(cfg) + self.gate_proj = nn.Linear( + cfg.hidden_dim, cfg.intermediate_dim, bias=False + ).to(dtype) + self.up_proj = nn.Linear( + cfg.hidden_dim, cfg.intermediate_dim, bias=False + ).to(dtype) + self.down_proj = nn.Linear( + cfg.intermediate_dim, cfg.hidden_dim, bias=False + ).to(dtype) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) + + +class TransformerBlock(nn.Module): + """Single Qwen3 transformer block: pre-norm attention + pre-norm MLP.""" + + def __init__(self, cfg: Qwen3Config): + super().__init__() + self.attn_norm = RMSNorm(cfg.hidden_dim, eps=cfg.rms_norm_eps) + self.attn = GQAAttention(cfg) + self.mlp_norm = RMSNorm(cfg.hidden_dim, eps=cfg.rms_norm_eps) + self.mlp = SwiGLUMLP(cfg) + + def forward( + self, + x: torch.Tensor, + rope_freqs: torch.Tensor, + mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + x = x + self.attn(self.attn_norm(x), rope_freqs, mask) + x = x + self.mlp(self.mlp_norm(x)) + return x + + +class Qwen3Model(nn.Module): + """Qwen3 causal language model (random weights, no real checkpoint). + + Uses a fixed seed for reproducible random initialization. + """ + + def __init__(self, cfg: Qwen3Config, seed: int = _DEFAULT_SEED): + super().__init__() + self.cfg = cfg + dtype = _get_dtype(cfg) + + torch.manual_seed(seed) + + self.embed = nn.Embedding(cfg.vocab_size, cfg.hidden_dim).to(dtype) + self.layers = nn.ModuleList( + [TransformerBlock(cfg) for _ in range(cfg.n_layers)] + ) + self.final_norm = RMSNorm(cfg.hidden_dim, eps=cfg.rms_norm_eps) + self.lm_head = nn.Linear(cfg.hidden_dim, cfg.vocab_size, bias=False).to( + dtype + ) + + # Precompute RoPE frequencies (cpu, moved to device in forward) + self.register_buffer( + "rope_freqs", + precompute_rope_freqs( + cfg.head_dim, cfg.max_seq_len, cfg.rope_theta + ), + persistent=False, + ) + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + """Forward pass. + + Args: + input_ids: (batch, seq) integer token IDs. + + Returns: + Logits tensor of shape (batch, seq, vocab_size). + """ + batch, seq = input_ids.shape + x = self.embed(input_ids) + + # Causal mask: upper-triangular -inf + mask = torch.full( + (seq, seq), float("-inf"), device=x.device, dtype=x.dtype + ) + mask = torch.triu(mask, diagonal=1) + mask = mask.unsqueeze(0).unsqueeze(0) # (1, 1, seq, seq) + + rope_freqs = self.rope_freqs.to(x.device) + + for layer in self.layers: + x = layer(x, rope_freqs, mask) + + x = self.final_norm(x) + return self.lm_head(x) diff --git a/examples/qwen3/test_harness.py b/examples/qwen3/test_harness.py new file mode 100644 index 00000000..34e078f9 --- /dev/null +++ b/examples/qwen3/test_harness.py @@ -0,0 +1,255 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Unit tests for the Qwen3 component harness. + +All GPU tests are skipped when CUDA is unavailable. +""" + +import pytest +import torch + +_CUDA = torch.cuda.is_available() +requires_cuda = pytest.mark.skipif(not _CUDA, reason="CUDA not available") + + +# --------------------------------------------------------------------------- +# Reference model tests +# --------------------------------------------------------------------------- + + +@requires_cuda +def test_ref_forward_shape(): + """Reference model produces logits of correct shape (1 layer, seq=128).""" + from .qwen3_config import Qwen3Config + from .qwen3_ref import Qwen3Model + + cfg = Qwen3Config(n_layers=1) + model = Qwen3Model(cfg, seed=42).cuda().half() + + batch, seq = 1, 128 + input_ids = torch.randint(0, cfg.vocab_size, (batch, seq), device="cuda") + + with torch.no_grad(): + logits = model(input_ids) + + assert logits.shape == ( + batch, + seq, + cfg.vocab_size, + ), f"Expected shape ({batch}, {seq}, {cfg.vocab_size}), got {logits.shape}" + assert logits.dtype == torch.float16 + + +@requires_cuda +def test_rmsnorm_unit_rms(): + """RMSNorm output has approximately unit RMS (within eps tolerance).""" + from .qwen3_ref import RMSNorm + + dim = 128 + eps = 1e-6 + norm = RMSNorm(dim, eps=eps).cuda() + + x = torch.randn(2, 64, dim, device="cuda", dtype=torch.float16) + with torch.no_grad(): + y = norm(x) + + # RMS of each vector should be close to 1 (since weight is ones) + rms = y.float().pow(2).mean(dim=-1).sqrt() + torch.testing.assert_close( + rms, + torch.ones_like(rms), + atol=0.05, + rtol=0.05, + msg="RMSNorm output should have approximately unit RMS", + ) + + +@requires_cuda +def test_attention_is_causal(): + """Causal attention zeroes future positions in attention weights. + + Verifies that the upper-triangular portion of the attention-weight + matrix is zero (future tokens cannot attend to past). + """ + import math + from .qwen3_config import Qwen3Config + from .qwen3_ref import GQAAttention, precompute_rope_freqs + + cfg = Qwen3Config(n_layers=1, n_q_heads=4, n_kv_heads=2, head_dim=32) + attn = GQAAttention(cfg).cuda().half() + rope_freqs = precompute_rope_freqs( + cfg.head_dim, cfg.max_seq_len, cfg.rope_theta + ).to("cuda") + + batch, seq = 1, 16 + x = torch.randn( + batch, seq, cfg.hidden_dim, device="cuda", dtype=torch.float16 + ) + + # Build causal mask + mask = torch.full( + (seq, seq), float("-inf"), device="cuda", dtype=torch.float16 + ) + mask = torch.triu(mask, diagonal=1).unsqueeze(0).unsqueeze(0) + + with torch.no_grad(): + # Manually compute attention weights to inspect them + q = ( + attn.q_proj(x) + .reshape(batch, seq, cfg.n_q_heads, cfg.head_dim) + .transpose(1, 2) + ) + k = ( + attn.k_proj(x) + .reshape(batch, seq, cfg.n_kv_heads, cfg.head_dim) + .transpose(1, 2) + ) + + from .qwen3_ref import apply_rope + + q, k = attn.qk_norm(q, k) + q = apply_rope(q, rope_freqs) + k = apply_rope(k, rope_freqs) + + if cfg.n_q_heads // cfg.n_kv_heads > 1: + k = k.repeat_interleave(cfg.n_q_heads // cfg.n_kv_heads, dim=1) + + scale = 1.0 / math.sqrt(cfg.head_dim) + scores = torch.matmul(q, k.transpose(-2, -1)) * scale + mask + weights = torch.softmax(scores.float(), dim=-1) + + # Upper triangle (future positions) must be zero + for h in range(weights.shape[1]): + upper = torch.triu(weights[0, h], diagonal=1) + assert upper.abs().max().item() < 1e-6, ( + f"Head {h}: future-position attention weight is non-zero " + f"(max={upper.abs().max().item():.2e})" + ) + + +@requires_cuda +def test_rope_applied_per_head(): + """RoPE modifies Q/K values (not a no-op).""" + from .qwen3_ref import apply_rope, precompute_rope_freqs + + head_dim = 64 + freqs = precompute_rope_freqs(head_dim, 256, theta=1e6).to("cuda") + x = torch.randn(1, 4, 32, head_dim, device="cuda", dtype=torch.float16) + + y = apply_rope(x, freqs) + + # RoPE should change the tensor + assert not torch.allclose( + x, y, atol=1e-4 + ), "RoPE had no effect on the tensor" + # Shape must be preserved + assert x.shape == y.shape + + +# --------------------------------------------------------------------------- +# Equivalence helper tests +# --------------------------------------------------------------------------- + + +def test_equiv_pass_identical(): + """assert_close passes for identical tensors.""" + from .equiv import assert_close + + t = torch.randn(4, 8, device="cpu", dtype=torch.float16) + assert_close(t, t.clone()) # should not raise + + +def test_equiv_fail_perturbed(): + """assert_close raises AssertionError on intentional mismatch.""" + from .equiv import assert_close + + t = torch.randn(4, 8, device="cpu", dtype=torch.float16) + perturbed = t + 10.0 # large perturbation + + with pytest.raises(AssertionError, match="not close"): + assert_close(perturbed, t, atol=1e-3, rtol=1e-3) + + +def test_get_dtype_invalid(): + """_get_dtype raises ValueError for invalid dtype strings.""" + from .qwen3_config import Qwen3Config + from .qwen3_ref import _get_dtype + + cfg = Qwen3Config(dtype="not_a_dtype") + with pytest.raises(ValueError, match="Invalid dtype"): + _get_dtype(cfg) + + +def test_equiv_shape_mismatch(): + """assert_close raises on shape mismatch.""" + from .equiv import assert_close + + a = torch.randn(4, 8, device="cpu") + b = torch.randn(4, 9, device="cpu") + + with pytest.raises(AssertionError, match="Shape mismatch"): + assert_close(a, b) + + +# --------------------------------------------------------------------------- +# Microbench helper tests +# --------------------------------------------------------------------------- + + +@requires_cuda +def test_microbench_returns_positive(): + """microbench returns dict with mean_us > 0 for a trivial matmul.""" + from .microbench import microbench + + a = torch.randn(256, 256, device="cuda", dtype=torch.float16) + b = torch.randn(256, 256, device="cuda", dtype=torch.float16) + + def fn(): + torch.matmul(a, b) + + result = microbench(fn, n_iters=10, use_cuda_graph=False, flush_l2=False) + + assert isinstance(result, dict) + assert "mean_us" in result + assert "std_us" in result + assert "n_iters" in result + assert ( + result["mean_us"] > 0 + ), f"Expected mean_us > 0, got {result['mean_us']}" + assert result["n_iters"] > 0 + + +@requires_cuda +def test_microbench_with_cuda_graph(): + """microbench works with CUDA graph capture.""" + from .microbench import microbench + + a = torch.randn(128, 128, device="cuda", dtype=torch.float16) + b = torch.randn(128, 128, device="cuda", dtype=torch.float16) + c = torch.empty(128, 128, device="cuda", dtype=torch.float16) + + def fn(): + torch.mm(a, b, out=c) + + result = microbench(fn, n_iters=10, use_cuda_graph=True, flush_l2=False) + + assert result["mean_us"] > 0 + assert result["n_iters"] > 0 + + +@requires_cuda +def test_microbench_with_flush_l2(): + """microbench works with L2 flush enabled.""" + from .microbench import microbench + + a = torch.randn(128, 128, device="cuda", dtype=torch.float16) + b = torch.randn(128, 128, device="cuda", dtype=torch.float16) + c = torch.empty(128, 128, device="cuda", dtype=torch.float16) + + def fn(): + torch.mm(a, b, out=c) + + result = microbench(fn, n_iters=5, use_cuda_graph=False, flush_l2=True) + assert result["mean_us"] > 0 + assert result["n_iters"] > 0