hrm-text : add prefix-LM attention - #1
Open
bolgacg wants to merge 1 commit into
Open
Conversation
HRM-Text models are trained as prefix language models: prompt tokens attend to each other in both directions and only generated tokens are causal. The GGUF already carries the prefix_lm flag; this makes the KQ mask honour it. For these models the mask is built without the causal check, so a token attends to every token of its own sequence that is in the KV cache or in the current ubatch. Generation appends one token per sequence per step, so generated tokens still see only earlier tokens, which is the training-time mask. llama_context::decode warns when a batch is split into ubatches, because prompt tokens in different ubatches cannot see each other (use -ub at least the longest prompt). llama-server disables prompt caching for these models: a cached prompt prefix was computed without the tokens that follow it. llama_model_is_prefix_lm() is added next to llama_model_is_diffusion() so tools can ask. Measured on DAISY (592 Danish quiz questions from the group that trained DFM-Mimir, their prompt and exact-match scorer, greedy, same Q8_0 GGUF): causal only 5.6, this change 8.3, the official transformers implementation 8.4. With the same prompt bytes the answers are word-identical to the transformers output on 56 of 60 questions through llama-server and 58 of 60 through a standalone driver.
1 task
Owner
|
Thanks for the submission. We wait for the PR to be accepted and then we can add this also |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
HRM-Text models are trained as prefix LMs: prompt tokens attend to each other in both directions and only generated tokens are causal. This branch already reads
prefix_lmintohparams.hrm_prefix_lm; this PR makes the mask honour it, 35 lines across six files.For these models
set_input_kq_maskbuilds the mask without the causal check, so a token attends to every token of its own sequence in the KV cache or the current ubatch. Generation appends one token per sequence per step, so generated tokens still see only the past, which reproduces the training-time mask. Doing it at the mask level keeps every context type and the server on the same behaviour. Three guards come with it:llama_context::decodewarns when a batch is split into ubatches (prompt tokens in different ubatches cannot see each other; use-ubat least the longest prompt).llama-serverdisables prompt caching for these models: a cached prefix was computed without the tokens that follow it.llama_model_is_prefix_lm()is added next tollama_model_is_diffusion()so tools can check.Validation against the official transformers implementation (prefix attention), same Q8_0 GGUF, greedy decoding, on DAISY, the 592-question Danish benchmark from the model's own group (server run with
--reasoning offso the rendered prompt matches the HF chat template):With identical prompt bytes the server's answers are word-identical to transformers on 93% of the 592 questions; the causal build agrees on 29%. On 1,000-token prompts the scores stay equal (65.9 both, 99% of answers identical), and on the group's reading benchmark (MultiWikiQA, 512 rows) the patched build scores 65.8 against the published 66.8. Scripts and logged outputs: github.com/bolgacg/daisy-tools (tools/prefix-run; runs in results/).
Not covered: speculative decoding (draft tokens of one sequence in one batch would see each other).