Mamba-3 SISO Long-Ragged Multi-Chunk Joint States Full VJP

Optimize a packed variable-length Mamba-3 SISO training operator on one CUDA device. You inherit a correct but weak eager PyTorch baseline and must make it as fast as possible while staying inside the numerical tolerances the verifier enforces.

Hard Constraints

Keep this function name and positional signature unchanged:

mamba3_siso_varlen_joint_states_full_vjp(
    q, k, v, adt, dt, trap, q_bias, k_bias, angles, d, z,
    cu_seqlens, dout, dfinal_angle, dfinal_ssm, dfinal_k, dfinal_v
)

Return exactly:

(output, final_angle_state, final_ssm_state, final_k_state, final_v_state,
 dq, dk, dv, dadt, ddt, dtrap, dq_bias, dk_bias, dangles, dd, dz)

Every returned value must be a concrete CUDA torch.Tensor. Do not mutate any input. Do not retrieve, import, vendor, or call an external Mamba implementation; do not inspect verifier files or processes, access clocks or timers, specialize by call order, or attempt held-out access.

Inputs Q, K, V, trap, angles, and Z use BF16. ADT, DT, both biases, and D use FP32. Output, final-angle, final-V, and the corresponding BF16 input gradients use BF16; final-SSM, final-K, and FP32 input gradients use FP32.

What You Have

The batch is one packed buffer, and cu_seqlens contains the cumulative boundaries of three to five independent sequences. State, causal attention, angle accumulation, and gradients reset at every boundary.

For each packed sequence, the reference does the following:

1. Expand the single Q/K head to the 32 value heads.
2. Compute theta = tanh(angles) * pi, then the per-sequence cumulative rotary angle cumsum(theta * dt) mod 2*pi.
3. Add q_bias and k_bias, and apply paired rotary embeddings to Q and K.
4. Compute p = sigmoid(trap), shifted next-token values dt_next and p_next (zero at the final token), then shifted_gamma = dt_next * (1 - p_next) and scale = dt * p + shifted_gamma.
5. Form the causal decayed QK interaction using ADT, scaled K, and V.
6. Add the per-head d * v skip term, subtract the shifted-QK correction, and apply the z * sigmoid(z) gate.
7. Return the final cumulative angle, SSM outer-product state, K state, and V state for every packed sequence.

The five nonzero cotangents contribute jointly to the gradients of all eleven differentiable inputs. Frozen shapes use GQA 1 -> 32, Q/K width 128, V width 64, rotary-angle width 32, chunk size 64, non-multiple ragged sequence lengths spanning 3--15 chunks, and three to five packed sequences.

The inherited solver is /app/methods/main/solver.py. Visible cases are in /app/problems/visible_cases.json, and python /app/selfcheck.py evaluates the complete visible panel. PyTorch, Triton, and locally authored helpers are available.

What You Submit

Submit the contents of /app/methods, with your final implementation at /app/methods/main/solver.py.

Record every attempted iteration as one row of a monotonically increasing v1, v2, ... table in /app/methods/experiment_log.md, and preserve that row's exact solver as the single file /app/methods/versions/vN/solver.py. A partial, interrupted, incorrect, or estimated iteration still needs its snapshot and must be marked invalid. Never reuse a label or delete a snapshot.

Afterwards restore the version you judge most likely to generalize into /app/methods/main/solver.py and run one final visible self-check. That final solver may duplicate a saved version; it is evaluated as an additional final entry.

How It Is Judged

Each case scores the frozen eager-baseline latency divided by your candidate latency, and the per-case speedups are combined by geometric mean. Any case that fails a correctness or protocol check zeroes the entire score.

Compilation and input generation are excluded from latency. All output, four final states, and eleven-input VJP work is included. The verifier checks every returned tensor, input immutability, fresh-call behavior, and timing stability, then re-runs your submission once on a sealed hidden panel to produce the final score.