feat(trainer): recurrent GDN forward matching vLLM decode for Qwen3.5

[hubert-marek]

Summary

The trainer computes GatedDeltaNet (Qwen3.5's linear attention) with FLA's chunked kernel, while vLLM generates tokens with a fused recurrent kernel. These are two different algorithms for the same function; their outputs diverge in the low bits and the divergence compounds through the recurrent state, so every sampled token carries a trainer/inference logprob mismatch that no amount of cu_seqlens plumbing removes.

This PR adds a model.gdn_recurrent_forward flag that runs the GDN forward through vLLM's own fused_sigmoid_gating_delta_rule_update kernel, with a chunked backward via FLA (the recurrent kernel has no backward — computing dg without materializing every state is an open problem). Gradients are bit-identical to the current training path; only forward activations move onto the generator's kernel.

• Import, not port. The forward calls vLLM's vendored kernel directly (lazy import) rather than copying its Triton, so parity tracks vLLM as it evolves instead of rotting against a frozen copy.
• Zero-state, unpaged. initial_state=None + ssm_state_indices=None makes vLLM's kernel recompute full packed sequences from a zero state — exactly the training case, no cache machinery.
• Off by default, rejected at config time with context parallelism.

Status / open items

This is a draft for early review. Known gaps, in order of when I'll hit them:

• GPU verification pending. The Triton path has not executed yet (developed on a CPU box). tests/unit/train/models/test_qwen3_5_gdn_parity.py is the gate.
• Activation-checkpointing interaction. The custom autograd.Function saves tensors in forward; under the default full checkpoint_wrapper this runs inside recompute and needs verification/fixing on GPU.
• Scope is the decode kernel only. The conv1d and vLLM's chunked prefill are still unaligned — a diagnostic test measures the conv gap. End-to-end mismatch_kl → 0 additionally needs the server run with an fp32 SSM cache (vllm_extra: mamba_ssm_cache_dtype = "float32", no code change) and, for full parity, a recurrent GDN prefill on the vLLM side.

Test plan

• pytest tests/unit/train/models/test_qwen3_5_gdn_parity.py -v -s on GPU (kernel parity, backward-equals-chunked, module wiring, conv diagnostic)
• Short Qwen3.5 RL run with the flag on + fp32 SSM cache, watching mismatch_kl/all/mean vs a flagged-off baseline

Co-Authored-By: Claude Opus 4.8 noreply@anthropic.com

Made with Cursor