mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 08:46:21 +08:00
28fb291db1
No behaviour change. `CONTRIBUTING.md` requires `ruff check --no-cache` and `ruff format --check --no-cache` to pass, and the SM90 branch predates that requirement: 72 findings, all in the files added here. Most of it is mechanical (`ruff format`), plus: - annotations for the helper signatures (`torch.Tensor` / `int` / `Generator` / return types); hooks are typed as `Callable[[nn.Module, tuple, object], None]` rather than left bare; - two `if`/`else` assignments become ternaries (SIM108), and the index read in `h20_real_requant_check.py` uses a context manager (SIM115). Also corrects the `sparse_mla_fwd` docstring, which still described the old `0 <= idx <= q_position` masking contract. Passing `is_causal=True` there is what produced the all-NaN layer 0: indices address the assembled `[swa cache | current swa | compressed]` buffer, so the row number is not a valid causal bound and the engine must pass False.
113 lines
4.5 KiB
Python
113 lines
4.5 KiB
Python
"""Check the torch fallback for `deep_select.interface.topk` against a reference.
|
|
|
|
Covers the parts of the contract the indexer and sampler actually depend on:
|
|
`end` truncation (scores past it are uninitialized garbage), OOB padding when a row has
|
|
fewer than `topk` valid entries, `output_idx_offset`, in-place writes into a non-contiguous
|
|
view, and descending order for `sorted_index=True`.
|
|
"""
|
|
|
|
import torch
|
|
|
|
from deepinfer.h20_ops.deep_select_sm90 import topk
|
|
|
|
|
|
def ref_topk(
|
|
scores: torch.Tensor,
|
|
k: int,
|
|
end: torch.Tensor | None = None,
|
|
offset: int | None = None,
|
|
fill: int = -1,
|
|
val_fill: float = float("-inf"),
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Reference: consider only [0, end) per row, take k largest, pad, then shift."""
|
|
rows = scores.shape[0]
|
|
idx_rows, val_rows = [], []
|
|
for r in range(rows):
|
|
n = int(end[r]) if end is not None else scores.shape[1]
|
|
row = scores[r, :n]
|
|
kk = min(k, n)
|
|
order = torch.argsort(row, descending=True)[:kk]
|
|
vals = row[order]
|
|
idxs = order.long()
|
|
if offset is not None:
|
|
idxs = idxs + int(offset[r])
|
|
pad = k - kk
|
|
if pad:
|
|
idxs = torch.cat([idxs, torch.full((pad,), fill, dtype=idxs.dtype, device=scores.device)])
|
|
vals = torch.cat([vals, torch.full((pad,), val_fill, dtype=vals.dtype, device=scores.device)])
|
|
idx_rows.append(idxs)
|
|
val_rows.append(vals)
|
|
return torch.stack(val_rows), torch.stack(idx_rows)
|
|
|
|
|
|
def main() -> None:
|
|
dev = "cuda"
|
|
torch.manual_seed(0)
|
|
rows, width, k = 6, 4096, 128
|
|
|
|
scores = torch.randn(rows, width, device=dev, dtype=torch.float32)
|
|
# Garbage past `end` that would win if the implementation ignored the bound.
|
|
end = torch.tensor([4096, 4096, 200, 128, 3, 50], dtype=torch.int32, device=dev)
|
|
for r, n in enumerate(end.tolist()):
|
|
scores[r, n:] = 1e6
|
|
|
|
vals, idxs = topk(scores, k, sorted=False, end=end, indices_type=torch.int32, return_value=True)
|
|
# Default `idx_oob_fill_value` is INT_MAX, so the reference must pad with the same value.
|
|
rv, ri = ref_topk(scores, k, end=end, fill=2147483647)
|
|
assert idxs.dtype == torch.int32
|
|
bad_rows = [r for r in range(rows) if not torch.equal(vals[r], rv[r]) or not torch.equal(idxs[r].long(), ri[r])]
|
|
print(f"[1] end truncation + OOB pad: rows mismatched = {bad_rows}")
|
|
assert not bad_rows
|
|
# Rows shorter than topk must be padded with the fill value, not with garbage indices.
|
|
short = (idxs[4] == 2147483647).sum().item()
|
|
print(f" row 4 (end=3) padded slots = {short} (expect {k - 3})")
|
|
assert short == k - 3
|
|
|
|
# Descending order for the valid prefix, which is what `sorted_index=True` promises.
|
|
lens = torch.clamp(end.long(), max=k)
|
|
for r in range(rows):
|
|
n = int(lens[r])
|
|
assert torch.equal(vals[r, :n], torch.sort(vals[r, :n], descending=True).values)
|
|
print("[2] descending order holds")
|
|
|
|
# In-place write into a non-contiguous view + per-row offset.
|
|
storage = torch.full((rows, k * 2), -7, dtype=torch.int32, device=dev)
|
|
view = storage[:, :k]
|
|
offset = torch.arange(rows, dtype=torch.int32, device=dev) * 1000
|
|
vals2, idxs2 = topk(
|
|
scores,
|
|
k,
|
|
end=end,
|
|
indices_type=torch.int32,
|
|
output_idx=view,
|
|
output_idx_offset=offset,
|
|
idx_oob_fill_value=-1,
|
|
return_value=False,
|
|
)
|
|
assert idxs2.data_ptr() == view.data_ptr()
|
|
assert vals2 is None
|
|
rv2, ri2 = ref_topk(scores, k, end=end, offset=offset, fill=-1)
|
|
print(f"[3] in-place view: match = {torch.equal(view.long(), ri2)}")
|
|
assert torch.equal(view.long(), ri2)
|
|
# The fill value must not be shifted by the offset.
|
|
assert int(view[4].min()) == -1
|
|
|
|
# No `end` at all takes the plain path.
|
|
v3, i3 = topk(scores, k, indices_type=torch.int64)
|
|
rv3, ri3 = ref_topk(scores, k)
|
|
print(f"[4] no end: values match = {torch.allclose(v3, rv3)}, indices match = {torch.equal(i3, ri3)}")
|
|
assert torch.allclose(v3, rv3) and torch.equal(i3, ri3)
|
|
|
|
# bf16 rows (the candidate-block path feeds bf16 block scores).
|
|
bs = torch.randn(3, 512, device=dev, dtype=torch.bfloat16)
|
|
vb, ib = topk(bs, 16, end=torch.tensor([512, 40, 8], dtype=torch.int32, device=dev), indices_type=torch.int32, return_value=False)
|
|
padded = int((ib[2] == 2147483647).sum())
|
|
print(f"[5] bf16: rows={tuple(ib.shape)} row2 padded={padded} (expect 8)")
|
|
assert padded == 16 - 8
|
|
|
|
print("PASS: deep_select torch fallback matches the documented contract")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|