biondizzle
3320abfe24
1. Compressor positional bias was being added to BOTH gate (softmax logit) AND KV content. Per paper eq. 9-12, position bias is only for the softmax logits (Z+B), NOT the KV content (C). Adding pb to kv_val corrupts every compressed KV entry with learned positional-bias content. Fixed in both CSA and HCA paths in compressor_reduce.cu. 2. SwiGLU clamp ordering: code was clamping silu(gate) instead of clamping raw gate before SiLU. Per paper §4.2.3: gate = clamp(gate, max=limit), then silu(clamp(gate)) * clamp(up). Fixed in moe.py (both unfused paths) and fused_swiglu.py (CuTeDSL kernel). shared_expert.py was already correct.
4.1 KiB
1. Possible bug: compressor positional bias is being added to KV content
In your dsv4/kernels/cuda/compressor_reduce.cu, the compressor appears to do this in both CSA and HCA paths:
g += pb;
kv_val += pb; // suspicious / wrong
The official compressor equations add positional bias only to the compression weights/logits Z + B, then use those weights to sum the raw projected KV content C. The bias is not added to the KV value itself. The paper defines compression as softmax over Z + B, followed by a weighted sum of C.
So this should be:
g += pb;
// do not add pb to kv_val
That bug would poison every compressed KV entry with learned positional-bias content. It may not fully explain the first token for a tiny prompt if SWA dominates, but it is absolutely wrong relative to the official architecture and will degrade CSA/HCA context quality. If your unit tests passed, they may have been comparing against a reference that made the same mistake or were too short to expose it.
2. Don’t use think_start as the canary here
In official thinking mode, the prompt formatter typically appends the assistant marker plus <think> before generation. That means decode step 0 is already inside the thinking span. The model should not necessarily emit think_start; a low think_start logit is not itself evidence that the model “failed to enter thinking mode.”
For this particular prompt, a high think_end logit can even be plausible because “The capital of France is” does not need much reasoning. Run the same current kernel with official chat mode, greedy decoding, no repetition penalty, no top-k/top-p, and compare first-token logits. The think_start = 1.77 observation is probably a misleading diagnostic.
3. Indexer parity still looks suspect
The official CSA/HCA details include RMSNorm on queries and compressed KV before attention, partial RoPE on the last 64 dims, sliding-window KV, and attention sink behavior.
For the CSA indexer, the official reference does more than a plain q/k dot product: indexer Q and compressed indexer K get the appropriate RoPE/rotation treatment, and the QK path is one of the explicitly FP4-QATed pieces. If your current indexer compressor is just producing compressed keys without the same rotate/RoPE/FP4 path, CSA top-k can select plausible-looking but wrong blocks. Again, probably not the first-token issue on a short prompt, but it will matter for any real context.
4. Check SwiGLU clamp ordering
The official behavior is effectively:
gate = clamp(gate, max=swiglu_limit)
up = clamp(up, -swiglu_limit, swiglu_limit)
out = silu(gate) * up
If your fused path clamps silu(gate) instead of clamping raw gate before SiLU, it is not equivalent. This is especially worth checking in both routed MoE and shared expert fused kernels, because a small-looking activation semantic mismatch repeats through every layer.
5. DEQUANT TO BF16 IN THIS ORDER JUST TO SEE WHAT HAPPENS (You are allowed to break the no bf16 rule for this because we can always revert back to previous commit)
The most suspicious surfaces to temporarily dequantize are, in order:
- lm head — FP4 lm head can directly flatten or reorder vocabulary logits.
- router gate — slight errors can change top-6 experts; wrong expert IDs are much worse than a small GEMM error.
- shared expert — official routed experts are the FP4 target; shared expert is a different sensitivity profile.
- attention q/kv/o projections and grouped output projection — these are not described as full FP4 QAT targets.
- compressor/indexer helper projections — only the CSA indexer QK path is explicitly FP4-QATed, not the whole compressor.
If a BF16/FP8 lm head alone makes Paris / . / answer-like tokens dominate again, you’ve found a high-leverage culprit. My money is on LM Head needing to be BF16
The fastest triage is basically: run thinking_mode=chat, greedy; switch only lm_head back to BF16/FP8; then switch router back; then patch the compressor bias-to-KV bug. If any one of those sharply separates the first-token distribution, you’ll know where to spend kernel time.