mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 16:56:21 +08:00
283ff893d5
Comparing final logits was impossible before: `ref_layer_dump.py` read `logits[0, -1]`, but the
head already collapsed the sequence axis (`x = x[:, -1]`), so logits are `[b, vocab]` and that
index silently stored a 0-dimensional scalar. Every earlier dump had an unusable
`final_row_logits`.
With that fixed, the comparison reframes the remaining drift. On the same 14 tokens:
prefill top-1 = 6127 (北京) margin(6127 - 666) = +0.256
reference top-1 = 6127 (北京) margin = +3.112
Our prefill picks the right token; the error only compresses the margin, by ~12x. But the
end-to-end run, which reaches that context via decode steps, picks 666 (`**`) instead. Same
context, different path, different answer -- so the remaining drift is in **decode**, not in the
per-128 quantization. The quantization floor is still 2.7e-2 and still irreducible
(`H20-DRIFT-ANALYSIS.md` has the DeepGEMM assertion that forces `gran_k == 128` on SM90), but it
is not what flips the token.
Two traps had to be handled to get numbers that mean anything:
- `execute_model_profile_run` runs a full forward + sampler on synthetic tokens. The layer dumps
were filtered by their prompt-length gate; logits are not, because their row count is the
number of *selected* tokens (1). The first attempt captured that step and produced cos = 0.34.
`debug_dump.profiling()` now suspends capture for the whole profile step.
- With `max_tokens=1` both a prompt and a decode step run, and only the prompt one is comparable
with a prefill-only reference. They are dumped as `logits_prompt` / `logits_decode`.
Also re-runs the layer comparison now that the expert-scale fix has landed: layer 0 went from
1.27e-1 to 2.88e-2 and layer 39 sits at 2.00e-1, so the error curve flattened by an order of
magnitude.
89 lines
3.6 KiB
Python
89 lines
3.6 KiB
Python
"""Probe which FP8 scale recipes DeepGEMM accepts on this GPU, via its own layout helper.
|
|
|
|
The SM90 path re-quantizes everything to fp32 per-128 because `get_gemm_recipe()` assumes Hopper
|
|
has nothing finer. The reference converter instead keeps an e8m0 scale per 32x32 block and folds
|
|
the relative offset into the FP8 values, which is lossless -- worth adopting only if the grouped
|
|
GEMM actually accepts it.
|
|
|
|
`transform_sf_into_required_layout` is the single place DeepGEMM validates a recipe (the engine
|
|
calls it for every MoE weight), so probing it answers the question directly and in milliseconds,
|
|
without materializing weights or launching a kernel.
|
|
|
|
python scripts/h20_recipe_probe.py
|
|
"""
|
|
|
|
import deep_gemm
|
|
import torch
|
|
|
|
E8M0 = torch.float8_e8m0fnu
|
|
|
|
|
|
def make_sf(m: int, k: int, block_m: int, block_k: int, num_groups: int, e8m0: bool) -> torch.Tensor:
|
|
"""Scales in the plainest layout: one value per (block_m x block_k) tile, per group."""
|
|
shape = (num_groups * (m // block_m), k // block_k)
|
|
if e8m0:
|
|
return torch.full(shape, 127, dtype=torch.uint8, device="cuda").view(E8M0)
|
|
return torch.ones(shape, dtype=torch.float32, device="cuda")
|
|
|
|
|
|
def probe(recipe_b: tuple[int, int], block_a: int, e8m0: bool, n: int, k: int, num_groups: int) -> str:
|
|
block_n, block_k = recipe_b
|
|
sf = make_sf(n, k, block_n, block_k, num_groups, e8m0)
|
|
try:
|
|
out = deep_gemm.transform_sf_into_required_layout(
|
|
sf,
|
|
n,
|
|
k,
|
|
recipe=(1, block_n, block_k),
|
|
num_groups=num_groups,
|
|
is_sfa=False,
|
|
disable_ue8m0_cast=e8m0,
|
|
)
|
|
got = tuple(out.shape)
|
|
# The helper silently re-lays-out whatever it is given, so acceptance alone is not proof;
|
|
# report the shape so an obviously wrong result is visible too.
|
|
return f"ACCEPTED out={got} dtype={out.dtype}"
|
|
except Exception as exc: # noqa: BLE001 - report whatever it raises
|
|
first = (str(exc).strip().splitlines() or [type(exc).__name__])[0]
|
|
return f"REJECTED {first[:100]}"
|
|
|
|
|
|
def main() -> None:
|
|
print(f"device: {torch.cuda.get_device_name(0)}")
|
|
print(f"deep_gemm: {deep_gemm.__version__}\n")
|
|
|
|
n, k, num_groups = 2304, 2560, 8
|
|
|
|
print("B-operand recipes (what expert weights use, is_sfa=False):")
|
|
for recipe, block_a, e8m0, label in [
|
|
((128, 128), 128, False, "fp32 per-128 (current SM90)"),
|
|
((128, 128), 128, True, "e8m0 per-128"),
|
|
((32, 32), 128, False, "fp32 per-32 (ref-converter grid)"),
|
|
((32, 32), 128, True, "e8m0 per-32 (ref-converter, lossless)"),
|
|
((1, 32), 128, True, "e8m0 per-row/32 (MXFP8)"),
|
|
((1, 128), 128, True, "e8m0 per-row/128"),
|
|
]:
|
|
print(f" {label:34} recipe={recipe} e8m0={e8m0:<5} -> {probe(recipe, block_a, e8m0, n, k, num_groups)}")
|
|
|
|
print("\nA-operand recipes (activations, is_sfa=True):")
|
|
for block_k, e8m0, label in [
|
|
(128, False, "fp32 per-128 (current SM90)"),
|
|
(32, True, "e8m0 per-32 (MXFP8)"),
|
|
(32, False, "fp32 per-32"),
|
|
]:
|
|
sf = make_sf(n, k, 1, block_k, 1, e8m0)
|
|
try:
|
|
out = deep_gemm.transform_sf_into_required_layout(
|
|
sf, n, k, recipe=(1, 1, block_k), num_groups=1, is_sfa=True, disable_ue8m0_cast=e8m0
|
|
)
|
|
print(f" {label:34} block_k={block_k:<4} e8m0={e8m0:<5} -> ACCEPTED out={tuple(out.shape)}")
|
|
except Exception as exc: # noqa: BLE001
|
|
first = (str(exc).strip().splitlines() or [type(exc).__name__])[0]
|
|
print(f" {label:34} block_k={block_k:<4} e8m0={e8m0:<5} -> REJECTED {first[:100]}")
|
|
|
|
print("\nPROBE_DONE")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|