The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Hey! I don't get how you want to run sdpa when it does not have attention sinks backed in it? You won't ever get good results.

ArthurZucker Thanks for your hard work.

Sorry everyone I trusted Claude too much

Sorry everyone I trusted Claude too much

人之常情，Keep going!

[For maintainers] Suggested jobs to run (before merge)

run-slow: deepseek_v4

Checked batch version!

Hey! I don't get how you want to run sdpa when it does not have attention sinks backed in it? You won't ever get good results.

Hi sorry, the sdpa test was just comparing the two mask forms apples-to-apples on a non-sink-aware backend, but anyway please ignore sdpa, the main point I want to make, is that the [S, S*top_k] mask is unnecessary; [S, compressed_len] is math-identical with better performance(both speed and mem).

Current (DeepseekV4CSACompressor.forward ):

gathered = flat_kv.index_select(0, flat_indices).view(batch, 1, -1, self.head_dim)
block_bias = gathered.new_full((batch, 1, seq_len, seq_len, top_k), float("-inf"))
allowed = torch.where(valid, gathered.new_zeros(()), gathered.new_full((), float("-inf")))
block_bias[:, 0, query_indices, query_indices, :] = allowed
return gathered, block_bias.view(batch, 1, seq_len, seq_len * top_k)

Alt:

safe_indices = torch.where(valid, top_k_indices, torch.full_like(top_k_indices, compressed_len))
block_bias = compressed_kv.new_full((batch, 1, seq_len, compressed_len + 1), float("-inf"))
block_bias.scatter_(-1, safe_indices.unsqueeze(1), 0.0)
return compressed_kv, block_bias[..., :compressed_len]

Each query's softmax has the same top_k non--inf entries either way (same indices into the compressed_kv axis, same dot-product values). Only the row width differs:
dense S*top_k mostly filled with -inf vs compressed_len with top_k nonzeros.

Regarding the issue i mentioned, its confirmed, not opened by me, flagged by another, but ive been looking at it, the paper and current implementation

A brief summary is this for the current forward cell block this is occuring:
Block: [E, F, G, PAD]
↓
Softmax over all 4 positions (including PAD)
↓
PAD contributes to compressed entry
↓
Polluted KV entry → incorrect attention

further details are in the issue itself, well v4 is after all a model unlike others, cant be helped that even after fixes issues happen

Seeing as theres also a bit of lack of diagrams in the very PR focused to solve v4 issues and i imagine we may need them, again, im attaching a few for everyones reference here too from the paper itself

Fig 1: CSA architecture

Fig 2: HCA architecture

This provides a more complete picture for whoever opens and uses this pr