[fix]Add logits_to_keep and shift_labels support for Qwen3-VL and Qwen3-VL-MoE (#1181)

# Summary

This PR adds `logits_to_keep` and `shift_labels` support for both
`Qwen3-VL` and `Qwen3-VL-MoE` in the Liger-patched forward path. The
change aligns the patched implementation with the expected Hugging Face
interface and enables selective logits materialization for long-context
inference.

# Testing Done

- `make test`
  - not fully green
- observed existing failures in `GRPO`, `fused_neighborhood_attention`,
and `gemma3` monkey patch tests
- `make test-convergence`
  - not fully green
- observed failure in
`test/convergence/bf16/test_mini_models_multimodal.py::test_mini_model_multimodal[mini_llama4-...]`
- `make checkstyle`
  - passed

Known limitation:

- The failed `make test` / `make test-convergence` cases above do not
directly exercise the `Qwen3-VL` or `Qwen3-VL-MoE` `logits_to_keep` /
`shift_labels` change in this PR

Author: luca-888

src/liger_kernel/transformers/model/qwen3_vl.py

@@ -33,6 +33,7 @@ def lce_forward(
     mm_token_type_ids: Optional[torch.IntTensor] = None,
     cache_position: Optional[torch.LongTensor] = None,
     second_per_grid_ts: Optional[torch.Tensor] = None,
+    logits_to_keep: Union[int, torch.Tensor] = 0,
     skip_logits: Optional[bool] = None,
     **kwargs,
 ) -> Union[Tuple, LigerQwen3VLCausalLMOutputWithPast]:
@@ -53,6 +54,9 @@ def lce_forward(
         The rope index difference between sequence length and multimodal rope.
     second_per_grid_ts (`torch.Tensor` of shape `(num_videos)`, *optional*):
         The time interval (in seconds) for each grid along the temporal dimension in the 3D position IDs.
+    logits_to_keep (`int` or `torch.Tensor`, *optional*):
+        If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
+        input tokens. If a `torch.Tensor`, it must contain the sequence indices to keep.
     Example:
     ```python
     >>> from PIL import Image
@@ -106,6 +110,8 @@ def lce_forward(
     )
 
     hidden_states = outputs[0]
+    slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+    kept_hidden_states = hidden_states[:, slice_indices, :]
 
     shift_labels = kwargs.pop("shift_labels", None)
     loss = None
@@ -121,7 +127,7 @@ def lce_forward(
 
     if skip_logits:
         result = LigerForCausalLMLoss(
-            hidden_states=hidden_states,
+            hidden_states=kept_hidden_states,
             lm_head_weight=self.lm_head.weight,
             labels=labels,
             shift_labels=shift_labels,
@@ -130,11 +136,17 @@ def lce_forward(
         )
         loss, _, token_accuracy, predicted_tokens = unpack_cross_entropy_result(result)
     else:
-        logits = self.lm_head(hidden_states)
+        logits = self.lm_head(kept_hidden_states)
 
         loss = None
-        if labels is not None:
-            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size)
+        if labels is not None or shift_labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                shift_labels=shift_labels,
+                vocab_size=self.config.text_config.vocab_size,
+                **kwargs,
+            )
 
     if not return_dict:
         output = (logits,) + outputs[1:]

src/liger_kernel/transformers/model/qwen3_vl_moe.py

@@ -34,11 +34,16 @@ def lce_forward(
     mm_token_type_ids: Optional[torch.IntTensor] = None,
     cache_position: Optional[torch.LongTensor] = None,
     second_per_grid_ts: Optional[torch.Tensor] = None,
+    logits_to_keep: Union[int, torch.Tensor] = 0,
     skip_logits: Optional[bool] = None,
     **kwargs,
 ) -> Union[Tuple, LigerQwen3VLMoeCausalLMOutputWithPast]:
     """
     Qwen3-VL-MoE forward with fused linear cross entropy support mirroring Qwen3-VL behaviour.
+
+    logits_to_keep (`int` or `torch.Tensor`, *optional*):
+        If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
+        input tokens. If a `torch.Tensor`, it must contain the sequence indices to keep.
     """
 
     output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
@@ -68,6 +73,8 @@ def lce_forward(
     )
 
     hidden_states = outputs[0]
+    slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+    kept_hidden_states = hidden_states[:, slice_indices, :]
 
     shift_labels = kwargs.pop("shift_labels", None)
     loss = None
@@ -83,7 +90,7 @@ def lce_forward(
 
     if skip_logits:
         result = LigerForCausalLMLoss(
-            hidden_states=hidden_states,
+            hidden_states=kept_hidden_states,
             lm_head_weight=self.lm_head.weight,
             labels=labels,
             shift_labels=shift_labels,
@@ -92,10 +99,16 @@ def lce_forward(
         )
         loss, _, token_accuracy, predicted_tokens = unpack_cross_entropy_result(result)
     else:
-        logits = self.lm_head(hidden_states)
-
-        if labels is not None:
-            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size)
+        logits = self.lm_head(kept_hidden_states)
+
+        if labels is not None or shift_labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                shift_labels=shift_labels,
+                vocab_size=self.config.text_config.vocab_size,
+                **kwargs,
+            )
 
     # Compute auxiliary load-balancing loss for MoE when requested
     aux_loss = None