-
Notifications
You must be signed in to change notification settings - Fork 250
feat(ds4): add default-off HIP GPU phase profiler #554
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,137 @@ | ||
| #!/usr/bin/env python3 | ||
| """Model-backed token parity regression for the public DS4 fused-verify flag.""" | ||
|
|
||
| from __future__ import annotations | ||
|
|
||
| import argparse | ||
| import json | ||
| import os | ||
| import re | ||
| import subprocess | ||
| import tempfile | ||
| import time | ||
| import urllib.request | ||
| from pathlib import Path | ||
|
|
||
|
|
||
| TOKEN_RE = re.compile(r"\[ds4-parity-tokens\] n=\d+ ids=\[([0-9 ]*)\]") | ||
|
|
||
|
|
||
| def wait_ready(port: int, proc: subprocess.Popen[bytes], timeout: float) -> None: | ||
| deadline = time.monotonic() + timeout | ||
| url = f"http://127.0.0.1:{port}/health" | ||
| while time.monotonic() < deadline: | ||
| if proc.poll() is not None: | ||
| raise RuntimeError(f"server exited before readiness: {proc.returncode}") | ||
| try: | ||
| with urllib.request.urlopen(url, timeout=2) as response: | ||
| if response.status == 200: | ||
| return | ||
| except OSError: | ||
| time.sleep(1) | ||
| raise TimeoutError(f"server did not become ready within {timeout:.0f}s") | ||
|
|
||
|
|
||
| def stop_server(proc: subprocess.Popen[bytes]) -> None: | ||
| if proc.poll() is not None: | ||
| return | ||
| proc.terminate() | ||
| try: | ||
| proc.wait(timeout=30) | ||
| except subprocess.TimeoutExpired: | ||
| proc.kill() | ||
| proc.wait(timeout=10) | ||
|
|
||
|
|
||
| def run_case(args: argparse.Namespace, fused_requested: bool, log_path: Path) -> list[int]: | ||
| env = os.environ.copy() | ||
| for name in ( | ||
| "DFLASH_DS4_FUSED_VERIFY", | ||
| "DFLASH_DS4_ALLOW_APPROX_FUSED_VERIFY", | ||
| "DFLASH_DS4_SEQ_VERIFY", | ||
| "DFLASH_DS4_DSPARK_DEBUG", | ||
| ): | ||
| env.pop(name, None) | ||
| env.update({ | ||
| "DFLASH_DS4_SPEC": "1", | ||
| "DFLASH_DS4_DRAFT": str(args.draft), | ||
| "DFLASH_DS4_SPEC_Q": "4", | ||
| "DFLASH_DS4_ADAPTIVE_WIDTH": "0", | ||
| "DFLASH_DS4_PARITY_TRACE": "1", | ||
| "LUCE_MMVQ_MAX_NCOLS": "4", | ||
| }) | ||
| if fused_requested: | ||
| env["DFLASH_DS4_FUSED_VERIFY"] = "1" | ||
|
|
||
| cmd = [ | ||
| str(args.server_bin), str(args.target), | ||
| "--host", "127.0.0.1", "--port", str(args.port), | ||
| "--max-ctx", str(args.max_ctx), "--chunk", "512", | ||
| "--target-device", "hip:0", "--ds4-fused-decode", | ||
| ] | ||
| with log_path.open("wb") as log: | ||
| proc = subprocess.Popen(cmd, env=env, stdout=log, stderr=subprocess.STDOUT) | ||
| try: | ||
| wait_ready(args.port, proc, args.startup_timeout) | ||
| body = json.dumps({ | ||
| "model": "dflash", | ||
| "messages": [{"role": "user", "content": args.prompt}], | ||
| "temperature": 0, | ||
| "seed": 1234, | ||
| "max_tokens": args.max_tokens, | ||
| "stream": False, | ||
| }).encode() | ||
| request = urllib.request.Request( | ||
| f"http://127.0.0.1:{args.port}/v1/chat/completions", | ||
| data=body, | ||
| headers={"Content-Type": "application/json"}, | ||
| ) | ||
| with urllib.request.urlopen(request, timeout=args.request_timeout) as response: | ||
| if response.status != 200: | ||
| raise RuntimeError(f"generation returned HTTP {response.status}") | ||
| json.load(response) | ||
| finally: | ||
| stop_server(proc) | ||
|
|
||
| matches = TOKEN_RE.findall(log_path.read_text(errors="replace")) | ||
| if not matches: | ||
| raise RuntimeError(f"token trace missing from {log_path}") | ||
| return [int(token) for token in matches[-1].split()] | ||
|
|
||
|
|
||
| def main() -> int: | ||
| parser = argparse.ArgumentParser() | ||
| parser.add_argument("--server-bin", required=True, type=Path) | ||
| parser.add_argument("--target", required=True, type=Path) | ||
| parser.add_argument("--draft", required=True, type=Path) | ||
| parser.add_argument("--port", type=int, default=18084) | ||
| parser.add_argument("--max-ctx", type=int, default=4096) | ||
| parser.add_argument("--max-tokens", type=int, default=32) | ||
| parser.add_argument("--startup-timeout", type=float, default=180) | ||
| parser.add_argument("--request-timeout", type=float, default=600) | ||
| parser.add_argument( | ||
| "--prompt", | ||
| default="Explain why a bicycle stays upright while moving.", | ||
| ) | ||
| args = parser.parse_args() | ||
| for path in (args.server_bin, args.target, args.draft): | ||
| if not path.is_file(): | ||
| parser.error(f"file not found: {path}") | ||
|
|
||
| with tempfile.TemporaryDirectory(prefix="ds4-fused-parity-") as tmp: | ||
| tmp_path = Path(tmp) | ||
| normal = run_case(args, False, tmp_path / "normal.log") | ||
| fused_flag = run_case(args, True, tmp_path / "fused-flag.log") | ||
|
|
||
| if normal != fused_flag: | ||
| limit = min(len(normal), len(fused_flag)) | ||
| first = next((i for i in range(limit) if normal[i] != fused_flag[i]), limit) | ||
| print(f"FAIL: first token mismatch at {first}: normal={normal[first:first+4]} " | ||
| f"fused_flag={fused_flag[first:first+4]}") | ||
| return 1 | ||
| print(f"PASS: {len(normal)} generated token IDs are identical") | ||
| return 0 | ||
|
|
||
|
|
||
| if __name__ == "__main__": | ||
| raise SystemExit(main()) |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -22,6 +22,7 @@ | |
| // DFLASH_DS4_FULL_SNAP=1 for A/B validation. | ||
|
|
||
| #include "deepseek4_dspark.h" | ||
| #include "deepseek4_gpu_profiler.h" | ||
| #include "deepseek4_internal.h" | ||
| #include "internal.h" | ||
| #include "common/dspark_head.h" | ||
|
|
@@ -39,6 +40,10 @@ | |
|
|
||
| namespace dflash::common { | ||
|
|
||
| namespace { | ||
| bool spec_env_flag(const char * name); | ||
| } | ||
|
|
||
| // ── DFlashTarget adapter over the DS4 target ──────────────────────────── | ||
| class DeepSeek4DFlashTarget : public DFlashTarget { | ||
| public: | ||
|
|
@@ -71,7 +76,16 @@ class DeepSeek4DFlashTarget : public DFlashTarget { | |
| const char * v = std::getenv("DFLASH_DS4_SEQ_VERIFY"); | ||
| return v && *v && *v != '0'; | ||
| }(); | ||
| const bool approximate_fused = | ||
| spec_env_flag("DFLASH_DS4_FUSED_VERIFY") && | ||
| spec_env_flag("DFLASH_DS4_ALLOW_APPROX_FUSED_VERIFY"); | ||
| const char * profile_mode = seq_verify ? "exact_sequential" : | ||
| (approximate_fused ? "approx_fused" : "exact"); | ||
| Ds4GpuProfiler verify_profiler( | ||
| spec_env_flag("DFLASH_DS4_GPU_PROFILE"), "verification", | ||
| profile_mode, n, base_pos); | ||
| if (seq_verify) { | ||
| verify_profiler.begin(Ds4GpuPhase::VerificationStep); | ||
| std::vector<int32_t> am_all; | ||
| std::vector<float> feat_all; | ||
| std::vector<float> logits_all; | ||
|
|
@@ -98,22 +112,27 @@ class DeepSeek4DFlashTarget : public DFlashTarget { | |
| last_tok = am_all.back(); | ||
| verify_n_ = n; | ||
| if (all_argmax) *all_argmax = std::move(am_all); | ||
| verify_profiler.end(Ds4GpuPhase::VerificationStep); | ||
| verify_profiler.emit(); | ||
| return true; | ||
| } | ||
| std::vector<int32_t> am; | ||
| // n==1 must take the dynamic (non-reuse) path: the reused decode graph | ||
| // skips the capture/all-logits hooks (backend HC), which this needs. | ||
| if (!deepseek4_dspark_verify_forward(backend_, device_, w_, cache_, capture_ids_, | ||
| embed_buf_.data(), tokens.data(), n, base_pos, am, | ||
| keep_logits_ ? &verify_logits_ : nullptr, | ||
| verify_features_, telemetry_, | ||
| /*allow_graph_reuse=*/n > 1)) { | ||
| verify_profiler.begin(Ds4GpuPhase::VerificationStep); | ||
| const bool verify_ok = deepseek4_dspark_verify_forward( | ||
| backend_, device_, w_, cache_, capture_ids_, embed_buf_.data(), tokens.data(), | ||
| n, base_pos, am, keep_logits_ ? &verify_logits_ : nullptr, | ||
| verify_features_, telemetry_, /*allow_graph_reuse=*/n > 1); | ||
| verify_profiler.end(Ds4GpuPhase::VerificationStep); | ||
| if (!verify_ok) { | ||
| return false; | ||
| } | ||
| if (am.empty()) return false; | ||
| last_tok = am.back(); | ||
| verify_n_ = n; | ||
| if (all_argmax) *all_argmax = std::move(am); | ||
| verify_profiler.emit(); | ||
| return true; | ||
| } | ||
|
|
||
|
|
@@ -426,13 +445,17 @@ bool run_deepseek4_dspark_spec_decode( | |
|
|
||
| const bool debug = spec_env_flag("DFLASH_DS4_DSPARK_DEBUG"); | ||
| const bool timing = spec_env_flag("DFLASH_DS4_TIMING"); | ||
| const bool parity_trace = spec_env_flag("DFLASH_DS4_PARITY_TRACE"); | ||
| const bool full_snap = spec_env_flag("DFLASH_DS4_FULL_SNAP"); | ||
| const bool seq_verify_mode = spec_env_flag("DFLASH_DS4_SEQ_VERIFY"); | ||
| // Laguna-style adaptive verify width: EWMA of accepted candidates, width = | ||
| // ewma + 2 (avg_commit << block means the wide tail is usually wasted). | ||
| // /tmp/ds4_awidth: 1 = on, 0 = off (default on). | ||
| bool adaptive_width = true; | ||
| if (std::FILE * f = std::fopen("/tmp/ds4_awidth", "r")) { | ||
| const char * adaptive_env = std::getenv("DFLASH_DS4_ADAPTIVE_WIDTH"); | ||
| if (adaptive_env && *adaptive_env) { | ||
| adaptive_width = *adaptive_env != '0'; | ||
| } else if (std::FILE * f = std::fopen("/tmp/ds4_awidth", "r")) { | ||
| int v = 1; | ||
| if (std::fscanf(f, "%d", &v) == 1) adaptive_width = (v != 0); | ||
| std::fclose(f); | ||
|
|
@@ -480,6 +503,7 @@ bool run_deepseek4_dspark_spec_decode( | |
|
|
||
| DeepSeek4DFlashTarget target(target_w, target_cache, backend, device, snap_backend, | ||
| drafter.capture_layer_ids, drafter.mask_token_id); | ||
| target.set_keep_logits(parity_trace); | ||
| DraftWeights dw = make_dspark_shim(drafter); | ||
| DeepSeek4SpecRollback rollback; | ||
| DeepSeek4StepTelemetry tel{}; | ||
|
|
@@ -680,6 +704,31 @@ bool run_deepseek4_dspark_spec_decode( | |
| } | ||
| tm_verify += spec_ms_since(t0); | ||
|
|
||
| if (parity_trace) { | ||
| std::vector<float> trace_logits; | ||
| if (target.read_verify_logits(q, trace_logits)) { | ||
| for (int row = 0; row < q; ++row) { | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. P3: When Prompt for AI agents |
||
| const float * logits = trace_logits.data() + (size_t) row * target_w.n_vocab; | ||
| int best = 0; | ||
| int second = -1; | ||
| for (int tok = 1; tok < target_w.n_vocab; ++tok) { | ||
| if (logits[tok] > logits[best]) { | ||
| second = best; | ||
| best = tok; | ||
| } else if (second < 0 || logits[tok] > logits[second]) { | ||
| second = tok; | ||
| } | ||
| } | ||
| const float second_logit = second >= 0 ? logits[second] : logits[best]; | ||
| std::fprintf(stderr, | ||
| "[ds4-parity-logits] step=%ld pos=%d row=%d best=%d val=%.9g " | ||
| "second=%d val2=%.9g margin=%.9g\n", | ||
| steps, pos, row, best, logits[best], second, second_logit, | ||
| logits[best] - second_logit); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| // Accept the longest matching prefix. accept counts the seed (slot 0) | ||
| // plus each candidate the target agrees with. | ||
| int accept = 1; | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
P3: This new inventory line won't be reproduced by the documented regeneration command (
grep -rE 'getenv("[A-Z0-9_]+")' server/src), since the actual reads go throughds4_env_flag(name)/spec_env_flag(name)wrappers rather than a literalgetenv("DFLASH_DS4_GPU_PROFILE")call. Consider noting this is manually curated or adjusting the generation instructions so the 'generated' section stays trustworthy.Prompt for AI agents