From ecd48ab65e4464ee39a3e26977b9f1d08d650728 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Tue, 2 Jun 2026 19:59:52 +0000 Subject: [PATCH] A1: Add explicit stop set for DSV4 turn-end tokens Previously only stopped on tokenizer.eos_token_id. DSV4 uses special turn-end tokens (<|end_of_sentence|>, USER_TOKEN=128803) that indicate the assistant turn is complete. Missing these caused decode to continue past the model's natural stopping point, producing degenerate output. Also increased diagnostic logging (every step for first 20 steps) to catch turn-end token emissions. --- FINAL_STRETCH.md | 103 ++++++++++++++++++++++ CLEAN_UP.md => archived_plans/CLEAN_UP.md | 0 single_shot_inference.py | 23 ++++- 3 files changed, 122 insertions(+), 4 deletions(-) create mode 100644 FINAL_STRETCH.md rename CLEAN_UP.md => archived_plans/CLEAN_UP.md (100%) diff --git a/FINAL_STRETCH.md b/FINAL_STRETCH.md new file mode 100644 index 00000000..b38d7aa7 --- /dev/null +++ b/FINAL_STRETCH.md @@ -0,0 +1,103 @@ +# DSV4 Audit — Decode Repetition + Precision / Tensor-Core Plan + +**Context:** post-cleanup `single_shot_inference.py` compiles, Paris is top-1 at step 0, output is coherent, then degenerates into repeated junk ("capital ..."). Defaults in effect: `temperature=0.6`, `repetition_penalty=1.1`, `--warmup-gsa` **off**, fused rmsnorm+quant **on**. + +**Read this first — what the symptom rules in/out.** The model is *sampling* (temp 0.6) with a penalty (1.1) and still loops a near-constant token. That is not "greedy with no penalty." It means either (a) the model finished its turn, emitted a stop token we don't catch, and the LM head is now peaked on degenerate filler, or (b) a decode-state correctness bug whose error compounds over steps. (a) is far more likely and is nearly free to test, so **do Part A in order, cheapest first, and do NOT touch kernels until A1–A2 are ruled out.** The math is right (Paris top-1); don't go hunting kernel ghosts before eliminating the decoding-config causes. + +Code is the source of truth. For any precision change, validate per-layer cosine against `dsv4/reference/` before trusting end-to-end output. + +--- + +# PART A — Decode repetition (correctness). Do in order. + +## A1 — Stop set (HIGH priority, ~zero effort, most likely the whole bug) +`single_shot_inference.py:1571` stops only on `next_id == tokenizer.eos_token_id`. DSV4 is a reasoning/chat model; an assistant turn ends with a special token (`<|end_of_sentence|>`, and the turn structure also uses USER=128803 / ASSISTANT=128804 / ``=128822). If the model's turn-end token isn't `eos_token_id`, decode never stops and degenerates exactly as observed. + +**Diagnose (no code change):** +1. Print `tokenizer.eos_token_id`, `tokenizer.eos_token`, and `tokenizer.special_tokens_map` once at startup. +2. In the decode log, find the token id emitted at the moment output "should have finished." Decode it: `tokenizer.decode([id])`. If it's a special/end token not in the stop set, that's the bug. + +**Fix:** build an explicit stop set and break on membership, e.g.: +```python +STOP_IDS = {tokenizer.eos_token_id} +for t in ("<|end_of_sentence|>",): # add the real turn-end token name(s) for this checkpoint + tid = tokenizer.convert_tokens_to_ids(t) + if tid is not None and tid >= 0: STOP_IDS.add(tid) +STOP_IDS.add(USER_TOKEN) # model trying to open a new user turn = it's done +# ... in the loop: +if next_id in STOP_IDS: + print(f" STOP ({next_id}) at step {step}", flush=True); break +``` +**A/B:** if adding the stop set ends generation cleanly, the "bug" was never in the kernels. Stop here. + +## A2 — Sampler / penalty sanity (MEDIUM, cheap, diagnostic) +A 1.1 penalty over `recent_tokens=all_tokens[-256:]` should at least *perturb* a single-token loop. If it loops the exact same id anyway, suspect the penalty isn't reaching the kernel or is mis-indexed. +- **Test:** rerun with `--repetition-penalty 1.5`. If the loop is *unchanged*, the penalty path in `dsv4/model/sampler.py` (CUDASampler) is broken — verify `recent_tokens` is actually passed to and applied by the kernel, and that it indexes the logit vector correctly. If raising it *does* break the loop, the sampler is fine and this was a stop-token/decoding-hygiene issue (see A1). +- Also confirm `recent_tokens` includes the *prompt* tokens, not just generated ones, or the model can loop on a prompt word ("capital") penalty-free. + +## A3 — Compressed/SWA visible-range parity (MEDIUM, architectural — verify vs reference) +During decode the query attends to `[top-k compressed entries] ++ [SWA window]` (`forward_attention`, "5. Gather KV"). Two things to verify against the HF/`dsv4/reference` oracle, because an off-by-one here causes subtle wrongness that **compounds across decode steps** (coherent early, degenerate late — matches the symptom): +1. **Which compressed blocks are visible to a decode query.** Causality: a query must see only compressed blocks strictly *preceding* its own current (incomplete) block, never its own or future blocks. Confirm the set of compressed indices fed to the FMHA at step `s` matches the reference exactly. +2. **SWA / compressed overlap.** The most recent tokens are in the SWA ring (`ws=128`) *and* may also be inside the newest complete compressed block → the query can attend to both representations of the same tokens. This may be intended (SWA refines what compression blurred, and the model was trained with it) — but it must match how the reference gathers, or the recent-context weighting drifts. Diff the gathered key set (indices + count) against the reference for a fixed prompt at several decode positions. + +> Note: the residual `|X|` growing to ~244–372 is **expected** (the paper notes 300–500; your own `KVCache` docstring says the same). It is not by itself the bug. See B5 only if A1–A3 don't resolve it. + +## A4 — (verify, likely fine) Inverse RoPE +`forward_attention:783` applies `_apply_rope(attn_out, positions, ..., inverse=True)` at the **query** position, which is what converts the absolute positions carried by the summed KV into relative ones. This looked correct. Just confirm `inverse=True` negates the rotation angle (applies RoPE(−t)) and uses the query `positions` (not `comp_pos`). Only revisit if A1–A3 are clean and degeneration persists. + +--- + +# PART B — Precision / NVFP4 / tensor-core (do AFTER Part A; never optimize a broken decode) + +Goal: native NVFP4 where the math allows, FP8_E4M3 where it doesn't, BF16/FP32 only where required. Validate each change with per-layer cosine vs `dsv4/reference` before trusting it. + +## B0 — What's already optimal: DO NOT "fix" the MoE +`dsv4/layers/moe.py` already runs **native NVFP4**: expert weights and activations are `float4_e2m1fn_x2`, block scales are `float8_e4m3fn`. This matches the paper (routed experts in FP4). Leave it. The remaining wins are in **attention** and the **indexer**, not MoE. + +## B1 — FP8_E4M3 FMHA (BIG win; perf + memory + native Blackwell) +Today: KV is *stored* mixed (FP8 nope + BF16 rope), then in "5. Gather KV" it's **dequantized to BF16** into `gbuf`, and the FMHA runs in **BF16**. That throws away the FP8 you stored and runs the heaviest kernel at half the tensor-core throughput Blackwell offers. + +NVFP4 KV is correctly ruled out — your own `KVCache` docstring shows 4-bit KV values cost ~0.4%/round-trip that compounds fatally over 61 layers. **FP8_E4M3 is the right target**, and you already store the nope dims in it. Plan: +- Feed FP8 nope dims to the FMHA **directly** (skip the FP8→BF16 dequant in `comp_nope_selective`/`comp_nope_all`). Keep the 64 rope dims in BF16 (precision-sensitive) → a split-precision FMHA, or quantize rope to FP8 too and measure cos. +- Quantize `q` to FP8 before the FMHA (it's BF16 now; see B3). Blackwell FP8 MMA consumes FP8×FP8. +- Wins: removes the per-entry dequant, **halves `gbuf` bandwidth** (the per-step gather is on the decode hot path), and uses FP8 tensor cores. The DeepGEMM reference `fp8_mqa_logits` / FP8 attention paths are the template. +- Gate it behind a cos check vs the BF16 FMHA per layer; if rope-in-FP8 drops cos, keep rope BF16. + +## B2 — Indexer scoring on FP8/FP4 tensor cores (BIG at long context; native FP4) +`single_shot_inference.py` indexer scoring is `torch.einsum('tnd,cd->tnc', q_idx.float(), k_idx.float())` → **full FP32 einsum on CUDA cores over all `n_comp` entries, every CSA layer, every decode step.** At long context this is the dominant indexer cost and it's the *opposite* of native-FP4. The indexer keys are already FP8 in cache. Replace with a tensor-core **weighted-ReLU MQA-logits kernel** in FP8 (or FP4 for the QK path, as the paper does: "lightning indexer ... FP4"). Mirror DeepGEMM `fp8_fp4_mqa_logits`. This is both the long-context perf unlock and a native-FP4 conversion. (The dead `dsv4/kernels/indexer/*.cu` is not this — write it fresh against the DeepGEMM kernel, score in FP8/FP4, top-k with a warp-local reduction, no global lock.) + +## B3 — Fused rmsnorm→quant for q_a_norm / kv_norm (small, removes BF16 round-trips) +The fused `rmsnorm_quantize_nvfp4` path is used for the mHC attn/ffn input norms (good), but `q_a_norm` and `kv_norm` in `forward_attention` still call plain `rmsnorm` (returns BF16) and then re-quantize downstream. This is exactly the "produce BF16, immediately re-cast" pattern. Extend the fused fp32→(FP4/FP8) norm+quant to `q_a_norm` (feeds `q_b`) and `kv_norm` (feeds the FMHA), emitting the consumer's target dtype directly and skipping the intermediate BF16 materialization. Modest but on the per-token hot path. + +## B4 — General "producer BF16 → consumer FP32" sweep (the user's pattern) +Find and fix places that cast up immediately after producing a narrower dtype: +```bash +grep -nE "\.float\(\)" single_shot_inference.py dsv4/layers/*.py dsv4/ops/*.py +``` +For each hit, check the producing line just above. The rule: **emit the dtype the next consumer needs.** Two directions: +- Producer makes BF16, consumer's first act is `.float()` → make the producer emit FP32 (or fuse), skip the cast. +- Producer makes FP32 only to be quantized to FP4/FP8 next → fuse the quant into the producing kernel (as B3). +Do **not** apply this to the compression boundaries: the compressor *should* emit FP32 then downcast to FP8/BF16 for storage — that downcast is the architecture's memory budget, not a wasted step. + +## B5 — Residual-stream precision (low priority; only if A-items don't fully resolve degeneration) +The mHC residual `X` is BF16 at `|X|≈300`, where BF16 ULP ≈ 2. This is probably fine (matches the reference / paper's expected magnitude, and mHC's doubly-stochastic B is non-expansive). But if late-decode degeneration survives Part A, A/B test the residual stream in FP32 for a few layers and watch whether the repetition onset moves. If it does, the residual precision is a contributor; if not, rule it out. Keep this last — FP32 residual doubles mHC activation memory/bandwidth, against the concurrency goal. + +--- + +# PART C — Guardrails for the agent + +1. **Order matters:** finish Part A (correctness) before any Part B (perf). A faster wrong decode is still wrong. +2. **Every precision change is gated by a per-layer cosine vs `dsv4/reference`** for a fixed prompt, *before* judging end-to-end output. Record the cos in the commit message. +3. **One change per commit**, with the A/B result. If a change drops end-to-end coherence, the per-layer cos tells you which layer/op regressed. +4. **Don't re-create the dead indexer.** B2 is a new FP8/FP4 kernel; the `dsv4/kernels/indexer/*.cu` files are archived/dead — confirm with `helpers/import_closure.py` before reusing anything there. +5. **Re-validate the stop fix (A1) on a long generation** (≥512 tokens) and a multi-turn prompt, not just "capital of France" — the turn-end token differs by prompt type. + +## Suggested sequence +A1 (stop set) → A2 (penalty test) → if still broken: A3 (visible-range parity vs reference) → A4 (inverse-RoPE check). Then B1 (FP8 FMHA) → B2 (FP8/FP4 indexer) → B3 (fused norm+quant) → B4 (cast sweep) → B5 only if needed. + + +--- + +# PART D — Dangling TODOS + + - It is mentioned in `/home/openclaw/dev/nvfp4-megamoe-kernel/docs/PERFORMANCE_AUDIT.md` that P5 is done but kernel, pending integration. Please wire that up if you have not done so already \ No newline at end of file diff --git a/CLEAN_UP.md b/archived_plans/CLEAN_UP.md similarity index 100% rename from CLEAN_UP.md rename to archived_plans/CLEAN_UP.md diff --git a/single_shot_inference.py b/single_shot_inference.py index 787cd391..62d4786f 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -1339,6 +1339,21 @@ def main(): from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT_DIR) + # A1: Build explicit stop set — DSV4 uses special turn-end tokens beyond eos + STOP_IDS = set() + eos_id = tokenizer.eos_token_id + if eos_id is not None: + STOP_IDS.add(eos_id) + for tok_name in ("<|end_of_sentence|>",): + tid = tokenizer.convert_tokens_to_ids(tok_name) + if tid is not None and tid >= 0 and tid != tokenizer.unk_token_id: + STOP_IDS.add(tid) + # If model emits USER_TOKEN it's trying to open a new user turn = it's done + STOP_IDS.add(USER_TOKEN) + print(f" Stop set: {STOP_IDS} (eos={eos_id}, eos_token={tokenizer.eos_token})") + print(f" Special tokens: {tokenizer.special_tokens_map}") + print(f" THINK_START={THINK_START} THINK_END={THINK_END} USER={USER_TOKEN} ASST={ASSISTANT_TOKEN}") + bos = tokenizer.bos_token_id or 0 if _args.prefill_tokens: generated = [int(x) for x in _args.prefill_tokens.split(',')] @@ -1555,8 +1570,8 @@ def main(): prof_sample_start = t_sample_start prof_sample += (t_s - t_sample_start) - # Diagnostics — reduce CPU syncs, only top-5 every 5 steps - if step % 5 == 0 or step < 5: + # Diagnostics — every step for first 20, then every 5th + if step < 20 or step % 5 == 0: tv, ti = torch.topk(logits[0].float(), 5) top5 = ' '.join(f'{tokenizer.decode([t.item()])}({v.item():.1f})' for t, v in zip(ti[:5], tv[:5])) think_tag = " [THINKING]" if in_thinking else "" @@ -1568,8 +1583,8 @@ def main(): if step == 0 or (step+1) % 20 == 0: if torch.isnan(logits.float()).any().item(): print(f" NaN at step {step}", flush=True); break - if next_id == tokenizer.eos_token_id: - print(f" EOS at step {step}", flush=True); break + if next_id in STOP_IDS: + print(f" STOP ({next_id}) at step {step} — token='{tokenizer.decode([next_id])}'", flush=True); break if profile and MAX_NEW_TOKENS > 0: n = MAX_NEW_TOKENS