wuyang
662 lines
24 KiB
Python
662 lines
24 KiB
Python
#!/usr/bin/env python3
|
|
"""Execute DeepSeek MLA's naive and weight-absorbed cache paths with real weights.
|
|
|
|
The probe deliberately separates three things:
|
|
|
|
1. DeepSeek-V2-Lite's official Hugging Face eager attention;
|
|
2. DeepSeek-V3's official pure-PyTorch ``naive`` and ``absorb`` reference paths;
|
|
3. FlashMLA's optimized-kernel architecture boundary.
|
|
|
|
The first official checkpoint shard contains the embedding, layer 0, and all
|
|
layer-1 attention weights. We execute layer 0, normalize the real layer-1 input,
|
|
map the unchanged V2-Lite attention tensors into the pinned V3 reference class,
|
|
then perform a 25-token prefill plus one-token incremental decode.
|
|
|
|
This is a correctness and cache-layout probe, not a serving benchmark.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import copy
|
|
import gc
|
|
import hashlib
|
|
import importlib.util
|
|
import json
|
|
import math
|
|
import platform
|
|
import subprocess
|
|
import sys
|
|
import types
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from safetensors import safe_open
|
|
from transformers import AutoTokenizer, DynamicCache
|
|
|
|
|
|
HF_REVISION = "604d5664dddd88a0433dbae533b7fe9472482de0"
|
|
V3_REVISION = "9b4e9788e4a3a731f7567338ed15d3ec549ce03b"
|
|
FLASHMLA_REVISION = "15f13e5030374295491c5ce31b02d7e63a7772c6"
|
|
CUTLASS_REVISION = "147f5673d0c1c3dcf66f78d677fd647e4a020219"
|
|
PROMPT = "用通俗的语言解释，为什么稀疏专家模型可以拥有很多参数，但每个 token 只使用其中一小部分。"
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--artifact-dir", type=Path, required=True)
|
|
parser.add_argument("--v3-repo", type=Path, required=True)
|
|
parser.add_argument("--flashmla-repo", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--device", default="cuda")
|
|
parser.add_argument("--cache-slots", type=int, default=32)
|
|
parser.add_argument("--captured-at", default=None)
|
|
return parser.parse_args()
|
|
|
|
|
|
def sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as handle:
|
|
for block in iter(lambda: handle.read(1024 * 1024), b""):
|
|
digest.update(block)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def git_revision(root: Path) -> str:
|
|
return subprocess.check_output(
|
|
["git", "-C", str(root), "rev-parse", "HEAD"], text=True
|
|
).strip()
|
|
|
|
|
|
def tensor_bytes(tensor: torch.Tensor) -> int:
|
|
return tensor.numel() * tensor.element_size()
|
|
|
|
|
|
def compare(left: torch.Tensor, right: torch.Tensor) -> dict[str, Any]:
|
|
if left.shape != right.shape:
|
|
raise ValueError(f"shape mismatch: {left.shape} != {right.shape}")
|
|
delta = (left.float() - right.float()).abs().flatten()
|
|
return {
|
|
"shape": list(left.shape),
|
|
"finite": bool(torch.isfinite(left).all() and torch.isfinite(right).all()),
|
|
"exact_fraction": (left == right).float().mean().item(),
|
|
"max_abs": delta.max().item(),
|
|
"mean_abs": delta.mean().item(),
|
|
"p50_abs": torch.quantile(delta, 0.50).item(),
|
|
"p90_abs": torch.quantile(delta, 0.90).item(),
|
|
"p99_abs": torch.quantile(delta, 0.99).item(),
|
|
}
|
|
|
|
|
|
def load_v2_modules(root: Path) -> tuple[Any, Any]:
|
|
package_name = "deepseek_v2_lite_absorb_official"
|
|
package = types.ModuleType(package_name)
|
|
package.__path__ = [str(root)]
|
|
sys.modules[package_name] = package
|
|
loaded = {}
|
|
for leaf in ("configuration_deepseek", "modeling_deepseek"):
|
|
name = f"{package_name}.{leaf}"
|
|
spec = importlib.util.spec_from_file_location(name, root / f"{leaf}.py")
|
|
if spec is None or spec.loader is None:
|
|
raise RuntimeError(f"cannot import official V2-Lite module: {leaf}")
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[name] = module
|
|
spec.loader.exec_module(module)
|
|
loaded[leaf] = module
|
|
return loaded["configuration_deepseek"], loaded["modeling_deepseek"]
|
|
|
|
|
|
def load_v3_module(root: Path) -> Any:
|
|
inference = root / "inference"
|
|
sys.path.insert(0, str(inference))
|
|
spec = importlib.util.spec_from_file_location(
|
|
"deepseek_v3_absorb_official", inference / "model.py"
|
|
)
|
|
if spec is None or spec.loader is None:
|
|
raise RuntimeError("cannot import official DeepSeek-V3 inference/model.py")
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[spec.name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def read_shard_state(
|
|
shard: Path, prefixes: tuple[str, ...], exact: tuple[str, ...] = ()
|
|
) -> dict[str, torch.Tensor]:
|
|
with safe_open(shard, framework="pt", device="cpu") as handle:
|
|
return {
|
|
key: handle.get_tensor(key)
|
|
for key in handle.keys()
|
|
if key in exact or any(key.startswith(prefix) for prefix in prefixes)
|
|
}
|
|
|
|
|
|
def load_real_layer1_input(
|
|
shard: Path,
|
|
modeling: Any,
|
|
config: Any,
|
|
tokenizer: Any,
|
|
device: torch.device,
|
|
) -> tuple[torch.Tensor, dict[str, torch.Tensor], list[int]]:
|
|
encoded = tokenizer(PROMPT, add_special_tokens=True, return_tensors="pt")
|
|
input_ids = encoded.input_ids
|
|
token_ids = input_ids[0].tolist()
|
|
if len(token_ids) != 26:
|
|
raise RuntimeError(f"prompt contract changed: expected 26 tokens, got {len(token_ids)}")
|
|
|
|
state = read_shard_state(
|
|
shard,
|
|
("model.layers.0.", "model.layers.1.self_attn."),
|
|
("model.embed_tokens.weight", "model.layers.1.input_layernorm.weight"),
|
|
)
|
|
embedding = state.pop("model.embed_tokens.weight")
|
|
hidden = F.embedding(input_ids, embedding).to(device)
|
|
|
|
layer0_prefix = "model.layers.0."
|
|
layer0_state = {
|
|
key[len(layer0_prefix) :]: value
|
|
for key, value in state.items()
|
|
if key.startswith(layer0_prefix)
|
|
}
|
|
with torch.device("meta"):
|
|
layer0 = modeling.DeepseekV2DecoderLayer(config, 0)
|
|
layer0.to_empty(device="cpu")
|
|
layer0.load_state_dict(layer0_state, strict=True, assign=True)
|
|
layer0 = layer0.to(device=device, dtype=torch.bfloat16).eval()
|
|
|
|
sequence = input_ids.shape[1]
|
|
position_ids = torch.arange(sequence, device=device).unsqueeze(0)
|
|
attention_mask = torch.ones_like(input_ids, device=device)
|
|
causal_mask = modeling._prepare_4d_causal_attention_mask(
|
|
attention_mask,
|
|
input_ids.shape,
|
|
hidden,
|
|
0,
|
|
)
|
|
with torch.inference_mode():
|
|
hidden = layer0(
|
|
hidden,
|
|
attention_mask=causal_mask,
|
|
position_ids=position_ids,
|
|
past_key_value=DynamicCache(),
|
|
use_cache=True,
|
|
)[0]
|
|
|
|
layer1_norm = modeling.DeepseekV2RMSNorm(
|
|
config.hidden_size, eps=config.rms_norm_eps
|
|
).to(device=device, dtype=torch.bfloat16)
|
|
with torch.no_grad():
|
|
layer1_norm.weight.copy_(
|
|
state["model.layers.1.input_layernorm.weight"].to(
|
|
device=device, dtype=torch.bfloat16
|
|
)
|
|
)
|
|
with torch.inference_mode():
|
|
normalized = layer1_norm(hidden)
|
|
|
|
attention_prefix = "model.layers.1.self_attn."
|
|
attention_state = {
|
|
key[len(attention_prefix) :]: value
|
|
for key, value in state.items()
|
|
if key.startswith(attention_prefix)
|
|
}
|
|
del layer0, layer0_state, embedding, hidden, layer1_norm, state
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
return normalized, attention_state, token_ids
|
|
|
|
|
|
def v3_args(v3: Any, config: Any, cache_slots: int) -> Any:
|
|
rope = config.rope_scaling
|
|
args = v3.ModelArgs(
|
|
max_batch_size=1,
|
|
max_seq_len=cache_slots,
|
|
dtype="bf16",
|
|
vocab_size=config.vocab_size,
|
|
dim=config.hidden_size,
|
|
n_layers=config.num_hidden_layers,
|
|
n_heads=config.num_attention_heads,
|
|
q_lora_rank=0,
|
|
kv_lora_rank=config.kv_lora_rank,
|
|
qk_nope_head_dim=config.qk_nope_head_dim,
|
|
qk_rope_head_dim=config.qk_rope_head_dim,
|
|
v_head_dim=config.v_head_dim,
|
|
original_seq_len=rope["original_max_position_embeddings"],
|
|
rope_theta=config.rope_theta,
|
|
rope_factor=rope["factor"],
|
|
beta_fast=rope["beta_fast"],
|
|
beta_slow=rope["beta_slow"],
|
|
mscale=rope["mscale"],
|
|
)
|
|
# The official class uses max_seq_len both to configure YaRN and to allocate
|
|
# cache buffers. Keep the workstation cache small, but retain the model's
|
|
# official 163,840-position YaRN contract for frequencies and scaling.
|
|
args.rope_reference_max_seq_len = config.max_position_embeddings
|
|
return args
|
|
|
|
|
|
def copy_v2_attention_weights(
|
|
module: Any,
|
|
state: dict[str, torch.Tensor],
|
|
device: torch.device,
|
|
dtype: torch.dtype,
|
|
) -> None:
|
|
mapping = {
|
|
"wq.weight": "q_proj.weight",
|
|
"wkv_a.weight": "kv_a_proj_with_mqa.weight",
|
|
"kv_norm.weight": "kv_a_layernorm.weight",
|
|
"wkv_b.weight": "kv_b_proj.weight",
|
|
"wo.weight": "o_proj.weight",
|
|
}
|
|
target = dict(module.named_parameters())
|
|
with torch.no_grad():
|
|
for destination, source in mapping.items():
|
|
target[destination].copy_(state[source].to(device=device, dtype=dtype))
|
|
|
|
|
|
def execute_v3_path(
|
|
v3: Any,
|
|
args: Any,
|
|
state: dict[str, torch.Tensor],
|
|
normalized: torch.Tensor,
|
|
device: torch.device,
|
|
dtype: torch.dtype,
|
|
implementation: str,
|
|
) -> tuple[torch.Tensor, dict[str, Any]]:
|
|
v3.attn_impl = implementation
|
|
v3.Linear.dtype = dtype
|
|
module = v3.MLA(args).to(device=device, dtype=dtype).eval()
|
|
copy_v2_attention_weights(module, state, device, dtype)
|
|
frequency_args = copy.copy(args)
|
|
frequency_args.max_seq_len = args.rope_reference_max_seq_len
|
|
frequencies = v3.precompute_freqs_cis(frequency_args).to(device)
|
|
if frequency_args.max_seq_len > args.original_seq_len:
|
|
mscale = 0.1 * args.mscale * math.log(args.rope_factor) + 1.0
|
|
module.softmax_scale = module.qk_head_dim ** -0.5 * mscale * mscale
|
|
prefill = normalized[:, :-1].to(dtype)
|
|
decode = normalized[:, -1:].to(dtype)
|
|
prefill_length = prefill.shape[1]
|
|
mask = torch.full(
|
|
(prefill_length, prefill_length),
|
|
float("-inf"),
|
|
device=device,
|
|
dtype=torch.float32,
|
|
).triu_(1)
|
|
with torch.inference_mode():
|
|
module(prefill, 0, frequencies[:prefill_length], mask)
|
|
output = module(
|
|
decode,
|
|
prefill_length,
|
|
frequencies[prefill_length : prefill_length + 1],
|
|
None,
|
|
)
|
|
|
|
if implementation == "naive":
|
|
active = {
|
|
"key": {
|
|
"shape": [1, prefill_length + 1, args.n_heads, args.qk_nope_head_dim + args.qk_rope_head_dim],
|
|
"bytes": (prefill_length + 1)
|
|
* args.n_heads
|
|
* (args.qk_nope_head_dim + args.qk_rope_head_dim)
|
|
* torch.empty((), dtype=dtype).element_size(),
|
|
},
|
|
"value": {
|
|
"shape": [1, prefill_length + 1, args.n_heads, args.v_head_dim],
|
|
"bytes": (prefill_length + 1)
|
|
* args.n_heads
|
|
* args.v_head_dim
|
|
* torch.empty((), dtype=dtype).element_size(),
|
|
},
|
|
}
|
|
allocated = {
|
|
"key_shape": list(module.k_cache.shape),
|
|
"key_bytes": tensor_bytes(module.k_cache),
|
|
"value_shape": list(module.v_cache.shape),
|
|
"value_bytes": tensor_bytes(module.v_cache),
|
|
}
|
|
else:
|
|
active = {
|
|
"latent": {
|
|
"shape": [1, prefill_length + 1, args.kv_lora_rank],
|
|
"bytes": (prefill_length + 1)
|
|
* args.kv_lora_rank
|
|
* torch.empty((), dtype=dtype).element_size(),
|
|
},
|
|
"rope": {
|
|
"shape": [1, prefill_length + 1, args.qk_rope_head_dim],
|
|
"bytes": (prefill_length + 1)
|
|
* args.qk_rope_head_dim
|
|
* torch.empty((), dtype=dtype).element_size(),
|
|
},
|
|
}
|
|
allocated = {
|
|
"latent_shape": list(module.kv_cache.shape),
|
|
"latent_bytes": tensor_bytes(module.kv_cache),
|
|
"rope_shape": list(module.pe_cache.shape),
|
|
"rope_bytes": tensor_bytes(module.pe_cache),
|
|
}
|
|
|
|
result = {
|
|
"implementation": implementation,
|
|
"dtype": str(dtype).replace("torch.", ""),
|
|
"active_cache": active,
|
|
"active_cache_bytes": sum(item["bytes"] for item in active.values()),
|
|
"allocated_cache": allocated,
|
|
"allocated_cache_bytes": sum(
|
|
value for key, value in allocated.items() if key.endswith("_bytes")
|
|
),
|
|
"decode_finite": bool(torch.isfinite(output).all()),
|
|
}
|
|
output = output.detach().cpu()
|
|
del module, frequencies, prefill, decode, mask
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
return output, result
|
|
|
|
|
|
def execute_hf_eager(
|
|
modeling: Any,
|
|
config: Any,
|
|
state: dict[str, torch.Tensor],
|
|
normalized: torch.Tensor,
|
|
device: torch.device,
|
|
) -> tuple[torch.Tensor, dict[str, Any]]:
|
|
# A standalone DynamicCache is zero-indexed; the weights still come from
|
|
# decoder layer 1, while this isolated attention module occupies cache slot 0.
|
|
attention = modeling.DeepseekV2Attention(config, layer_idx=0)
|
|
attention.load_state_dict(state, strict=True, assign=True)
|
|
attention = attention.to(device=device, dtype=torch.bfloat16).eval()
|
|
prefill = normalized[:, :-1]
|
|
decode = normalized[:, -1:]
|
|
prefill_length = prefill.shape[1]
|
|
prefill_mask = modeling._prepare_4d_causal_attention_mask(
|
|
torch.ones((1, prefill_length), dtype=torch.long, device=device),
|
|
(1, prefill_length),
|
|
prefill,
|
|
0,
|
|
)
|
|
cache = DynamicCache()
|
|
with torch.inference_mode():
|
|
attention(
|
|
prefill,
|
|
attention_mask=prefill_mask,
|
|
position_ids=torch.arange(prefill_length, device=device).unsqueeze(0),
|
|
past_key_value=cache,
|
|
use_cache=True,
|
|
)
|
|
output = attention(
|
|
decode,
|
|
attention_mask=torch.zeros(
|
|
(1, 1, 1, prefill_length + 1),
|
|
device=device,
|
|
dtype=decode.dtype,
|
|
),
|
|
position_ids=torch.tensor([[prefill_length]], device=device),
|
|
past_key_value=cache,
|
|
use_cache=True,
|
|
)[0]
|
|
key = cache.key_cache[0]
|
|
value = cache.value_cache[0]
|
|
result = {
|
|
"implementation": "huggingface_eager",
|
|
"dtype": "bfloat16",
|
|
"key_shape": list(key.shape),
|
|
"value_shape": list(value.shape),
|
|
"active_cache_bytes": tensor_bytes(key) + tensor_bytes(value),
|
|
"decode_finite": bool(torch.isfinite(output).all()),
|
|
}
|
|
output = output.detach().cpu()
|
|
del attention, cache, key, value, prefill, decode, prefill_mask
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
return output, result
|
|
|
|
|
|
def nvidia_smi() -> dict[str, str]:
|
|
fields = ["name", "driver_version", "memory.total", "compute_cap"]
|
|
output = subprocess.check_output(
|
|
[
|
|
"nvidia-smi",
|
|
f"--query-gpu={','.join(fields)}",
|
|
"--format=csv,noheader,nounits",
|
|
],
|
|
text=True,
|
|
).strip()
|
|
return dict(zip(fields, [item.strip() for item in output.split(",")], strict=True))
|
|
|
|
|
|
def flashmla_boundary(root: Path, device: torch.device) -> dict[str, Any]:
|
|
readme = (root / "README.md").read_text()
|
|
setup = (root / "setup.py").read_text()
|
|
dense_header = (root / "csrc/api/dense_decode.h").read_text()
|
|
required = [
|
|
"| Dense Decoding | SM90 | MQA | BF16 |",
|
|
"| Sparse Decoding | SM90 & SM100 | MQA | FP8",
|
|
'"arch=compute_100f,code=sm_100f"',
|
|
'"arch=compute_90a,code=sm_90a"',
|
|
"Dense decode MLA is only supported on SM90a architecture",
|
|
]
|
|
combined = "\n".join((readme, setup, dense_header))
|
|
missing = [needle for needle in required if needle not in combined]
|
|
if missing:
|
|
raise RuntimeError(f"FlashMLA source contract changed: {missing}")
|
|
capability = torch.cuda.get_device_capability(device)
|
|
return {
|
|
"official_revision": FLASHMLA_REVISION,
|
|
"cutlass_revision": CUTLASS_REVISION,
|
|
"official_support_matrix": [
|
|
{"kernel": "dense_decode", "architectures": ["SM90"], "mode": "MQA", "cache": "BF16"},
|
|
{"kernel": "sparse_decode", "architectures": ["SM90", "SM100"], "mode": "MQA", "cache": "FP8"},
|
|
{"kernel": "dense_prefill", "architectures": ["SM100"], "mode": "MHA", "cache": None},
|
|
{"kernel": "sparse_prefill", "architectures": ["SM90", "SM100"], "mode": "MQA", "cache": None},
|
|
],
|
|
"compiled_gencode": ["sm_100f", "sm_90a"],
|
|
"dense_decode_runtime_guard": "Dense decode MLA is only supported on SM90a architecture",
|
|
"local_architecture": f"SM{capability[0]}{capability[1]}",
|
|
"optimized_kernel_executed": False,
|
|
"why_not": "The pinned official source does not list or generate SM120 kernels; dense decode additionally rejects non-SM90a devices.",
|
|
"build_attempts": [
|
|
{
|
|
"environment": "host CUDA 12.8 / g++ 15.2",
|
|
"result": "failed before compilation",
|
|
"evidence": "CUDA 12.8 rejects host compiler versions newer than g++ 13; g++-13 is not installed.",
|
|
},
|
|
{
|
|
"environment": "isolated CUDA 13.0.2 / torch 2.11.0+cu130 / g++ 13.3",
|
|
"result": "failed during official extension compilation",
|
|
"evidence": "csrc/api.cpp could not find cuda/std/utility; no wheel or runtime kernel was produced.",
|
|
},
|
|
],
|
|
"source_sha256": {
|
|
"readme": sha256(root / "README.md"),
|
|
"setup": sha256(root / "setup.py"),
|
|
"dense_decode_header": sha256(root / "csrc/api/dense_decode.h"),
|
|
},
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
artifact_root = args.artifact_dir.resolve()
|
|
v3_root = args.v3_repo.resolve()
|
|
flashmla_root = args.flashmla_repo.resolve()
|
|
shard = artifact_root / "model-00001-of-000004.safetensors"
|
|
required = [
|
|
shard,
|
|
artifact_root / "config.json",
|
|
artifact_root / "configuration_deepseek.py",
|
|
artifact_root / "modeling_deepseek.py",
|
|
artifact_root / "tokenizer.json",
|
|
v3_root / "inference/model.py",
|
|
flashmla_root / "README.md",
|
|
flashmla_root / "setup.py",
|
|
flashmla_root / "csrc/api/dense_decode.h",
|
|
]
|
|
missing = [str(path) for path in required if not path.exists()]
|
|
if missing:
|
|
raise FileNotFoundError(f"missing required official artifacts: {missing}")
|
|
if git_revision(v3_root) != V3_REVISION:
|
|
raise RuntimeError("DeepSeek-V3 revision does not match the probe contract")
|
|
if git_revision(flashmla_root) != FLASHMLA_REVISION:
|
|
raise RuntimeError("FlashMLA revision does not match the probe contract")
|
|
if args.cache_slots < 26:
|
|
raise ValueError("cache-slots must be at least 26")
|
|
device = torch.device(args.device)
|
|
if device.type == "cuda" and not torch.cuda.is_available():
|
|
raise RuntimeError("CUDA requested but unavailable")
|
|
|
|
torch.manual_seed(0)
|
|
configuration, modeling = load_v2_modules(artifact_root)
|
|
config = configuration.DeepseekV2Config.from_pretrained(artifact_root)
|
|
config._attn_implementation = "eager"
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
artifact_root,
|
|
trust_remote_code=True,
|
|
local_files_only=True,
|
|
)
|
|
v3 = load_v3_module(v3_root)
|
|
|
|
normalized, attention_state, token_ids = load_real_layer1_input(
|
|
shard, modeling, config, tokenizer, device
|
|
)
|
|
hf_output, hf_result = execute_hf_eager(
|
|
modeling, config, attention_state, normalized, device
|
|
)
|
|
model_args = v3_args(v3, config, args.cache_slots)
|
|
naive_bf16, naive_result = execute_v3_path(
|
|
v3,
|
|
model_args,
|
|
attention_state,
|
|
normalized,
|
|
device,
|
|
torch.bfloat16,
|
|
"naive",
|
|
)
|
|
absorb_bf16, absorb_result = execute_v3_path(
|
|
v3,
|
|
model_args,
|
|
attention_state,
|
|
normalized,
|
|
device,
|
|
torch.bfloat16,
|
|
"absorb",
|
|
)
|
|
naive_fp32, naive_fp32_result = execute_v3_path(
|
|
v3,
|
|
model_args,
|
|
attention_state,
|
|
normalized,
|
|
device,
|
|
torch.float32,
|
|
"naive",
|
|
)
|
|
absorb_fp32, absorb_fp32_result = execute_v3_path(
|
|
v3,
|
|
model_args,
|
|
attention_state,
|
|
normalized,
|
|
device,
|
|
torch.float32,
|
|
"absorb",
|
|
)
|
|
|
|
naive_bytes = naive_result["active_cache_bytes"]
|
|
absorb_bytes = absorb_result["active_cache_bytes"]
|
|
result = {
|
|
"schema_version": 1,
|
|
"captured_at": args.captured_at or datetime.now(timezone.utc).isoformat(),
|
|
"evidence_identity": "X / official V2-Lite weights executed through official HF eager and V3 naive/absorb reference paths",
|
|
"boundary": {
|
|
"correctness_probe": True,
|
|
"serving_benchmark": False,
|
|
"full_model_generation": False,
|
|
"optimized_flashmla_kernel_executed": False,
|
|
"scope": "one real layer-1 attention input; 25-token prefill plus one-token decode",
|
|
},
|
|
"provenance": {
|
|
"huggingface_model": "deepseek-ai/DeepSeek-V2-Lite",
|
|
"huggingface_revision": HF_REVISION,
|
|
"deepseek_v3_repository": "deepseek-ai/DeepSeek-V3",
|
|
"deepseek_v3_revision": V3_REVISION,
|
|
"flashmla_repository": "deepseek-ai/FlashMLA",
|
|
"flashmla_revision": FLASHMLA_REVISION,
|
|
"sha256": {
|
|
"v2_config": sha256(artifact_root / "config.json"),
|
|
"v2_modeling": sha256(artifact_root / "modeling_deepseek.py"),
|
|
"v2_shard_1": sha256(shard),
|
|
"v3_model": sha256(v3_root / "inference/model.py"),
|
|
},
|
|
},
|
|
"environment": {
|
|
"python": platform.python_version(),
|
|
"platform": platform.platform(),
|
|
"libc": list(platform.libc_ver()),
|
|
"torch": torch.__version__,
|
|
"torch_cuda": torch.version.cuda,
|
|
"transformers": __import__("transformers").__version__,
|
|
"safetensors": __import__("safetensors").__version__,
|
|
"device": str(device),
|
|
"cuda_capability": list(torch.cuda.get_device_capability(device)),
|
|
"nvidia_smi": nvidia_smi(),
|
|
},
|
|
"input": {
|
|
"prompt": PROMPT,
|
|
"token_ids": token_ids,
|
|
"tokens": tokenizer.convert_ids_to_tokens(token_ids),
|
|
"sequence_tokens": len(token_ids),
|
|
"prefill_tokens": len(token_ids) - 1,
|
|
"decode_tokens": 1,
|
|
"layer0_executed": True,
|
|
"layer1_input_normalized": True,
|
|
},
|
|
"configuration": {
|
|
"hidden": config.hidden_size,
|
|
"attention_heads": config.num_attention_heads,
|
|
"qk_nope_head_dim": config.qk_nope_head_dim,
|
|
"qk_rope_head_dim": config.qk_rope_head_dim,
|
|
"v_head_dim": config.v_head_dim,
|
|
"kv_lora_rank": config.kv_lora_rank,
|
|
"cache_slots": args.cache_slots,
|
|
},
|
|
"execution": {
|
|
"huggingface_eager": hf_result,
|
|
"v3_naive_bfloat16": naive_result,
|
|
"v3_absorb_bfloat16": absorb_result,
|
|
"v3_naive_float32": naive_fp32_result,
|
|
"v3_absorb_float32": absorb_fp32_result,
|
|
},
|
|
"correctness": {
|
|
"hf_eager_vs_v3_naive_bfloat16": compare(hf_output, naive_bf16),
|
|
"v3_naive_vs_absorb_bfloat16": compare(naive_bf16, absorb_bf16),
|
|
"hf_eager_vs_v3_absorb_bfloat16": compare(hf_output, absorb_bf16),
|
|
"v3_naive_vs_absorb_float32": compare(naive_fp32, absorb_fp32),
|
|
},
|
|
"cache_accounting": {
|
|
"active_tokens": len(token_ids),
|
|
"dtype": "BF16",
|
|
"naive_active_bytes": naive_bytes,
|
|
"absorb_active_bytes": absorb_bytes,
|
|
"naive_elements_per_token": config.num_attention_heads
|
|
* (
|
|
config.qk_nope_head_dim
|
|
+ config.qk_rope_head_dim
|
|
+ config.v_head_dim
|
|
),
|
|
"absorb_elements_per_token": config.kv_lora_rank
|
|
+ config.qk_rope_head_dim,
|
|
"naive_over_absorb_ratio": naive_bytes / absorb_bytes,
|
|
"absorb_reduction_vs_naive": 1 - absorb_bytes / naive_bytes,
|
|
},
|
|
"flashmla_boundary": flashmla_boundary(flashmla_root, device),
|
|
}
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
args.output.write_text(json.dumps(result, ensure_ascii=False, indent=2) + "\n")
|
|
print(json.dumps(result, ensure_ascii=False, indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|