nvfp4-megamoe-kernel/FINAL_STRETCH.md

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 / </think>=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.:

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:

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 (Fuse mHC pre_block + RMSNorm into a single op) is done but kernel, pending integration. Please wire that up if you have not done so already