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.
89 lines
3.1 KiB
Python
89 lines
3.1 KiB
Python
"""Reference next-token prediction from the official V4.1 inference code.
|
|
|
|
`generate.py` always wraps text in a chat template, so its outputs are not directly comparable
|
|
with a raw-token harness. This script takes explicit token ids and prints the greedy next token
|
|
(plus the top-5), which makes the comparison exact regardless of prompt formatting.
|
|
|
|
torchrun --nproc-per-node 8 ref_next_token.py --ids 0,128803,2056,... --ckpt-path ...
|
|
|
|
Writes `<out>.json` with the ids, the chosen token and the top-5 so a harness can diff it.
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
|
|
from safetensors.torch import load_model # noqa: E402
|
|
|
|
from model import ModelArgs, Transformer # noqa: E402
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--ckpt-path", required=True)
|
|
parser.add_argument("--config", required=True)
|
|
parser.add_argument("--ids", required=True, help="comma-separated prompt token ids")
|
|
parser.add_argument("--out", default="/data/deepinfer-h20/ref/next_token.json")
|
|
parser.add_argument("--topk", type=int, default=5)
|
|
args = parser.parse_args()
|
|
|
|
ids = [int(x) for x in args.ids.replace(" ", "").split(",") if x]
|
|
|
|
rank = int(os.getenv("RANK", "0"))
|
|
local_rank = int(os.getenv("LOCAL_RANK", "0"))
|
|
world_size = int(os.getenv("WORLD_SIZE", "1"))
|
|
torch.cuda.set_device(local_rank)
|
|
if world_size > 1:
|
|
dist.init_process_group("nccl")
|
|
|
|
torch.set_default_dtype(torch.bfloat16)
|
|
torch.set_num_threads(8)
|
|
with open(args.config) as f:
|
|
model_args = ModelArgs(**json.load(f))
|
|
model_args.temperature = 0.0
|
|
model_args.max_batch_size = 1
|
|
model_args.max_seq_len = 4096
|
|
|
|
# The engram hash table is built from the tokenizer's compressed token map, so it is not
|
|
# optional even for a text-only forward.
|
|
from transformers import AutoTokenizer
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(args.ckpt_path)
|
|
with torch.device("cuda"):
|
|
model = Transformer(model_args, tokenizer)
|
|
load_model(model, os.path.join(args.ckpt_path, f"model{rank}-mp{world_size}.safetensors"))
|
|
torch.set_default_device("cuda")
|
|
|
|
tokens = torch.tensor([ids], dtype=torch.long)
|
|
with torch.inference_mode():
|
|
# `Transformer.forward` returns (output_ids, logits, main_hidden).
|
|
output_ids, logits, _ = model.forward(tokens, 0)
|
|
|
|
row = logits[0, -1].float()
|
|
top = torch.topk(row, args.topk)
|
|
assert int(top.indices[0]) == int(output_ids[0, -1]), "argmax disagrees with the sampled id"
|
|
|
|
payload = {
|
|
"prompt_ids": ids,
|
|
"next_token": int(top.indices[0]),
|
|
"top_tokens": [int(x) for x in top.indices],
|
|
"top_logits": [round(float(x), 4) for x in top.values],
|
|
}
|
|
if rank == 0:
|
|
print(json.dumps(payload, ensure_ascii=False), flush=True)
|
|
os.makedirs(os.path.dirname(args.out), exist_ok=True)
|
|
with open(args.out, "w") as f:
|
|
json.dump(payload, f, ensure_ascii=False, indent=2)
|
|
if world_size > 1:
|
|
dist.destroy_process_group()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|