Skip to content

hrm-text : add prefix-LM attention - #1

Open
bolgacg wants to merge 1 commit into
noctrex:hrm-text-supportfrom
bolgacg:hrm-text-prefix-lm
Open

hrm-text : add prefix-LM attention#1
bolgacg wants to merge 1 commit into
noctrex:hrm-text-supportfrom
bolgacg:hrm-text-prefix-lm

Conversation

@bolgacg

@bolgacg bolgacg commented Sep 6, 2026

Copy link
Copy Markdown

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_lm into hparams.hrm_prefix_lm; this PR makes the mask honour it, 35 lines across six files.

For these models set_input_kq_mask builds 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::decode warns when a batch is split into ubatches (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 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 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 off so the rendered prompt matches the HF chat template):

run exact match (%)
this branch, causal only 5.6
this branch plus this PR 8.3
official transformers 8.4
reported in the Mimir paper 9.6

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).

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.
@noctrex

noctrex commented Sep 6, 2026

Copy link
Copy Markdown
Owner

Thanks for the submission. We wait for the PR to be accepted and then we can add this also

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants