mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-30 15:57:20 +08:00
5e2838aa5e
The last blocker for `fused_sample` was the random stream: `T.rng_init`/`T.rng_rand` lower to
`curand_init`/`curand` with `curandStatePhilox4_32_10_t`, and the seeding is not documented. Its
mixing function is public but `curand_init` folds the seed, the subsequence and the offset into the
counter in a way that has to be measured rather than derived.
Measured (`scripts/h20_rng_dump.py` dumps the real stream), the layout is:
key = (seed & 0xFFFFFFFF, seed >> 32)
ctr = (offset & 0xFFFFFFFF, offset >> 32, subsequence & 0xFFFFFFFF, subsequence >> 32)
with draws taking the block's four words in order before `ctr.x` advances. The subsequence sits in
the third and fourth counter words, not the first -- a cuRAND subsequence spans 2^64 blocks. That
detail is invisible at subsequence 0 and wrong everywhere else, which is why the check covers both.
`scripts/h20_philox_check.py` compares our kernel against the dumped reference over five seeds
(including 2^64-1) and four subsequences (0, 1, 2048, 4097 -- the last two are the actual
`sample_positions` from a run): 20/20 streams bit-identical.
One trap worth noting: the multipliers and the second Weyl constant all exceed int32, so referencing
them as Python ints inside the kernel silently wraps them negative. They are materialised as uint32
tiles instead.
This does not yet change any output -- `fused_sample` is still the TileLang kernel. It removes the
reason it had to be, which is the point.
69 lines
2.4 KiB
Python
69 lines
2.4 KiB
Python
"""Check the Triton Philox against TileLang's `T.rng_rand()` stream.
|
|
|
|
The reference is not derived from documentation -- it is dumped from the real kernel
|
|
(`scripts/h20_rng_dump.py`), because `curand_init`'s seeding and subsequence layout are not
|
|
documented. This script re-dumps and compares, so it stays honest if the container's CUDA version
|
|
changes the seeding.
|
|
|
|
Both `subsequence=0` and a non-zero one are covered on purpose: the subsequence lands in the third
|
|
counter word, which is invisible when it is zero.
|
|
|
|
Exit code is non-zero on any mismatch, so it can gate the sampler port.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import sys
|
|
|
|
import torch
|
|
|
|
from deepinfer.model_executor.layers.tilelang_ops.triton.philox import philox_stream_tl
|
|
|
|
sys.path.insert(0, "scripts")
|
|
from h20_rng_dump import rng_dump_tl # noqa: E402 (path set above so the script is a library here)
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--n", type=int, default=16)
|
|
parser.add_argument("--seeds", default="0,1,42,12345,18446744073709551615")
|
|
parser.add_argument("--subseqs", default="0,1,2048,4097")
|
|
args = parser.parse_args()
|
|
|
|
n = args.n
|
|
seeds = [int(x) for x in args.seeds.split(",")]
|
|
subseqs = [int(x) for x in args.subseqs.split(",")]
|
|
|
|
reference = rng_dump_tl(n)
|
|
candidate = philox_stream_tl()
|
|
|
|
failures = 0
|
|
checked = 0
|
|
for seed in seeds:
|
|
for seq in subseqs:
|
|
seed_t = torch.tensor([seed], dtype=torch.uint64, device="cuda")
|
|
seq_t = torch.tensor([seq], dtype=torch.uint64, device="cuda")
|
|
|
|
want = torch.zeros(n, dtype=torch.uint32, device="cuda")
|
|
reference(seed_t, seq_t, want)
|
|
|
|
got = torch.zeros(n, dtype=torch.uint32, device="cuda")
|
|
candidate(seed_t, seq_t, got)
|
|
|
|
checked += 1
|
|
if not torch.equal(want, got):
|
|
failures += 1
|
|
want_cpu = want.cpu().tolist()
|
|
got_cpu = got.cpu().tolist()
|
|
first = next(i for i, (a, b) in enumerate(zip(want_cpu, got_cpu, strict=True)) if a != b)
|
|
print(f"MISMATCH seed={seed} seq={seq}: first diff at {first}, want {want_cpu[first]} got {got_cpu[first]}")
|
|
|
|
print(f"\n{checked - failures}/{checked} streams bit-identical")
|
|
if failures:
|
|
raise SystemExit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|