Apply max_length truncation on HF assistant-mask path

[guan404ming]

Purpose

Truncate text-only inputs on the HF assistant-mask path to max_length, which it previously ignored, making --seq-length a no-op for text models using HF masks. Multimodal inputs stay untruncated since vLLM re-tokenizes them from messages.

Tests

test_preprocess_batch_truncation_hf_mask_path → passed. Fails before the fix, passes after; test_preprocess_batch_multimodal still passes.

Checklist

I have filled in:

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan/results, such as providing test command and pasting the results.
• (Optional) The necessary documentation update.
• I (a human) have written or reviewed the code in this pr to the best of my ability.

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 5d8c319c-760c-424f-a3b6-97adb9ea21f5

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

In _get_input_ids_loss_mask, the HF assistant-token-mask path now conditionally builds truncation_kwargs: empty for ProcessorMixin processors (to avoid multimodal retokenization desync), and {max_length, truncation: True} for text-only processors. A new integration test verifies that _preprocess_batch truncates both input_ids and loss_mask when assistant_pattern=None.

Changes

HF assistant-mask path truncation

Layer / File(s)	Summary
Conditional truncation kwargs and integration test src/speculators/data_generation/preprocessing.py, tests/integration/datagen/test_preprocessing.py	_get_input_ids_loss_mask selects empty or populated truncation_kwargs based on whether the processor is a ProcessorMixin, then expands them into apply_chat_template. test_preprocess_batch_truncation_hf_mask_path confirms input_ids and loss_mask are both capped at max_length on the HF mask path.

Possibly related PRs

• vllm-project/speculators#495: Directly overlaps — refactored preprocessing to use ProcessorLike and originally implemented _get_input_ids_loss_mask with the HF assistant-masking logic that this PR extends with conditional truncation.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 75.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title directly and specifically describes the main change: applying max_length truncation to the HF assistant-mask path, which aligns with the core fix in the changeset.
Description check	✅ Passed	The description clearly explains the purpose of the fix, the specific issue it addresses, and includes test results demonstrating the fix works correctly.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (1)

tests/integration/datagen/test_preprocessing.py (1)

672-688: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add a _supports_assistant_mask skip guard to avoid a spurious failure.

This test selects the HF mask path (assistant_pattern=None) but only guards on chat_template. If the processor has a chat template yet does not emit a functional assistant_masks/assistant_mask, _get_input_ids_loss_mask raises and _preprocess_batch swallows it (continue), leaving results["input_ids"] empty. The subsequent assert len(results["input_ids"]) == 1 then fails rather than skipping. The existing test_preprocess_batch_uses_hf_assistant_mask already guards for this case; mirror it here.

💚 Proposed guard

     if not hasattr(processor, "apply_chat_template") or processor.chat_template is None:
         pytest.skip("Processor does not support chat templates")
+
+    if not _supports_assistant_mask(processor):
+        pytest.skip("Processor does not support assistant token mask")

As per path instructions: "Ensure PyTest tests are clear, comprehensive, and cover edge cases ... Check that new code paths introduced in the PR are covered."

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/integration/datagen/test_preprocessing.py` around lines 672 - 688, The
test currently only guards on the presence of a chat_template but does not
verify that the processor supports assistant_mask functionality. When
assistant_pattern is set to None to use the HF mask path, the processor must
emit a functional assistant_mask, otherwise _get_input_ids_loss_mask will fail
silently and leave results["input_ids"] empty, causing the assertion to fail
instead of skipping. Add a second skip guard that checks
_supports_assistant_mask (mirroring the existing guard pattern used in
test_preprocess_batch_uses_hf_assistant_mask) to ensure the processor supports
both chat_template and assistant_mask before proceeding with the test.

Source: Path instructions

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@tests/integration/datagen/test_preprocessing.py`:
- Around line 672-688: The test currently only guards on the presence of a
chat_template but does not verify that the processor supports assistant_mask
functionality. When assistant_pattern is set to None to use the HF mask path,
the processor must emit a functional assistant_mask, otherwise
_get_input_ids_loss_mask will fail silently and leave results["input_ids"]
empty, causing the assertion to fail instead of skipping. Add a second skip
guard that checks _supports_assistant_mask (mirroring the existing guard pattern
used in test_preprocess_batch_uses_hf_assistant_mask) to ensure the processor
supports both chat_template and assistant_mask before proceeding with the test.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: a24dfaed-5c01-4564-98b8-0842de5d543a

📥 Commits

Reviewing files that changed from the base of the PR and between 03ba008 and 0d646ca.

📒 Files selected for processing (2)

• src/speculators/data_generation/preprocessing.py
• tests/integration/datagen/test_preprocessing.py

[fynnsu]

Multimodal inputs stay untruncated since vLLM re-tokenizes them from messages.

It seems like in the fallback case we are currently truncating the messages?

[guan404ming]

Thanks for catching that. Yes, the fallback path truncates MM input_ids via processor_kwargs, which desyncs from vLLM's re-tokenization of full messages (the prompt_token_ids != token_ids check in extract_output would reject it).

I think that's a latent bug. Gate it the same way so only the text-only branch truncates is better. Should I fold this fix into the current PR, or keep it scoped and open a separate one?

[WindChimeRan]

#652

I think these are related. Will double check the behavior consistency.

[orestis-z]

LGTM. Clean fix for a real bug — the HF assistant-mask path was silently ignoring max_length. Agree with @coderabbitai’s suggestion to add the _supports_assistant_mask skip guard to the new test. Recommend approving once that’s addressed.

🤖 Generated with Claude Code using the /pr-review skill

[WindChimeRan]

This PR is on the legacy HF path, which will be superseded by #794 . We want single-source of the truth from vllm.

[speculatorsbot]

LGTM. Correct fix — the ProcessorMixin gate is consistent with the fallback path, and the regression test validates both input_ids and loss_mask truncation. Outstanding items from earlier reviews (the _supports_assistant_mask skip guard and the #794 supersession question) are the right things to resolve before merging. Recommend approving once addressed.

🤖 Generated with Claude Code using the /pr-review skill

[shanjiaz]

Thanks for putting up this fix! We're deprecating this path way tho. Please take a look at #794 and let us know if you have other concerns/recommendations.