Lfm2: thread seq_idx through ShortConv for packed/varlen inputs

[ChangyiYang]

What this PR fixes

Lfm2ShortConv.cuda_kernels_forward calls causal_conv1d_fn without
seq_idx. When the model is run on packed/varlen inputs — e.g. multiple
samples flattened into a single sequence with boundaries marked via
cu_seqlens (derived from position_ids), as in any sequence-packed
training pipeline — the 1-D conv slides across sample boundaries and the
conv state of one sample leaks into the next.

Attention is already isolated per-sample (FlashAttention / cu_seqlens),
but conv state is not. This causes a per-token logprob divergence
between packed-and-unpacked forwards of the same sample, observable as
a higher k3_kl when comparing the actor's packed forward against an
inference engine's per-sample forward.

This is the standard pattern already used by bamba, mamba2,
qwen3_next, qwen3_5(_moe), falcon_h1, zamba2, nemotron_h,
granitemoehybrid, …: pass a per-token int tensor seq_idx (segment
i labels its own tokens with i) to causal_conv1d_fn; the CUDA
kernel resets conv state at every boundary.

Changes

• Lfm2ShortConv.cuda_kernels_forward and Lfm2ShortConv.forward accept
  a new seq_idx: torch.IntTensor | None = None kwarg and forward it to
  causal_conv1d_fn.
• Lfm2DecoderLayer.forward passes seq_idx=kwargs.get("seq_idx") to
  self.conv(...) so callers can supply it through the model forward
  kwargs (same convention as qwen3_next).
• Backward-compatible: seq_idx=None (the default) preserves the prior
  behaviour exactly; the causal_conv1d_fn kernel treats None as "no
  boundaries". No change for unpacked / single-sequence inference.
• LFM2-MoE inherits via class Lfm2MoeShortConv(Lfm2ShortConv): pass,
  so the change benefits LFM2-MoE too.

Note on modular conversion

The change to the parent class body in modular_lfm2.py did not
propagate into the generated modeling_lfm2_moe.py when running
utils/modular_model_converter.py locally (the subclass body in
modular_lfm2_moe.py is pass). I applied the same change manually in
modeling_lfm2_moe.py to keep modeling and modular intent in sync.
Happy to instead extend modular_lfm2_moe.py with explicit method
overrides if reviewers prefer that, or to dig into the converter.

Caller usage (example)

# Build seq_idx for a flattened batch with cu_seqlens (LongTensor[N+1]):
seq_idx = torch.zeros(total_tokens, dtype=torch.int32, device=device)
for i in range(len(cu_seqlens) - 1):
    seq_idx[cu_seqlens[i]:cu_seqlens[i+1]] = i
outputs = model(input_ids=flat, position_ids=pos_ids, seq_idx=seq_idx)

🤖 Generated with Claude Code

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: lfm2, lfm2_moe

[github-actions]

CI Dashboard: View test results in Grafana

[Rocketknight1]

Yes, this looks good! It matches Qwen, and I trust the Qwen logic because it has a lot of usage

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.