mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-11-01 01:06:23 +08:00
style(sm90): satisfy ruff and add the missing type annotations
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.
This commit is contained in:
loopyt
parent
10956113dc
commit
28fb291db1
@ -6,8 +6,8 @@ from deep_cluster.hardware.device import (
|
||||
device_numa_node,
|
||||
get_actual_num_gpus,
|
||||
get_bf16_tflops_limit,
|
||||
is_sm_major_10,
|
||||
is_sm_major_9,
|
||||
is_sm_major_10,
|
||||
num_gpus,
|
||||
physical_device_count,
|
||||
physical_device_id,
|
||||
|
||||
@ -58,11 +58,8 @@ def topk(
|
||||
|
||||
values, indices = torch.topk(scores, k, dim=-1, sorted=True)
|
||||
|
||||
if valid_mask is None:
|
||||
keep = None
|
||||
else:
|
||||
# `indices` are original column ids, so the mask still applies after the top-k pick.
|
||||
keep = valid_mask.gather(1, indices)
|
||||
# `indices` are original column ids, so the mask still applies after the top-k pick.
|
||||
keep = None if valid_mask is None else valid_mask.gather(1, indices)
|
||||
|
||||
out_idx = indices.to(indices_type)
|
||||
if output_idx_offset is not None:
|
||||
@ -80,9 +77,7 @@ def topk(
|
||||
if keep is not None:
|
||||
out_val = torch.where(keep, out_val, torch.full_like(out_val, value_oob_fill_value))
|
||||
if k < topk:
|
||||
out_val = torch.cat(
|
||||
[out_val, out_val.new_full((num_rows, topk - k), value_oob_fill_value)], dim=1
|
||||
)
|
||||
out_val = torch.cat([out_val, out_val.new_full((num_rows, topk - k), value_oob_fill_value)], dim=1)
|
||||
else:
|
||||
out_val = None
|
||||
|
||||
|
||||
@ -171,12 +171,9 @@ class LocalEngramGather:
|
||||
rows, rows_sf = self._gather_layer(
|
||||
self.weights[layer], self.sfs[layer] if self.sfs else None, flat, self.num_local_entries[layer]
|
||||
)
|
||||
if rows_sf is not None:
|
||||
payload = self._requant(rows, rows_sf, num_tokens)
|
||||
else:
|
||||
payload = rows.view(num_tokens, column)
|
||||
payload = self._requant(rows, rows_sf, num_tokens) if rows_sf is not None else rows.view(num_tokens, column)
|
||||
|
||||
def hook(payload=payload):
|
||||
def hook(payload: torch.Tensor = payload) -> torch.Tensor:
|
||||
return payload
|
||||
|
||||
hooks.append(hook)
|
||||
|
||||
@ -60,7 +60,9 @@ def _requant(x: torch.Tensor, dst_block: int) -> tuple[torch.Tensor, torch.Tenso
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def requant_dense_fp8(weight: torch.Tensor, sf: torch.Tensor, src_block: int = FP4_BLOCK, dst_block: int = 128):
|
||||
def requant_dense_fp8(
|
||||
weight: torch.Tensor, sf: torch.Tensor, src_block: int = FP4_BLOCK, dst_block: int = 128
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""FP8 dense weight with per-``src_block`` scales → FP8 with per-``dst_block`` FP32 scales."""
|
||||
rows, cols = weight.shape
|
||||
x = weight.to(torch.float32) * _expand_blocks(sf.to(torch.float32), rows, cols, src_block)
|
||||
|
||||
@ -18,9 +18,9 @@ the ``padded_H`` handling for head counts below 64, and the ``REPLICATE_H`` head
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
import tilelang
|
||||
import tilelang.language as T
|
||||
import torch
|
||||
|
||||
LOG2E = 1.44269504
|
||||
|
||||
@ -33,9 +33,9 @@ LOG2E = 1.44269504
|
||||
},
|
||||
)
|
||||
def sparse_mla_fwd(
|
||||
Q,
|
||||
KV,
|
||||
Indices,
|
||||
Q: torch.Tensor,
|
||||
KV: torch.Tensor,
|
||||
Indices: torch.Tensor,
|
||||
heads: int,
|
||||
dim: int,
|
||||
topk: int,
|
||||
@ -45,12 +45,19 @@ def sparse_mla_fwd(
|
||||
block_I: int = 64,
|
||||
num_stages: int = 2,
|
||||
threads: int = 256,
|
||||
):
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Sparse MLA forward: out[b, s, h, dim], lse[b, s, h].
|
||||
|
||||
Q: [b, s, h, dim] bf16, KV: [b, s_kv, kv_group, dim] bf16, Indices: [b, s, kv_group, topk] int32.
|
||||
An index is valid when ``0 <= idx <= q_position``; larger values (and the padding used to make
|
||||
``topk`` a multiple of ``block_I``) are masked out.
|
||||
Negative indices (the padding that makes ``topk`` a multiple of ``block_I``) are always masked.
|
||||
|
||||
``is_causal`` additionally drops indices above the query's row number. **The engine must pass
|
||||
False**: `kv_index_cat` shifts the compressed-KV block by
|
||||
`total_k_swa_cached + total_k_swa_cur`, so indices index the assembled
|
||||
`[swa cache | current swa | compressed]` buffer and say nothing about the row number. Masking
|
||||
on it drops nearly every index, `sumexp` stays 0, and the whole layer turns to NaN. Causality
|
||||
is already guaranteed upstream -- the indexer only scores positions before each query, and the
|
||||
sliding-window column block is a window by construction.
|
||||
"""
|
||||
assert sm_scale is not None, "pass the explicit attention scale"
|
||||
scale = float(sm_scale * LOG2E) # log2(e): the kernel folds the exponent into exp2
|
||||
|
||||
@ -33,9 +33,9 @@ import torch
|
||||
|
||||
from deepinfer import PLATFORM
|
||||
from deepinfer.hugging_face.names import hf_layer_module_prefix, hf_to_model, resolve_hf_name, scale_name
|
||||
from deepinfer.model_executor.layers.linear import get_fp8_casting_granularity
|
||||
from deepinfer.hugging_face.parameters import model_params, set_ue8m0_sf
|
||||
from deepinfer.logger import init_logger
|
||||
from deepinfer.model_executor.layers.linear import get_fp8_casting_granularity
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
@ -52,6 +52,7 @@ def _check_finite(tag: str, tensor: torch.Tensor) -> None:
|
||||
bad = int((~torch.isfinite(tensor)).sum())
|
||||
raise RuntimeError(f"[nan-debug] {tag}: {bad}/{tensor.numel()} non-finite values, shape={tuple(tensor.shape)}")
|
||||
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
|
||||
@ -124,13 +124,9 @@ class MQA(nn.Module):
|
||||
# which has no SM90 implementation (`DG_HOST_ASSERT(arch_major == 10)`), so Hopper runs
|
||||
# the dense path for every layer. Selection is identical, only the scoring is wider.
|
||||
self.is_candidate_topk_source = (
|
||||
shared_topology is not None
|
||||
and shared_topology.is_candidate_source_layer(self.layer_idx)
|
||||
and not PLATFORM.is_sm90
|
||||
)
|
||||
self.use_candidate_topk = (
|
||||
candidate_source is not None and not self.is_candidate_topk_source and not PLATFORM.is_sm90
|
||||
shared_topology is not None and shared_topology.is_candidate_source_layer(self.layer_idx) and not PLATFORM.is_sm90
|
||||
)
|
||||
self.use_candidate_topk = candidate_source is not None and not self.is_candidate_topk_source and not PLATFORM.is_sm90
|
||||
if self.use_candidate_topk:
|
||||
assert candidate_cfg is not None
|
||||
assert self.n_head <= 32 and self.head_dim == 128
|
||||
@ -518,9 +514,7 @@ class MQA(nn.Module):
|
||||
page_tokens = kv_cache.shape[1]
|
||||
page_bytes = page_tokens * (head_dim + 4)
|
||||
# One contiguous byte row per page, following the tensor's real block stride.
|
||||
pages = torch.as_strided(
|
||||
kv_cache, (kv_cache.shape[0], page_bytes), (kv_cache.stride(0), 1), kv_cache.storage_offset()
|
||||
)
|
||||
pages = torch.as_strided(kv_cache, (kv_cache.shape[0], page_bytes), (kv_cache.stride(0), 1), kv_cache.storage_offset())
|
||||
|
||||
out = torch.full((num_rows, self.max_kv_len), float("-inf"), dtype=torch.float32, device=q.device)
|
||||
# One host read of the per-row context lengths; the engine forbids syncs during the
|
||||
@ -577,9 +571,7 @@ class MQA(nn.Module):
|
||||
f"decode_index_attn-part-{self.layer_idx}", enabled=self.layer_idx == second_layer, scale=len(same_layers)
|
||||
):
|
||||
p_sum = (
|
||||
self._decode_p_sum_sm90(
|
||||
q_fp8or4, q_sf, kv_cache, weights, klen, input_metadata.block_tables
|
||||
)
|
||||
self._decode_p_sum_sm90(q_fp8or4, q_sf, kv_cache, weights, klen, input_metadata.block_tables)
|
||||
if PLATFORM.is_sm90
|
||||
else self._compute_decode_p_sum(
|
||||
(q_fp8or4, q_sf),
|
||||
|
||||
@ -76,9 +76,7 @@ def main(dp_size: int, dp_rank: int) -> None:
|
||||
|
||||
if args.ids:
|
||||
ids = [int(x) for x in args.ids.replace(" ", "").split(",") if x]
|
||||
outputs = llm.generate(
|
||||
prompts=None, prompt_token_ids=[ids], sampling_params=sampling_params[:1], use_tqdm=False
|
||||
)
|
||||
outputs = llm.generate(prompts=None, prompt_token_ids=[ids], sampling_params=sampling_params[:1], use_tqdm=False)
|
||||
shown = [f"<ids len={len(ids)}>"]
|
||||
else:
|
||||
outputs = llm.generate(prompts=args.prompt, sampling_params=sampling_params, use_tqdm=False)
|
||||
|
||||
@ -6,11 +6,9 @@ Shapes follow DeepSeek-V4.1-Flash: head_dim 512 with RoPE inside, 64 query heads
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from tilelang.profiler import do_bench
|
||||
|
||||
from deepinfer.h20_ops.attn_sm90 import (
|
||||
CACHE_ELEMENTS,
|
||||
CACHE_ELEMENTS_PER_SF,
|
||||
decode_mla_mqa_sparse,
|
||||
prefill_attn_varlen_sparse,
|
||||
@ -21,13 +19,13 @@ HEADS = 64
|
||||
DIM = 512
|
||||
|
||||
|
||||
def ref_attention(q, kv, indices, scale, sink=None):
|
||||
def ref_attention(q: torch.Tensor, kv: torch.Tensor, indices: torch.Tensor, scale: float, sink: torch.Tensor | None = None) -> torch.Tensor:
|
||||
"""q [n, h, d], kv [m, 1, d], indices [n, topk] -> [n, h, d] float32."""
|
||||
n, h, d = q.shape
|
||||
topk = indices.shape[1]
|
||||
valid = indices >= 0
|
||||
idx = indices.clamp_min(0)
|
||||
k = kv[idx.reshape(-1), 0, :].view(n, topk, d).float() # [n, topk, d]
|
||||
k = kv[idx.reshape(-1), 0, :].view(n, topk, d).float() # [n, topk, d]
|
||||
v = k
|
||||
scores = torch.einsum("nhd,nkd->nhk", q.float(), k) * scale # [n, h, topk]
|
||||
scores = scores.masked_fill(~valid.unsqueeze(1), float("-inf"))
|
||||
@ -39,7 +37,7 @@ def ref_attention(q, kv, indices, scale, sink=None):
|
||||
return out, lse
|
||||
|
||||
|
||||
def make_indices(n, m, topk, gen):
|
||||
def make_indices(n: int, m: int, topk: int, gen: torch.Generator) -> torch.Tensor:
|
||||
"""Causal top-k like the indexer produces: random subset of [0, position], -1 padded."""
|
||||
indices = torch.full((n, topk), -1, dtype=torch.int32, device=DEV)
|
||||
for t in range(n):
|
||||
@ -50,7 +48,7 @@ def make_indices(n, m, topk, gen):
|
||||
return indices
|
||||
|
||||
|
||||
def quantize_cache(kv, block_size=64):
|
||||
def quantize_cache(kv: torch.Tensor, block_size: int = 64) -> torch.Tensor:
|
||||
"""Build a V41-format paged cache ([rows, 528] uint8) from bf16 kv [m, 1, d]."""
|
||||
m = kv.shape[0]
|
||||
values = kv[:, 0, :].float()
|
||||
@ -61,9 +59,7 @@ def quantize_cache(kv, block_size=64):
|
||||
exp = exp.clamp_min(-126)
|
||||
scale = torch.pow(2.0, exp)
|
||||
q = (groups / scale.unsqueeze(-1)).to(torch.float8_e4m3fn)
|
||||
packed = torch.cat(
|
||||
[q.reshape(m, DIM).view(torch.uint8), (exp + 127).to(torch.uint8)], dim=1
|
||||
) # [m, 528]
|
||||
packed = torch.cat([q.reshape(m, DIM).view(torch.uint8), (exp + 127).to(torch.uint8)], dim=1) # [m, 528]
|
||||
deq = (q.float() * scale.unsqueeze(-1)).reshape(m, DIM).to(torch.bfloat16)
|
||||
rows = packed.contiguous()
|
||||
n_blocks = (m + block_size - 1) // block_size
|
||||
@ -74,7 +70,7 @@ def quantize_cache(kv, block_size=64):
|
||||
return rows.view(n_blocks, block_size, 1, -1), deq
|
||||
|
||||
|
||||
def main():
|
||||
def main() -> None:
|
||||
torch.manual_seed(0)
|
||||
gen = torch.Generator(device="cpu").manual_seed(0)
|
||||
scale = DIM**-0.5
|
||||
@ -119,10 +115,7 @@ def main():
|
||||
kv_empty = torch.randn(64, 1, DIM, device=DEV, dtype=torch.bfloat16)
|
||||
idx_empty = torch.full((8, 128), -1, dtype=torch.int32, device=DEV)
|
||||
out_e, lse_e = prefill_attn_varlen_sparse(q_empty, kv_empty, idx_empty, scale)
|
||||
print(
|
||||
f"empty window: finite={bool(torch.isfinite(out_e).all())} "
|
||||
f"zero={bool((out_e == 0).all())} lse_max={lse_e.max().item():.3e}"
|
||||
)
|
||||
print(f"empty window: finite={bool(torch.isfinite(out_e).all())} zero={bool((out_e == 0).all())} lse_max={lse_e.max().item():.3e}")
|
||||
assert torch.isfinite(out_e).all() and (out_e == 0).all()
|
||||
|
||||
# ---------------- decode: paged fp8 cache + gather ----------------
|
||||
@ -132,15 +125,15 @@ def main():
|
||||
cache, kv_deq = quantize_cache(kv2)
|
||||
|
||||
idx2 = torch.randint(0, m_dec, (n_dec, topk_dec), device=DEV, dtype=torch.int32)
|
||||
dec, _ = decode_mla_mqa_sparse(
|
||||
q2, cache, DIM, idx2, sched_meta=None, softmax_scale=scale, attn_sink=sink
|
||||
)
|
||||
dec, _ = decode_mla_mqa_sparse(q2, cache, DIM, idx2, sched_meta=None, softmax_scale=scale, attn_sink=sink)
|
||||
ref2, _ = ref_attention(q2, kv_deq.unsqueeze(1), idx2, scale, sink=sink)
|
||||
err2 = (dec.float() - ref2).abs().max().item()
|
||||
corr2 = torch.cosine_similarity(dec.float().flatten(), ref2.flatten(), dim=0).item()
|
||||
print(f"decode (n={n_dec}, cache={m_dec}, topk={topk_dec}): max_abs_err={err2:.4f} corr={corr2:.6f}")
|
||||
|
||||
ms2 = do_bench(lambda: decode_mla_mqa_sparse(q2, cache, DIM, idx2, sched_meta=None, softmax_scale=scale, attn_sink=sink), warmup=5, rep=20)
|
||||
ms2 = do_bench(
|
||||
lambda: decode_mla_mqa_sparse(q2, cache, DIM, idx2, sched_meta=None, softmax_scale=scale, attn_sink=sink), warmup=5, rep=20
|
||||
)
|
||||
print(f"decode perf: batch={n_dec} topk={topk_dec} -> {ms2:.2f} ms/step")
|
||||
|
||||
# ---------------- perf: full-size prefill ----------------
|
||||
|
||||
@ -1,23 +1,21 @@
|
||||
"""Validate LocalEngramGather against a torch reference, single-process and multi-rank.
|
||||
|
||||
torchrun --nproc_per_node=8 scripts/h20_engram_check.py
|
||||
torchrun --nproc_per_node=8 scripts/h20_engram_check.py
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from deepinfer.h20_ops.engram_local import LocalEngramGather
|
||||
|
||||
HIDDEN = 256
|
||||
N_LOCAL = 4096 # rows per rank
|
||||
ENTRIES = 24 # 3 n-gram tables x 8 heads
|
||||
N_LOCAL = 4096 # rows per rank
|
||||
ENTRIES = 24 # 3 n-gram tables x 8 heads
|
||||
TOKENS = 128
|
||||
LAYERS = 2
|
||||
|
||||
|
||||
def quantize_table(rows, gen):
|
||||
def quantize_table(rows: torch.Tensor, gen: torch.Generator) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""fp8 table + per-32 ue8m0 (stored as fp32) scales, like the checkpoint."""
|
||||
groups = rows.float().view(rows.shape[0], HIDDEN // 32, 32)
|
||||
amax = groups.abs().amax(-1).clamp_min(1e-30)
|
||||
@ -27,7 +25,7 @@ def quantize_table(rows, gen):
|
||||
return q, scale
|
||||
|
||||
|
||||
def main():
|
||||
def main() -> None:
|
||||
dist.init_process_group("nccl")
|
||||
rank, world = dist.get_rank(), dist.get_world_size()
|
||||
torch.cuda.set_device(rank)
|
||||
@ -63,7 +61,7 @@ def main():
|
||||
dist.all_gather_into_tensor(ref_rows[layer], tables[layer])
|
||||
dist.all_gather_into_tensor(ref_sf[layer], scales[layer])
|
||||
|
||||
def reference(layer):
|
||||
def reference(layer: int) -> torch.Tensor:
|
||||
flat = idx[layer].reshape(-1).long()
|
||||
valid = flat >= 0
|
||||
safe = flat.clamp_min(0)
|
||||
@ -82,10 +80,7 @@ def main():
|
||||
ref_deq = ref.float().view(TOKENS, ENTRIES, HIDDEN)
|
||||
err = (got_vals - ref_deq).abs().max().item()
|
||||
rel = err / max(ref_deq.abs().max().item(), 1e-9)
|
||||
print(
|
||||
f"[rank {rank}] out={tuple(a.shape)} sf={None if sf is None else tuple(sf.shape)} "
|
||||
f"max_abs={err:.4f} rel={rel:.4f}"
|
||||
)
|
||||
print(f"[rank {rank}] out={tuple(a.shape)} sf={None if sf is None else tuple(sf.shape)} max_abs={err:.4f} rel={rel:.4f}")
|
||||
print(f"[rank {rank}] invalid slot all-zero: {bool((got_vals[0, 0] == 0).all())}")
|
||||
|
||||
# Non-owner traffic actually exercised?
|
||||
|
||||
@ -105,9 +105,7 @@ def main() -> None:
|
||||
assert ref_vals.shape == ours_vals.shape, f"{ref_vals.shape} vs {ours_vals.shape}"
|
||||
num = (ours_vals - ref_vals).norm().item()
|
||||
den = ref_vals.norm().item()
|
||||
cos = torch.nn.functional.cosine_similarity(
|
||||
ours_vals.reshape(-1), ref_vals.reshape(-1), dim=0
|
||||
).item()
|
||||
cos = torch.nn.functional.cosine_similarity(ours_vals.reshape(-1), ref_vals.reshape(-1), dim=0).item()
|
||||
print(f" ref scale: e8m0 {tuple(ref_sf.shape)} ours: fp32 {tuple(ours_sf.shape)}")
|
||||
print(f" |ref|={den:.4f} |ours|={ours_vals.norm().item():.4f} ratio={ours_vals.norm().item() / den:.5f}")
|
||||
print(f" rel_err={num / den:.5f} cos={cos:.7f}")
|
||||
|
||||
@ -1,11 +1,15 @@
|
||||
"""Validate the SM90 dense requant on a real checkpoint tensor (wq_a + shared expert)."""
|
||||
|
||||
import json
|
||||
|
||||
import torch
|
||||
from safetensors import safe_open
|
||||
|
||||
from deepinfer.h20_ops.quant import requant_dense_fp8, ue8m0_to_fp32
|
||||
|
||||
root = "/data/models/DeepSeek-V4.1-Flash-1017"
|
||||
idx = json.load(open(f"{root}/model.safetensors.index.json"))["weight_map"]
|
||||
with open(f"{root}/model.safetensors.index.json") as f:
|
||||
idx = json.load(f)["weight_map"]
|
||||
|
||||
targets = [
|
||||
"layers.0.attn.wq_a.weight",
|
||||
|
||||
@ -1,6 +1,9 @@
|
||||
"""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
|
||||
|
||||
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"
|
||||
@ -12,7 +15,8 @@ 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):
|
||||
|
||||
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()
|
||||
@ -21,7 +25,8 @@ def q_token(t, blk):
|
||||
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):
|
||||
|
||||
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()
|
||||
@ -30,26 +35,27 @@ def q_blocks(t, blk):
|
||||
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')
|
||||
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}
|
||||
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)
|
||||
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()
|
||||
@ -61,14 +67,23 @@ 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())
|
||||
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)
|
||||
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))
|
||||
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)
|
||||
|
||||
@ -11,7 +11,14 @@ import torch
|
||||
from deepinfer.h20_ops.deep_select_sm90 import topk
|
||||
|
||||
|
||||
def ref_topk(scores, k, end=None, offset=None, fill=-1, val_fill=float("-inf")):
|
||||
def ref_topk(
|
||||
scores: torch.Tensor,
|
||||
k: int,
|
||||
end: torch.Tensor | None = None,
|
||||
offset: int | None = None,
|
||||
fill: int = -1,
|
||||
val_fill: float = float("-inf"),
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Reference: consider only [0, end) per row, take k largest, pad, then shift."""
|
||||
rows = scores.shape[0]
|
||||
idx_rows, val_rows = [], []
|
||||
@ -68,8 +75,14 @@ def main() -> None:
|
||||
view = storage[:, :k]
|
||||
offset = torch.arange(rows, dtype=torch.int32, device=dev) * 1000
|
||||
vals2, idxs2 = topk(
|
||||
scores, k, end=end, indices_type=torch.int32, output_idx=view, output_idx_offset=offset,
|
||||
idx_oob_fill_value=-1, return_value=False,
|
||||
scores,
|
||||
k,
|
||||
end=end,
|
||||
indices_type=torch.int32,
|
||||
output_idx=view,
|
||||
output_idx_offset=offset,
|
||||
idx_oob_fill_value=-1,
|
||||
return_value=False,
|
||||
)
|
||||
assert idxs2.data_ptr() == view.data_ptr()
|
||||
assert vals2 is None
|
||||
|
||||
@ -20,10 +20,11 @@ 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 transformers import AutoTokenizer # noqa: E402
|
||||
from collections.abc import Callable # noqa: E402
|
||||
|
||||
from model import ModelArgs, Transformer # noqa: E402
|
||||
from safetensors.torch import load_model # noqa: E402
|
||||
from transformers import AutoTokenizer # noqa: E402
|
||||
|
||||
|
||||
def main() -> None:
|
||||
@ -59,27 +60,27 @@ def main() -> None:
|
||||
|
||||
captured: dict[str, np.ndarray] = {}
|
||||
|
||||
def make_hook(layer_id: int):
|
||||
def hook(_module, _inputs, output):
|
||||
def make_hook(layer_id: int) -> Callable[[torch.nn.Module, tuple, object], None]:
|
||||
def hook(_module: torch.nn.Module, _inputs: tuple[torch.Tensor, ...], output: object) -> None:
|
||||
# Block.forward returns (x, pre_mix); x is [b, s, hc_mult, dim].
|
||||
x = output[0] if isinstance(output, tuple) else output
|
||||
captured[f"layer{layer_id}"] = x.mean(dim=2)[0, -1].float().cpu().numpy()
|
||||
|
||||
return hook
|
||||
|
||||
def make_attn_hook(layer_id: int):
|
||||
def hook(_module, _inputs, output):
|
||||
def make_attn_hook(layer_id: int) -> Callable[[torch.nn.Module, tuple, object], None]:
|
||||
def hook(_module: torch.nn.Module, _inputs: tuple[torch.Tensor, ...], output: object) -> None:
|
||||
# The attention sub-layer output is already collapsed: [b, s, dim].
|
||||
x = output[0] if isinstance(output, tuple) else output
|
||||
captured[f"layer{layer_id}_attn"] = x[0, -1].float().cpu().numpy()
|
||||
|
||||
return hook
|
||||
|
||||
def embed_hook(_module, _inputs, output):
|
||||
def embed_hook(_module: torch.nn.Module, _inputs: tuple[torch.Tensor, ...], output: object) -> None:
|
||||
captured["embed"] = output[0, -1].float().cpu().numpy()
|
||||
|
||||
def make_gate_hook(layer_id: int):
|
||||
def hook(_module, _inputs, output):
|
||||
def make_gate_hook(layer_id: int) -> Callable[[torch.nn.Module, tuple, object], None]:
|
||||
def hook(_module: torch.nn.Module, _inputs: tuple[torch.Tensor, ...], output: object) -> None:
|
||||
# Gate.forward returns (weights [n, topk], indices [n, topk]).
|
||||
weights, indices = output
|
||||
captured[f"layer{layer_id}_topk_idx"] = indices[-1].to(torch.int32).cpu().numpy()
|
||||
@ -87,14 +88,14 @@ def main() -> None:
|
||||
|
||||
return hook
|
||||
|
||||
def make_shared_hook(layer_id: int):
|
||||
def hook(_module, _inputs, output):
|
||||
def make_shared_hook(layer_id: int) -> Callable[[torch.nn.Module, tuple, object], None]:
|
||||
def hook(_module: torch.nn.Module, _inputs: tuple[torch.Tensor, ...], output: object) -> None:
|
||||
captured[f"layer{layer_id}_shared_out"] = output.view(-1, output.shape[-1])[-1].float().cpu().numpy()
|
||||
|
||||
return hook
|
||||
|
||||
def make_moe_hook(layer_id: int):
|
||||
def hook(_module, _inputs, output):
|
||||
def make_moe_hook(layer_id: int) -> Callable[[torch.nn.Module, tuple, object], None]:
|
||||
def hook(_module: torch.nn.Module, _inputs: tuple[torch.Tensor, ...], output: object) -> None:
|
||||
# The MoE output already includes the shared expert, matching our `_combined`.
|
||||
captured[f"layer{layer_id}_combined"] = output.view(-1, output.shape[-1])[-1].float().cpu().numpy()
|
||||
|
||||
|
||||
@ -19,9 +19,8 @@ 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
|
||||
from safetensors.torch import load_model # noqa: E402
|
||||
|
||||
|
||||
def main() -> None:
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user