Same shape, different layout: four inputs to keep in a Triton reproducer

[wbavon]

A FlagOS community note for developers writing or integrating Triton operators.

shape=(7, 17), dtype=float32 does not fully describe a kernel input. These four tensors have the same shape, but ask the implementation to read memory differently:

Input construction	Strides	Storage offset	What it exercises
randn(7, 17)	(17, 1)	0	Contiguous input with an odd dimension
randn(17, 7).t()	(1, 7)	0	Transposed indexing
randn(7, 35)[:, 1::2]	(35, 2)	1	A stepped slice starting inside storage
randn(1, 17).expand(7, 17)	(0, 1)	0	Multiple logical rows reading the same storage

A subtle way to lose a reproduction is to call .clone() or .contiguous() before the candidate operator. PyTorch documents that the default clone preserves strides for dense inputs, but converts a non-dense stepped slice to a dense layout. The values can remain correct while the memory-access condition that triggered the problem disappears. Build the views on the target device, and record metadata at the call site.

Here is a small correctness probe using the direct flag_gems.add entry point, as used by FlagGems' public tests. Run it in a fresh process in a supported FlagGems environment, without globally enabling FlagGems dispatch. This is a diagnostic template, not a report of new backend test results.

import json

import torch
import triton
import flag_gems

torch.manual_seed(0)
kw = {"device": flag_gems.device, "dtype": torch.float32}

def randn(*shape):
    return torch.randn(shape, **kw)

cases = {
    "contiguous": randn(7, 17),
    "transpose": randn(17, 7).t(),
    "slice_offset": randn(7, 35)[:, 1::2],
    "expand": randn(1, 17).expand(7, 17),
}
bias = randn(17)  # Broadcast along the first dimension.

print(json.dumps({
    "torch": torch.__version__,
    "triton": triton.__version__,
    "flag_gems": getattr(flag_gems, "__version__", "unknown"),
    "backend": flag_gems.vendor_name,
    "device": str(flag_gems.device),
}))

failed = []
for name, x in cases.items():
    print(json.dumps({"case": name, "shape": list(x.shape),
                      "stride": list(x.stride()),
                      "storage_offset": x.storage_offset(),
                      "dtype": str(x.dtype)}))
    try:
        reference = torch.add(x, bias, alpha=0.5)
        result = flag_gems.add(x, bias, alpha=0.5)
        torch.testing.assert_close(result, reference, rtol=1e-5, atol=1e-6)
        print(name, "PASS")
    except Exception as error:
        failed.append(name)
        print(name, type(error).__name__, str(error))

raise SystemExit(1 if failed else 0)

Keep expanded inputs read-only: their elements may alias the same storage. This probe uses an out-of-place operation. It checks each input against its own native reference; it is not a performance comparison between layouts. The float32 tolerances are specific to this small probe, and should not become a blanket accuracy policy for other operators or dtypes.

When sharing a failure, include the failing construction, complete output, exact hardware model, driver/runtime versions, and installed package or source revisions. A source checkout should include its Git commit; triton.__version__ alone does not identify a compiler fork. Save the smallest reproduction in the relevant GitHub issue. Passing these four cases is a useful starting point, not proof of full operator or model compatibility.

Which input condition has caught a bug in your own kernels: a transpose, a nonzero offset, or a zero stride? A small public reproducer is more useful than a screenshot of the error.

Sources: FlagGems add tests, PyTorch clone semantics, and PyTorch expand semantics. For measurement after correctness, see the earlier timing note.

For ongoing discussion, FlagOS Community on Slack has a #community channel. The material and issue discussions are public; joining Slack is optional.