mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-30 15:57:20 +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.
92 lines
3.8 KiB
Python
92 lines
3.8 KiB
Python
"""SM90 sanity checks: DeepGEMM FP8 1D2D (fp32 scales) + the H20 requant helpers."""
|
|
|
|
import deep_gemm
|
|
import torch
|
|
|
|
from deepinfer.h20_ops.quant import _unpack_fp4, requant_dense_fp8, requant_expert_fp4, ue8m0_to_fp32
|
|
|
|
torch.manual_seed(0)
|
|
dev = "cuda"
|
|
print("cap", torch.cuda.get_device_capability(0))
|
|
|
|
# ---------------- 1. DeepGEMM fp8/fp8 1D2D GEMM with fp32 per-128 scales on SM90 ----
|
|
M, N, K = 256, 512, 1024
|
|
blk = 128
|
|
x = torch.randn(M, K, device=dev, dtype=torch.bfloat16)
|
|
w = torch.randn(N, K, device=dev, dtype=torch.bfloat16)
|
|
|
|
|
|
def q_token(t: torch.Tensor, blk: int) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""A operand: per-token 1x128 blocks -> sf (rows, cols/128) fp32."""
|
|
r, c = t.shape
|
|
tf = t.float()
|
|
amax = tf.abs().view(r, c // blk, blk).amax(dim=2)
|
|
s = (amax / 448.0).clamp_min(1e-30)
|
|
q = (tf.view(r, c // blk, blk) / s[:, :, None]).to(torch.float8_e4m3fn)
|
|
return q.view(r, c), s.contiguous()
|
|
|
|
|
|
def q_blocks(t: torch.Tensor, blk: int) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""B operand: 128x128 blocks -> sf (rows/128, cols/128) fp32."""
|
|
r, c = t.shape
|
|
tf = t.float()
|
|
amax = tf.abs().view(r // blk, blk, c // blk, blk).amax(dim=(1, 3))
|
|
s = (amax / 448.0).clamp_min(1e-30)
|
|
q = (tf.view(r // blk, blk, c // blk, blk) / s[:, None, :, None]).to(torch.float8_e4m3fn)
|
|
return q.view(r, c), s.contiguous()
|
|
|
|
|
|
xq, xsf = q_token(x, blk)
|
|
wq, wsf = q_blocks(w, blk)
|
|
out = torch.empty(M, N, device=dev, dtype=torch.bfloat16)
|
|
deep_gemm.fp8_fp4_gemm_nt((xq, xsf), (wq, wsf), out, None, disable_ue8m0_cast=True)
|
|
print("recipe default (1,128,128) OK")
|
|
ref = x @ w.t()
|
|
err = (out.float() - ref).abs().max().item()
|
|
rel = err / ref.abs().max().item()
|
|
print(f"gemm 1D2D fp32-sf on SM90: max_abs_err={err:.4f} rel={rel:.5f}")
|
|
|
|
# ---------------- 2. per-token activation quant via the runtime recipe ----------
|
|
params = {"fmt": "e4m3", "num_per_channels": 128, "round_sf": True, "use_packed_ue8m0": False, "use_tma_aligned_col_major_sf": True}
|
|
import tile_kernels
|
|
|
|
a8, asf = tile_kernels.quant.per_token_cast(x, **params)
|
|
print("per_token_cast ->", a8.shape, a8.dtype, asf.shape, asf.dtype)
|
|
|
|
# ---------------- 3. expert FP4 -> FP8 requant round-trip ----------------------
|
|
n_out, k = 256, 512
|
|
vals = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, 0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], device=dev)
|
|
codes = torch.randint(0, 16, (n_out, k), device=dev)
|
|
ref_x = vals[codes]
|
|
packed = ((codes[..., 1::2] << 4) | codes[..., 0::2]).to(torch.uint8).view(torch.int8).contiguous()
|
|
assert torch.allclose(_unpack_fp4(packed), ref_x, atol=0, rtol=0), "fp4 unpack mismatch"
|
|
|
|
exp = torch.randint(118, 132, (n_out, k // 32), device=dev)
|
|
sf_e8m0 = (exp.to(torch.int32)).to(torch.uint8).view(torch.float8_e8m0fnu)
|
|
x_true = ref_x * ue8m0_to_fp32(sf_e8m0).repeat_interleave(32, dim=1)
|
|
|
|
w8, sf = requant_expert_fp4(packed, ue8m0_to_fp32(sf_e8m0), out_block=128)
|
|
xr = w8.to(torch.float32) * sf.t().repeat_interleave(128, 0).repeat_interleave(128, 1)
|
|
print(
|
|
"expert requant: fp8",
|
|
tuple(w8.shape),
|
|
"sf",
|
|
tuple(sf.shape),
|
|
"max_abs_err",
|
|
(xr - x_true).abs().max().item(),
|
|
"rel",
|
|
((xr - x_true).abs().max() / x_true.abs().max()).item(),
|
|
)
|
|
|
|
# ---------------- 4. dense 32->128 requant ------------------------------------
|
|
w_ref = torch.randn(256, 512, device=dev, dtype=torch.bfloat16) / 8
|
|
w8d, sfd = q_blocks(w_ref, 32)
|
|
sf32 = ue8m0_to_fp32(
|
|
(torch.randint(118, 130, (256 // 32, 512 // 32), device=dev).to(torch.int32)).to(torch.uint8).view(torch.float8_e8m0fnu)
|
|
)
|
|
src = w8d.float() * sf32.repeat_interleave(32, 0).repeat_interleave(32, 1)
|
|
w8n, sfn = requant_dense_fp8(w8d, sf32, dst_block=128)
|
|
rec = w8n.float() * sfn.repeat_interleave(128, 0).repeat_interleave(128, 1)
|
|
print("dense 32->128 requant: sf", tuple(sfn.shape), "rel_err", ((rec - src).abs().max() / src.abs().max()).item())
|
|
print("ALL_CHECKS_DONE")
|