mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-11-01 01:06:23 +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.
78 lines
2.8 KiB
Python
78 lines
2.8 KiB
Python
"""Dump TileLang's `T.rng_rand()` stream so a Triton implementation can be checked against it.
|
|
|
|
`T.rng_init(seed, seq, off, generator)` lowers to `curand_init(seed, seq, off, &state)` and
|
|
`T.rng_rand()` to `curand(&state)`, with the generator defaulting to `curandStatePhilox4_32_10_t`.
|
|
That is cuRAND's Philox4x32-10, whose mixing function is public but whose `curand_init` seeds and
|
|
fast-forwards the counter in a way that is not documented. So the stream is measured here rather
|
|
than derived, and `h20_philox_check.py` compares against it.
|
|
|
|
Prints, for each seed: the first `n` values, which is enough to pin both the counter progression and
|
|
the key schedule, since any mismatch in either shows up in the first few outputs.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
|
|
import tilelang
|
|
import tilelang.language as T
|
|
import torch
|
|
|
|
|
|
@tilelang.jit(
|
|
pass_configs={
|
|
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
|
|
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
|
|
},
|
|
)
|
|
def rng_dump_tl(n: int) -> object:
|
|
@T.prim_func
|
|
def _rng_dump(
|
|
seed: T.Tensor((1,), T.uint64),
|
|
seq: T.Tensor((1,), T.uint64),
|
|
out: T.Tensor((n,), T.uint32),
|
|
) -> None:
|
|
with T.Kernel(1, threads=32) as _:
|
|
T.rng_init(seed[0], seq[0])
|
|
for i in T.serial(n):
|
|
# Every thread advances its own state, which is per-thread; only lane 0 writes, so
|
|
# the recorded sequence is exactly that lane's.
|
|
value = T.rng_rand()
|
|
if T.get_thread_binding() == 0:
|
|
out[i] = value
|
|
|
|
return _rng_dump
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--n", type=int, default=16)
|
|
parser.add_argument("--seeds", default="0,1,42,12345")
|
|
parser.add_argument("--out", default="")
|
|
args = parser.parse_args()
|
|
|
|
n = args.n
|
|
seeds = [int(x) for x in args.seeds.split(",")]
|
|
kernel = rng_dump_tl(n)
|
|
payload: dict[str, list[int]] = {}
|
|
for seed in seeds:
|
|
seed_t = torch.tensor([seed], dtype=torch.uint64, device="cuda")
|
|
# Subsequence 0 is what the sampling kernel uses: `T.rng_init(seed, seq)` passes the row's
|
|
# sample position as `seq`, so every distinct position is its own stream.
|
|
for seq in (0, 1):
|
|
out = torch.zeros(n, dtype=torch.uint32, device="cuda")
|
|
kernel(seed_t, torch.tensor([seq], dtype=torch.uint64, device="cuda"), out)
|
|
key = f"{seed}:{seq}"
|
|
payload[key] = [int(x) for x in out.cpu()]
|
|
print(f"seed={seed} seq={seq}: {payload[key]}", flush=True)
|
|
|
|
if args.out:
|
|
with open(args.out, "w") as f:
|
|
json.dump(payload, f, indent=2)
|
|
print(f"wrote {args.out}", flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|