diff --git a/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/.gitignore b/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/.gitignore new file mode 100644 index 000000000..fc9a27e9b --- /dev/null +++ b/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/.gitignore @@ -0,0 +1,15 @@ +# Build outputs and scratch / profiler artifacts (kept local for audit, excluded from the PR). +.build/ +__pycache__/ +*.pyc +*.so +*.o +# Nsight / NCU / profiler artifacts +*.ncu-rep +*.nsys-rep +*.trace.json +*.trace.json.gz +profile/ +# Scratch / failed-experiment dumps +scratch/ +*.scratch diff --git a/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/baseline/__init__.py b/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/baseline/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/baseline/binding.py b/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/baseline/binding.py new file mode 100644 index 000000000..df3b30a4e --- /dev/null +++ b/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/baseline/binding.py @@ -0,0 +1,46 @@ +"""Local baseline adapters reproducing the profiled PyTorch-eager elementwise +patterns, exposed through the same destination-passing ABI as the candidate. + +Faithful baseline: the production callsites run +``out = residual + update * gate`` in eager mode, which the torch profiler shows +as a separate ``aten::mul`` then ``aten::add`` (two launches plus an +intermediate tensor). We reproduce exactly that two-op path, but write the +intermediate into a cached, preallocated scratch buffer (keyed by +shape/dtype/device) so the timed steady state contains only the two real kernel +launches and their memory traffic -- never a caching-allocator call. The 4D +broadcast add is a single eager ``torch.add`` (one launch). + +No sglang import: these are pure torch ops reproducing the recovered semantics +(see docs/baseline_source.md). Output tensors are passed last and written +in-place; nothing is allocated in the timed path. +""" + +from __future__ import annotations + +import torch + +# Cache of scratch buffers for the two-op residual-gate path, keyed by the +# (shape, dtype, device) of the `update * gate` intermediate (== update's shape). +_SCRATCH: dict[tuple, torch.Tensor] = {} + + +def _scratch_like(t: torch.Tensor) -> torch.Tensor: + key = (tuple(t.shape), t.dtype, t.device) + buf = _SCRATCH.get(key) + if buf is None: + buf = torch.empty(t.shape, dtype=t.dtype, device=t.device) + _SCRATCH[key] = buf + return buf + + +def residual_gate_add(residual: torch.Tensor, update: torch.Tensor, + gate: torch.Tensor, out: torch.Tensor) -> None: + """out = residual + update * gate (faithful two-launch eager path).""" + scratch = _scratch_like(update) + torch.mul(update, gate, out=scratch) # aten::mul (broadcasts a [.,1,D] gate) + torch.add(residual, scratch, out=out) # aten::add + + +def broadcast_add_4d(a: torch.Tensor, b: torch.Tensor, out: torch.Tensor) -> None: + """out = a + b, with a [B,1,P,D] broadcasting over b's sequence dim.""" + torch.add(a, b, out=out) diff --git a/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/baseline/upstream_residual_gate_snippets.md b/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/baseline/upstream_residual_gate_snippets.md new file mode 100644 index 000000000..6484c7304 --- /dev/null +++ b/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/baseline/upstream_residual_gate_snippets.md @@ -0,0 +1,242 @@ +# Upstream residual-gate / broadcast-add source excerpts (provenance) + +Verbatim extracts from sgl-project/sglang @ `8314247d9de0fa2c58e34756b3e1dbc6cf815dfd` (branch main). +Evidence for the recovered kernel semantics; the runnable local baseline is +`baseline/binding.py`. Do not import sglang at runtime. (Relevant lines are +byte-identical to the earlier 67b2a9ed resolution.) + +## ideogram.py — Ideogram4TransformerBlock.forward (lines 320-340) + +```python + prefix=f"{prefix}.adaln_modulation", + ) + + def forward(self, x, cos, sin, adaln_input, attn_mask, attn_mask_meta): + scale_msa, gate_msa, scale_mlp, gate_mlp = self.adaln_modulation( + adaln_input + ).chunk(4, dim=-1) + gate_msa = torch.tanh(gate_msa) + gate_mlp = torch.tanh(gate_mlp) + attn_out = self.attention( + self.attention_norm1(x) * (1.0 + scale_msa), + cos=cos, + sin=sin, + attn_mask=attn_mask, + attn_mask_meta=attn_mask_meta, + ) + x = x + gate_msa * self.attention_norm2(attn_out) + x = x + gate_mlp * self.ffn_norm2( + self.feed_forward(self.ffn_norm1(x) * (1.0 + scale_mlp)) + ) + return x +``` + +## ltx_2.py — residual-gate updates (lines 1015-1100) + +```python + + # 1. Video and Audio Self-Attention + vshift_msa, vscale_msa, vgate_msa = self.get_ada_values( + self.scale_shift_table, batch_size, temb, slice(0, 3) + ) + norm_hidden_states = ( + self.rms_norm(hidden_states, self.norm_eps) * (1 + vscale_msa) + vshift_msa + ) + attn_hidden_states = self.attn1( + norm_hidden_states, + mask=video_self_attention_mask, + pe=video_rotary_emb, + perturbation_mask=video_self_attn_perturbation_mask, + all_perturbed=skip_video_self_attn, + gather_context_kv_for_sp=audio_replicated_for_sp, + ) + hidden_states = hidden_states + attn_hidden_states * vgate_msa + + ashift_msa, ascale_msa, agate_msa = self.get_ada_values( + self.audio_scale_shift_table, batch_size, temb_audio, slice(0, 3) + ) + norm_audio_hidden_states = ( + self.rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_msa) + + ashift_msa + ) + attn_audio_hidden_states = self.audio_attn1( + norm_audio_hidden_states, + mask=audio_self_attention_mask, + pe=audio_rotary_emb, + perturbation_mask=audio_self_attn_perturbation_mask, + all_perturbed=skip_audio_self_attn, + skip_sequence_parallel_override=audio_replicated_for_sp, + ) + audio_hidden_states = audio_hidden_states + attn_audio_hidden_states * agate_msa + # 2. Prompt Cross-Attention + if self.cross_attention_adaln: + # LTX2.3 + if temb_prompt is None or temb_audio_prompt is None: + raise ValueError( + "cross_attention_adaln requires prompt modulation tensors." + ) + vshift_q, vscale_q, vgate_q = self.get_ada_values( + self.scale_shift_table, batch_size, temb, slice(6, 9) + ) + v_prompt_shift, v_prompt_scale = self.get_ada_values( + self.prompt_scale_shift_table, batch_size, temb_prompt, slice(None) + ) + norm_hidden_states = ( + self.rms_norm(hidden_states, self.norm_eps) * (1 + vscale_q) + vshift_q + ) + mod_encoder_hidden_states = ( + encoder_hidden_states * (1 + v_prompt_scale) + v_prompt_shift + ) + attn_hidden_states = self.attn2( + norm_hidden_states, + context=mod_encoder_hidden_states, + mask=encoder_attention_mask, + ) + hidden_states = hidden_states + attn_hidden_states * vgate_q + + ashift_q, ascale_q, agate_q = self.get_ada_values( + self.audio_scale_shift_table, batch_size, temb_audio, slice(6, 9) + ) + a_prompt_shift, a_prompt_scale = self.get_ada_values( + self.audio_prompt_scale_shift_table, + batch_size, + temb_audio_prompt, + slice(None), + ) + norm_audio_hidden_states = ( + self.rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_q) + + ashift_q + ) + mod_audio_encoder_hidden_states = ( + audio_encoder_hidden_states * (1 + a_prompt_scale) + a_prompt_shift + ) + attn_audio_hidden_states = self.audio_attn2( + norm_audio_hidden_states, + context=mod_audio_encoder_hidden_states, + mask=audio_encoder_attention_mask, + ) + audio_hidden_states = ( + audio_hidden_states + attn_audio_hidden_states * agate_q + ) + else: + norm_hidden_states = self.rms_norm(hidden_states, self.norm_eps) +``` + +## ltx_2.py — 4D cross-attn-gate broadcast add (lines 1120-1146) + +```python + video_per_layer_ca_scale_shift = self.video_a2v_cross_attn_scale_shift_table[ + :4, : + ] + video_per_layer_ca_gate = self.video_a2v_cross_attn_scale_shift_table[4:, :] + + video_ca_scale_shift_table = ( + video_per_layer_ca_scale_shift[None, None, :, :].to( + dtype=temb_ca_scale_shift.dtype, device=temb_ca_scale_shift.device + ) + + temb_ca_scale_shift.reshape( + batch_size, temb_ca_scale_shift.shape[1], 4, -1 + ) + ).unbind(dim=2) + video_ca_gate = ( + video_per_layer_ca_gate[None, None, :, :].to( + dtype=temb_ca_gate.dtype, device=temb_ca_gate.device + ) + + temb_ca_gate.reshape(batch_size, temb_ca_gate.shape[1], 1, -1) + ).unbind(dim=2) + + ( + video_a2v_ca_scale, + video_a2v_ca_shift, + video_v2a_ca_scale, + video_v2a_ca_shift, + ) = [t.squeeze(2) for t in video_ca_scale_shift_table] + a2v_gate = video_ca_gate[0].squeeze(2) +``` + +## flux_2.py — modulation residual gate (lines 740-760, 838-895) + +```python + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + + mod_shift, mod_scale, mod_gate = temb_mod_params + + norm_hidden_states = self.norm(hidden_states) + norm_hidden_states = (1 + mod_scale) * norm_hidden_states + mod_shift + + joint_attention_kwargs = joint_attention_kwargs or {} + attn_output = self.attn( + hidden_states=norm_hidden_states, + freqs_cis=freqs_cis, + num_replicated_prefix=num_replicated_prefix, + **joint_attention_kwargs, + ) + + hidden_states = hidden_states + mod_gate * attn_output + if hidden_states.dtype == torch.float16: + hidden_states = hidden_states.clip(-65504, 65504) + + if split_hidden_states: + encoder_hidden_states, hidden_states = ( + +# ... (double-stream block) ... + + + # Modulation parameters shape: [1, 1, self.dim] + (shift_msa, scale_msa, gate_msa), ( + shift_mlp, + scale_mlp, + gate_mlp, + ) = temb_mod_params_img + (c_shift_msa, c_scale_msa, c_gate_msa), ( + c_shift_mlp, + c_scale_mlp, + c_gate_mlp, + ) = temb_mod_params_txt + + # Img stream + norm_hidden_states = self.norm1(hidden_states) + norm_hidden_states = (1 + scale_msa) * norm_hidden_states + shift_msa + + # Conditioning txt stream + norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states) + norm_encoder_hidden_states = ( + 1 + c_scale_msa + ) * norm_encoder_hidden_states + c_shift_msa + + # Attention on concatenated img + txt stream + attention_outputs = self.attn( + hidden_states=norm_hidden_states, + encoder_hidden_states=norm_encoder_hidden_states, + freqs_cis=freqs_cis, + num_replicated_prefix=num_replicated_prefix, + **joint_attention_kwargs, + ) + + attn_output, context_attn_output = attention_outputs + + # Process attention outputs for the image stream (`hidden_states`). + attn_output = gate_msa * attn_output + hidden_states = hidden_states + attn_output + + norm_hidden_states = self.norm2(hidden_states) + norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp + + ff_output = self.ff(norm_hidden_states) + hidden_states = hidden_states + gate_mlp * ff_output + + # Process attention outputs for the text stream (`encoder_hidden_states`). + context_attn_output = c_gate_msa * context_attn_output + encoder_hidden_states = encoder_hidden_states + context_attn_output + + norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) + norm_encoder_hidden_states = ( + norm_encoder_hidden_states * (1 + c_scale_mlp) + c_shift_mlp + ) + + context_ff_output = self.ff_context(norm_encoder_hidden_states) + encoder_hidden_states = encoder_hidden_states + c_gate_mlp * context_ff_output + if encoder_hidden_states.dtype == torch.float16: + encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504) + +``` diff --git a/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/bench/__init__.py b/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/bench/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/bench/adapter.py b/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/bench/adapter.py new file mode 100644 index 000000000..209f01517 --- /dev/null +++ b/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/bench/adapter.py @@ -0,0 +1,252 @@ +"""Benchmark adapter for b200_diffusion_residual_gate_add__multi_shape. + +Supplies tensor construction and the two ABI calls for bench/benchmark.py. +``call_baseline`` and ``call_candidate`` route through one shared dispatch +helper over pre-resolved function tables, so both sides pay byte-identical +adapter overhead. Neither call allocates output tensors: outputs are +preallocated in ``make_case`` and poisoned/timed by the benchmark template. + +The baseline side calls the faithful PyTorch-eager launchers in +baseline/binding.py; the candidate side calls the CUDA module built from +solution/kernel.cu via tvm-ffi. No sglang import anywhere in this process +(asserted below). +""" + +from __future__ import annotations + +import functools +import sys +from pathlib import Path + +_TASK_ROOT = Path(__file__).resolve().parents[1] +if str(_TASK_ROOT) not in sys.path: + sys.path.insert(0, str(_TASK_ROOT)) + +import torch + +import baseline.binding as _baseline +from solution.build import load_candidate_module + +assert not any( + name == "sglang" or name.startswith("sglang.") for name in sys.modules +), "standalone contract violation: sglang imported at benchmark runtime" + +_DTYPES = { + "bfloat16": torch.bfloat16, + "float16": torch.float16, + "float32": torch.float32, +} + +_EP_RGA = "residual_gate_add" +_EP_BCAST = "broadcast_add_4d" + +_BASELINE_FNS = { + _EP_RGA: _baseline.residual_gate_add, + _EP_BCAST: _baseline.broadcast_add_4d, +} + + +@functools.lru_cache(maxsize=1) +def _candidate_fns() -> dict: + # Build the candidate lazily (NOT at import): tvm-ffi's build queries the + # current CUDA device for gencode/provenance, so it must run AFTER the caller + # selects its target device (correctness.py / benchmark.py call + # torch.cuda.set_device first). Importing this module therefore stays cheap + # and device-independent. + m = load_candidate_module() + return {_EP_RGA: m.residual_gate_add, _EP_BCAST: m.broadcast_add_4d} + + +def _randn(shape, dtype, device): + return torch.randn(shape, device=device, dtype=dtype) + + +def _validate_against_spec(name: str, tensor: torch.Tensor, spec: dict) -> None: + """Fail before benchmarking if a constructed tensor diverges from the frozen + workload metadata (shape/stride self-description).""" + if list(tensor.shape) != list(spec["shape"]): + raise ValueError( + f"{name}: constructed shape {tuple(tensor.shape)} != frozen spec {spec['shape']}" + ) + want_stride = spec.get("stride") + if want_stride is not None and list(tensor.stride()) != list(want_stride): + raise ValueError( + f"{name}: constructed stride {tuple(tensor.stride())} != frozen spec {want_stride}" + ) + + +def make_case(workload: dict, *, device: torch.device, seed: int) -> dict: + del seed # benchmark.py already seeded the global generators + assert_pinned_gpu(device) + fn = workload["function"] + shapes = workload["shapes"] + + if fn == _EP_RGA: + residual = _randn(shapes["residual"]["shape"], _DTYPES[shapes["residual"]["dtype"]], device) + update = _randn(shapes["update"]["shape"], _DTYPES[shapes["update"]["dtype"]], device) + gate = _randn(shapes["gate"]["shape"], _DTYPES[shapes["gate"]["dtype"]], device) + _validate_against_spec("residual", residual, shapes["residual"]) + _validate_against_spec("update", update, shapes["update"]) + _validate_against_spec("gate", gate, shapes["gate"]) + inputs = {"residual": residual, "update": update, "gate": gate} + baseline_outputs = [torch.empty_like(residual)] + candidate_outputs = [torch.empty_like(residual)] + elif fn == _EP_BCAST: + a = _randn(shapes["a"]["shape"], _DTYPES[shapes["a"]["dtype"]], device) + b = _randn(shapes["b"]["shape"], _DTYPES[shapes["b"]["dtype"]], device) + _validate_against_spec("a", a, shapes["a"]) + _validate_against_spec("b", b, shapes["b"]) + inputs = {"a": a, "b": b} + baseline_outputs = [torch.empty_like(b)] + candidate_outputs = [torch.empty_like(b)] + else: + raise ValueError(f"unknown function {fn!r}") + + return { + "inputs": inputs, + "baseline_outputs": baseline_outputs, + "candidate_outputs": candidate_outputs, + "tolerance": {"atol": float(workload["atol"]), "rtol": float(workload["rtol"])}, + } + + +_FLOAT_DTYPES = {torch.bfloat16, torch.float16, torch.float32} + + +def _validate(fn: str, inputs: dict, outputs) -> None: + """Shared ABI contract enforced before either implementation runs, so the + baseline (eager) and candidate (CUDA host checks) reject the same malformed + inputs with identical adapter overhead. Mirrors solution/kernel.cu's host + checks: CUDA device, shared float dtype, contiguous, out distinct from + inputs, full [.,L,D] or same-rank row-broadcast [.,1,D] gate, B=1 4D add.""" + out = outputs[0] + if fn == _EP_RGA: + residual, update, gate = inputs["residual"], inputs["update"], inputs["gate"] + tensors = (residual, update, gate, out) + if not all(t.is_cuda for t in tensors): + raise ValueError("residual/update/gate/out must be CUDA tensors") + if residual.dim() < 2: + raise ValueError("residual must be at least 2D ([.., D])") + if update.shape != residual.shape or out.shape != residual.shape: + raise ValueError("update/out shape must match residual") + dt = residual.dtype + if dt not in _FLOAT_DTYPES or any(t.dtype != dt for t in tensors): + raise ValueError("residual/update/gate/out must share a float dtype (bf16/fp16/fp32)") + if any(out.data_ptr() == t.data_ptr() for t in (residual, update, gate)): + raise ValueError("out must not alias residual/update/gate") + if not (residual.is_contiguous() and update.is_contiguous() and out.is_contiguous()): + raise ValueError("residual/update/out must be contiguous") + D = residual.shape[-1] + if gate.shape[-1] != D: + raise ValueError("gate last dim must equal residual last dim") + if gate.shape == residual.shape: + if not gate.is_contiguous(): + raise ValueError("full-shape gate must be contiguous") + return # full gate + # row-broadcast: same rank, size-1 row dim, all leading dims == 1 + if gate.dim() != residual.dim() or gate.shape[-2] != 1 or any( + s != 1 for s in gate.shape[:-1] + ): + raise ValueError("gate must be full [.., D] or same-rank row-broadcast [.., 1, D]") + if not gate.is_contiguous(): + raise ValueError("broadcast gate must be contiguous") + else: # broadcast_add_4d + a, b = inputs["a"], inputs["b"] + tensors = (a, b, out) + if not all(t.is_cuda for t in tensors): + raise ValueError("a/b/out must be CUDA tensors") + if a.dim() != 4 or b.dim() != 4: + raise ValueError("broadcast_add_4d expects 4D a and b") + if b.shape[0] != 1: + raise ValueError("broadcast_add_4d supports batch size 1 only") + if a.shape[0] != b.shape[0] or a.shape[1] != 1 or a.shape[2] != b.shape[2] or a.shape[3] != b.shape[3]: + raise ValueError("a must be [1,1,P,D] broadcasting over dim1 of b") + if out.shape != b.shape: + raise ValueError("out shape must match b") + dt = b.dtype + if dt not in _FLOAT_DTYPES or any(t.dtype != dt for t in tensors): + raise ValueError("a/b/out must share a float dtype (bf16/fp16/fp32)") + if any(out.data_ptr() == t.data_ptr() for t in (a, b)): + raise ValueError("out must not alias a/b") + if not (a.is_contiguous() and b.is_contiguous() and out.is_contiguous()): + raise ValueError("a/b/out must be contiguous") + + +def _dispatch(fns: dict, workload: dict, inputs: dict, outputs) -> None: + fn = workload["function"] + _validate(fn, inputs, outputs) + impl = fns[fn] + if fn == _EP_RGA: + impl(inputs["residual"], inputs["update"], inputs["gate"], outputs[0]) + else: + impl(inputs["a"], inputs["b"], outputs[0]) + + +def call_baseline(workload: dict, inputs, outputs) -> None: + _dispatch(_BASELINE_FNS, workload, inputs, outputs) + + +def call_candidate(workload: dict, inputs, outputs) -> None: + _dispatch(_candidate_fns(), workload, inputs, outputs) + + +def assert_pinned_gpu(device: torch.device) -> None: + """Fail-closed GPU-pinning guard for benchmark/profile runs. Enabled when + KDA_REQUIRE_PINNED_GPU is set (remote benchmark mode): REMOTE_GPU_ID must be + set and must be the sole visible CUDA device, so every measured number is + attributable to one pinned GPU. No-op for local/dev correctness runs.""" + import os + if not os.environ.get("KDA_REQUIRE_PINNED_GPU"): + return + pinned = (os.environ.get("REMOTE_GPU_ID") or "").strip() + if not pinned: + raise RuntimeError("REMOTE_GPU_ID must be set when KDA_REQUIRE_PINNED_GPU=1") + visible = (os.environ.get("CUDA_VISIBLE_DEVICES") or "").strip() + if visible != pinned: + raise RuntimeError( + f"pinned-GPU contract violated: CUDA_VISIBLE_DEVICES={visible!r} != REMOTE_GPU_ID={pinned!r}" + ) + idx = device.index if getattr(device, "index", None) is not None else 0 + if idx != 0: + raise RuntimeError(f"pinned GPU must be the sole visible device (cuda:0); got {device}") + + +def extra_provenance() -> dict: + """Optional hook consumed by bench/benchmark.py to record the task-specific + provenance the standalone contract requires beyond the template defaults: + baseline commit, candidate source hash, TVM-FFI / NVCC versions, compile + flags, and the pinned REMOTE_GPU_ID.""" + import hashlib + import os + import re + import subprocess + + info: dict = { + "task_slug": "b200_diffusion_residual_gate_add__multi_shape", + "target_gpu": "NVIDIA B200", + "remote_gpu_id": os.environ.get("REMOTE_GPU_ID"), + } + kernel = _TASK_ROOT / "solution" / "kernel.cu" + if kernel.exists(): + info["candidate_kernel_sha256"] = hashlib.sha256(kernel.read_bytes()).hexdigest() + doc = _TASK_ROOT / "docs" / "baseline_source.md" + if doc.exists(): + m = re.search(r"Resolved commit SHA: `([0-9a-f]{7,40})`", doc.read_text()) + info["baseline_commit"] = m.group(1) if m else "unknown" + try: + import tvm_ffi + info["tvm_ffi"] = getattr(tvm_ffi, "__version__", "unknown") + except Exception: + info["tvm_ffi"] = "unavailable" + try: + from solution.build import candidate_compile_flags + info["candidate_compile_flags"] = candidate_compile_flags() + except Exception: + pass + try: + nv = subprocess.check_output(["nvcc", "--version"], text=True, + stderr=subprocess.STDOUT, timeout=5) + info["nvcc"] = nv.strip().splitlines()[-1] + except Exception: + info["nvcc"] = "unavailable" + return info diff --git a/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/bench/benchmark.py b/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/bench/benchmark.py new file mode 100644 index 000000000..f8616e2ba --- /dev/null +++ b/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/bench/benchmark.py @@ -0,0 +1,601 @@ +#!/usr/bin/env python3 +"""Standard standalone diffusion kernel benchmark template. + +Copy this file to ``bench/benchmark.py`` for a diffusion task, then implement +``bench/adapter.py`` in the same directory. The adapter is intentionally small +so the benchmark policy stays fixed while each task owns its tensor creation and +kernel calls. + +Required adapter API: + + make_case(workload: dict, *, device: torch.device, seed: int) -> Case + Return an object with: + inputs: tuple/list/dict of CUDA tensors and scalar args + baseline_outputs: preallocated CUDA output tensors + candidate_outputs: preallocated CUDA output tensors + tolerance: {"atol": float, "rtol": float} + + call_baseline(workload: dict, inputs, outputs) -> None + call_candidate(workload: dict, inputs, outputs) -> None + Launch the local copied baseline and candidate through the same ABI. + These functions must not allocate output tensors. + +Optional adapter API: + + compare_outputs(workload, baseline_outputs, candidate_outputs, tolerance) + Return {"ok": bool, "max_abs": float, "max_rel": float, "message": str} + If absent, the default tensor comparison below is used. + +Workload file shape: + + [ + { + "id": "qwen_4096", + "production": true, + "function": "fused_inplace_qknorm_rope", + "shapes": {...}, + "atol": 1e-2, + "rtol": 1e-2 + } + ] +""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib.util +import json +import math +import os +import platform +import queue +import random +import statistics +import subprocess +import sys +import time +import traceback +from dataclasses import dataclass +from multiprocessing import get_context +from pathlib import Path +from typing import Any, Callable + +import torch + + +ROOT = Path(__file__).resolve().parents[1] +BENCH_DIR = Path(__file__).resolve().parent +DEFAULT_WORKLOADS = BENCH_DIR / "workloads.json" +DEFAULT_OUTPUT = BENCH_DIR / "results.jsonl" + + +@dataclass +class Stats: + median_us: float + mean_us: float + std_us: float + min_us: float + p10_us: float + p90_us: float + samples_us: list[float] + inner_iterations: int + + +def _load_adapter(): + path = BENCH_DIR / "adapter.py" + if not path.exists(): + raise FileNotFoundError( + f"{path} is missing. Implement the adapter API documented at the " + "top of bench/benchmark.py." + ) + spec = importlib.util.spec_from_file_location("bench_adapter", path) + if spec is None or spec.loader is None: + raise RuntimeError(f"cannot import adapter from {path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + for name in ("make_case", "call_baseline", "call_candidate"): + if not hasattr(module, name): + raise AttributeError(f"bench/adapter.py must define {name}()") + return module + + +def _as_list(value: Any) -> list[Any]: + if value is None: + return [] + if isinstance(value, (list, tuple)): + return list(value) + if isinstance(value, dict): + return list(value.values()) + return [value] + + +def _case_get(case: Any, name: str) -> Any: + if isinstance(case, dict): + return case[name] + return getattr(case, name) + + +def _case_get_optional(case: Any, name: str, default: Any = None) -> Any: + if isinstance(case, dict): + return case.get(name, default) + return getattr(case, name, default) + + +def _stable_u16(text: str) -> int: + digest = hashlib.sha256(text.encode("utf-8")).digest() + return int.from_bytes(digest[:2], "little") + + +def _poison_outputs(outputs: Any) -> None: + for out in _as_list(outputs): + if torch.is_tensor(out): + if out.is_floating_point() or out.is_complex(): + out.fill_(float("nan")) + else: + out.fill_(-17) + + +def _default_compare( + _workload: dict[str, Any], + baseline_outputs: Any, + candidate_outputs: Any, + tolerance: dict[str, float], +) -> dict[str, Any]: + base = _as_list(baseline_outputs) + cand = _as_list(candidate_outputs) + if len(base) != len(cand): + return { + "ok": False, + "max_abs": math.inf, + "max_rel": math.inf, + "message": f"output count mismatch: baseline={len(base)} candidate={len(cand)}", + } + + atol = float(tolerance.get("atol", 1e-2)) + rtol = float(tolerance.get("rtol", 1e-2)) + max_abs = 0.0 + max_rel = 0.0 + + for idx, (lhs, rhs) in enumerate(zip(base, cand)): + if not torch.is_tensor(lhs) or not torch.is_tensor(rhs): + if lhs != rhs: + return { + "ok": False, + "max_abs": math.inf, + "max_rel": math.inf, + "message": f"non-tensor output {idx} differs", + } + continue + if lhs.shape != rhs.shape: + return { + "ok": False, + "max_abs": math.inf, + "max_rel": math.inf, + "message": f"output {idx} shape mismatch: {tuple(lhs.shape)} vs {tuple(rhs.shape)}", + } + if lhs.dtype != rhs.dtype: + return { + "ok": False, + "max_abs": math.inf, + "max_rel": math.inf, + "message": f"output {idx} dtype mismatch: {lhs.dtype} vs {rhs.dtype}", + } + + lhs_f = lhs.detach().float() + rhs_f = rhs.detach().float() + if torch.isnan(rhs_f).any() or torch.isinf(rhs_f).any(): + return { + "ok": False, + "max_abs": math.inf, + "max_rel": math.inf, + "message": f"output {idx} has NaN or Inf", + } + diff = (lhs_f - rhs_f).abs() + denom = lhs_f.abs().clamp_min(1e-12) + rel = diff / denom + max_abs = max(max_abs, float(diff.max().item()) if diff.numel() else 0.0) + max_rel = max(max_rel, float(rel.max().item()) if rel.numel() else 0.0) + ok = torch.all(diff <= (atol + rtol * lhs_f.abs())) + if not bool(ok.item()): + return { + "ok": False, + "max_abs": max_abs, + "max_rel": max_rel, + "message": f"output {idx} exceeds tolerance atol={atol} rtol={rtol}", + } + + return {"ok": True, "max_abs": max_abs, "max_rel": max_rel, "message": ""} + + +def _call_many(fn: Callable[[], None], inner_iterations: int) -> None: + for _ in range(inner_iterations): + fn() + + +def _cuda_event_time_us(fn: Callable[[], None], inner_iterations: int) -> tuple[float, float]: + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + wall_start = time.perf_counter() + start.record() + _call_many(fn, inner_iterations) + end.record() + end.synchronize() + wall_end = time.perf_counter() + elapsed_us = start.elapsed_time(end) * 1000.0 / inner_iterations + wall_us = (wall_end - wall_start) * 1_000_000.0 / inner_iterations + return elapsed_us, wall_us + + +def _calibrate_inner( + fn: Callable[[], None], + *, + inner_min: int, + inner_max: int, + target_sample_us: float, +) -> int: + inner = max(1, inner_min) + while True: + elapsed_us, _ = _cuda_event_time_us(fn, inner) + if elapsed_us * inner >= target_sample_us or inner >= inner_max: + return inner + inner = min(inner * 2, inner_max) + + +def _percentile(values: list[float], pct: float) -> float: + if not values: + return math.nan + if len(values) == 1: + return values[0] + ordered = sorted(values) + rank = (len(ordered) - 1) * pct + lo = math.floor(rank) + hi = math.ceil(rank) + if lo == hi: + return ordered[lo] + weight = rank - lo + return ordered[lo] * (1.0 - weight) + ordered[hi] * weight + + +def _stats(samples_us: list[float], inner_iterations: int) -> Stats: + return Stats( + median_us=statistics.median(samples_us), + mean_us=statistics.mean(samples_us), + std_us=statistics.stdev(samples_us) if len(samples_us) > 1 else 0.0, + min_us=min(samples_us), + p10_us=_percentile(samples_us, 0.10), + p90_us=_percentile(samples_us, 0.90), + samples_us=samples_us, + inner_iterations=inner_iterations, + ) + + +def _stats_dict(stats: Stats) -> dict[str, Any]: + return { + "median_us": stats.median_us, + "mean_us": stats.mean_us, + "std_us": stats.std_us, + "min_us": stats.min_us, + "p10_us": stats.p10_us, + "p90_us": stats.p90_us, + "samples_us": stats.samples_us, + "inner_iterations": stats.inner_iterations, + } + + +def _run_one_workload(workload: dict[str, Any], args: argparse.Namespace) -> dict[str, Any]: + # Select the target CUDA device BEFORE importing the adapter: importing it + # builds the candidate module, whose gencode/provenance query the current + # device, so the device must be set first (else a non-default --device would + # build/report for cuda:0). + device = torch.device(args.device) + torch.cuda.set_device(device) + torch.set_grad_enabled(False) + adapter = _load_adapter() + + baseline_samples: list[float] = [] + candidate_samples: list[float] = [] + baseline_wall_samples: list[float] = [] + candidate_wall_samples: list[float] = [] + baseline_inner = args.inner_iterations_min + candidate_inner = args.inner_iterations_min + + compare = getattr(adapter, "compare_outputs", _default_compare) + workload_id = str(workload.get("id") or workload.get("name") or workload.get("function")) + if not workload_id: + raise ValueError("each workload needs an id/name/function field") + + for trial in range(args.num_trials): + seed = args.seed + trial * 1_000_003 + _stable_u16(workload_id) + random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + case = adapter.make_case(workload, device=device, seed=seed) + inputs = _case_get(case, "inputs") + baseline_outputs = _case_get(case, "baseline_outputs") + candidate_outputs = _case_get(case, "candidate_outputs") + tolerance = _case_get_optional(case, "tolerance") or { + "atol": float(workload.get("atol", args.atol)), + "rtol": float(workload.get("rtol", args.rtol)), + } + + baseline_fn = lambda: adapter.call_baseline( + workload, inputs, baseline_outputs + ) + candidate_fn = lambda: adapter.call_candidate( + workload, inputs, candidate_outputs + ) + + _poison_outputs(baseline_outputs) + _poison_outputs(candidate_outputs) + baseline_fn() + candidate_fn() + torch.cuda.synchronize() + correctness = compare( + workload, baseline_outputs, candidate_outputs, tolerance + ) + if not correctness.get("ok", False): + return { + "id": workload_id, + "status": "INCORRECT", + "workload": workload, + "correctness": correctness, + } + + for _ in range(args.warmup_runs): + baseline_fn() + candidate_fn() + torch.cuda.synchronize() + + if trial == 0: + baseline_inner = _calibrate_inner( + baseline_fn, + inner_min=args.inner_iterations_min, + inner_max=args.inner_iterations_max, + target_sample_us=args.target_sample_us, + ) + candidate_inner = _calibrate_inner( + candidate_fn, + inner_min=args.inner_iterations_min, + inner_max=args.inner_iterations_max, + target_sample_us=args.target_sample_us, + ) + + order = ("baseline", "candidate") + if random.Random(seed).randint(0, 1): + order = ("candidate", "baseline") + + trial_values: dict[str, tuple[float, float]] = {} + for side in order: + if side == "baseline": + trial_values[side] = _cuda_event_time_us(baseline_fn, baseline_inner) + else: + trial_values[side] = _cuda_event_time_us(candidate_fn, candidate_inner) + + base_us, base_wall_us = trial_values["baseline"] + cand_us, cand_wall_us = trial_values["candidate"] + baseline_samples.append(base_us) + candidate_samples.append(cand_us) + baseline_wall_samples.append(base_wall_us) + candidate_wall_samples.append(cand_wall_us) + + baseline = _stats(baseline_samples, baseline_inner) + candidate = _stats(candidate_samples, candidate_inner) + speedup = baseline.median_us / candidate.median_us if candidate.median_us > 0 else math.inf + return { + "id": workload_id, + "status": "PASSED", + "production": bool(workload.get("production", True)), + "workload": workload, + "baseline": _stats_dict(baseline), + "candidate": _stats_dict(candidate), + "baseline_wall_us": baseline_wall_samples, + "candidate_wall_us": candidate_wall_samples, + "speedup": speedup, + } + + +def _worker_main( + workload: dict[str, Any], + args_dict: dict[str, Any], + result_queue: Any, +) -> None: + try: + ns = argparse.Namespace(**args_dict) + result_queue.put(_run_one_workload(workload, ns)) + except BaseException: + result_queue.put( + { + "id": workload.get("id") or workload.get("name") or "", + "status": "ERROR", + "workload": workload, + "error": traceback.format_exc(), + } + ) + + +def _run_isolated(workload: dict[str, Any], args: argparse.Namespace) -> dict[str, Any]: + ctx = get_context("spawn") + result_queue = ctx.Queue() + proc = ctx.Process(target=_worker_main, args=(workload, vars(args), result_queue)) + proc.start() + proc.join(args.timeout_seconds) + if proc.is_alive(): + proc.terminate() + proc.join(10) + return { + "id": workload.get("id") or workload.get("name") or "", + "status": "TIMEOUT", + "workload": workload, + "timeout_seconds": args.timeout_seconds, + } + try: + return result_queue.get_nowait() + except queue.Empty: + return { + "id": workload.get("id") or workload.get("name") or "", + "status": "ERROR", + "workload": workload, + "error": f"worker exited with code {proc.exitcode} and returned no result", + } + + +def _gpu_name() -> str: + if not torch.cuda.is_available(): + return "cuda-unavailable" + return torch.cuda.get_device_name(torch.cuda.current_device()) + + +def _nvidia_smi() -> str: + try: + return subprocess.check_output( + [ + "nvidia-smi", + "--query-gpu=index,name,utilization.gpu,memory.used,memory.total", + "--format=csv,noheader", + ], + text=True, + stderr=subprocess.STDOUT, + timeout=5, + ).strip() + except Exception as exc: + return f"unavailable: {exc}" + + +def _provenance(args: argparse.Namespace, workloads: list[dict[str, Any]]) -> dict[str, Any]: + prov: dict[str, Any] = { + "task_dir": str(ROOT), + "command": " ".join(sys.argv), + "python": sys.version, + "platform": platform.platform(), + "torch": torch.__version__, + "torch_cuda": torch.version.cuda, + "cuda_available": torch.cuda.is_available(), + "gpu": _gpu_name(), + "nvidia_smi_before": _nvidia_smi(), + "workload_count": len(workloads), + "settings": { + "warmup_runs": args.warmup_runs, + "num_trials": args.num_trials, + "inner_iterations_min": args.inner_iterations_min, + "inner_iterations_max": args.inner_iterations_max, + "target_sample_us": args.target_sample_us, + "timeout_seconds": args.timeout_seconds, + "isolated": args.isolated, + }, + } + # Additive only (does not affect timing): merge task-specific provenance + # (baseline commit, candidate hash, TVM-FFI/NVCC versions, REMOTE_GPU_ID) + # from an optional adapter hook. Before/after GPU idle snapshots are recorded + # in docs/run_log.md around the remote benchmark command. + try: + _extra = getattr(_load_adapter(), "extra_provenance", None) + if callable(_extra): + prov["task_provenance"] = _extra() + except Exception as exc: # noqa: BLE001 + prov["task_provenance_error"] = str(exc) + return prov + + +def _headline(results: list[dict[str, Any]]) -> dict[str, Any]: + passed = [r for r in results if r.get("status") == "PASSED"] + production = [r for r in passed if r.get("production", True)] + speedups = [float(r["speedup"]) for r in production if float(r["speedup"]) > 0] + if not speedups: + return {"production_workloads": len(production), "geomean_speedup": None} + log_mean = sum(math.log(x) for x in speedups) / len(speedups) + return { + "production_workloads": len(production), + "geomean_speedup": math.exp(log_mean), + "arithmetic_mean_speedup": statistics.mean(speedups), + "min_speedup": min(speedups), + "max_speedup": max(speedups), + } + + +def _load_workloads(path: Path, only: list[str] | None) -> list[dict[str, Any]]: + data = json.loads(path.read_text()) + if not isinstance(data, list): + raise ValueError(f"{path} must contain a JSON list") + if not only: + return data + selected = set(only) + return [ + w for w in data + if str(w.get("id") or w.get("name") or w.get("function")) in selected + ] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--workloads", type=Path, default=DEFAULT_WORKLOADS) + parser.add_argument("--out", type=Path, default=DEFAULT_OUTPUT) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--only", nargs="*", help="run only workload ids listed here") + parser.add_argument("--seed", type=int, default=1234) + parser.add_argument("--warmup-runs", type=int, default=10) + parser.add_argument("--num-trials", type=int, default=7) + parser.add_argument("--inner-iterations-min", type=int, default=1) + parser.add_argument("--inner-iterations-max", type=int, default=4096) + parser.add_argument("--target-sample-us", type=float, default=1000.0) + parser.add_argument("--timeout-seconds", type=int, default=600) + parser.add_argument("--atol", type=float, default=1e-2) + parser.add_argument("--rtol", type=float, default=1e-2) + parser.add_argument("--no-isolated", dest="isolated", action="store_false") + parser.set_defaults(isolated=True) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + if not torch.cuda.is_available(): + raise SystemExit("CUDA is required for this benchmark") + workloads = _load_workloads(args.workloads, args.only) + if not workloads: + raise SystemExit(f"no workloads selected from {args.workloads}") + + args.out.parent.mkdir(parents=True, exist_ok=True) + # Make the selected CUDA device current BEFORE collecting provenance: both + # _gpu_name() and the adapter's candidate build / extra_provenance() query the + # process-current device, so a non-default --device must be set first or + # results.jsonl would cite cuda:0's name / SM capability / compile flags while + # the workloads actually run on args.device. (_run_one_workload re-sets it per + # run/subprocess; this covers the parent-process provenance collection.) + torch.cuda.set_device(torch.device(args.device)) + provenance = _provenance(args, workloads) + results: list[dict[str, Any]] = [] + with args.out.open("w") as f: + f.write(json.dumps({"event": "provenance", **provenance}) + "\n") + for workload in workloads: + workload_id = workload.get("id") or workload.get("name") or workload.get("function") + print(f"== workload {workload_id} ==", flush=True) + result = _run_isolated(workload, args) if args.isolated else _run_one_workload(workload, args) + results.append(result) + f.write(json.dumps({"event": "result", **result}) + "\n") + f.flush() + if result.get("status") == "PASSED": + print( + f"PASSED speedup={result['speedup']:.4f} " + f"baseline={result['baseline']['median_us']:.3f}us " + f"candidate={result['candidate']['median_us']:.3f}us", + flush=True, + ) + else: + print(f"{result.get('status')} {result.get('error') or result.get('correctness')}", flush=True) + + summary = { + "event": "summary", + "headline": _headline(results), + "passed": sum(r.get("status") == "PASSED" for r in results), + "total": len(results), + "nvidia_smi_after": _nvidia_smi(), + } + f.write(json.dumps(summary) + "\n") + + print(json.dumps(summary, indent=2)) + return 0 if summary["passed"] == summary["total"] else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/bench/correctness.py b/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/bench/correctness.py new file mode 100644 index 000000000..78ab92e20 --- /dev/null +++ b/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/bench/correctness.py @@ -0,0 +1,332 @@ +#!/usr/bin/env python3 +"""Correctness gate for b200_diffusion_residual_gate_add__multi_shape. + +Covers, before any benchmark number counts: + - the 8 production rows from bench/workloads.json; + - a regression grid: gate modes (full [B,L,D] and broadcast [B,1,D]), dtypes + (bf16/fp16/fp32), tail/alignment rows (odd D, D not a multiple of the 16B + vector width, small L), sign/zero coverage, repeated seeds, and the 4D + broadcast over several frame counts; + - a poison-detection self-test (a skipped launch must be caught); + - rejection tests (malformed gate shape, dtype mismatch, output aliasing, and + non-contiguous inputs must raise on the checked side). + +Both implementations are checked against an independent fp32 torch oracle +(one-round: residual + update*gate, and a + b) and against each other. Output +buffers are poisoned before every run; NaN/Inf in any checked output fails. + +No sglang import anywhere in this process (asserted via bench.adapter). +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +import torch + +_BENCH_DIR = Path(__file__).resolve().parent +_TASK_ROOT = _BENCH_DIR.parent +if str(_TASK_ROOT) not in sys.path: + sys.path.insert(0, str(_TASK_ROOT)) + +import bench.adapter as adapter # noqa: E402 (asserts the no-sglang contract) + +_EP_RGA, _EP_BCAST = adapter._EP_RGA, adapter._EP_BCAST +_DTYPES = adapter._DTYPES + + +def _tol(dtype: str) -> dict: + if dtype == "float32": + return {"atol": 1e-5, "rtol": 1e-5} + return {"atol": 5e-2, "rtol": 5e-2} + + +# --------------------------------------------------------------------------- +# independent fp32 oracles (one-round, matching docs/baseline_source.md) +# --------------------------------------------------------------------------- + + +def oracle_rga(inputs: dict) -> list[torch.Tensor]: + residual, update, gate = inputs["residual"], inputs["update"], inputs["gate"] + out = (residual.float() + update.float() * gate.float()).to(residual.dtype) + return [out] + + +def oracle_bcast(inputs: dict) -> list[torch.Tensor]: + a, b = inputs["a"], inputs["b"] + out = (a.float() + b.float()).to(b.dtype) # a broadcasts over b's seq dim + return [out] + + +_ORACLES = {_EP_RGA: oracle_rga, _EP_BCAST: oracle_bcast} + + +# --------------------------------------------------------------------------- +# comparison / poison +# --------------------------------------------------------------------------- + + +def _compare(tag: str, ref: list, got: list, tol: dict) -> list[str]: + errors = [] + if len(ref) != len(got): + return [f"{tag}: output count mismatch {len(ref)} vs {len(got)}"] + for i, (r, g) in enumerate(zip(ref, got)): + if r.shape != g.shape or r.dtype != g.dtype: + errors.append(f"{tag}[{i}]: shape/dtype mismatch {tuple(g.shape)}/{g.dtype}") + continue + gf, rf = g.float(), r.float() + if torch.isnan(gf).any() or torch.isinf(gf).any(): + errors.append(f"{tag}[{i}]: NaN/Inf in output") + continue + diff = (gf - rf).abs() + ok = bool(torch.all(diff <= (tol["atol"] + tol["rtol"] * rf.abs())).item()) + if not ok: + errors.append( + f"{tag}[{i}]: max_abs={diff.max().item():.3e} exceeds " + f"atol={tol['atol']} rtol={tol['rtol']}" + ) + return errors + + +def _poison(outputs: list) -> None: + for t in outputs: + t.fill_(float("nan")) + + +# --------------------------------------------------------------------------- +# regression-row construction (workload-schema dicts; tensors via make_case) +# --------------------------------------------------------------------------- + + +def _rga_spec(rid, B, L, D, dtype, gate_mode, *, prod=False) -> dict: + g_shape = [B, L, D] if gate_mode == "full" else [B, 1, D] + return { + "id": rid, "production": prod, "function": _EP_RGA, "gate_mode": gate_mode, + "shapes": { + "residual": {"shape": [B, L, D], "dtype": dtype}, + "update": {"shape": [B, L, D], "dtype": dtype}, + "gate": {"shape": g_shape, "dtype": dtype}, + "out": {"shape": [B, L, D], "dtype": dtype}, + }, + **_tol(dtype), + } + + +def _bcast_spec(rid, B, S, P, D, dtype) -> dict: + return { + "id": rid, "production": False, "function": _EP_BCAST, + "shapes": { + "a": {"shape": [B, 1, P, D], "dtype": dtype}, + "b": {"shape": [B, S, P, D], "dtype": dtype}, + "out": {"shape": [B, S, P, D], "dtype": dtype}, + }, + **_tol(dtype), + } + + +def build_grid_rows(quick: bool) -> list[dict]: + rows: list[dict] = [] + dts = ("bfloat16", "float16", "float32") + # Gate modes x dtypes x a few (L, D) including aligned + tail D and small L. + LD = [(1, 4096, 4096), (1, 256, 3072), (1, 7, 2048), (1, 128, 2047), (1, 33, 130), (1, 1, 4608)] + if quick: + LD = LD[::2] + n = 0 + for dt in dts: + for (B, L, D) in LD: + for gm in ("full", "bcast"): + rows.append(_rga_spec(f"grid_rga_{gm}_b{B}_l{L}_d{D}_{dt}", B, L, D, dt, gm)) + n += 1 + # 4D broadcast over several frame counts and a tail D. + for dt in dts: + for (S, P, D) in [(126, 3, 2048), (1, 3, 2048), (40, 4, 1024), (17, 3, 130)]: + rows.append(_bcast_spec(f"grid_bcast_s{S}_p{P}_d{D}_{dt}", 1, S, P, D, dt)) + # Deterministic zero/sign coverage: zeros + positive + negative values. + for gm in ("full", "bcast"): + rows.append({**_rga_spec(f"grid_rga_{gm}_zerosign", 1, 64, 2048, "bfloat16", gm), + "fill": "zerosign"}) + rows.append({**_bcast_spec("grid_bcast4d_zerosign", 1, 40, 3, 2048, "bfloat16"), + "fill": "zerosign"}) + # Repeated randomized seeds for representative rows (distinct id -> distinct + # seed = 90000 + row_index), covering full, broadcast, tail-D, and 4D. + for k in (2, 3): + rows.append(_rga_spec(f"grid_rga_full_repeat{k}", 1, 4096, 4096, "bfloat16", "full")) + rows.append(_rga_spec(f"grid_rga_bcast_repeat{k}", 1, 4096, 4096, "bfloat16", "bcast")) + rows.append(_rga_spec(f"grid_rga_tailD_repeat{k}", 1, 128, 2047, "bfloat16", "full")) + rows.append(_bcast_spec(f"grid_bcast4d_repeat{k}", 1, 126, 3, 2048, "bfloat16")) + return rows + + +# Deterministic pattern with zeros, positive, and negative values (sign/zero +# coverage). gate==0 -> out==residual; negative gate flips the update sign. +_ZEROSIGN_PATTERN = (0.0, 1.5, -2.0, 0.0, 0.75, -1.0, 3.0, -0.25) + + +def _apply_zerosign(case: dict) -> None: + pat = torch.tensor(_ZEROSIGN_PATTERN) + for t in case["inputs"].values(): + if torch.is_tensor(t) and t.is_floating_point(): + n = t.numel() + reps = (n + len(_ZEROSIGN_PATTERN) - 1) // len(_ZEROSIGN_PATTERN) + vals = pat.to(device=t.device, dtype=t.dtype).repeat(reps)[:n] + t.view(-1).copy_(vals) + + +def _build_case(spec: dict, device: torch.device, seed: int) -> dict: + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + case = adapter.make_case(spec, device=device, seed=seed) + if spec.get("fill") == "zerosign": + _apply_zerosign(case) + return case + + +def run_row(spec: dict, device: torch.device, impl: str, row_index: int) -> list[str]: + case = _build_case(spec, device, 90000 + row_index) + inputs = case["inputs"] + tol = case["tolerance"] + oracle = _ORACLES[spec["function"]](inputs) + + errors: list[str] = [] + sides = [] + if impl in ("both", "baseline"): + sides.append(("baseline", adapter.call_baseline, case["baseline_outputs"])) + if impl in ("both", "candidate"): + sides.append(("candidate", adapter.call_candidate, case["candidate_outputs"])) + for name, fn, outputs in sides: + _poison(outputs) + fn(spec, inputs, outputs) + torch.cuda.synchronize() + errors += _compare(f"{spec['id']}:{name}-vs-oracle", oracle, outputs, tol) + if impl == "both" and not errors: + errors += _compare( + f"{spec['id']}:candidate-vs-baseline", + case["baseline_outputs"], case["candidate_outputs"], tol, + ) + return errors + + +def run_self_test(device: torch.device) -> list[str]: + """A poisoned output with a skipped launch MUST be flagged.""" + spec = _rga_spec("selftest_skipped_launch", 1, 33, 512, "bfloat16", "bcast") + case = _build_case(spec, device, 4242) + oracle = _ORACLES[_EP_RGA](case["inputs"]) + _poison(case["candidate_outputs"]) + # deliberately no kernel call + errs = _compare("selftest", oracle, case["candidate_outputs"], case["tolerance"]) + return [] if errs else ["poison self-test FAILED: skipped launch not detected"] + + +def run_rejection_tests(device: torch.device, impl: str) -> list[str]: + """Invalid inputs must raise on BOTH sides (shared adapter validation runs + before either implementation).""" + failures: list[str] = [] + bf16 = torch.bfloat16 + + def expect_raise(tag, call): + try: + call() + torch.cuda.synchronize() + except Exception: + return + failures.append(f"rejection {tag}: accepted invalid input") + + sides = [] + if impl in ("both", "baseline"): + sides.append(("baseline", adapter.call_baseline)) + if impl in ("both", "candidate"): + sides.append(("candidate", adapter.call_candidate)) + + res = torch.randn(1, 64, 2048, device=device, dtype=bf16) + upd = torch.randn(1, 64, 2048, device=device, dtype=bf16) + out = torch.empty_like(res) + nc = torch.randn(1, 64, 4096, device=device, dtype=bf16)[:, :, ::2] # non-contiguous + rga, bca = {"function": _EP_RGA}, {"function": _EP_BCAST} + + cases = [ + # tag, workload, inputs, outputs + ("rga-bad-gate-2d", rga, + {"residual": res, "update": upd, "gate": torch.randn(1, 2048, device=device, dtype=bf16)}, [out]), + ("rga-gate-wrong-D", rga, + {"residual": res, "update": upd, "gate": torch.randn(1, 1, 1024, device=device, dtype=bf16)}, [out]), + ("rga-dtype-mismatch", rga, + {"residual": res, "update": upd, "gate": torch.randn(1, 1, 2048, device=device, dtype=torch.float32)}, [out]), + ("rga-gate-leaddim-not1", rga, + {"residual": res, "update": upd, "gate": torch.randn(2, 1, 2048, device=device, dtype=bf16)}, [out]), + ("rga-alias-out", rga, + {"residual": res, "update": upd, "gate": torch.randn(1, 1, 2048, device=device, dtype=bf16)}, [res]), + ("rga-noncontig", rga, + {"residual": nc, "update": torch.empty(1, 64, 2048, device=device, dtype=bf16), + "gate": torch.randn(1, 1, 2048, device=device, dtype=bf16)}, + [torch.empty(1, 64, 2048, device=device, dtype=bf16)]), + ("rga-full-gate-noncontig", rga, + {"residual": res, "update": upd, + "gate": torch.randn(1, 64, 4096, device=device, dtype=bf16)[:, :, ::2]}, [out]), + ("bcast-batch-gt1", bca, + {"a": torch.randn(2, 1, 3, 8, device=device, dtype=bf16), + "b": torch.randn(2, 5, 3, 8, device=device, dtype=bf16)}, + [torch.empty(2, 5, 3, 8, device=device, dtype=bf16)]), + ] + for side, call in sides: + for tag, wl, inp, outs in cases: + expect_raise(f"{tag}:{side}", lambda c=call, w=wl, i=inp, o=outs: c(w, i, o)) + return failures + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--impl", choices=("both", "baseline", "candidate"), default="both") + parser.add_argument("--rows", choices=("all", "grid", "production"), default="all") + parser.add_argument("--quick", action="store_true", help="subsample the regression grid") + parser.add_argument("--max-failures", type=int, default=25) + parser.add_argument("--report", default=None, help="optional JSON report path") + args = parser.parse_args() + + device = torch.device(args.device) + torch.cuda.set_device(device) + torch.set_grad_enabled(False) + + rows: list[dict] = [] + if args.rows in ("all", "grid"): + rows += build_grid_rows(args.quick) + if args.rows in ("all", "production"): + rows += json.loads((_BENCH_DIR / "workloads.json").read_text()) + + failures: list[str] = [] + failures += run_self_test(device) + failures += run_rejection_tests(device, args.impl) + + ran = 0 + for i, spec in enumerate(rows): + if len(failures) >= args.max_failures: + failures.append(f"... stopping early after {args.max_failures} failures") + break + try: + failures += run_row(spec, device, args.impl, i) + except Exception as exc: # noqa: BLE001 - any crash is a row failure + failures.append(f"{spec['id']}: EXCEPTION {type(exc).__name__}: {exc}") + ran += 1 + + summary = { + "impl": args.impl, "rows_requested": len(rows), "rows_ran": ran, + "failures": failures, "ok": not failures, "device": str(device), + "gpu": torch.cuda.get_device_name(device), "torch": torch.__version__, + } + if args.report: + Path(args.report).write_text(json.dumps(summary, indent=2) + "\n") + + print(f"[correctness] impl={args.impl} rows={ran}/{len(rows)} " + f"failures={len(failures)} gpu={summary['gpu']}") + for f in failures: + print(f" FAIL {f}") + if not failures: + print("[correctness] PASS") + return 0 if not failures else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/bench/workloads.json b/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/bench/workloads.json new file mode 100644 index 000000000..ad8b03100 --- /dev/null +++ b/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/bench/workloads.json @@ -0,0 +1,128 @@ +[ + { + "id": "ltx2_full_s8160_c4096", + "production": true, + "preset": "ltx2-hq", + "function": "residual_gate_add", + "gate_mode": "full", + "shapes": { + "residual": {"shape": [1, 8160, 4096], "dtype": "bfloat16", "stride": [33423360, 4096, 1]}, + "update": {"shape": [1, 8160, 4096], "dtype": "bfloat16", "stride": [33423360, 4096, 1]}, + "gate": {"shape": [1, 8160, 4096], "dtype": "bfloat16", "stride": [33423360, 4096, 1]}, + "out": {"shape": [1, 8160, 4096], "dtype": "bfloat16"} + }, + "seed": 1901, + "atol": 0.05, + "rtol": 0.05 + }, + { + "id": "ltx2_bcast_s32640_c4096", + "production": true, + "preset": "ltx2-hq", + "function": "residual_gate_add", + "gate_mode": "bcast", + "shapes": { + "residual": {"shape": [1, 32640, 4096], "dtype": "bfloat16", "stride": [133693440, 4096, 1]}, + "update": {"shape": [1, 32640, 4096], "dtype": "bfloat16", "stride": [133693440, 4096, 1]}, + "gate": {"shape": [1, 1, 4096], "dtype": "bfloat16", "stride": [4096, 4096, 1]}, + "out": {"shape": [1, 32640, 4096], "dtype": "bfloat16"} + }, + "seed": 1902, + "atol": 0.05, + "rtol": 0.05 + }, + { + "id": "ltx2_full_s126_c2048", + "production": true, + "preset": "ltx2-hq", + "function": "residual_gate_add", + "gate_mode": "full", + "shapes": { + "residual": {"shape": [1, 126, 2048], "dtype": "bfloat16", "stride": [258048, 2048, 1]}, + "update": {"shape": [1, 126, 2048], "dtype": "bfloat16", "stride": [258048, 2048, 1]}, + "gate": {"shape": [1, 126, 2048], "dtype": "bfloat16", "stride": [258048, 2048, 1]}, + "out": {"shape": [1, 126, 2048], "dtype": "bfloat16"} + }, + "seed": 1903, + "atol": 0.05, + "rtol": 0.05 + }, + { + "id": "ideogram4_bcast_s4096_c4608", + "production": true, + "preset": "ideogram4-fp8", + "function": "residual_gate_add", + "gate_mode": "bcast", + "shapes": { + "residual": {"shape": [1, 4096, 4608], "dtype": "bfloat16", "stride": [18874368, 4608, 1]}, + "update": {"shape": [1, 4096, 4608], "dtype": "bfloat16", "stride": [18874368, 4608, 1]}, + "gate": {"shape": [1, 1, 4608], "dtype": "bfloat16", "stride": [4608, 4608, 1]}, + "out": {"shape": [1, 4096, 4608], "dtype": "bfloat16"} + }, + "seed": 1904, + "atol": 0.05, + "rtol": 0.05 + }, + { + "id": "flux2_bcast_s4608_c3072", + "production": true, + "preset": "flux2-klein-base", + "function": "residual_gate_add", + "gate_mode": "bcast", + "shapes": { + "residual": {"shape": [1, 4608, 3072], "dtype": "bfloat16", "stride": [14155776, 3072, 1]}, + "update": {"shape": [1, 4608, 3072], "dtype": "bfloat16", "stride": [14155776, 3072, 1]}, + "gate": {"shape": [1, 1, 3072], "dtype": "bfloat16", "stride": [3072, 3072, 1]}, + "out": {"shape": [1, 4608, 3072], "dtype": "bfloat16"} + }, + "seed": 1905, + "atol": 0.05, + "rtol": 0.05 + }, + { + "id": "flux2_bcast_s4096_c3072", + "production": true, + "preset": "flux2-klein-base", + "function": "residual_gate_add", + "gate_mode": "bcast", + "shapes": { + "residual": {"shape": [1, 4096, 3072], "dtype": "bfloat16", "stride": [12582912, 3072, 1]}, + "update": {"shape": [1, 4096, 3072], "dtype": "bfloat16", "stride": [12582912, 3072, 1]}, + "gate": {"shape": [1, 1, 3072], "dtype": "bfloat16", "stride": [3072, 3072, 1]}, + "out": {"shape": [1, 4096, 3072], "dtype": "bfloat16"} + }, + "seed": 1906, + "atol": 0.05, + "rtol": 0.05 + }, + { + "id": "flux2_bcast_s512_c3072", + "production": true, + "preset": "flux2-klein-base", + "function": "residual_gate_add", + "gate_mode": "bcast", + "shapes": { + "residual": {"shape": [1, 512, 3072], "dtype": "bfloat16", "stride": [1572864, 3072, 1]}, + "update": {"shape": [1, 512, 3072], "dtype": "bfloat16", "stride": [1572864, 3072, 1]}, + "gate": {"shape": [1, 1, 3072], "dtype": "bfloat16", "stride": [3072, 3072, 1]}, + "out": {"shape": [1, 512, 3072], "dtype": "bfloat16"} + }, + "seed": 1907, + "atol": 0.05, + "rtol": 0.05 + }, + { + "id": "ltx2_broadcast_add_4d_s126_p3_c2048", + "production": true, + "preset": "ltx2-hq", + "function": "broadcast_add_4d", + "shapes": { + "a": {"shape": [1, 1, 3, 2048], "dtype": "bfloat16", "stride": [6144, 6144, 2048, 1]}, + "b": {"shape": [1, 126, 3, 2048], "dtype": "bfloat16", "stride": [774144, 6144, 2048, 1]}, + "out": {"shape": [1, 126, 3, 2048], "dtype": "bfloat16"} + }, + "seed": 1908, + "atol": 0.05, + "rtol": 0.05 + } +] diff --git a/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/config.toml b/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/config.toml index 576104eda..b1c9ac52e 100644 --- a/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/config.toml +++ b/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/config.toml @@ -12,8 +12,21 @@ entry_points = [ [build] language = "cuda" destination_passing_style = true -baseline_entry_point = "baseline/kernel.cu::residual_gate_add_baseline" -candidate_entry_point = "solution/kernel.cu::residual_gate_add_candidate" +# The upstream "kernel" is PyTorch-eager elementwise math (residual + update*gate +# and a 4D broadcast add), not a CUDA/Triton kernel. Per the standalone contract +# (and the b200_diffusion_fuse_scale_shift precedent for non-CUDA upstreams), the +# baseline is exposed through local destination-passing Python launchers in +# baseline/binding.py with the same ABI as the candidate (outputs passed last); +# the candidate is the fused CUDA module from solution/kernel.cu. One entry per +# public op. (Reconciled from the generator's singular kernel.cu default to the binding.py baseline.) +baseline_entry_points = [ + "baseline/binding.py::residual_gate_add", + "baseline/binding.py::broadcast_add_4d", +] +candidate_entry_points = [ + "solution/kernel.cu::residual_gate_add", + "solution/kernel.cu::broadcast_add_4d", +] [benchmark] warmup_runs = 10 diff --git a/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/docs/baseline_source.md b/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/docs/baseline_source.md new file mode 100644 index 000000000..661135360 --- /dev/null +++ b/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/docs/baseline_source.md @@ -0,0 +1,91 @@ +# Baseline Source Lineage + +## Upstream provenance + +- Repository: `https://github.com/sgl-project/sglang` +- Branch: `main` +- Resolved commit SHA: `8314247d9de0fa2c58e34756b3e1dbc6cf815dfd` +- Resolution method: `git ls-remote https://github.com/sgl-project/sglang.git refs/heads/main` +- Resolution time: 2026-06-25 (UTC), at baseline-recovery time for this task. +- Fetch method: raw file fetch at the resolved commit + (`https://raw.githubusercontent.com/sgl-project/sglang/8314247d9de0fa2c58e34756b3e1dbc6cf815dfd/python/sglang/multimodal_gen/runtime/models/dits/.py`). +- Pin note: an earlier resolution during this task returned + `67b2a9ed0cfba8ec625d3f26548e502646fd914d`; upstream `main` advanced to + `8314247d…` shortly after. The residual-gate (`residual + update*gate`) and the + LTX 4D broadcast-add expressions in `ltx_2.py` / `ideogram.py` / `flux_2.py` + are byte-identical between the two commits (verified by diffing the relevant + lines), so the recovered semantics are unaffected. This task pins to + `8314247d…` as the baseline-recovery commit. + +## Copied / referenced upstream files (at the resolved commit) + +| Upstream path | Role | Relevant lines | +|---|---|---| +| `python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py` | LTX-2 DiT blocks: residual-gate updates and the 4D cross-attn-gate broadcast add | `residual+update*gate` at 1031, 1048, 1073, 1097, 1198, 1222; 4D broadcast at 1126-1140 | +| `python/sglang/multimodal_gen/runtime/models/dits/ideogram.py` | Ideogram4 DiT block: `x = x + gate * norm(sublayer(x))` | 327-337 | +| `python/sglang/multimodal_gen/runtime/models/dits/flux_2.py` | FLUX.2 single/double-stream blocks: modulation residual gate | 755, 873, 880, 883, 892 | + +The verbatim relevant excerpts (with line numbers) are preserved in +`baseline/upstream_residual_gate_snippets.md` for provenance/evidence. These full +model files are NOT copied wholesale into `baseline/` because they contain large +amounts of unrelated model code; the standalone baseline is the elementwise +expression itself (see "Baseline implementation" below). + +## K — Kernel semantics and callsite contract (recovered) + +Two elementwise operator families, all bf16 in production: + +- `residual_gate_add`: `out = residual + update * gate` (the multiply is commutative; + upstream writes both `residual + update * gate` and `residual + gate * update`). + - LTX-2: `hidden_states = hidden_states + attn_hidden_states * vgate_msa` (gate is `[B,1,D]` or `[B,L,D]`). + - Ideogram4: `x = x + gate_msa * self.attention_norm2(attn_out)` and `x = x + gate_mlp * self.ffn_norm2(...)` (gate `[B,1,D]`, tanh-gated upstream; the gate tensor is passed in already-computed). + - FLUX.2: `hidden_states = hidden_states + mod_gate * attn_output`, `+ gate_mlp * ff_output`, `+ c_gate_mlp * context_ff_output` (gate `[B,1,D]`). + - `gate` appears in two shapes: full `[B,L,D]` (equal to `residual`) and row-broadcast `[B,1,D]`. +- `broadcast_add_4d`: LTX-2 cross-attn gate/scale-shift table construction, + `out = per_layer_table[None, None, :, :] + temb.reshape(B, S, P, D)`, i.e. + `a=[1,1,P,D]` broadcasts over the sequence dim against `b=[1,S,P,D]`. + Production row: `a=[1,1,3,2048]`, `b=[1,126,3,2048]`, `out=[1,126,3,2048]`; + `out[0,s,p,d] = a[0,0,p,d] + b[0,s,p,d]`. + +### Supported shape contract: batch size 1 (B=1) + +Every frozen production row is B=1, so the standalone kernel ABI deliberately +supports only B=1: `residual_gate_add` accepts a full gate `[1,L,D]` or a +row-broadcast gate `[1,1,D]`; `broadcast_add_4d` accepts `a=[1,1,P,D]`, +`b=[1,S,P,D]`. A true batched broadcast (B>1 — e.g. gate `[B,1,D]` or +`b=[B,S,P,D]` with B>1, which eager PyTorch would broadcast per batch) is OUT OF +SCOPE and is rejected on BOTH sides: the candidate via its CUDA host checks, and +the eager baseline via the shared `bench/adapter.py::_validate` (which runs before +either implementation). This keeps the ABI symmetric and reward-hack-resistant; +the `rga-gate-leaddim-not1` and `bcast-batch-gt1` rejection tests cover it. +Supporting B>1 would add per-vector batch index arithmetic to the +bandwidth-critical path for a shape absent from the frozen workload, so it is +documented as out-of-scope rather than implemented. + +These callsites are PyTorch-eager elementwise ops. The torch profiler (see +`../../docs/sglang_recent_diffusion_b200_profile_audit_2026-06-24.md`) shows them +as separate `aten::mul` then `aten::add` launches plus a hot 4D broadcast add. + +## R — Correctness oracle and baseline path + +- Oracle (fp32, one-round): `residual_gate_add` -> `(residual.float() + update.float() * gate.float()).to(out_dtype)`; + `broadcast_add_4d` -> `(a.float() + b.float()).to(out_dtype)` with `a` broadcast over the seq dim. +- Baseline implementation (faithful eager): `baseline/binding.py` reproduces the + profiled two-op eager path (`update*gate` into a cached preallocated scratch buffer, + then `+residual` into the output) for `residual_gate_add`, and a single eager broadcast + add for `broadcast_add_4d`. Exposed through a destination-passing ABI identical to the + candidate (output passed last, current CUDA stream, no allocation in the timed steady + state). No `sglang` import at correctness/benchmark runtime. + +## W — Workload shape set and benchmark methodology + +- The 8 frozen production rows are in `bench/workloads.json`, matching the + `diffusion_residual_gate_add__multi_shape` section of + `../../docs/diffusion_benchmark_shape_coverage.md` (source of truth). +- Benchmark methodology follows `../../docs/standalone_diffusion_benchmark.md` + and starts `bench/benchmark.py` from `../../docs/standalone_diffusion_benchmark_template.py`. + +## Local edits + +- None to upstream source (the relevant expressions are reproduced in `baseline/binding.py` + as the local baseline; upstream excerpts are preserved verbatim for provenance only). diff --git a/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/docs/benchmark_method.md b/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/docs/benchmark_method.md new file mode 100644 index 000000000..564826752 --- /dev/null +++ b/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/docs/benchmark_method.md @@ -0,0 +1,72 @@ +# Benchmark Method + +## Harness +`bench/benchmark.py` starts from the repository standalone timing harness +(`../../docs/standalone_diffusion_benchmark_template.py`). The timing/scoring +logic is unchanged; the only deltas are (1) an additive provenance merge in +`_provenance` (it calls the optional `adapter.extra_provenance()` hook), which +does not affect measurement, and (2) `torch.cuda.set_device(args.device)` is +called before the adapter is imported in `_run_one_workload`, so a non-default +`--device` builds/reports the candidate for the intended GPU. It provides +CUDA-event timing, interleaved A/B sampling per trial, inner-loop +amplification toward `target_sample_us`, fresh random inputs per trial, +preallocated + poisoned outputs, the equal-weight geometric-mean headline +(`baseline_median_us / candidate_median_us`), and per-row +median/mean/std/min/p10/p90 with full provenance in `results.jsonl`. + +The task-specific `bench/adapter.py` supplies only `make_case`, `call_baseline`, +and `call_candidate`, routed through one shared `_dispatch` so both sides pay +byte-identical adapter overhead. Outputs are preallocated in `make_case`; nothing +is allocated in the timed path. `bench/adapter.py` asserts no `sglang` import at +runtime. + +## ABI +Both sides use a destination-passing ABI with the output tensor passed last: +- `residual_gate_add(residual, update, gate, out)` +- `broadcast_add_4d(a, b, out)` + +Both ops are batch-1 only (every production row is B=1); a true B>1 broadcast is +deliberately out of scope and rejected symmetrically on both sides (see +`docs/baseline_source.md`). + +Baseline (`baseline/binding.py`) is the faithful PyTorch-eager path: +`residual_gate_add` runs `torch.mul(update, gate, out=scratch)` then +`torch.add(residual, scratch, out=out)` — the profiled two-launch `mul`+`add` +pattern — with a cached, preallocated scratch buffer so the timed steady state +contains only the two real kernel launches; `broadcast_add_4d` is a single eager +`torch.add`. Candidate (`solution/kernel.cu`, built by `solution/build.py` via +`tvm_ffi.cpp.load`) fuses each into a single CUDA pass. Every launch uses +`at::cuda::getCurrentCUDAStream()`. + +## Compile flags (symmetric) +- Candidate CUDA: `-std=c++17 -O3` and the device's native gencode + (`-gencode=arch=compute_100,code=sm_100` on B200). No `--use_fast_math` (the + eager baseline does not use fast math, so the candidate must not either). +- Baseline: pure torch eager ops; no extra compile flags. Both sides are invoked + through one Python adapter dispatch with matched overhead. + +## Workloads +`bench/workloads.json` is the frozen source of truth (8 production rows matching +the `diffusion_residual_gate_add__multi_shape` section of +`../../docs/diffusion_benchmark_shape_coverage.md`). Workloads, tolerances, +scoring, and timing rules are frozen before tuning; changing any of them requires +deleting old results and remeasuring both baseline and candidate. + +## Numerics / tolerances +Candidate accumulates in fp32 and rounds once. Correctness compares candidate vs +an fp32 one-round oracle and vs the baseline, within bf16/fp16 `atol=rtol=5e-2` +and fp32 `atol=rtol=1e-5`, with NaN/Inf rejection. + +## Reporting (per the contract) +Headline: equal-weight geomean over all 8 production rows. Secondary diagnostics +(a residual-gate-only geomean and a call-count-weighted view) and the per-row +speed-of-light/roofline table with NCU-derived bounds are reported in +`docs/results.md` (the 4D broadcast row has little fusion headroom and is +reported with its speed-of-light ceiling). + +## Remote / GPU +All correctness, benchmark, profiler, and NCU runs execute on B200 (host +`ion-b200`). An idle GPU is selected via `nvidia-smi`, its id exported as +`REMOTE_GPU_ID`, and reused across baseline/candidate/correctness/benchmark/ +profile; host/id/model and before/after idle state are recorded in +`docs/results.md`/`docs/run_log.md`. diff --git a/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/docs/results.md b/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/docs/results.md new file mode 100644 index 000000000..9cdd64d4d --- /dev/null +++ b/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/docs/results.md @@ -0,0 +1,147 @@ +# Results — b200_diffusion_residual_gate_add__multi_shape + +## Conclusion: GO +The fused candidate is correctness-clean on B200 and faster than the faithful +PyTorch-eager production path on every production row. Headline equal-weight +geomean **2.193x** (all rows >= 1.18x). NCU + roofline confirm the result is +bandwidth/launch-bound as expected for elementwise traffic, and the analyze pass +(task3/task7/task10) accepts the baseline/numerics/layout and ranks remaining +optimization as modest, conditional upside — so the generic vectorized candidate +is accepted as the final implementation for this task. + +Scope: the ABI is batch-1 (every frozen production row is B=1). Full gate +`[1,L,D]`, row-broadcast gate `[1,1,D]`, and 4D `a=[1,1,P,D]`/`b=[1,S,P,D]` are +supported; a true B>1 broadcast is deliberately out of scope and is rejected +symmetrically on both the candidate and the eager baseline (shared +`bench/adapter.py::_validate`), with `rga-gate-leaddim-not1` and `bcast-batch-gt1` +rejection tests. See `docs/baseline_source.md` for the contract rationale. + +## Environment +- Host `ion-b200` (`innomatrix-us-adc-smb200-0003`), container `sglang_bbuf` + (`lmsysorg/sglang:dev`); torch 2.11.0+cu130, CUDA 13.0, tvm-ffi 0.1.9, nvcc 13.0. +- NVIDIA B200 (192 GB HBM3e, 148 SMs, ~8 TB/s nominal; NCU sustained ref ~7.2 TB/s). +- **All final evidence is on ONE pinned idle GPU: physical GPU 7** + (`REMOTE_GPU_ID=7 CUDA_VISIBLE_DEVICES=7`, fail-closed guard + `KDA_REQUIRE_PINNED_GPU=1`); correctness, benchmark, and NCU were all collected + on GPU 7, idle before/after (0%/0MiB). See `docs/run_log.md`. (Earlier rounds + also ran on GPU 2/GPU 0; those are superseded by this unified GPU-7 chain.) +- Baseline source: sgl-project/sglang `main` @ `8314247d9de0fa2c58e34756b3e1dbc6cf815dfd` + (`docs/baseline_source.md`). Candidate `solution/kernel.cu` sha256: the Round-4 + benchmark was measured on `a450f863…`; subsequent review fixes are host-side and + compute-neutral on the timed single-GPU path — R5 added a `out.ndim()==4` reject + (`27f67c5f…`), R7 added same-device validation + a `CUDAGuard` (no-op when the + tensor device is already current) giving the current source `a6ab9c86…` (R6 was + adapter-only, kernel unchanged). The device kernels and the valid-input timed + path are identical across these, so the measured geomean stands; correctness was + re-confirmed 67/67 on each rebuild (see `docs/run_log.md`). R8 was harness-only + (`bench/benchmark.py` now sets the selected CUDA device before collecting + provenance, so a non-default `--device` records the right GPU/SM/flags); the + candidate hash is unchanged and both gates were re-run fresh on idle GPU 7 + (67/67, geomean 2.186x). + +## Final commands (all on GPU 7) +```bash +export KDA_REQUIRE_PINNED_GPU=1 REMOTE_GPU_ID=7 CUDA_VISIBLE_DEVICES=7 +python bench/correctness.py --impl both --rows all --report /tmp/rga_correctness_final.json +python bench/benchmark.py --out bench/results.jsonl +ncu --set basic --launch-skip 6 --launch-count 1 --target-processes all -o /tmp/ncu_ python /tmp/rga_profile.py +``` + +## Correctness (AC-3 / AC-4) +`bench/correctness.py --impl both --rows all` (GPU 7): **67/67 PASS**. Covers the 8 +production rows (candidate vs fp32 one-round oracle AND vs faithful eager baseline, +bf16 atol=rtol=5e-2), the regression grid (full/broadcast gate; bf16/fp16/fp32; +odd-D / non-vec-aligned-D / small-L tails; deterministic zero/sign rows; repeated +randomized seeds; 4D over multiple frame counts), a poison self-test, and both-side +rejection (full-gate-noncontig, bad-gate-2d, gate-leaddim-not1, dtype-mismatch, +alias, non-contiguous, 4D batch>1). + +## Performance (AC-5) — candidate vs faithful eager two-op baseline (GPU 7) +Baseline = the profiled production path (`torch.mul(update,gate,out=scratch)` then +`torch.add(residual,scratch,out=out)`, two launches + one temp + two dispatches; +single `torch.add` for the 4D row). Candidate = one fused CUDA pass. CUDA-event +median per call (matched ratio 1.0). + +| Workload | gate | speedup | baseline us | candidate us | +|---|---|---:|---:|---:| +| ltx2_full_s8160_c4096 | full | 1.5977 | 66.180 | 41.421 | +| ltx2_bcast_s32640_c4096 | bcast | 2.9781 | 420.216 | 141.104 | +| ltx2_full_s126_c2048 | full | 1.7535 | 10.590 | 6.039 | +| ideogram4_bcast_s4096_c4608 | bcast | 2.8348 | 64.616 | 22.793 | +| flux2_bcast_s4608_c3072 | bcast | 3.2779 | 47.492 | 14.489 | +| flux2_bcast_s4096_c3072 | bcast | 3.3286 | 41.456 | 12.455 | +| flux2_bcast_s512_c3072 | bcast | 1.7542 | 19.622 | 11.186 | +| ltx2_broadcast_add_4d | - | 1.1805 | 11.716 | 9.925 | + +Headline and secondary views (the row mix weights broadcast-gate cases heavily, so +secondary views are reported): +- **All-8 equal-weight geomean: 2.193x** (the contract headline). Min 1.18x, max 3.33x. +- Residual-gate-only (7 rows) geomean: **2.396x**. +- Call-count-weighted geomean (documented profile call-counts, 6-row subset, + 82,134 calls): **1.676x** — lower because the highest-frequency rows are the tiny + full-gate `[1,126,2048]` (33,123 calls, 1.75x) and the 4D add (13,392 calls, + 1.18x), which have the smallest per-call speedups. + +The win is "fused single CUDA kernel vs the faithful eager two-op production path" +— it removes one kernel launch, the intermediate temp's full write+read, and one +Python dispatch; it is not a single-kernel-vs-single-kernel algorithmic speedup. +(Run-to-run the headline geomean is stable: 2.199x on GPU 2, 2.193x on GPU 7, and +2.186x on a fresh GPU-7 re-run after the Round-8 benchmark-provenance fix — all +within run-to-run noise, per-row deltas < 1%.) + +## Roofline / speed-of-light (AC-7), GPU 7 +Candidate byte model (bf16, 2 B/elem): full gate ~8 B/elem (r+u+g read, out write); +broadcast gate ~6 B/elem (r+u read, out write; gate cached); 4D add ~4 B/elem +(b read, out write; a cached). Achieved BW from the CUDA-event median above; +DRAM%/SM% from NCU (`--set basic`, GPU 7) for the three profiled rows. + +| Workload | elems | cand B/elem | cand us | achieved GB/s | % of ~8 TB/s | NCU DRAM% | NCU SM% | named bound | +|---|---:|---:|---:|---:|---:|---:|---:|---| +| ltx2_full_s8160_c4096 | 33.42M | 8 | 41.42 | 6455 | 80.7 | 72.5 | 55.7 | DRAM-bound (near roofline) | +| ltx2_bcast_s32640_c4096 | 133.69M | 6 | 141.10 | 5685 | 71.1 | 59.5 | 69.7 | SM/occupancy-leaning (modulo+ldg; occ 50.5%) | +| ltx2_full_s126_c2048 | 0.258M | 8 | 6.04 | 342 | 4.3 | 3.8 | 3.1 | launch/grid-bound (grid 126 < 148 SMs) | +| ideogram4_bcast_s4096_c4608 | 18.87M | 6 | 22.79 | 4968 | 62.1 | - | - | DRAM-bound (bandwidth) | +| flux2_bcast_s4608_c3072 | 14.16M | 6 | 14.49 | 5862 | 73.3 | - | - | DRAM-bound (bandwidth) | +| flux2_bcast_s4096_c3072 | 12.58M | 6 | 12.46 | 6062 | 75.8 | - | - | DRAM-bound (bandwidth) | +| flux2_bcast_s512_c3072 | 1.57M | 6 | 11.19 | 844 | 10.5 | - | - | launch/occupancy-leaning (small) | +| ltx2_broadcast_add_4d | 0.774M | 4 | 9.93 | 312 | 3.9 | - | - | launch/low-fusion (inferred, not NCU-profiled) | + +NCU bound interpretation (task10, Codex-confirmed; GPU 7 `.ncu-rep` retained under +the ignored remote `/tmp`): +- Large full-gate (`ltx2_full_s8160`): DRAM-bound near roofline (NCU 72.5% DRAM, + "DRAM bottleneck"; ~6.46 TB/s effective at benchmark speed). +- Large broadcast-gate (`ltx2_bcast_s32640`): SM/occupancy-leaning, DRAM headroom + (NCU SM 69.7% > DRAM 59.5%, occupancy 50.5%) — the per-vector `v % row_vec` + modulo + `__ldg` indexing is the plausible instruction-side limiter. +- Small (`ltx2_full_s126`): launch/grid-bound (NCU: grid 126 blocks < 148 SMs, + occupancy 12.7%) — the win is collapsing two launches into one. +- 4D add: launch/low-fusion bound, **inferred** from the byte model + 1.18x speedup + (this row was not separately NCU-profiled). + +## Analyze pass (task3 / task7 / task10, via Codex) +- task3 — baseline faithfulness ACCEPT, fp32 one-round numerics + bf16 5e-2 oracle + ACCEPT, out-of-place + contiguous-only alias/layout ACCEPT. No unfair/illusory + win once framed as fused-vs-eager (above) with secondary views. +- task7 — ranked optimization directions: (1) highest-value = remove the + broadcast-gate per-vector modulo (bcast_big SM-leaning), "modest but real, + moderate risk, worth a bounded experiment if more work is allowed"; (2) full-gate + large already ~roofline (low headroom); (3) small rows launch-bound (per-row CTA + split risks hurting large rows unless separately dispatched); (4) 4D low priority. +- task10 — confirmed the per-row bound characterization above (4D marked inferred). + +## Optimization decision (task7 -> task8) +Accept the current generic grid-stride fused candidate as the final implementation +for this task. Rationale (evidence-backed): the success bar is decisively met +(67/67 correctness; positive geomean 2.193x on every row; bounds explained by +NCU/roofline); the one ranked edit (broadcast-gate modulo / occupancy) is, per NCU +and the analyze pass, modest and conditional upside for rows that already win +2.83-3.33x, and the natural implementation (row-shaped launch) carries multi-shape +load-imbalance risk; and a single generic fused kernel is explicitly within the +plan's lower-bound path boundary. No `docs/dispatch.md` is written because no +shape-bucket specialization landed. + +### Ranked future optimization (not pursued for this deliverable) +1. Broadcast-gate: eliminate the per-vector `v % row_vec` modulo (grid-stride + row-offset tracking, or a one-block-per-row launch) to lift `ltx2_bcast_s32640` + from SM/occupancy-leaning (NCU SM 69.7%, occ 50.5%) toward its DRAM roofline. + Requires re-running strict-pinned correctness + benchmark after the edit. diff --git a/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/docs/run_log.md b/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/docs/run_log.md new file mode 100644 index 000000000..926d1cf53 --- /dev/null +++ b/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/docs/run_log.md @@ -0,0 +1,170 @@ +# Remote B200 Run Log + +## Environment (AC-6) +- Host: `ion-b200` (`innomatrix-us-adc-smb200-0003`), user `sglang-omni`. +- Container: `sglang_bbuf` (`lmsysorg/sglang:dev`). +- Toolchain: torch `2.11.0+cu130`, torch.version.cuda `13.0`, tvm-ffi `0.1.9`, + nvcc `cuda_13.0.r13.0` (V13.0.88). +- GPU model: NVIDIA B200 (192 GB HBM3e, 148 SMs). +- Pinned GPU: `REMOTE_GPU_ID=2`, `CUDA_VISIBLE_DEVICES=2` (sole visible device -> + `cuda:0`). The fail-closed pin guard (`bench/adapter.assert_pinned_gpu`, + enabled via `KDA_REQUIRE_PINNED_GPU=1`) was active for the benchmark run. +- Idle evidence for GPU 2 (`index, util%, mem_used MiB`): + - Before: `2, 0, 0` + - After: `2, 0, 0` (post-run summary snapshot: `2, 0%, 4 MiB`) + - GPUs 3-6 were busy (~95%) on unrelated jobs throughout; GPU 1 had standing + memory occupancy. GPU 2 remained uncontended. + +## Workspace +- Task folder streamed (tar over ssh, excluding `.humanize`/`.build`/`.git`) into + `/home/sglang-omni/bbuf/kda_tasks/residual_gate_add` inside the container. +- Local source commit: `70ceefac4` (working branch + `kda/k19_b200_diffusion_residual_gate_add__multi_shape-20260625-083922-58057`). + +## Commands +```bash +# build (tvm-ffi) +docker exec -e CUDA_VISIBLE_DEVICES=2 sglang_bbuf bash -lc \ + 'cd && python -c "from solution.build import load_candidate_module; load_candidate_module()"' +# correctness gate +docker exec -e CUDA_VISIBLE_DEVICES=2 sglang_bbuf bash -lc \ + 'cd && python bench/correctness.py --impl both --rows all' +# frozen benchmark (pin guard active) +docker exec -e KDA_REQUIRE_PINNED_GPU=1 -e REMOTE_GPU_ID=2 -e CUDA_VISIBLE_DEVICES=2 sglang_bbuf bash -lc \ + 'cd && python bench/benchmark.py --out bench/results.jsonl' +``` + +## Build (AC-2/AC-3) +- `solution/kernel.cu` compiled cleanly via `tvm_ffi.cpp.load`; both exports + (`residual_gate_add`, `broadcast_add_4d`) present. No `--use_fast_math`. + +## Correctness (AC-3 / AC-4) — PASS +- `bench/correctness.py --impl both --rows all`: **56/56 rows, 0 failures**, on + NVIDIA B200. Covers the 8 production rows (candidate vs fp32 one-round oracle + AND candidate vs faithful eager baseline, bf16 atol=rtol=5e-2), the regression + grid (full/broadcast gate; bf16/fp16/fp32; odd-D / non-vec-aligned-D / small-L + tails; 4D over multiple frame counts; sign/zero), the poison self-test, and + both-side rejection tests (incl. `rga-full-gate-noncontig`, `rga-bad-gate-2d`, + `rga-gate-leaddim-not1`, dtype mismatch, aliasing, non-contiguous, 4D batch>1). + +## Benchmark (AC-5) — candidate vs faithful eager baseline, geomean 2.199x +All 8 production rows PASSED (matched ratio 1.0). Baseline = profiled two-launch +eager `mul`+`add`; candidate = single fused CUDA pass. + +| Workload | gate | speedup | baseline (us) | candidate (us) | +|---|---|---:|---:|---:| +| ltx2_full_s8160_c4096 | full | 1.6016 | 66.328 | 41.413 | +| ltx2_bcast_s32640_c4096 | bcast | 2.9638 | 419.032 | 141.384 | +| ltx2_full_s126_c2048 | full | 1.7970 | 17.661 | 9.828 | +| ideogram4_bcast_s4096_c4608 | bcast | 2.8252 | 64.350 | 22.777 | +| flux2_bcast_s4608_c3072 | bcast | 3.2763 | 47.488 | 14.495 | +| flux2_bcast_s4096_c3072 | bcast | 3.3356 | 41.568 | 12.462 | +| flux2_bcast_s512_c3072 | bcast | 1.7261 | 19.546 | 11.324 | +| ltx2_broadcast_add_4d_s126_p3_c2048 | - | 1.2028 | 11.510 | 9.569 | + +Headline: **geomean 2.1990x**, arithmetic mean 2.3411x, min 1.2028x, max 3.3356x. + +## Provenance (in `bench/results.jsonl`, AC-5) +- baseline_commit `8314247d9de0fa2c58e34756b3e1dbc6cf815dfd` +- candidate_kernel_sha256 `a450f8635e11791f…` +- remote_gpu_id `2`; tvm_ffi `0.1.9`; nvcc `cuda_13.0` +- gpu `NVIDIA B200`; torch `2.11.0+cu130`; torch_cuda `13.0` + +## Preliminary roofline sanity (AC-7 — full per-row table + NCU is the next round) +- `ltx2_bcast_s32640_c4096` (133.69M bf16 elems): candidate streams residual+update + (read) + out (write) ~ 3 x 133.69M x 2B = ~802 MB in 141.384 us -> + ~5.67 TB/s achieved (~71% of ~8 TB/s HBM3e). Bandwidth-bound, near roofline. + Eager baseline (~5 full passes incl. temp write/read + 2 launches + 2 dispatches) + ~ 1337 MB in 419.032 us -> ~3.19 TB/s; the candidate's win is the eliminated + temp traffic + one launch + one dispatch. +- Full per-row speed-of-light/bytes/%peak table and the NCU bound interpretation + are in `docs/results.md` (collected in the round below). + +## Round 3 — strict-pinned reruns + NCU (GPU 0, idle) +- GPU 0 idle before: `0, NVIDIA B200, 0%, 0 MiB`; after: idle. +- Correctness rerun (strict pin: `KDA_REQUIRE_PINNED_GPU=1 REMOTE_GPU_ID=0 + CUDA_VISIBLE_DEVICES=0 python bench/correctness.py --impl both --rows all + --report /tmp/rga_correctness_final.json`): **67/67 PASS** (adds AC-4 zero/sign + rows + repeated-seed rows over the R2 56/56). The pin guard was active. +- NCU (`ncu --set basic --launch-skip 6 --launch-count 1`, GPU 0) on the three + representative rows: + | row | DRAM % | Compute(SM) % | NCU duration | occupancy | grid | NCU bound | + |---|---:|---:|---:|---:|---|---| + | ltx2_full_s8160_c4096 (full) | 72.09 | 55.80 | 51.3us | 60.5% | 16320x256 | DRAM bottleneck | + | ltx2_bcast_s32640_c4096 (bcast) | 59.53 | 69.77 | 195.7us | 50.6% | 65280x256 | SM/instruction-leaning | + | ltx2_full_s126_c2048 (full) | 3.72 | 3.14 | 6.4us | 12.8% | 126x256 | grid < 148 SMs (launch-bound) | +- (R3 NCU was captured to stdout via `--set basic`, not saved as `.ncu-rep`; the + scratch harness is `/tmp/rga_profile.py`. The Round-4 reruns below save the raw + `.ncu-rep` files.) + +## Round 4 — UNIFIED final evidence on one pinned GPU (GPU 7, idle) +Per the plan's AC-6 (one `REMOTE_GPU_ID` reused across correctness/benchmark/ +profile), all final evidence is collected on a single physical GPU (7), idle +before and after (`7, NVIDIA B200, 0%, 0 MiB` both), with the fail-closed pin env +`KDA_REQUIRE_PINNED_GPU=1 REMOTE_GPU_ID=7 CUDA_VISIBLE_DEVICES=7` on every command. +- Correctness: `bench/correctness.py --impl both --rows all` -> **67/67 PASS** + (report `/tmp/rga_correctness_final.json`). +- Benchmark: `bench/benchmark.py --out bench/results.jsonl` -> 8/8 rows PASS, + geomean **2.193x**; `results.jsonl` provenance `remote_gpu_id=7`, + baseline_commit 8314247d, candidate sha256 a450f863. +- NCU (`-o /tmp/ncu_`, GPU 7): full_big DRAM 72.49% / SM 55.73%; bcast_big + DRAM 59.45% / SM 69.68% (occ 50.5%); full_small DRAM 3.77% / SM 3.06% + (grid 126 < 148 SMs). Raw `.ncu-rep` saved under the ignored remote `/tmp` + (`/tmp/ncu_full_big.ncu-rep`, `/tmp/ncu_bcast_big.ncu-rep`, + `/tmp/ncu_full_small.ncu-rep`); not staged for the PR. +- Per-row roofline table and bound interpretation: `docs/results.md`. + +## Round 5 — review-fix re-validation (compute-neutral) +Two Codex code-review findings fixed: P2 (`bench/benchmark.py` now calls +`torch.cuda.set_device` before importing/building the adapter) and P3 +(`solution/kernel.cu` `broadcast_add_4d` now requires `out.ndim()==4`). Both are +host-side only — the device kernels and the timed path for valid inputs are +byte-identical — so the geomean 2.193x measured in Round 4 stands. +- New candidate `solution/kernel.cu` sha256: `27f67c5ff283d87b8d56b3c55f8caef749e73c845a8667e063f0b8d479a6c749` + (Round-4 benchmark provenance recorded the pre-fix `a450f863…`; the only kernel + delta is the added host-side rank reject, which no valid/production row triggers). +- Correctness re-confirmed on the clean rebuild: `bench/correctness.py --impl both + --rows all` = **67/67 PASS** on B200 (GPU 1; report `/tmp/rga_correctness_r5.json`). +- Benchmark re-run deferred: at fix time all 8 B200 GPUs were busy (util 49-96%) or + had standing memory occupancy, so per the idle-GPU rule no timed numbers were + collected on a non-idle card. The Round-4 geomean stands (compute-neutral edits); + a refresh on an idle GPU is optional and would reproduce the same numbers. + +## Round 7 — review fix: multi-GPU device guard (compute-neutral) +Codex review P2: the exported kernels launched on the process-current device's +stream rather than the input tensors' device. Fixed in `solution/kernel.cu`: both +`residual_gate_add` and `broadcast_add_4d` now validate all operands share one CUDA +device id and install a `c10::cuda::CUDAGuard(dev)` + take that device's stream +before launching. On a single visible GPU (the pinned runs) the guard is a no-op +(the tensor device is already current), so the timed path and the Round-4 geomean +2.193x are unchanged. +- New candidate `solution/kernel.cu` sha256: `a6ab9c86e7ecf2c056d00c568e8895d0fbbf9dcac1494c78a71d569500725fc8`. +- Rebuilt on B200 (GPU 1); `bench/correctness.py --impl both --rows all` = **67/67 + PASS** (report `/tmp/rga_correctness_r7.json`). Benchmark re-run deferred (no + fully-idle GPU; compute-neutral on single-GPU). + +## Round 8 — review fix: benchmark provenance device-order (re-validated on GPU 7) +Codex review P2: `bench/benchmark.py::main()` collected provenance +(`_provenance` -> `_gpu_name()` + adapter `extra_provenance()`/candidate build) +BEFORE setting the CUDA device, so a non-default `--device` would record cuda:0's +name/SM/compile-flags while workloads ran on `args.device`. Fixed: `main()` now +calls `torch.cuda.set_device(torch.device(args.device))` before `_provenance()`, +mirroring the R5 `_run_one_workload` device-first pattern. Harness-only — the +candidate `solution/kernel.cu` is unchanged (hash still +`a6ab9c86…`), and on a pinned single-visible-GPU run the current device already +equals `args.device`, so the fix is compute-neutral. +- An idle GPU was available this round, so both gates were re-run fresh on the + unified-evidence card (physical **GPU 7**, strict pin + `KDA_REQUIRE_PINNED_GPU=1 REMOTE_GPU_ID=7 CUDA_VISIBLE_DEVICES=7`): + - `bench/correctness.py --impl both --rows all` -> **67/67 PASS** + (report `/tmp/rga_correctness_r8.json`). + - `bench/benchmark.py --out bench/results.jsonl` -> **8/8 PASS**, geomean + **2.186x** (min 1.186x, max 3.277x) — reproduces the Round-4 2.193x within + run-to-run noise (per-row deltas < 1%). `results.jsonl` provenance now records + `gpu=NVIDIA B200`, `remote_gpu_id=7`, `candidate_kernel_sha256=a6ab9c86…`, + `baseline_commit=8314247d…` (1 provenance + 8 result + 1 summary events). + - Idle evidence GPU 7: before `7, NVIDIA B200, 0%, 0 MiB`. After my run returned, + an unrelated job landed on GPU 7 (post-run snapshot `7, …, 96%, ~42 GB`); the + measured times matching Round-4's idle-GPU numbers within noise confirm the + measurement window itself was uncontended (heavy contention would have inflated + and de-stabilized the per-row times, which it did not). diff --git a/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/solution/__init__.py b/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/solution/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/solution/build.py b/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/solution/build.py new file mode 100644 index 000000000..632e0e6d6 --- /dev/null +++ b/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/solution/build.py @@ -0,0 +1,60 @@ +"""Build and load the candidate CUDA module via tvm-ffi. + +Compiles solution/kernel.cu into a tvm-ffi module exposing the two +destination-passing entry points (``residual_gate_add``, ``broadcast_add_4d``). +Uses ``tvm_ffi.cpp.load`` directly (the same underlying builder the production +jit stack uses) plus torch's include and library paths for +``at::cuda::getCurrentCUDAStream``. No sglang import. + +Compile flags (recorded in docs/benchmark_method.md): ``-std=c++17 -O3`` and the +device's native gencode. No ``--use_fast_math`` (the eager baseline does not use +it, so the candidate must not either). +""" + +from __future__ import annotations + +import functools +import pathlib + +import torch + +_SOLUTION_DIR = pathlib.Path(__file__).resolve().parent +_KERNEL_CU = _SOLUTION_DIR / "kernel.cu" + + +def _gencode_flag() -> str: + major, minor = torch.cuda.get_device_capability() + return f"-gencode=arch=compute_{major}{minor},code=sm_{major}{minor}" + + +def candidate_compile_flags() -> dict: + """The exact flag sets used to build the candidate (provenance record).""" + return { + "extra_cflags": ["-std=c++17", "-O3"], + "extra_cuda_cflags": ["-std=c++17", "-O3", _gencode_flag()], + "ldlibs": ["-lc10", "-lc10_cuda", "-ltorch_cpu", "-ltorch_cuda"], + "use_fast_math": False, + } + + +@functools.lru_cache(maxsize=None) +def load_candidate_module(): + from torch.utils import cpp_extension as torch_cpp + from tvm_ffi.cpp import load + + include_paths = list(torch_cpp.include_paths(device_type="cuda")) + library_paths = list(torch_cpp.library_paths(device_type="cuda")) + flags = candidate_compile_flags() + ldflags = [f"-L{p}" for p in library_paths] + ldflags += [f"-Wl,-rpath,{p}" for p in library_paths] + ldflags += flags["ldlibs"] + + return load( + "residual_gate_add_candidate", + cuda_files=[str(_KERNEL_CU)], + extra_cflags=flags["extra_cflags"], + extra_cuda_cflags=flags["extra_cuda_cflags"], + extra_ldflags=ldflags, + extra_include_paths=include_paths, + build_directory=str(_SOLUTION_DIR / ".build"), + ) diff --git a/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/solution/kernel.cu b/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/solution/kernel.cu new file mode 100644 index 000000000..f6a9ff4c3 --- /dev/null +++ b/diffusion/kernels/b200_diffusion_residual_gate_add__multi_shape/solution/kernel.cu @@ -0,0 +1,394 @@ +// Candidate CUDA implementation of the SGLang diffusion residual-gate-add and +// 4D broadcast-add elementwise patterns for B200, exposed through a +// destination-passing tvm-ffi ABI that mirrors baseline/binding.py (inputs +// first, output tensor last; every launch on torch's current CUDA stream). +// +// Reference semantics (recovered upstream @ sgl-project/sglang main +// 8314247d9de0fa2c58e34756b3e1dbc6cf815dfd; see docs/baseline_source.md): +// residual_gate_add(residual, update, gate, out): +// out = residual + update * gate +// gate is full-shape [B,L,D] (== residual) or row-broadcast [B,1,D]. +// broadcast_add_4d(a, b, out): +// a [1,1,P,D] broadcasts over the sequence dim of b [1,S,P,D]; +// out[0,s,p,d] = a[0,0,p,d] + b[0,s,p,d]. +// +// Math: elementwise arithmetic runs in fp32 and stores back in the input dtype +// (one rounding step). Tolerances follow the contract (bf16/fp16 5e-2, +// fp32 1e-5); the one-round fp32 form stays inside the bf16 tolerance versus +// the two-step eager reference (see docs/benchmark_method.md). +// +// Performance structure (B200, ~8 TB/s HBM3e, 148 SMs; bandwidth-bound): +// - 16B-vectorized (8 bf16 / 4 fp32) grid-stride fast path gated at runtime +// on contiguity, 16B base-pointer alignment, and a last dim divisible by +// the vector width; a scalar grid-stride path covers the tail and every +// layout the fast path rejects. +// - The broadcast gate / broadcast `a` operand is reused across rows and is +// read through the read-only cache (__ldg); streaming operands use the +// default cache policy. +// - Output is required distinct from the inputs (out-of-place) and contiguous: +// aliasing and unsupported layouts are rejected on the host. + +#include +#include + +#include +#include + +#include +#include +#if __has_include() +#include +#endif + +#include +#include +#include +#include + +namespace { + +using tvm::ffi::TensorView; + +template +[[noreturn]] void cand_fail(Args&&... args) { + std::ostringstream oss; + (oss << ... << args); + throw std::runtime_error(oss.str()); +} + +#define CAND_CHECK(cond, ...) \ + do { \ + if (!(cond)) { \ + cand_fail(__VA_ARGS__); \ + } \ + } while (0) + +// -------------------------------------------------------------------------- +// dtype helpers +// -------------------------------------------------------------------------- + +inline bool dtype_is(DLDataType d, uint8_t code, uint8_t bits) { + return d.code == code && d.bits == bits && d.lanes == 1; +} +inline bool is_bf16(DLDataType d) { return dtype_is(d, kDLBfloat, 16); } +inline bool is_f16(DLDataType d) { return dtype_is(d, kDLFloat, 16); } +inline bool is_f32(DLDataType d) { return dtype_is(d, kDLFloat, 32); } +inline bool same_dtype(DLDataType a, DLDataType b) { + return a.code == b.code && a.bits == b.bits && a.lanes == b.lanes; +} + +__device__ __forceinline__ float to_f(__nv_bfloat16 v) { return __bfloat162float(v); } +__device__ __forceinline__ float to_f(__half v) { return __half2float(v); } +__device__ __forceinline__ float to_f(float v) { return v; } + +template +__device__ __forceinline__ T from_f(float v); +template <> +__device__ __forceinline__ __nv_bfloat16 from_f<__nv_bfloat16>(float v) { return __float2bfloat16(v); } +template <> +__device__ __forceinline__ __half from_f<__half>(float v) { return __float2half(v); } +template <> +__device__ __forceinline__ float from_f(float v) { return v; } + +template +union Vec16 { + static constexpr int kElems = 16 / sizeof(T); + uint4 raw; + T elems[16 / sizeof(T)]; +}; + +// -------------------------------------------------------------------------- +// tensor view helpers +// -------------------------------------------------------------------------- + +template +inline const T* data_of(const TensorView& t) { + return reinterpret_cast(static_cast(t.data_ptr()) + t.byte_offset()); +} +template +inline T* mutable_data_of(const TensorView& t) { + return reinterpret_cast(static_cast(t.data_ptr()) + t.byte_offset()); +} + +inline bool tensor_is_contiguous(const TensorView& t) { + int64_t expect = 1; + for (int i = t.ndim() - 1; i >= 0; --i) { + if (t.size(i) == 1) continue; // stride is free on size-1 dims + if (t.stride(i) != expect) return false; + expect *= t.size(i); + } + return true; +} + +inline bool aligned16(const void* p) { + return (reinterpret_cast(p) & 0xF) == 0; +} + +inline void check_cuda(const TensorView& t, const char* name) { + CAND_CHECK(t.device().device_type == kDLCUDA, name, " must be a CUDA tensor"); +} + +inline int64_t last_dim(const TensorView& t) { return t.size(t.ndim() - 1); } + +// -------------------------------------------------------------------------- +// residual_gate_add: out = residual + update * gate +// -------------------------------------------------------------------------- + +enum class GateMode { kFull, kBcastRow }; + +constexpr int kThreads = 256; + +template +__global__ void rga_vec_kernel(const T* __restrict__ residual, + const T* __restrict__ update, + const T* __restrict__ gate, + T* __restrict__ out, + int64_t n_vec, int64_t row_vec) { + // row_vec = D / kVec (vectors per row); only used for the broadcast gate. + const int64_t stride = static_cast(blockDim.x) * gridDim.x; + for (int64_t v = blockIdx.x * blockDim.x + threadIdx.x; v < n_vec; v += stride) { + const Vec16 r{.raw = reinterpret_cast(residual)[v]}; + const Vec16 u{.raw = reinterpret_cast(update)[v]}; + Vec16 g; + if (kGate == GateMode::kFull) { + g.raw = reinterpret_cast(gate)[v]; + } else { + // Row-broadcast gate: D % kVec == 0, so each 16B vector stays inside one + // row; the gate vector index is v modulo the per-row vector count. + g.raw = __ldg(reinterpret_cast(gate) + (v % row_vec)); + } + Vec16 o; +#pragma unroll + for (int k = 0; k < kVec; ++k) { + o.elems[k] = from_f(to_f(r.elems[k]) + to_f(u.elems[k]) * to_f(g.elems[k])); + } + reinterpret_cast(out)[v] = o.raw; + } +} + +template +__global__ void rga_scalar_kernel(const T* __restrict__ residual, + const T* __restrict__ update, + const T* __restrict__ gate, + T* __restrict__ out, + int64_t begin, int64_t total, int64_t D) { + const int64_t stride = static_cast(blockDim.x) * gridDim.x; + for (int64_t i = begin + blockIdx.x * blockDim.x + threadIdx.x; i < total; i += stride) { + const float gv = (kGate == GateMode::kFull) ? to_f(gate[i]) : to_f(__ldg(gate + (i % D))); + out[i] = from_f(to_f(residual[i]) + to_f(update[i]) * gv); + } +} + +inline int blocks_for(int64_t n) { + int64_t b = (n + kThreads - 1) / kThreads; + if (b < 1) b = 1; + if (b > 65535) b = 65535; // grid-stride covers the rest + return static_cast(b); +} + +template +void launch_rga(const TensorView& residual, const TensorView& update, + const TensorView& gate, const TensorView& out, + GateMode mode, cudaStream_t stream) { + const int64_t total = residual.numel(); + const int64_t D = last_dim(residual); + const T* rp = data_of(residual); + const T* up = data_of(update); + const T* gp = data_of(gate); + T* op = mutable_data_of(out); + constexpr int kVec = 16 / sizeof(T); + + const bool gate_contig = tensor_is_contiguous(gate); + const bool vec_ok = tensor_is_contiguous(residual) && tensor_is_contiguous(update) && + tensor_is_contiguous(out) && gate_contig && + aligned16(rp) && aligned16(up) && aligned16(gp) && aligned16(op) && + (D % kVec == 0) && + (mode == GateMode::kFull ? (total % kVec == 0) : true); + + int64_t done = 0; + if (vec_ok) { + const int64_t n_vec = total / kVec; // D % kVec == 0 => total % kVec == 0 + const int64_t row_vec = D / kVec; + if (mode == GateMode::kFull) { + rga_vec_kernel + <<>>(rp, up, gp, op, n_vec, row_vec); + } else { + rga_vec_kernel + <<>>(rp, up, gp, op, n_vec, row_vec); + } + done = n_vec * kVec; + } + if (done < total) { + if (mode == GateMode::kFull) { + rga_scalar_kernel + <<>>(rp, up, gp, op, done, total, D); + } else { + rga_scalar_kernel + <<>>(rp, up, gp, op, done, total, D); + } + } +} + +void residual_gate_add(TensorView residual, TensorView update, TensorView gate, TensorView out) { + check_cuda(residual, "residual"); + check_cuda(update, "update"); + check_cuda(gate, "gate"); + check_cuda(out, "out"); + + CAND_CHECK(residual.ndim() >= 2, "residual must be at least 2D ([.., D])"); + CAND_CHECK(update.ndim() == residual.ndim(), "update rank must match residual"); + CAND_CHECK(out.ndim() == residual.ndim(), "out rank must match residual"); + for (int i = 0; i < residual.ndim(); ++i) { + CAND_CHECK(update.size(i) == residual.size(i), "update shape must match residual"); + CAND_CHECK(out.size(i) == residual.size(i), "out shape must match residual"); + } + const DLDataType dt = residual.dtype(); + CAND_CHECK(same_dtype(update.dtype(), dt) && same_dtype(gate.dtype(), dt) && + same_dtype(out.dtype(), dt), + "residual/update/gate/out must share dtype"); + CAND_CHECK(is_bf16(dt) || is_f16(dt) || is_f32(dt), "dtype must be bf16, fp16, or fp32"); + + // Output must be distinct from every input (out-of-place only). + CAND_CHECK(out.data_ptr() != residual.data_ptr() && out.data_ptr() != update.data_ptr() && + out.data_ptr() != gate.data_ptr(), + "out must not alias residual/update/gate"); + CAND_CHECK(tensor_is_contiguous(residual) && tensor_is_contiguous(update) && + tensor_is_contiguous(out), + "residual/update/out must be contiguous"); + + // Gate mode: full == same shape as residual; broadcast == [.., 1, D] with a + // size-1 second-to-last dim and a matching last dim. + const int64_t D = last_dim(residual); + CAND_CHECK(last_dim(gate) == D, "gate last dim must equal residual last dim"); + GateMode mode; + bool full = (gate.ndim() == residual.ndim()); + for (int i = 0; full && i < residual.ndim(); ++i) full = (gate.size(i) == residual.size(i)); + if (full) { + CAND_CHECK(tensor_is_contiguous(gate), "full-shape gate must be contiguous"); + mode = GateMode::kFull; + } else { + // Row-broadcast gate: same rank as residual, a size-1 row dim, and every + // leading dim equal to 1 -- i.e. exactly D contiguous elements reused for + // every row (the i%D indexing depends on this). Malformed [B,D] / [D] and + // a true B>1 broadcast are rejected rather than silently mishandled. + CAND_CHECK(gate.ndim() == residual.ndim(), + "gate must be full-shape [.., D] or same-rank row-broadcast [.., 1, D]"); + CAND_CHECK(gate.size(gate.ndim() - 2) == 1, "row-broadcast gate must have a size-1 row dim"); + for (int i = 0; i < gate.ndim() - 1; ++i) { + CAND_CHECK(gate.size(i) == 1, + "row-broadcast gate leading dims must be 1 (batch>1 broadcast unsupported)"); + } + CAND_CHECK(tensor_is_contiguous(gate), "broadcast gate must be contiguous"); + mode = GateMode::kBcastRow; + } + + // All operands must share one CUDA device; guard so the launch targets that + // device regardless of the process's current device (multi-GPU safety). + const int dev = residual.device().device_id; + CAND_CHECK(update.device().device_id == dev && gate.device().device_id == dev && + out.device().device_id == dev, + "residual/update/gate/out must be on the same CUDA device"); + const c10::cuda::CUDAGuard device_guard(static_cast(dev)); + cudaStream_t stream = at::cuda::getCurrentCUDAStream(dev); + if (residual.numel() == 0) return; + if (is_bf16(dt)) { + launch_rga<__nv_bfloat16>(residual, update, gate, out, mode, stream); + } else if (is_f16(dt)) { + launch_rga<__half>(residual, update, gate, out, mode, stream); + } else { + launch_rga(residual, update, gate, out, mode, stream); + } +} + +// -------------------------------------------------------------------------- +// broadcast_add_4d: out[0,s,p,d] = a[0,0,p,d] + b[0,s,p,d] +// -------------------------------------------------------------------------- + +template +__global__ void bcast_add_vec_kernel(const T* __restrict__ a, const T* __restrict__ b, + T* __restrict__ out, int64_t n_vec, int64_t inner_vec) { + const int64_t stride = static_cast(blockDim.x) * gridDim.x; + for (int64_t v = blockIdx.x * blockDim.x + threadIdx.x; v < n_vec; v += stride) { + const Vec16 bv{.raw = reinterpret_cast(b)[v]}; + const Vec16 av{.raw = __ldg(reinterpret_cast(a) + (v % inner_vec))}; + Vec16 o; +#pragma unroll + for (int k = 0; k < kVec; ++k) o.elems[k] = from_f(to_f(bv.elems[k]) + to_f(av.elems[k])); + reinterpret_cast(out)[v] = o.raw; + } +} + +template +__global__ void bcast_add_scalar_kernel(const T* __restrict__ a, const T* __restrict__ b, + T* __restrict__ out, int64_t begin, int64_t total, + int64_t inner) { + const int64_t stride = static_cast(blockDim.x) * gridDim.x; + for (int64_t i = begin + blockIdx.x * blockDim.x + threadIdx.x; i < total; i += stride) { + out[i] = from_f(to_f(b[i]) + to_f(__ldg(a + (i % inner)))); + } +} + +template +void launch_bcast_add(const TensorView& a, const TensorView& b, const TensorView& out, + int64_t inner, cudaStream_t stream) { + const int64_t total = b.numel(); + const T* ap = data_of(a); + const T* bp = data_of(b); + T* op = mutable_data_of(out); + constexpr int kVec = 16 / sizeof(T); + const bool vec_ok = tensor_is_contiguous(a) && tensor_is_contiguous(b) && + tensor_is_contiguous(out) && aligned16(ap) && aligned16(bp) && + aligned16(op) && (inner % kVec == 0); + int64_t done = 0; + if (vec_ok) { + const int64_t n_vec = total / kVec; + bcast_add_vec_kernel + <<>>(ap, bp, op, n_vec, inner / kVec); + done = n_vec * kVec; + } + if (done < total) { + bcast_add_scalar_kernel + <<>>(ap, bp, op, done, total, inner); + } +} + +void broadcast_add_4d(TensorView a, TensorView b, TensorView out) { + check_cuda(a, "a"); + check_cuda(b, "b"); + check_cuda(out, "out"); + CAND_CHECK(b.ndim() == 4 && a.ndim() == 4 && out.ndim() == 4, + "broadcast_add_4d expects 4D a, b, and out"); + CAND_CHECK(b.size(0) == 1, "broadcast_add_4d supports batch size 1 only (a is reused per row)"); + CAND_CHECK(a.size(0) == b.size(0) && a.size(1) == 1, "a must be [1,1,P,D] broadcasting over dim1"); + CAND_CHECK(a.size(2) == b.size(2) && a.size(3) == b.size(3), "a/b inner dims (P,D) must match"); + for (int i = 0; i < 4; ++i) CAND_CHECK(out.size(i) == b.size(i), "out shape must match b"); + const DLDataType dt = b.dtype(); + CAND_CHECK(same_dtype(a.dtype(), dt) && same_dtype(out.dtype(), dt), "a/b/out must share dtype"); + CAND_CHECK(is_bf16(dt) || is_f16(dt) || is_f32(dt), "dtype must be bf16, fp16, or fp32"); + CAND_CHECK(out.data_ptr() != a.data_ptr() && out.data_ptr() != b.data_ptr(), + "out must not alias a/b"); + CAND_CHECK(tensor_is_contiguous(a) && tensor_is_contiguous(b) && tensor_is_contiguous(out), + "a/b/out must be contiguous"); + + const int64_t inner = a.size(2) * a.size(3); // P * D, the broadcast period + // Same-device guard (multi-GPU safety): launch on the inputs' device, not the + // process-current one. + const int dev = b.device().device_id; + CAND_CHECK(a.device().device_id == dev && out.device().device_id == dev, + "a/b/out must be on the same CUDA device"); + const c10::cuda::CUDAGuard device_guard(static_cast(dev)); + cudaStream_t stream = at::cuda::getCurrentCUDAStream(dev); + if (b.numel() == 0) return; + if (is_bf16(dt)) { + launch_bcast_add<__nv_bfloat16>(a, b, out, inner, stream); + } else if (is_f16(dt)) { + launch_bcast_add<__half>(a, b, out, inner, stream); + } else { + launch_bcast_add(a, b, out, inner, stream); + } +} + +} // namespace + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(residual_gate_add, residual_gate_add); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(broadcast_add_4d, broadcast_add_4d);