HyenaCascade.sequential_forward silently truncates multi-token inputs after prefill (`u = u[:, -1]`) — wrong outputs, corrupted state, no error or warning

[dhnihaoya]

Summary

After prefill, HyenaCascade.forward routes every subsequent call to sequential_forward, which begins with:

if len(u.shape) > 2:
    u = u[:, -1]

If a caller passes several tokens at once (seqlen > 1) with inference state — e.g. a chunk/block forward with initial state, which is exactly what speculative-decoding verification needs — the first seqlen - 1 tokens are silently discarded:

• the call returns a single position instead of seqlen positions, with no error and no warning;
• the returned position is wrong (the conv/recurrent window never saw the dropped tokens);
• the layer state advances by 1 token instead of seqlen, corrupting every subsequent decode step.

The failure is completely silent, so downstream code keeps running on wrong logits.

Affected code and versions

• vortex/model/model.py, HyenaCascade.forward L224–229 (routing to sequential_forward once layer_idx is in inference_params.fir_state_dict):
  

  vortex/vortex/model/model.py

  Lines 224 to 229 in 8b00afe

  	def forward(self, u, inference_params=None, padding_mask=None, *args, **kwargs):
  	if inference_params is not None and self.layer_idx in inference_params.fir_state_dict.keys():
  	return self.sequential_forward(u, inference_params)
  	else:
  	return self.parallel_forward(u, inference_params, padding_mask)

• vortex/model/model.py, HyenaCascade.sequential_forward L328, truncation at L332–333:
  

  vortex/vortex/model/model.py

  Lines 328 to 334 in 8b00afe

  	def sequential_forward(self, u, inference_params):
  	if self.data_dtype is None:
  	self.data_dtype = u.dtype
  	if len(u.shape) > 2:
  	u = u[:, -1]

Reproduced on:

• main @ 8b00afebeac745d1f31e7e2788f0e0e39fa47637 (HEAD as of 2026-09-02)
• PyPI vtx==1.1.0 (latest release as of 2026-09-02); its vortex/model/model.py is byte-identical to main at the lines above

All Hyena layer types (HCS/HCM FIR and HCL IIR) are affected — the truncation sits in the shared entry of sequential_forward, before any FIR/IIR branching.

Minimal reproduction

Self-contained, no checkpoint required (random-weight HCS-style HyenaCascade, hidden 16, short FIR-3 + inner FIR-7), runs on CPU in ~1 s (torch 2.7.1; also runs unchanged on GPU hosts):

import copy
import sys
import types
import warnings

import torch

# CPU-only shim: vortex.ops.*_interface use @triton.autotune at module level, which
# raises RuntimeError (not ImportError, so engine.py's guard misses it) when no GPU
# driver is active. These Triton kernels are opt-in (use_hc*_kernel=False); stubbing
# them to None mirrors the engine's own fallback. The repro only uses F.conv1d paths.
if not torch.cuda.is_available():
    for mod_name, symbols in {
        "vortex.ops.hcs_interface": ("hcs_conv",),
        "vortex.ops.hcm_interface": ("hcm_fft_conv",),
        "vortex.ops.hcl_interface": ("hcl_fft_conv",),
    }.items():
        stub = types.ModuleType(mod_name)
        for sym in symbols:
            setattr(stub, sym, None)
        sys.modules.setdefault(mod_name, stub)

from vortex.model.cache import HyenaCascadeFIRInferenceParams
from vortex.model.model import HyenaCascade


class dotdict(dict):
    __getattr__ = dict.get
    __setattr__ = dict.__setitem__


def build_layer(seed=0):
    torch.manual_seed(seed)
    hidden = 16
    config = dotdict(
        hidden_size=hidden,
        num_filters=hidden,
        num_attention_heads=2,
        state_size=16,
        short_filter_length=3,
        short_filter_bias=True,
        interleave=False,
        column_split_hyena=False,
    )
    layer = HyenaCascade(config, layer_idx=0,
                         hyena_filter_groups=hidden, fir_inner_filter_length=7)
    layer.eval()
    return layer


torch.manual_seed(1234)
layer = build_layer()
B, L_PREFILL, GAMMA = 1, 8, 4
H = layer.hidden_size
u_prefill = torch.randn(B, L_PREFILL, 3 * H)
u_chunk = torch.randn(B, GAMMA, 3 * H)
u_next = torch.randn(B, 1, 3 * H)

# 1) prefill (parallel_forward populates FIR inference state)
ip0 = HyenaCascadeFIRInferenceParams()
with torch.no_grad():
    _, ip0 = layer(u_prefill, ip0)

