[DeepSeekV4] Compressor does not seem to account for padding tokens when forming compressed KV blocks

[WKQ9411]

Description

I noticed a potential padding-related issue in the current DeepSeekV4 implementation.

The sliding-window attention path appears to correctly receive the user-provided attention_mask through create_sliding_window_causal_mask, so padded key/value tokens can be masked out in the normal attention path.

However, the HCA / CSA compressor path seems to form compressed KV entries from the projected kv / gate tensors without using the input attention_mask.

From my reading, the compressor first projects hidden states into kv and gate, then stores/chunks them through the cache layer, and finally performs softmax-gated pooling to produce compressed KV entries. This means padding tokens may still participate in compressed KV construction.

For example, with right padding:

attention_mask = [1, 1, 1, 0]
compress_rate = 4

the compressor may form one compressed block from:

[token_0, token_1, token_2, pad]

Even though the pad token is masked in the normal sliding-window attention mask, it may still contribute to the compressed KV entry.

Question

If padding is intended to be supported, should the compressor receive attention_mask and use it when forming compressed KV blocks?

Thanks!

[ArjunSrivastava1]

huh? another deepseekv4 issue?
i understand ya, but its a bit unclear rn
well, i suppose you should read through these too #45820, and #45758

these mention how there were problems in the query mask and how causality is not maintained, basically query q[t] where t -> limit could still retrieve or see info from out of the limit
which, was fixed just yesterday fully

ig it should for the attention masking, but ur point remains a bit vague, specify it further like, which stages do u refer to and where the values are supposed to be etc etc

[WKQ9411]

huh? another deepseekv4 issue?
i understand ya, but its a bit unclear rn
well, i suppose you should read through these too #45820, and #45758

these mention how there were problems in the query mask and how causality is not maintained, basically query q[t] where t -> limit could still retrieve or see info from out of the limit
which, was fixed just yesterday fully

ig it should for the attention masking, but ur point remains a bit vague, specify it further like, which stages do u refer to and where the values are supposed to be etc etc

Thanks, I understand that #45820 / #45758 fixed the query-to-compressed-slot visibility issue after compressed entries are created.

My concern is about an earlier stage: source-token-to-compressed-entry construction, especially in a padded batch.

For example, suppose:

compress_rate = 4

sample 0 attention_mask = [1, 1, 1, 1]
sample 0 tokens          = [A, B, C, D]

sample 1 attention_mask = [1, 1, 1, 0]
sample 1 tokens          = [E, F, G, PAD]

Since the batch tensor has shape "[2, 4]", the compressor may chunk both samples by physical positions and form:

sample 0 compressed_entry_0 = compress([A, B, C, D])
sample 1 compressed_entry_0 = compress([E, F, G, PAD])

The "block_bias" logic can control whether a query is allowed to attend to "compressed_entry_0", but it does not seem to prevent "PAD" from participating in the softmax-gated pooling that creates "sample 1 compressed_entry_0".

So I am not referring to whether a query can see future/out-of-limit compressed slots. I am asking whether padding tokens should be excluded when constructing compressed KV entries, or whether padded batches are intentionally unsupported for the compressor path.

[ArjunSrivastava1]

@Rocketknight1, i take it this has been looked into as well then?

comparing with the same pr #45892 , its isnt really mentioned there, i was studying the diagrams for attentions in the original paper just now, and yeah, this is a concern along with that other one, even more so cus this is a very short while ago

@WKQ9411 as for ur question... what i found from the paper is that padding is, indeed and explicitly mentioned

1. Padding is acceptable with <10% overhead
2. When compression spans ranks, padding entries are intentionally produced and placed at the tail
3. Block padding for cache alignment is explicitly endorsed (multiples of lcm(m, m'))

which means, the padded batches are meant to be supported, theyre used to ensure a lot of things - see above three, however for the latter part; theyre also to be excluded and not accounted for the compressed portion

will have to run some correctness checks on the implementation now ig

If padding is intended to be supported, should the compressor receive attention_mask and use it when forming compressed > KV blocks?

Yes, it is

[WKQ9411]

Thanks for checking the paper. That matches my concern.

[ArjunSrivastava1]

Confirmed and is a rather srs bug
went down hunting in the transformer and model files, 3 matches of compressor and it doesnt indeed, pass the attention mask as needed for csa - thats, the main feature for csa too, it does provide pad(only one match in the entire codebase)

file: modeling_deepseek_v4.py
lines of codes 816, 817, it doesnt handle the attention masking for padded cells at all
further reading up on things? well, the compressor interface might have to be changed a bit for the forward func, i can look into fixes but will prefer a maintianer cus honestly? touching something like on my own has chances of needing even further repairs, i can fix it up but not worth the risk

escalating to @ArthurZucker , seems we got another thing to add to the pr, i will mention the same in there

[puneetdixit200]

I’d like to investigate this. I can create a minimal padding-aware generation test for DeepSeekV4, inspect the compressed KV block formation, and patch the compressor to respect attention_mask semantics. Please assign this to me.