Promote latest kernel artifacts to main

README.md

@@ -1,9 +1,69 @@
-# flashrt/transformer-layout-primitives
+# transformer-layout-primitives
 
-This repository is a compatibility mirror for older `kernels` clients
-that resolve repositories through the default Hugging Face model repo API.
+Generic BF16 layout, RoPE, and text-state primitives for transformer pipelines.
 
-Canonical Kernel Hub repo: https://huggingface.co/kernels/flashrt/transformer-layout-primitives
+Hub repo: `flashrt/transformer-layout-primitives`
 
-Do not edit this mirror by hand. It is generated from the Kernel Hub
-`vN` branches and contains the same `build/**` artifacts.
+## Public API
+
+- `fill_neginf_bf16(dst) -> dst`
+- `add_bias_bf16_(data, bias) -> data`
+- `repeat_interleave_heads_bf16(src, repeat) -> bf16`
+- `text_gather_bf16(src, batch, seq) -> bf16`
+- `text_scatter_bf16(dst, src, batch, seq) -> dst`
+- `rope_rotate_half_bf16_(x, cos, sin) -> x`
+- `qk_rmsnorm_rope_bf16_(qk, weight, cos, sin, eps=1e-6) -> qk`
+- `qk_pair_rmsnorm_rope_bf16(q, k, q_weight, k_weight, cos, sin, eps=1e-6) -> (q, k)`
+- `gather_rows_bf16(src, row_indices, out=None) -> bf16`
+- `scatter_rows_bf16(src, row_indices, rows, out=None) -> bf16`
+
+The package is intentionally model-neutral. It exposes Tensor APIs for common
+transformer integration gaps: head repeat for GQA/MQA, first/last token gather
+and scatter, bias add, RoPE rotate-half, and fused Q/K RMSNorm+RoPE. The pair
+API handles different Q and KV head counts in one launch.
+
+## Example
+
+```python
+from kernels import get_kernel
+import torch
+
+ops = get_kernel("flashrt/transformer-layout-primitives", version=1)
+
+q = torch.randn((128, 32, 128), device="cuda", dtype=torch.bfloat16)
+weight = torch.ones((128,), device="cuda", dtype=torch.bfloat16)
+cos = torch.randn((128, 128), device="cuda", dtype=torch.bfloat16)
+sin = torch.randn((128, 128), device="cuda", dtype=torch.bfloat16)
+ops.qk_rmsnorm_rope_bf16_(q, weight, cos, sin)
+
+k = torch.randn((128, 8, 128), device="cuda", dtype=torch.bfloat16)
+q, k = ops.qk_pair_rmsnorm_rope_bf16(
+    q, k, weight, weight, cos, sin
+)
+```
+
+## Shape contract
+
+- All tensors are contiguous CUDA BF16 tensors.
+- `repeat_interleave_heads_bf16`: `src` is `(seq, heads, head_dim)`.
+- `text_gather_bf16`: `src` is flattened `(batch * seq, dim)` and returns
+  first and last token rows as `(2 * batch, dim)`.
+- `text_scatter_bf16`: writes `(2 * batch, dim)` rows back to first and last
+  positions in `(batch * seq, dim)`.
+- RoPE functions use rotate-half layout with `cos/sin` shaped `(seq, head_dim)`
+  or `(rows, head_dim)`.
+- `qk_pair_rmsnorm_rope_bf16` accepts Q `(rows, q_heads, head_dim)` and K
+  `(rows, kv_heads, head_dim)`. `head_dim` must be even and in `[8, 256]`.
+  Q and K may have different head counts.
+- `gather_rows_bf16` and `scatter_rows_bf16` use contiguous CUDA int64 row
+  indices. Scatter indices must be unique.
+- Indexed-row validation includes the Cosmos3-Edge production layout
+  `128 -> 60` rows with hidden size `2048`, including exact CUDA Graph replay.
+
+## Validation
+
+Correctness is tested against PyTorch BF16/FP32 reference formulas with exact
+checks for pure layout operations and strict BF16 tolerances for math ops. The
+pair API must additionally match the two staged single-tensor kernels exactly.
+
+See `benchmarks/RESULTS.md` for current local RTX 5090 source benchmark data.

