mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 00:36:20 +08:00
feat(triton): port the sampler, bit-exact on the draw and honest about the rest
`fused_sample` for the configuration the engine builds: 4096 columns after the force_topk cut,
sampling on, rejection off. With one tile the baseline's two-level search collapses -- the coarse
offset *is* the token id -- so this implements log-softmax, the top-k/top-p cut as a threshold plus
the tie rule, renormalisation, a fixed-point cumulative distribution and one Philox draw.
The draw matches cuRAND bit for bit (that is what the previous commit established). The rest does
not all match, and the split is measured rather than claimed:
engine's configuration (top_p=0.99, 2644 kept columns) tokens identical on every row
narrow cuts (top_k=128, top_p=0.5) ~90% of rows
Kept sets match column for column and log-probabilities agree to ~1e-5 in every case. What differs
is the token, and the cause is a discontinuity in the baseline rather than a defect here: the draw
is `d % total` with `total = sum(floor(p_i * 2^31))`, and the baseline's rejection limit
`(2^32-1) // total * total` steps from `2 * total` to `1 * total` exactly at `total == 2^31` -- which
`total` sits next to, because the probabilities sum to one. On the near side almost no draw is
rejected, on the far side about half are, and each rejection moves to the next draw. A few units of
rounding error in `total`, which is a reduction order that cannot match, therefore change the token
outright. Clamping `total` under 2^31 does not help: the baseline crosses in either direction too,
so it only relocates the disagreement (30 rows instead of 11).
Four bugs found on the way, all worth recording because each looked like something else:
- Triton evaluates `//`, `%` and `<=` on **uint32 operands with signed arithmetic**, so 4294967295
behaves as -1. That made `limit` come out as `2^32 - total`, flipping the rejection decision.
Fixed by doing the whole cumulative distribution in int64.
- Nested runtime `if`s produced a threshold that was not the one assigned inside the branch; the
cuts are now combined with `tl.where` instead.
- The threshold was taken from the probability array and compared against logits. The sort has to be
of the logits (equivalent, since the map is monotone -- the mass is recomputed from the sorted
logits rather than sorting twice).
- Renormalising by division rather than through the log domain, as the baseline does, moved the
token agreement from 76% to 91% before the remaining gap was understood.
Also gives the A/B harness a fourth case element: an agreement floor for outputs whose answer is a
random draw, so the rate is reported rather than compared for equality -- a comparison that can pass
without testing anything would be worse than none.
This commit is contained in:
loopyt
parent
5e2838aa5e
commit
afbfdb6c7a
@ -0,0 +1,258 @@
|
||||
"""Triton version of `fused_sample`, for the configuration the engine actually builds.
|
||||
|
||||
Measured on a real run (`DEEPINFER_TRACE_KERNEL_ARGS=1`), that is `fused_sample_tl(4096, True,
|
||||
False)`: the logits have already been cut to `force_topk` = 4096 columns, sampling is on, and
|
||||
rejection sampling is off. `vocab_size` fits in one tile there, so `N_LSE_CHUNKS`,
|
||||
`N_FILTER_CHUNKS` and `N_SAMPLE_CHUNKS` are all 1 -- which removes the two-level search entirely:
|
||||
with one chunk the coarse offset *is* the token id. Other shapes are refused rather than
|
||||
approximated, because the multi-chunk path would need the slot scan this does not implement.
|
||||
|
||||
The algorithm, matching the baseline step for step:
|
||||
|
||||
1. temperature scaling and a numerically-stable log-softmax over the row;
|
||||
2. a top-k / top-p cut, expressed as a threshold logit plus a tie rule (keep every logit above the
|
||||
threshold, then the lowest-numbered columns whose logit *equals* it, until the cut count is
|
||||
reached) -- this is what the baseline's radix select plus its tie scan computes;
|
||||
3. renormalisation over the kept set, which is the baseline's `lse0 -= log(mass)`;
|
||||
4. a fixed-point cumulative distribution at a 2^31 scale and one draw from the Philox stream
|
||||
(`triton/philox.py`), taking the first column whose cumulative mass exceeds the draw.
|
||||
|
||||
Step 4 is why the Philox port mattered: without a matching stream the sampled token would differ on
|
||||
every non-greedy call, and `temperatures` is 1.0 in practice, so every call is non-greedy.
|
||||
|
||||
Deterministic agreement is exact where it can be: the kept set (which columns survive the cut)
|
||||
matches the baseline column for column, and the returned log-probabilities agree to about 1e-5
|
||||
relative -- that is the row sum having a different reduction tree.
|
||||
|
||||
The drawn token is a different matter, and it is worth being precise about why, because the cause is
|
||||
not obvious and I got it wrong twice before measuring it:
|
||||
|
||||
r = draw % total, total = sum(floor(p_i * 2^31))
|
||||
token = first column whose cumulative mass exceeds r
|
||||
|
||||
The draw itself is bit-identical to cuRAND's (see `triton/philox.py`). But the baseline's rejection
|
||||
limit is `(2^32 - 1) // total * total`, and that expression steps from `2 * total` down to
|
||||
`1 * total` exactly when `total` crosses 2^31 -- which it sits right next to, because the
|
||||
probabilities sum to one. On the high side almost no draw is rejected; on the low side about half
|
||||
are, and each rejection moves to the next draw, so the token changes outright. A few units of
|
||||
rounding error in `total` therefore flip the token about a tenth of the time, and since `total`
|
||||
depends on a reduction order that cannot match, this is not something the port can fix.
|
||||
|
||||
Clamping `total` to stay under 2^31 does not help: the baseline's own total crosses too, in either
|
||||
direction, so clamping only moves the disagreement to other rows.
|
||||
|
||||
What that leaves is a measured split, reported by the A/B harness rather than asserted:
|
||||
|
||||
* the engine's configuration (`top_p=0.99`, no top-k, 2644 kept columns): tokens match on every row;
|
||||
* narrow cuts (`top_k=128`, `top_p=0.5`): about 90% of rows match, the rest are a different draw
|
||||
from a distribution that agrees to 1e-5.
|
||||
|
||||
So the sampler is a correct sampler, and its randomness does not line up with the baseline's on
|
||||
narrow cuts. Anything reading this file should not read that as "the kernels disagree".
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from deepinfer.model_executor.layers.tilelang_ops.triton.philox import draw_at
|
||||
|
||||
# The fixed-point scale the baseline uses for the cumulative distribution: 2^31, i.e. one bit short
|
||||
# of overflowing uint32, which leaves room for the total mass to land on the boundary.
|
||||
_SCALE_BITS = 31
|
||||
_SCALE = float(1 << _SCALE_BITS)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _fused_sample_kernel(
|
||||
LOGITS,
|
||||
TEMPERATURES,
|
||||
TOP_PS,
|
||||
TOP_KS,
|
||||
SAMPLE_POSITIONS,
|
||||
SEQ_SEEDS,
|
||||
LOGPROBS,
|
||||
CANDIDATE,
|
||||
logits_stride,
|
||||
logprobs_stride,
|
||||
seed_offset,
|
||||
V: tl.constexpr,
|
||||
BLOCK: tl.constexpr,
|
||||
SCALE: tl.constexpr,
|
||||
):
|
||||
row = tl.program_id(0)
|
||||
cols = tl.arange(0, BLOCK)
|
||||
live = cols < V
|
||||
|
||||
inv_temperature = 1.0 / tl.load(TEMPERATURES + row)
|
||||
x = tl.load(LOGITS + row * logits_stride + cols, mask=live, other=float("-inf"))
|
||||
|
||||
# Log-softmax with the temperature folded into the exponent. `other=-inf` above makes the padding
|
||||
# contribute nothing to the sum without a second mask.
|
||||
max_v = tl.max(x, axis=0)
|
||||
exp_v = tl.where(live, tl.exp((x - max_v) * inv_temperature), 0.0)
|
||||
total = tl.sum(exp_v, axis=0)
|
||||
prob = tl.where(live, exp_v / total, 0.0)
|
||||
|
||||
top_p = tl.load(TOP_PS + row)
|
||||
top_k = tl.load(TOP_KS + row)
|
||||
top_p_active = top_p < 1.0
|
||||
top_k_active = (top_k >= 0) & (top_k < V)
|
||||
|
||||
# The cut count: how many columns survive. `V` means "no cut", so taking a minimum over the
|
||||
# active cuts needs no branching -- and this deliberately avoids nested runtime `if`s, which
|
||||
# produced a wrong threshold: the value read outside the branch was not the one assigned inside.
|
||||
count_p = tl.full((), V, tl.int32)
|
||||
count_k = tl.full((), V, tl.int32)
|
||||
threshold = tl.full((), float("-inf"), tl.float32)
|
||||
|
||||
if top_p_active or top_k_active:
|
||||
# A descending sort gives both cuts at once, and it has to sort the *logits*: the threshold
|
||||
# is compared against logits below, so taking it from the probabilities would compare two
|
||||
# different quantities and silently drop every column.
|
||||
#
|
||||
# Sorting logits still yields the probability order, since the map from logit to probability
|
||||
# is monotone, so the cumulative mass is recomputed from the sorted logits rather than
|
||||
# sorting twice and hoping the two permutations agree.
|
||||
sorted_x = tl.sort(x, descending=True)
|
||||
sorted_p = tl.exp((sorted_x - max_v) * inv_temperature) / total
|
||||
|
||||
# The cumulative mass is compared in fixed point at the same 2^31 scale the baseline uses
|
||||
# (`ps = cast(top_p * SCALE) + 1`), not in float: with floats, a column whose cumulative mass
|
||||
# lands within an ulp of the target is kept on one side and dropped on the other, which
|
||||
# showed up as a one-column difference in the kept set. The sum needs int64 -- 4096 columns
|
||||
# at 2^31 each exceeds int32.
|
||||
fixed = (sorted_p * SCALE).to(tl.int64)
|
||||
cdf = tl.cumsum(fixed, axis=0)
|
||||
ps = (top_p * SCALE).to(tl.int64) + 1
|
||||
|
||||
# First index whose cumulative mass reaches top_p, and the k-th largest for top_k. Each is
|
||||
# only consulted when its cut is active, so an inactive cut is pinned to V.
|
||||
count_p = tl.minimum(tl.sum((cdf < ps).to(tl.int32), axis=0) + 1, V)
|
||||
count_p = tl.where(top_p_active, count_p, V)
|
||||
count_k = tl.where(top_k_active, tl.minimum(top_k, V), V)
|
||||
|
||||
keep_count = tl.minimum(count_p, count_k)
|
||||
# The threshold logit is the last column the cut keeps.
|
||||
selected = tl.minimum(keep_count - 1, BLOCK - 1)
|
||||
threshold = tl.sum(tl.where(tl.arange(0, BLOCK) == selected, sorted_x, 0.0), axis=0)
|
||||
|
||||
# Tie rule, as the baseline applies it: keep everything strictly above the threshold, then
|
||||
# the lowest-numbered columns equal to it until the count is reached. Real logits almost
|
||||
# never repeat, but the rule decides the answer when they do.
|
||||
above = live & (x > threshold)
|
||||
num_above = tl.sum(above.to(tl.int32), axis=0)
|
||||
is_tie = live & (x == threshold)
|
||||
# Rank within the ties by column order, which the cumulative count gives directly.
|
||||
tie_rank = tl.cumsum(is_tie.to(tl.int32), axis=0)
|
||||
keep = above | (is_tie & (tie_rank <= keep_count - num_above))
|
||||
else:
|
||||
keep = live
|
||||
|
||||
# Renormalise over the kept set through the log domain, as the baseline does (`lse0 -= log(mass)`
|
||||
# then `exp((logit + lse1) * inv_t + lse0)`), rather than dividing the probabilities. The two are
|
||||
# the same arithmetic but not the same rounding, and the total below lands on a discontinuity
|
||||
# (see `limit`), so matching the baseline's rounding here is worth the extra `log`.
|
||||
kept_prob = tl.where(keep, prob, 0.0)
|
||||
mass = tl.sum(kept_prob, axis=0)
|
||||
log_norm = -tl.log(total) - tl.log(mass)
|
||||
log_prob = tl.where(keep, (x - max_v) * inv_temperature + log_norm, float("-inf"))
|
||||
|
||||
normalised = tl.where(keep, tl.exp(log_prob), 0.0)
|
||||
tl.store(LOGPROBS + row * logprobs_stride + cols, log_prob, mask=live)
|
||||
# Fixed-point cumulative distribution, then one draw, all in int64.
|
||||
#
|
||||
# The scale is 2^31, which is one bit too large for the arithmetic to be safe in uint32: Triton
|
||||
# evaluates `//`, `%` and `<=` on uint32 operands with *signed* arithmetic, so values above 2^31
|
||||
# behave as negative. That broke the draw in two ways at once -- `limit` came out as
|
||||
# `2^32 - total`, flipping the rejection decision and moving to a different draw entirely, and
|
||||
# the search compared a wrapped cumulative against the target. int64 sidesteps both.
|
||||
fixed = (normalised * SCALE).to(tl.int64)
|
||||
cumulative = tl.cumsum(fixed, axis=0)
|
||||
|
||||
# Not clamped to 2^31 even though the probabilities sum to one: the baseline's total crosses
|
||||
# that value too (its own rounding goes either way), so clamping just moves the disagreement to
|
||||
# different rows instead of removing it. See the note in the module docstring.
|
||||
total_mass = tl.max(cumulative, axis=0)
|
||||
|
||||
seed = tl.load(SEQ_SEEDS + row).to(tl.uint64) + seed_offset
|
||||
# The same position derivation as the baseline: three streams per sampling position, and the
|
||||
# first one is the sample itself.
|
||||
subseq = (tl.load(SAMPLE_POSITIONS + row).to(tl.uint64) * 3).to(tl.uint64)
|
||||
|
||||
k0 = (seed & 0xFFFFFFFF).to(tl.uint32)
|
||||
k1 = (seed >> 32).to(tl.uint32)
|
||||
sub_lo = (subseq & 0xFFFFFFFF).to(tl.uint32)
|
||||
sub_hi = (subseq >> 32).to(tl.uint32)
|
||||
|
||||
# The baseline rejects draws above the largest exact multiple of the total, so the residual is
|
||||
# uniform rather than biased, and consumes one draw per rejection. With a total near 2^31 that
|
||||
# rejects about half the time, so this loops a couple of times on average.
|
||||
max_draw = tl.full((), 0xFFFFFFFF, tl.int64)
|
||||
limit = (max_draw // total_mass) * total_mass
|
||||
n = tl.zeros((), tl.uint32)
|
||||
draw = draw_at(k0, k1, sub_lo, sub_hi, n)
|
||||
while draw >= limit:
|
||||
n = n + 1
|
||||
draw = draw_at(k0, k1, sub_lo, sub_hi, n)
|
||||
|
||||
target = tl.where(total_mass == 0, tl.zeros((), tl.int64), draw.to(tl.int64) % total_mass)
|
||||
# First column whose cumulative mass exceeds the draw; `<=` counts the columns that do not.
|
||||
token = tl.sum((cumulative <= target).to(tl.int32), axis=0)
|
||||
token = tl.minimum(token, V - 1)
|
||||
tl.store(CANDIDATE + row, token)
|
||||
|
||||
|
||||
def fused_sample_tl(vocab_size: int, do_sample: bool, do_reject_sampling: bool) -> Any:
|
||||
"""Build the sampler for one configuration."""
|
||||
assert vocab_size >= 32, "fused_sample requires vocab_size >= 32"
|
||||
assert not do_reject_sampling or do_sample, "reject sampling requires sampling"
|
||||
assert not do_reject_sampling, "rejection sampling (the MTP verify path) is not ported yet"
|
||||
assert do_sample, "the greedy path is a different kernel in the baseline and is not ported"
|
||||
block = triton.next_power_of_2(vocab_size)
|
||||
|
||||
def launch(
|
||||
logits: torch.Tensor,
|
||||
temperatures: torch.Tensor,
|
||||
top_p: torch.Tensor,
|
||||
top_k: torch.Tensor,
|
||||
sample_positions: torch.Tensor,
|
||||
seq_seeds: torch.Tensor,
|
||||
is_verify: torch.Tensor | None,
|
||||
logprobs: torch.Tensor,
|
||||
candidate_token_ids: torch.Tensor,
|
||||
last_nextn_token_logprobs: torch.Tensor | None,
|
||||
mtp_info_read_rows: torch.Tensor | None,
|
||||
last_next_n_token_ids: torch.Tensor | None,
|
||||
accept_result_verify: torch.Tensor | None,
|
||||
draft_selected_logprobs: torch.Tensor | None,
|
||||
accept_probs: torch.Tensor | None,
|
||||
seed_offset: int,
|
||||
lse_workspace: torch.Tensor,
|
||||
) -> None:
|
||||
del is_verify, last_nextn_token_logprobs, mtp_info_read_rows, last_next_n_token_ids
|
||||
del accept_result_verify, draft_selected_logprobs, accept_probs, lse_workspace
|
||||
total_q = logits.shape[0]
|
||||
if total_q == 0:
|
||||
return
|
||||
_fused_sample_kernel[(total_q,)](
|
||||
logits,
|
||||
temperatures,
|
||||
top_p,
|
||||
top_k,
|
||||
sample_positions,
|
||||
seq_seeds.view(torch.int64),
|
||||
logprobs,
|
||||
candidate_token_ids,
|
||||
logits.stride(0),
|
||||
logprobs.stride(0),
|
||||
seed_offset,
|
||||
V=vocab_size,
|
||||
BLOCK=block,
|
||||
SCALE=_SCALE,
|
||||
num_warps=max(1, min(8, block // 512)),
|
||||
)
|
||||
|
||||
return launch
|
||||
@ -50,6 +50,20 @@ def philox4x32_10(c0, c1, c2, c3, k0, k1):
|
||||
return c0, c1, c2, c3
|
||||
|
||||
|
||||
@triton.jit
|
||||
def draw_at(k0, k1, sub_lo, sub_hi, n):
|
||||
"""The `n`-th draw of a stream, addressed directly.
|
||||
|
||||
`curand()` returns the four words of a block in order before advancing the counter, so draw `n`
|
||||
is word `n % 4` of block `n // 4`. Addressing it directly (rather than keeping a running state)
|
||||
lets a rejection loop re-draw with a scalar index, which a tuple cannot be indexed by at runtime.
|
||||
"""
|
||||
block = n // 4
|
||||
word = n % 4
|
||||
w0, w1, w2, w3 = philox4x32_10(block, tl.zeros((), tl.uint32), sub_lo, sub_hi, k0, k1)
|
||||
return tl.where(word == 0, w0, tl.where(word == 1, w1, tl.where(word == 2, w2, w3)))
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _philox_stream_kernel(
|
||||
SEED,
|
||||
|
||||
@ -25,6 +25,7 @@ from deepinfer.model_executor.layers.tilelang_ops.compress import (
|
||||
compress_and_update_state_prefill,
|
||||
)
|
||||
from deepinfer.model_executor.layers.tilelang_ops.flash_mla_utils import convert_kv_index
|
||||
from deepinfer.model_executor.layers.tilelang_ops.fused_sample import fused_sample
|
||||
from deepinfer.model_executor.layers.tilelang_ops.fused_transform import (
|
||||
CastArgs,
|
||||
KVCacheOutArgs,
|
||||
@ -44,6 +45,7 @@ from deepinfer.model_executor.layers.tilelang_ops.sampler import map_lookup
|
||||
from deepinfer.model_executor.layers.tilelang_ops.store_kv_cache import store_kv_cache
|
||||
from deepinfer.model_executor.layers.tilelang_ops.store_mtp_info import store_mtp_info
|
||||
from deepinfer.model_executor.layers.tilelang_ops.swa_index import expand_swa_index
|
||||
from deepinfer.model_executor.sampling_inputs import LogitsTransformInputs, SamplingKernelInputs
|
||||
|
||||
|
||||
def _cases_gather() -> list[tuple[str, bool, Callable[[], Any]]]:
|
||||
@ -888,6 +890,65 @@ def _cases_compress_decode() -> list[tuple[str, bool, Callable[[], Any]]]:
|
||||
]
|
||||
|
||||
|
||||
def _cases_fused_sample() -> list[tuple[str, bool, Callable[[], Any]]]:
|
||||
"""Sampler output: log-probabilities and the drawn token.
|
||||
|
||||
The comparison is not `exact`: a different reduction tree shifts the row sum by an ulp, which
|
||||
propagates into the fixed-point cumulative mass, where a column sitting exactly on a draw
|
||||
boundary can flip. So log-probabilities are compared with a tolerance and the token count is
|
||||
checked as an agreement rate -- the honest standard for a kernel whose answer is a random draw.
|
||||
"""
|
||||
|
||||
def make(vocab: int, rows: int, top_p: float, top_k: int):
|
||||
def run() -> Any:
|
||||
torch.manual_seed(0)
|
||||
logits = torch.randn((rows, vocab), device="cuda", dtype=torch.float32) * 2.0
|
||||
temperatures = torch.ones(rows, device="cuda", dtype=torch.float32)
|
||||
top_ps = torch.full((rows,), top_p, device="cuda", dtype=torch.float32)
|
||||
top_ks = torch.full((rows,), top_k, device="cuda", dtype=torch.int32)
|
||||
positions = torch.arange(1000, 1000 + rows, device="cuda", dtype=torch.int32)
|
||||
# torch has no uint64 arange, so build in int64 and reinterpret: the seeds are a bit
|
||||
# pattern, and they exceed 2^63 as uint64.
|
||||
seeds = (torch.arange(rows, device="cuda", dtype=torch.int64) * 7919 + 12345).view(torch.uint64)
|
||||
logprobs, token, *_ = fused_sample(
|
||||
logits,
|
||||
SamplingKernelInputs(
|
||||
logits=LogitsTransformInputs(temperatures, top_ps, top_ks),
|
||||
seq_seeds=seeds,
|
||||
sample_positions=positions,
|
||||
),
|
||||
)
|
||||
return logprobs, token
|
||||
|
||||
return run
|
||||
|
||||
# The last element is the token agreement floor, and it is a measured number, not a guess.
|
||||
#
|
||||
# The log-probabilities are compared with a tolerance and the kept set has to match exactly; the
|
||||
# token cannot, because the draw is `d % total` and the baseline's rejection limit steps at
|
||||
# `total == 2^31`, which `total` sits next to. A few units of rounding error in `total` therefore
|
||||
# move to a different draw. The engine's own configuration is unaffected (its totals stay on one
|
||||
# side), so it is held at 0.9 and passes with every row matching; the narrow cuts sit either
|
||||
# side of that step on 10-13% of rows, so their floor is 0.85.
|
||||
return [
|
||||
# The engine's own configuration: 4096 columns after the force_topk cut, top_p=0.99, no top_k.
|
||||
("engine_cfg", False, make(4096, 8, 0.99, -1), 0.9),
|
||||
("engine_cfg_batch32", False, make(4096, 32, 0.99, -1), 0.9),
|
||||
# top_k only, and both cuts at once.
|
||||
# A narrow cut concentrates the mass in few columns, so the fixed-point cumulative values are
|
||||
# coarse relative to the spacing between them and a boundary lands on a different side more
|
||||
# often. 128 rows rather than 8, because with 8 a single flip is 12.5% and the rate says
|
||||
# nothing.
|
||||
("top_k_128", False, make(4096, 128, 1.0, 128), 0.85),
|
||||
("top_k_1", False, make(4096, 32, 1.0, 1), 0.85),
|
||||
("top_p_and_k", False, make(4096, 128, 0.9, 256), 0.85),
|
||||
("no_cut", False, make(4096, 128, 1.0, -1), 0.85),
|
||||
# Sharp logits make the distribution peaky, so a small numeric difference cannot move the
|
||||
# draw; flat ones put many columns on similar mass and are the harder case.
|
||||
("top_p_0.5", False, make(4096, 128, 0.5, -1), 0.85),
|
||||
]
|
||||
|
||||
|
||||
CASES: dict[str, Callable[[], list[tuple[str, bool, Callable[[], Any]]]]] = {
|
||||
"store_kv_cache_tl": _cases_store_kv_cache,
|
||||
"load_paged_kv_cache_tl": _cases_load_paged_kv_cache,
|
||||
@ -899,6 +960,7 @@ CASES: dict[str, Callable[[], list[tuple[str, bool, Callable[[], Any]]]]] = {
|
||||
"prepare_v4_decode_metadata_tl": _cases_prepare_v4_decode_metadata,
|
||||
"compress_and_update_state_prefill_tl": _cases_compress_prefill,
|
||||
"compress_and_update_state_decode_tl": _cases_compress_decode,
|
||||
"fused_sample_tl": _cases_fused_sample,
|
||||
"copy_block_tables_tl": _cases_copy_block_tables,
|
||||
"gather_tl": _cases_gather,
|
||||
"gather_input_ids_zos_tl": _cases_gather_input_ids,
|
||||
|
||||
@ -68,8 +68,38 @@ _DTYPE_TOL: dict[torch.dtype, tuple[float, float]] = {
|
||||
}
|
||||
|
||||
|
||||
def compare(name: str, expected: Any, actual: Any, *, exact: bool, rtol: float | None = None, atol: float | None = None) -> bool:
|
||||
def compare(
|
||||
name: str,
|
||||
expected: Any,
|
||||
actual: Any,
|
||||
*,
|
||||
exact: bool,
|
||||
rtol: float | None = None,
|
||||
atol: float | None = None,
|
||||
agreement: float | None = None,
|
||||
) -> bool:
|
||||
"""Report whether one case agrees; returns True on success."""
|
||||
if (
|
||||
agreement is not None
|
||||
and isinstance(expected, torch.Tensor)
|
||||
and expected.dtype
|
||||
in (
|
||||
torch.int32,
|
||||
torch.int64,
|
||||
)
|
||||
):
|
||||
# Random-draw outputs: require a high agreement rate rather than equality, and say what the
|
||||
# rate was so a regression that halves it cannot hide behind the threshold.
|
||||
same = int((expected == actual).sum())
|
||||
total = expected.numel()
|
||||
rate = same / max(total, 1)
|
||||
if rate >= agreement:
|
||||
if rate < 1.0:
|
||||
print(f" note {name}: {total - same}/{total} differ (random draw on a boundary)")
|
||||
return True
|
||||
print(f" FAIL {name}: agreement {rate:.4f} < {agreement} ({total - same}/{total})")
|
||||
return False
|
||||
|
||||
if isinstance(expected, torch.Tensor):
|
||||
if expected.dtype != actual.dtype:
|
||||
print(f" FAIL {name}: dtype {expected.dtype} vs {actual.dtype}")
|
||||
@ -115,7 +145,7 @@ def compare(name: str, expected: Any, actual: Any, *, exact: bool, rtol: float |
|
||||
return False
|
||||
ok = True
|
||||
for i, (e, a) in enumerate(zip(expected, actual, strict=True)):
|
||||
ok &= compare(f"{name}[{i}]", e, a, exact=exact, rtol=rtol, atol=atol)
|
||||
ok &= compare(f"{name}[{i}]", e, a, exact=exact, rtol=rtol, atol=atol, agreement=agreement)
|
||||
return ok
|
||||
|
||||
same = expected == actual
|
||||
@ -146,7 +176,13 @@ def main() -> None:
|
||||
failures = 0
|
||||
for kernel in selected:
|
||||
print(f"[{kernel}]")
|
||||
for case_name, exact, fn in registry[kernel]():
|
||||
for case in registry[kernel]():
|
||||
# A case is `(name, exact, fn)`, optionally with a fourth element: the fraction of
|
||||
# entries allowed to differ. Only for outputs whose answer is a random draw, where the
|
||||
# last bit of a cumulative mass decides which side of a boundary the token falls on;
|
||||
# everything else is compared exactly or within a dtype tolerance.
|
||||
case_name, exact, fn = case[:3]
|
||||
agreement = case[3] if len(case) > 3 else None
|
||||
try:
|
||||
expected, actual = run_both(kernel, fn)
|
||||
except Exception as exc: # noqa: BLE001 - a raise on one backend is itself a result
|
||||
@ -154,7 +190,7 @@ def main() -> None:
|
||||
print(f" FAIL {case_name}: {type(exc).__name__}: {lines[0][:160] if lines else ''}")
|
||||
failures += 1
|
||||
continue
|
||||
if not compare(case_name, expected, actual, exact=exact):
|
||||
if not compare(case_name, expected, actual, exact=exact, agreement=agreement):
|
||||
failures += 1
|
||||
elif args.verbose:
|
||||
print(f" ok {case_name}")
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user