mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-11-01 01:06:23 +08:00
1e04aea7ea
DeepInfer targets Blackwell: the vendored FlashMLA fork only ships SM100 kernels (`CUDA_ARCHS = ["sm_100a", "sm_103a"]`, every entry point asserts `arch.is_sm100f()`), and several DeepGEMM attention APIs plus the whole quantization recipe assume FP4 tensor cores, packed UE8M0 scales and per-32 blocks. None of that exists on Hopper, so this branch adds an SM90 backend and gets an 8x H20 node to produce tokens end to end. Scope and status: the pipeline runs (prefill, decode, sampling) and matches the reference on the first token, but numerics still diverge from position ~12 onward, so the output is not yet correct. H20-DEBUG-PLAN.md tracks the remaining work with the exact reproduction case; see H20-SM90-FINDINGS.md for the analysis. Quantization - `Platform.sf_pack_factor` separates two meanings that `fp8_sf_type` conflated: the GEMM scale type (FP32 on SM90) and the number of scales packed per element (1, instead of 4 UE8M0 exponents per int32). KV caches keep the ue8m0 byte layout. - Requantize FP4 experts and ue8m0 dense weights to FP8 E4M3 with FP32 per-128 scales while loading (`hugging_face/weights.py`, `h20_ops/quant.py`), so the runtime sees a plain FP8 checkpoint: `mxfp4` is disabled at runtime, expert weights are repacked as 2D, and `use_fp4_sparse_kv` is off since the FP4 cache kernels do not exist on Hopper. - GEMM recipes switch to `(1, 128, 128)` with `disable_ue8m0_cast=True`, which is the only shape SM90 DeepGEMM accepts. Attention - `h20_ops/sparse_mla_sm90.py` adapts TileLang's `examples/deepseek_v32` sparse MLA kernel (which is Hopper-targeted) to the V4.1 layout: RoPE lives inside the 512-wide latent, so the V3.2 `dim`/`tail` split is removed and QK/PV are both 512. Measured on H20: 104 TFLOP/s, bf16-noise accuracy. It takes no causal mask, because the indices are already causal -- deriving causality from the row number drops every index once the KV buffer is offset, which produced an all-NaN layer 0. - `h20_ops/attn_sm90.py` re-implements `prefill_attn_varlen_sparse` and `decode_mla_mqa_sparse` with the same signatures as `attention_backend.py`, and registers them through `patch_imports`, so call sites are untouched and the FlashMLA-only module is never imported. - Softmax normalisation clamps `sumexp`, so a padded row with no valid index yields zeros instead of 0/0 and cannot poison the whole hidden state. Engram - `h20_ops/engram_local.py` replaces `deep_ep.EngramBuffer.fetch`, which needs NCCL GIN (RDMA) and is unavailable on a single node, with a distributed gather over NVLink built from two `all_to_all`s. Row ownership follows the C++ layout (`owner = idx // num_local`). Verified against a full reference table on 8 ranks. Other Blackwell assumptions removed - DeepSelect's top-k kernel faults at launch on SM90, so `h20_ops/deep_select_sm90.py` reproduces its contract with `torch.topk` (end truncation, OOB fill, in-place output views). The shared-layer candidate-block optimisation is disabled because `fp8_fp4_sparse_mqa_logits` has no SM90 implementation, and the dense MQA-logits path is used instead; its fp32 output is narrowed for the block-max kernel. - Decode-time index scoring uses `_decode_p_sum_sm90`: Hopper's paged kernel hard-requires 64-token pages while the cache uses 128, and the pages are strided views that cannot be split without a copy. - `mega_mhc` / `mega_moe`, the FlashAttention-4 CuTe DSL kernels and the fused norm+rope+cast kernel are all SM100-only and take a Hopper fallback. - `allow_cuda_sync()` is used for the few places that genuinely need a host read (all-to-all sizing, the ViT fallback, the NaN diagnostic) instead of disabling the sync guard globally. Vendored submodules (not part of this commit; submodule gitlinks are unchanged) - `vendor/DeepSelect/setup.py`: re-enable the `sm_90a` gencode that is commented out, otherwise top-k fails with "no kernel image". Scripts: `h20_attn_check.py`, `h20_engram_check.py`, `h20_topk_check.py`, `h20_smoke_check.py`, `h20_real_requant_check.py` are unit checks that run on 8x H20; `h20_ask.py` is a one-prompt harness and `ref_next_token.py` produces the official reference next-token for arbitrary ids, so both sides can be compared on identical input.
77 lines
3.6 KiB
Python
77 lines
3.6 KiB
Python
"""SM90 sanity checks: DeepGEMM FP8 1D2D (fp32 scales) + the H20 requant helpers."""
|
|
import torch, deep_gemm
|
|
from deepinfer.h20_ops.quant import ue8m0_to_fp32, requant_dense_fp8, requant_expert_fp4, _unpack_fp4
|
|
|
|
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, blk):
|
|
"""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, blk):
|
|
"""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., .5, 1., 1.5, 2., 3., 4., 6., 0., -.5, -1., -1.5, -2., -3., -4., -6.], 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")
|