tests/test_transformer_layout_primitives.py

@@ -0,0 +1,512 @@
+#!/usr/bin/env python3
+"""Correctness tests for transformer-layout-primitives."""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import os
+import sys
+from pathlib import Path
+
+import torch
+
+ROOT = Path(__file__).resolve().parents[2]
+PACKAGE = ROOT / "transformer-layout-primitives"
+REGISTRATION_INCLUDE = ROOT.parent / "kernels" / "kernel-builder" / "src" / "pyproject" / "templates" / "torch"
+
+
+class SourceOps:
+    def __init__(self, namespace: str) -> None:
+        self.ops = getattr(torch.ops, namespace)
+
+    def fill_neginf_bf16(self, dst):
+        self.ops.fill_neginf_bf16(dst)
+        return dst
+
+    def add_bias_bf16_(self, data, bias):
+        self.ops.add_bias_bf16_(data, bias)
+        return data
+
+    def repeat_interleave_heads_bf16(self, src, repeat, out=None):
+        if out is None:
+            out = torch.empty((src.shape[0], src.shape[1] * repeat, src.shape[2]), device=src.device, dtype=src.dtype)
+        self.ops.repeat_interleave_heads_bf16(src, int(repeat), out)
+        return out
+
+    def text_gather_bf16(self, src, batch, seq, out=None):
+        if out is None:
+            out = torch.empty((2 * batch, src.shape[1]), device=src.device, dtype=src.dtype)
+        self.ops.text_gather_bf16(src, int(batch), int(seq), out)
+        return out
+
+    def text_scatter_bf16(self, dst, src, batch, seq):
+        self.ops.text_scatter_bf16(dst, src, int(batch), int(seq))
+        return dst
+
+    def rope_rotate_half_bf16_(self, x, cos, sin):
+        self.ops.rope_rotate_half_bf16_(x, cos, sin)
+        return x
+
+    def qk_rmsnorm_rope_bf16_(self, qk, weight, cos, sin, eps=1e-6):
+        self.ops.qk_rmsnorm_rope_bf16_(qk, weight, cos, sin, float(eps))
+        return qk
+
+    def qk_pair_rmsnorm_rope_bf16(
+        self, q, k, q_weight, k_weight, cos, sin, eps=1e-6, q_out=None, k_out=None
+    ):
+        if q_out is None:
+            q_out = torch.empty_like(q)
+        if k_out is None:
+            k_out = torch.empty_like(k)
+        self.ops.qk_pair_rmsnorm_rope_bf16(
+            q, k, q_weight, k_weight, cos, sin, float(eps), q_out, k_out
+        )
+        return q_out, k_out
+
+    def gather_rows_bf16(self, src, row_indices, out=None):
+        if out is None:
+            out = torch.empty(
+                (row_indices.numel(), src.shape[1]), device=src.device, dtype=src.dtype
+            )
+        self.ops.gather_rows_bf16(src, row_indices, out)
+        return out
+
+    def scatter_rows_bf16(self, src, row_indices, destination_rows, out=None):
+        if out is None:
+            out = torch.zeros(
+                (destination_rows, src.shape[1]), device=src.device, dtype=src.dtype
+            )
+        self.ops.scatter_rows_bf16(src, row_indices, out)
+        return out
+
+
+def _arch_list() -> str:
+    major, minor = torch.cuda.get_device_capability(0)
+    return "12.0a" if major >= 12 else f"{major}.{minor}"
+
+
+def load_source_ops() -> SourceOps:
+    from torch.utils.cpp_extension import load
+
+    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", _arch_list())
+    namespace = "transformer_layout_primitives_source_test"
+    load(
+        name=namespace,
+        sources=[
+            str(PACKAGE / "torch-ext" / "torch_binding.cpp"),
+            str(PACKAGE / "csrc" / "transformer_layout_primitives.cu"),
+        ],
+        extra_include_paths=[str(PACKAGE / "csrc"), str(REGISTRATION_INCLUDE)],
+        extra_cflags=["-O3", "-DCUDA_KERNEL"],
+        extra_cuda_cflags=[
+            "-O3",
+            "--expt-relaxed-constexpr",
+            "-DCUDA_KERNEL",
+            "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
+            "-U__CUDA_NO_BFLOAT16_OPERATORS__",
+            "-U__CUDA_NO_BFLOAT162_OPERATORS__",
+        ],
+        is_python_module=False,
+        verbose=False,
+    )
+    return SourceOps(namespace)
+
+
+def load_installed_ops(artifact: str | None):
+    if artifact:
+        sys.path.insert(0, artifact)
+    try:
+        return importlib.import_module("transformer_layout_primitives")
+    finally:
+        if artifact:
+            sys.path.remove(artifact)
+
+
+def rotate_half_ref(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    half = x.shape[-1] // 2
+    lo = x[..., :half].float()
+    hi = x[..., half:].float()
+    c = cos[:, None, :half].float()
+    s = sin[:, None, :half].float()
+    return torch.cat([lo * c - hi * s, hi * c + lo * s], dim=-1).to(torch.bfloat16)
+
+
+def qk_rmsnorm_rope_ref(
+    qk: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float = 1e-6,
+) -> torch.Tensor:
+    rms = torch.rsqrt((qk.float() * qk.float()).mean(dim=-1, keepdim=True) + eps)
+    normed = (qk.float() * rms * weight.float()).to(torch.bfloat16)
+    return rotate_half_ref(normed, cos, sin)
+
+
+def qk_pair_rmsnorm_rope_ref(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    eps: float = 1e-6,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    return (
+        qk_rmsnorm_rope_ref(q, q_weight, cos, sin, eps),
+        qk_rmsnorm_rope_ref(k, k_weight, cos, sin, eps),
+    )
+
+
+def metrics(got: torch.Tensor, ref: torch.Tensor) -> tuple[float, float, float]:
+    diff = (got.float() - ref.float()).abs()
+    cos = torch.nn.functional.cosine_similarity(got.float().flatten(), ref.float().flatten(), dim=0).item()
+    return float(diff.max().item()), float(diff.mean().item()), float(cos)
+
+
+def bf16_max_ulp(got: torch.Tensor, ref: torch.Tensor) -> int:
+    def ordered_bits(value: torch.Tensor) -> torch.Tensor:
+        bits = value.contiguous().view(torch.int16).to(torch.int32) & 0xFFFF
+        magnitude = bits & 0x7FFF
+        return torch.where(
+            (bits & 0x8000) != 0,
+            0x8000 - magnitude,
+            0x8000 + magnitude,
+        )
+
+    return int(
+        (ordered_bits(got) - ordered_bits(ref)).abs().max().item()
+    )
+
+
+def assert_close(
+    name: str,
+    got: torch.Tensor,
+    ref: torch.Tensor,
+    atol: float,
+    cos_min: float,
+    rtol: float = 0.0,
+) -> None:
+    max_abs, mean_abs, cos = metrics(got, ref)
+    diff = (got.float() - ref.float()).abs()
+    p99_abs = float(torch.quantile(diff, 0.99).item())
+    tolerance = atol + rtol * ref.float().abs()
+    violations = int((diff > tolerance).sum().item())
+    max_ulp = bf16_max_ulp(got, ref)
+    print(
+        f"{name}: max_abs={max_abs:.6f} p99_abs={p99_abs:.6f} "
+        f"mean_abs={mean_abs:.6e} max_ulp={max_ulp} cosine={cos:.8f} "
+        f"violations={violations} rtol={rtol} atol={atol}"
+    )
+    if violations or cos < cos_min:
+        raise AssertionError(
+            f"{name} failed: max_abs={max_abs} p99_abs={p99_abs} "
+            f"max_ulp={max_ulp} cosine={cos} violations={violations}"
+        )
+
+
+def expect_runtime_error(name: str, fn) -> None:
+    try:
+        fn()
+    except RuntimeError:
+        print(f"{name}: rejected")
+        return
+    raise AssertionError(f"{name}: expected RuntimeError")
+
+
+def run(ops, mode: str) -> int:
+    torch.manual_seed(31)
+    count = 0
+    layout_shapes = [(2, 5, 128), (3, 49, 256)] if mode == "smoke" else [
+        (1, 1, 128),
+        (2, 5, 128),
+        (3, 49, 256),
+        (4, 256, 1024),
+        (2, 2520, 2048),
+    ]
+    for batch, seq, dim in layout_shapes:
+        x = torch.randn((batch * seq, dim), device="cuda", dtype=torch.bfloat16)
+
+        dst = torch.empty_like(x)
+        ops.fill_neginf_bf16(dst)
+        ref_neginf = torch.full_like(x, -1e30)
+        torch.testing.assert_close(dst.float().cpu(), ref_neginf.float().cpu(), rtol=0, atol=0)
+        count += 1
+
+        bias = torch.randn((dim,), device="cuda", dtype=torch.bfloat16)
+        got = x.clone()
+        ops.add_bias_bf16_(got, bias)
+        ref = (x.float() + bias.float()).to(torch.bfloat16)
+        torch.testing.assert_close(got.cpu(), ref.cpu(), rtol=0, atol=0)
+        count += 1
+
+        gathered = ops.text_gather_bf16(x, batch, seq)
+        ref_gather = torch.stack([x[b * seq + offset] for b in range(batch) for offset in (0, seq - 1)], dim=0)
+        torch.testing.assert_close(gathered.cpu(), ref_gather.cpu(), rtol=0, atol=0)
+        count += 1
+
+        scattered = torch.zeros_like(x)
+        ops.text_scatter_bf16(scattered, gathered, batch, seq)
+        ref_scatter = torch.zeros_like(x)
+        for b in range(batch):
+            ref_scatter[b * seq] = gathered[2 * b]
+            ref_scatter[b * seq + seq - 1] = gathered[2 * b + 1]
+        torch.testing.assert_close(scattered.cpu(), ref_scatter.cpu(), rtol=0, atol=0)
+        count += 1
+
+    cpu_rng_state = torch.random.get_rng_state()
+    cuda_rng_state = torch.cuda.get_rng_state()
+    indexed_shapes = [(17, 5, 64), (277, 51, 128)] if mode == "smoke" else [
+        (1, 1, 8),
+        (17, 5, 64),
+        (277, 51, 128),
+        (128, 60, 2048),
+        (2520, 105, 1152),
+        (5070, 257, 4096),
+    ]
+    for source_rows, selected_rows, hidden in indexed_shapes:
+        src = torch.randn(
+            (source_rows, hidden), device="cuda", dtype=torch.bfloat16
+        )
+        indices = torch.randperm(source_rows, device="cuda", dtype=torch.int64)[
+            :selected_rows
+        ].contiguous()
+        got = ops.gather_rows_bf16(src, indices)
+        ref = src.index_select(0, indices)
+        torch.testing.assert_close(got, ref, rtol=0, atol=0)
+        count += 1
+
+        destination_rows = source_rows + 3
+        got = ops.scatter_rows_bf16(ref, indices, destination_rows)
+        ref_scatter = torch.zeros(
+            (destination_rows, hidden), device="cuda", dtype=torch.bfloat16
+        )
+        ref_scatter.index_copy_(0, indices, ref)
+        torch.testing.assert_close(got, ref_scatter, rtol=0, atol=0)
+        count += 1
+    torch.random.set_rng_state(cpu_rng_state)
+    torch.cuda.set_rng_state(cuda_rng_state)
+
+    repeat_shapes = [(17, 4, 64, 2), (128, 8, 128, 4)] if mode == "smoke" else [
+        (1, 1, 64, 8),
+        (17, 4, 64, 2),
+        (128, 8, 128, 4),
+        (2520, 8, 128, 4),
+    ]
+    for seq, heads, dim, repeat in repeat_shapes:
+        src = torch.randn((seq, heads, dim), device="cuda", dtype=torch.bfloat16)
+        got = ops.repeat_interleave_heads_bf16(src, repeat)
+        ref = src.repeat_interleave(repeat, dim=1)
+        torch.testing.assert_close(got.cpu(), ref.cpu(), rtol=0, atol=0)
+        count += 1
+
+    rope_shapes = [(17, 4, 64), (128, 8, 128)] if mode == "smoke" else [
+        (1, 1, 64),
+        (17, 4, 64),
+        (128, 8, 128),
+        (2520, 32, 128),
+    ]
+    for seq, heads, dim in rope_shapes:
+        x = torch.randn((seq, heads, dim), device="cuda", dtype=torch.bfloat16)
+        cos = torch.randn((seq, dim), device="cuda", dtype=torch.bfloat16)
+        sin = torch.randn((seq, dim), device="cuda", dtype=torch.bfloat16)
+        got = x.clone()
+        ops.rope_rotate_half_bf16_(got, cos, sin)
+        ref = rotate_half_ref(x, cos, sin)
+        assert_close(f"rope seq={seq} heads={heads} dim={dim}", got, ref, atol=0, cos_min=0.999999)
+        count += 1
+
+        weight = torch.randn((dim,), device="cuda", dtype=torch.bfloat16)
+        got = x.clone()
+        ops.qk_rmsnorm_rope_bf16_(got, weight, cos, sin)
+        ref = qk_rmsnorm_rope_ref(x, weight, cos, sin)
+        # The CUDA kernel is bitwise-gated against the established staged
+        # native path below. PyTorch eager may use a different FP32 reduction
+        # order, so this semantic comparison uses BF16 elementwise tolerances.
+        assert_close(
+            f"qk_rmsnorm_rope seq={seq} heads={heads} dim={dim}",
+            got,
+            ref,
+            atol=0.015625,
+            rtol=0.02,
+            cos_min=0.999999,
+        )
+        count += 1
+
+    pair_shapes = [(17, 16, 8, 128), (49, 16, 16, 72)] if mode == "smoke" else [
+        (1, 16, 8, 128),
+        (17, 16, 8, 128),
+        (49, 16, 16, 72),
+        (51, 16, 16, 80),
+        (65, 32, 8, 128),
+        (277, 16, 8, 128),
+        (512, 24, 24, 128),
+        (2520, 24, 24, 128),
+        (5070, 24, 24, 128),
+    ]
+    for rows, q_heads, k_heads, dim in pair_shapes:
+        q = torch.randn((rows, q_heads, dim), device="cuda", dtype=torch.bfloat16)
+        k = torch.randn((rows, k_heads, dim), device="cuda", dtype=torch.bfloat16)
+        q_weight = torch.randn((dim,), device="cuda", dtype=torch.bfloat16)
+        k_weight = torch.randn((dim,), device="cuda", dtype=torch.bfloat16)
+        angles = torch.randn((rows, dim // 2), device="cuda", dtype=torch.float32)
+        cos_half = angles.cos().to(torch.bfloat16)
+        sin_half = angles.sin().to(torch.bfloat16)
+        cos = torch.cat((cos_half, cos_half), dim=-1)
+        sin = torch.cat((sin_half, sin_half), dim=-1)
+        got_q, got_k = ops.qk_pair_rmsnorm_rope_bf16(
+            q, k, q_weight, k_weight, cos, sin
+        )
+        staged_q = q.clone()
+        staged_k = k.clone()
+        ops.qk_rmsnorm_rope_bf16_(staged_q, q_weight, cos, sin)
+        ops.qk_rmsnorm_rope_bf16_(staged_k, k_weight, cos, sin)
+        torch.testing.assert_close(got_q, staged_q, rtol=0, atol=0)
+        torch.testing.assert_close(got_k, staged_k, rtol=0, atol=0)
+        ref_q, ref_k = qk_pair_rmsnorm_rope_ref(
+            q, k, q_weight, k_weight, cos, sin
+        )
+        label = f"qk_pair rows={rows} qh={q_heads} kh={k_heads} dim={dim}"
+        # The fused path is already required to be bitwise equal to the
+        # established staged native kernels above. Torch 2.11's eager
+        # reduction order can differ by up to 0.03125 in BF16 at large rows.
+        assert_close(
+            f"{label}/q", got_q, ref_q, atol=0.015625, rtol=0.02, cos_min=0.999999
+        )
+        assert_close(
+            f"{label}/k", got_k, ref_k, atol=0.015625, rtol=0.02, cos_min=0.999999
+        )
+        count += 4
+
+    q = torch.randn((17, 4, 64), device="cuda", dtype=torch.bfloat16)
+    k = torch.randn((17, 2, 64), device="cuda", dtype=torch.bfloat16)
+    weight = torch.ones((64,), device="cuda", dtype=torch.bfloat16)
+    cos = torch.ones((17, 64), device="cuda", dtype=torch.bfloat16)
+    sin = torch.zeros_like(cos)
+    expect_runtime_error(
+        "qk_pair mismatched rows",
+        lambda: ops.qk_pair_rmsnorm_rope_bf16(
+            q, k[:-1].contiguous(), weight, weight, cos, sin
+        ),
+    )
+    expect_runtime_error(
+        "qk_pair invalid head_dim",
+        lambda: ops.qk_pair_rmsnorm_rope_bf16(
+            q[:, :, :-1].contiguous(),
+            k[:, :, :-1].contiguous(),
+            weight[:-1].contiguous(),
+            weight[:-1].contiguous(),
+            cos[:, :-1].contiguous(),
+            sin[:, :-1].contiguous(),
+        ),
+    )
+    expect_runtime_error(
+        "qk_pair noncontiguous",
+        lambda: ops.qk_pair_rmsnorm_rope_bf16(
+            q.transpose(0, 1), k, weight, weight, cos, sin
+        ),
+    )
+    count += 3
+
+    return count
+
+
+def run_compile_default_eps(ops) -> int:
+    seq, heads, dim = 17, 4, 64
+    x = torch.randn((seq, heads, dim), device="cuda", dtype=torch.bfloat16)
+    weight = torch.randn((dim,), device="cuda", dtype=torch.bfloat16)
+    cos = torch.randn((seq, dim), device="cuda", dtype=torch.bfloat16)
+    sin = torch.randn((seq, dim), device="cuda", dtype=torch.bfloat16)
+    ref = qk_rmsnorm_rope_ref(x, weight, cos, sin)
+
+    def invoke(qk, rms_weight, rope_cos, rope_sin):
+        return ops.qk_rmsnorm_rope_bf16_(qk, rms_weight, rope_cos, rope_sin)
+
+    compiled = torch.compile(invoke, fullgraph=True)
+    got = x.clone()
+    compiled(got, weight, cos, sin)
+    assert_close("qk_rmsnorm_rope compile default eps", got, ref, atol=0.015625, cos_min=0.999999)
+
+    k = torch.randn((seq, 2, dim), device="cuda", dtype=torch.bfloat16)
+    k_weight = torch.randn((dim,), device="cuda", dtype=torch.bfloat16)
+    pair_ref = qk_pair_rmsnorm_rope_ref(x, k, weight, k_weight, cos, sin)
+
+    def invoke_pair(q, key, qw, kw, rope_cos, rope_sin):
+        return ops.qk_pair_rmsnorm_rope_bf16(q, key, qw, kw, rope_cos, rope_sin)
+
+    compiled_pair = torch.compile(invoke_pair, fullgraph=True)
+    got_q, got_k = compiled_pair(x, k, weight, k_weight, cos, sin)
+    assert_close("qk_pair compile default eps/q", got_q, pair_ref[0], atol=0.015625, cos_min=0.999999)
+    assert_close("qk_pair compile default eps/k", got_k, pair_ref[1], atol=0.015625, cos_min=0.999999)
+    return 3
+
+
+def run_indexed_compile(ops) -> int:
+    src = torch.randn((277, 128), device="cuda", dtype=torch.bfloat16)
+    indices = torch.randperm(277, device="cuda", dtype=torch.int64)[:51].contiguous()
+
+    def invoke(src, indices):
+        gathered = ops.gather_rows_bf16(src, indices)
+        return ops.scatter_rows_bf16(gathered, indices, src.shape[0])
+
+    eager = invoke(src, indices)
+    compiled = torch.compile(invoke, fullgraph=True)(src, indices)
+    torch.testing.assert_close(compiled, eager, rtol=0, atol=0)
+    print("PASS gather/scatter torch.compile fullgraph")
+    return 1
+
+
+def run_cosmos_edge_indexed_graph(ops) -> int:
+    source_rows, selected_rows, hidden = 128, 60, 2048
+    src = torch.randn((source_rows, hidden), device="cuda", dtype=torch.bfloat16)
+    indices = torch.randperm(
+        source_rows, device="cuda", dtype=torch.int64
+    )[:selected_rows].contiguous()
+    gathered = torch.empty(
+        (selected_rows, hidden), device="cuda", dtype=torch.bfloat16
+    )
+    scattered = torch.zeros(
+        (source_rows, hidden), device="cuda", dtype=torch.bfloat16
+    )
+
+    graph = torch.cuda.CUDAGraph()
+    torch.cuda.synchronize()
+    with torch.cuda.graph(graph):
+        ops.gather_rows_bf16(src, indices, out=gathered)
+        ops.scatter_rows_bf16(gathered, indices, source_rows, out=scattered)
+    graph.replay()
+    torch.cuda.synchronize()
+    expected_gathered = src.index_select(0, indices)
+    expected_scattered = torch.zeros_like(scattered)
+    expected_scattered.index_copy_(0, indices, expected_gathered)
+    torch.testing.assert_close(gathered, expected_gathered, rtol=0.0, atol=0.0)
+    torch.testing.assert_close(scattered, expected_scattered, rtol=0.0, atol=0.0)
+    first = (gathered.clone(), scattered.clone())
+    graph.replay()
+    torch.cuda.synchronize()
+    torch.testing.assert_close(gathered, first[0], rtol=0.0, atol=0.0)
+    torch.testing.assert_close(scattered, first[1], rtol=0.0, atol=0.0)
+    print("PASS Cosmos3-Edge gather/scatter CUDA Graph replay")
+    return 2
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="source")
+    parser.add_argument("--artifact", default=None)
+    parser.add_argument("--mode", choices=["smoke", "full"], default="smoke")
+    args = parser.parse_args()
+    ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact)
+    count = run(ops, args.mode)
+    count += run_indexed_compile(ops)
+    if args.mode == "full":
+        count += run_cosmos_edge_indexed_graph(ops)
+    if args.backend == "installed":
+        count += run_compile_default_eps(ops)
+    print(f"transformer-layout-primitives {args.backend} {args.mode}: passed {count}/{count}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())