# 2) reference: step the gamma=4 tokens one at a time (correct chunk semantics)
ip_ref = copy.deepcopy(ip0)
ys = []
with torch.no_grad():
    for i in range(GAMMA):
        y_i, ip_ref = layer(u_chunk[:, i:i + 1], ip_ref)
        ys.append(y_i)
y_ref = torch.cat(ys, dim=1)                       # [B, 4, H]

# 3) under test: one forward with all 4 tokens at once (block-verification call)
ip_bug = copy.deepcopy(ip0)
with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")
    with torch.no_grad():
        y_bug, ip_bug = layer(u_chunk, ip_bug)     # silently truncated

# 4) control: feed only the last token after prefill
with torch.no_grad():
    y_ctrl, _ = layer(u_chunk[:, -1:], copy.deepcopy(ip0))

# 5) state pollution: decode one more token from each final state
with torch.no_grad():
    y_next_ref, _ = layer(u_next, copy.deepcopy(ip_ref))
    y_next_bug, _ = layer(u_next, ip_bug)

print("expected shape:", tuple(y_ref.shape), " actual shape:", tuple(y_bug.shape))
print("warnings raised:", len(caught))
print("chunk output == last-token-only output:", torch.equal(y_bug, y_ctrl))
print("max|chunk_out - true last-position| :", (y_bug - y_ref[:, -1:]).abs().max().item())
print("max|next-token out (buggy vs ref)|  :", (y_next_bug - y_next_ref).abs().max().item())

Observed output (CPU, fp32):

expected shape: (1, 4, 16)  actual shape: (1, 1, 16)
warnings raised: 0
chunk output == last-token-only output: True
max|chunk_out - true last-position| : 30.00904083251953
max|next-token out (buggy vs ref)|  : 24.278169631958008

Expected vs actual behavior

Expected — either of:

1. a correct chunk forward: [B, γ, H] outputs for all γ positions, with the layer state advanced by γ tokens; or
2. a loud failure (raise) telling the caller that sequential_forward only supports seqlen = 1.

Actual:

• returns [B, 1, H] — γ−1 positions silently missing, zero warnings;
• the single output is bitwise identical to feeding only the last token, i.e. the first γ−1 tokens have zero effect on output and state;
• that surviving position is itself wrong (max abs diff 30.0 above) because the conv window is missing the dropped tokens;
• the FIR state advances by 1 token instead of γ, so all subsequent decode steps are corrupted (next-token max abs diff 24.3 above).

Why this matters

Vortex is the inference stack of Evo 2 (StripedHyena2). Speculative decoding requires verifying γ draft tokens in a single forward with initial inference state (block verification). Any caller who does this through the public API lands in sequential_forward and gets silently wrong logits — there is no signal that anything went wrong unless the caller independently cross-checks against per-token decoding.

We hit this while building EvSpark, lossless (DSpark-style) speculative decoding for Evo 2 7B; we disclose this exact pitfall in §3.1 of our paper and currently maintain an internal chunk-forward-with-initial-state path to work around it. A guard upstream would turn this silent correctness bug into an explicit, debuggable error for everyone.

Note that vortex's own Generator never triggers the multi-token case — decode feeds one token per step (x = x[:, -1:], vortex/model/generation.py L169) and prefill goes through parallel_forward — so raising on seqlen > 1 cannot break any first-party flow.

Suggested minimal guard

Fail loudly on the unsupported path (diff against main @ 8b00afe):

--- a/vortex/model/model.py
+++ b/vortex/model/model.py
@@ -328,9 +328,20 @@ class HyenaCascade(nn.Module):
     def sequential_forward(self, u, inference_params):
         if self.data_dtype is None:
             self.data_dtype = u.dtype
 
-        if len(u.shape) > 2:
+        if len(u.shape) > 2 and u.shape[1] > 1:
+            raise ValueError(
+                "HyenaCascade.sequential_forward received a multi-token input "
+                f"(seqlen={u.shape[1]}) after prefill, but chunk forward with "
+                "initial state is not implemented. Previously this input was "
+                "silently truncated to its last token (`u = u[:, -1]`), "
+                f"discarding {u.shape[1] - 1} token(s), returning a wrong "
+                "output for the surviving position, and corrupting the "
+                "recurrent/conv state for all subsequent tokens."
+            )
+        if len(u.shape) > 2:
             u = u[:, -1]
 
         z_pre, fir_state = self.engine.step_fir(

(A warnings.warn variant would preserve backward compatibility, but since every seqlen > 1 result on this path is guaranteed wrong, raising seems safer.)

Longer term, a proper chunk forward with initial state (parallel FIR conv spliced with the cached window; IIR zero-state FFT conv plus initial-state decay term) would enable speculative-decoding-style workloads directly — happy to discuss or contribute. In the meantime, the guard above is the minimal fix.

Happy to open a PR with the guard if that helps.