mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 08:46:21 +08:00
4a5606ad2d
Groundwork for replacing TileLang with Triton on the H20 path, plus the first ported kernel. The scope is measured rather than guessed. `DEEPINFER_TRACE_KERNELS=<path>` proxies every `*_tl` factory in the dispatcher and records its first use; one full text-generation run on 8xH20 says exactly 24 kernels execute. `H20-TRITON-PLAN.md` lists them, in dependency order. `DEEPINFER_OPS_BACKEND` selects the backend package (unset keeps the platform default, so nothing changes by omission) and `DEEPINFER_OPS_FALLBACK` names individual kernels that must stay on the default. Kernels the selected backend does not export at all are taken from the default with a one-line notice, which lets a partial port run end to end the whole way through instead of only at the end. That coexistence is the reason for a parallel backend rather than an in-place rewrite: the TileLang kernel is a validated baseline, so each Triton kernel can be diffed against it on real shapes through the same public interface function, which is a much stronger check than comparing against a fresh Python reference. `scripts/h20_ops_ab.py` does that swap and reports per case. First kernel: `gather_tl`, all 9 cases bitwise identical (zeros filled for out-of-range indices, int8/bfloat16 payloads, strided rows, int32/int64 indices, empty and single-row inputs).
54 lines
2.4 KiB
Python
54 lines
2.4 KiB
Python
"""Per-kernel A/B cases: run the public interface function once per backend and compare.
|
|
|
|
Each entry maps a kernel name to a builder returning ``(case_name, exact, callable)`` triples.
|
|
``exact=True`` demands bitwise equality, which is right for index and metadata kernels; floating
|
|
point kernels are compared with a tolerance and are labelled ``exact=False``.
|
|
|
|
Shapes and dtypes reuse the ones the engine actually passes, taken from the TileLang test matrices
|
|
in ``cuda/*.py`` and from the traced inference run.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Callable
|
|
from typing import Any
|
|
|
|
import torch
|
|
|
|
from deepinfer.model_executor.layers.tilelang_ops.gather import gather
|
|
|
|
|
|
def _cases_gather() -> list[tuple[str, bool, Callable[[], Any]]]:
|
|
"""Row gather across the widths the engine uses: metadata rows, heads, and hidden states."""
|
|
|
|
def make(n: int, rows: int, d: int, dtype: torch.dtype, index_dtype: torch.dtype, strided: bool):
|
|
def run() -> Any:
|
|
torch.manual_seed(0)
|
|
base = (
|
|
torch.randn(n, d, device="cuda").to(dtype)
|
|
if dtype.is_floating_point
|
|
else torch.randint(0, 100, (n, d), device="cuda", dtype=dtype)
|
|
)
|
|
x = base[:, ::1] if strided else base
|
|
index = torch.randint(-3, n + 3, (rows,), device="cuda", dtype=index_dtype)
|
|
return gather(x, index)
|
|
|
|
return run
|
|
|
|
return [
|
|
("n1k_m233_d128_int32_idx", True, make(1024, 233, 128, torch.int8, torch.int32, False)),
|
|
("n1k_m1_d4_int8", True, make(1024, 1, 4, torch.int8, torch.int32, False)),
|
|
("n64_m256_d528_bytes", True, make(64, 256, 528, torch.int8, torch.int32, False)),
|
|
("n512_m300_d4096_bytes", True, make(512, 300, 4096, torch.int8, torch.int32, False)),
|
|
("n512_m300_d7168_int64_idx", True, make(512, 300, 7168, torch.int8, torch.int64, False)),
|
|
("n256_m256_d4096_bf16", True, make(256, 256, 4096, torch.bfloat16, torch.int32, False)),
|
|
("n256_m0_d128_empty", True, make(256, 0, 128, torch.int8, torch.int32, False)),
|
|
("n1_m17_d128_single_row", True, make(1, 17, 128, torch.int8, torch.int32, False)),
|
|
("n300_m233_d1025_odd", True, make(300, 233, 1025, torch.int8, torch.int32, False)),
|
|
]
|
|
|
|
|
|
CASES: dict[str, Callable[[], list[tuple[str, bool, Callable[[], Any]]]]] = {
|
|
"gather_tl": _cases_gather,
|
|
}
|