align speculative Mamba state with sequential decode

[McZyWu]

Summary

Align the NPU speculative Mamba state update path with sequential single-token decode.

This change fixes numerical inconsistencies in the recurrent gated delta rule and causal Conv1D kernels when multiple MTP tokens are verified in one invocation.

Problem

MTP target verification processes multiple draft tokens in a single kernel invocation, while normal decode processes one token per invocation.

The two paths previously crossed different state persistence boundaries:

• The recurrent gated delta rule kernel kept FP32 intermediate states between speculative steps, while sequential decode persisted and reloaded BF16 state after every token.
• The multi-token causal Conv1D path forced parts of the computation to FP16 and used a different accumulation and persistence order from the single-token path.

As a result, processing the same token sequence with MTP and sequential decode could produce different intermediate states. The differences accumulated across decode steps and could affect subsequent token selection.

Changes

Recurrent gated delta rule

• Round the recurrent state through BF16 between speculative tokens.
• Reload the persisted state before advancing the next recurrence step.
• Match the per-token state transition used by sequential single-token decode.

Causal Conv1D

• Preserve the input and history tensor dtype instead of forcing FP16.
• Perform convolution multiply-accumulate and bias addition in FP32.
• Persist the BF16 convolution result before applying SiLU.
• Align the single-token and multi-token update paths to use the same operation order and precision boundaries.

Validation

Added focused BF16 consistency tests:

• Compare one 4-token recurrent gated delta rule invocation with four sequential single-token invocations.
• Compare all recurrent intermediate states with rtol=0 and atol=0.
• Compare one 4-token causal Conv1D update with four sequential single-token updates.
• Compare Conv1D outputs and final convolution state with rtol=0 and atol=0.

Both kernel paths now produce exactly matching outputs and states for the covered cases.

Scope

This change guarantees consistency between speculative multi-token execution and sequential execution of the same NPU kernels.

It does not claim full-model bitwise equivalence between MTP-enabled and non-MTP inference. Other model layers, different recurrent kernel implementations, and the M=4 versus M=1 execution shapes may still introduce numerical differences.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.