Add unfused SBHD compressed sparse attention - #6400
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
Signed-off-by: Deyu Fu <deyuf@nvidia.com>
c463cbe to
10ae87b
Compare
|
/ok to test 10ae87b |
|
/claude strict-review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| kv_t = kv_full.permute(1, 0, 2) | ||
|
|
||
| safe_indices = topk_indices.clamp(min=0).long() # [b, sq, topk] | ||
| safe_indices_exp = safe_indices.unsqueeze(-1).expand(-1, -1, -1, hn) # [b, sq, topk, hn] | ||
| # [b, n_kv, hn] -> [b, 1, n_kv, hn] -> gather -> [b, sq, topk, hn] | ||
| kv_gathered = torch.gather( | ||
| kv_t.unsqueeze(1).expand(-1, sq, -1, -1), dim=2, index=safe_indices_exp | ||
| ) |
There was a problem hiding this comment.
[CRITICAL Performance] torch.gather over an expanded input materializes a full [b, sq, n_kv, hn] gradient buffer in backward.
The forward is fine — expand is stride-0 and gather reads it strided. The backward is not: gather_backward is grad.new_zeros(self.sizes()).scatter_add_(dim, index, grad), and self here is the logical [b, sq, n_kv, hn] expanded view. So every backward through this kernel allocates a dense tensor of b · sq · n_kv · hn elements.
For a realistic CSA-4 layer at sq=4096, window=128, ratio=4 → n_compressed=1024, n_kv=5120, v_head_dim=128, b=1: 4096 · 5120 · 128 · 2 B ≈ 5.4 GB of transient allocation per layer per backward (×4 if the graph ever holds it in fp32). This will OOM long before the fp32 scores tensor does, and it scales as sq · n_kv rather than sq · topk.
The fix is already used elsewhere in this file — _compute_unfused_csa_non_compressed_lse (lines 292-308) flattens to batch-major and uses index_select, whose backward is an index_add into [b * n_kv, hn]. Applying the same shape here:
n_kv = kv_full.size(0)
# [n_kv, b, hn] -> [b, n_kv, hn] -> [b * n_kv, hn]
kv_flat = kv_full.permute(1, 0, 2).reshape(b * n_kv, hn)
batch_offsets = (
torch.arange(b, device=kv_full.device, dtype=torch.int64) * n_kv
).view(b, 1, 1)
safe_indices = topk_indices.clamp(min=0).long() + batch_offsets # [b, sq, topk]
# backward is an index_add into [b * n_kv, hn], not a [b, sq, n_kv, hn] zeros tensor
kv_gathered = kv_flat.index_select(0, safe_indices.reshape(-1)).view(b, sq, topk, hn)Everything downstream (kv_g, the two einsums, the invalid mask) is unchanged.
There was a problem hiding this comment.
Thanks — fixed in d02991d. The sparse-attention gather now flattens batch and KV position to [b * n_kv, hn], applies per-batch offsets, and uses index_select, so backward accumulates into the compact KV buffer instead of materializing [b, sq, n_kv, hn]. I added a CUDA regression test covering b=2, repeated indices, and -1 padding, with forward and query/KV/sink gradients checked against a loop oracle; it also explicitly rejects torch.gather. The exact-head full functional CI passed.
| weights = torch.softmax(score, dim=1, dtype=torch.float32).to(kv.dtype) | ||
| kv = (kv * weights).sum(dim=1) # [n_compressed, b, head_dim] | ||
|
|
||
| kv = self.norm(kv.to(x.dtype)) |
There was a problem hiding this comment.
[IMPORTANT Implementation] Pooling weights are computed in FP32 and then immediately thrown back to BF16 before the weighted sum.
torch.softmax(..., dtype=torch.float32) buys FP32 normalization, but .to(kv.dtype) discards it before the product and the reduction. For compress_ratio == 128 (non-overlapping, coff = 1) each weight is ~1/128 ≈ 0.0078; BF16's 8-bit mantissa gives ~0.4% relative error per weight, and the kv * weights product is also formed in BF16. sum(dim=1) does accumulate in FP32 internally on CUDA, but that does not recover precision already lost in the operands. This is the value path feeding the compressed KV that the whole sparse-attention denominator depends on, so the error propagates into both the student top-k and the teacher target.
Keep the reduction in FP32 and cast once:
weights = torch.softmax(score, dim=1, dtype=torch.float32)
kv = (kv.float() * weights).sum(dim=1).to(x.dtype) # [n_compressed, b, head_dim]
kv = self.norm(kv)This also removes the no-op kv.to(x.dtype) on line 484 (kv is already x.dtype there today).
There was a problem hiding this comment.
Thanks — fixed in d02991d. Compressor pooling now keeps the softmax weights, kv.float() * weights, and the reduction in FP32, then casts the pooled result once to x.dtype before normalization. The new BF16-input test checks the forward result and both KV/score gradients exactly against an FP32 pooling oracle. The exact-head full functional CI passed.
| csa_window_size: int = 128 | ||
| """Sliding window size for compressed sparse attention.""" | ||
|
|
||
| csa_compress_ratios: Optional[List[int]] = None | ||
| """Per-layer compress ratios, e.g. [0, 0, 4, 128, 4, 128, ...].""" | ||
|
|
||
| csa_compress_rotary_base: float = 40000.0 | ||
| """RoPE base for compressed KV positions in compressed sparse attention.""" | ||
|
|
||
| csa_dense_mode: bool = False |
There was a problem hiding this comment.
[IMPORTANT Compatibility] Two of these four new public config fields have no runtime reader in this PR.
Grepping the whole tree after this change:
csa_window_size— read atcsa.py:716. ✅csa_dense_mode— read atcsa.py:748. ✅csa_compress_ratios— read only intest_attention_variant_csa.py. Production code takescompress_ratioas an explicitCompressedSparseAttention.__init__argument and never indexes this list. TheCompressedSparseAttentiondocstring claims it does (see separate comment).csa_compress_rotary_base— zero readers anywhere, production or test. The rotary module is injected via therotary_pos_embargument, which this PR never constructs.
TransformerConfig is public API: once a field ships, its name, type, and default are effectively frozen, and 40000.0 in particular is a numeric default that a future consumer either honors or silently contradicts. Shipping it dead means the first PR that actually wires the rotary base inherits a name and default nobody validated against the DSv4 reference.
Please either move csa_compress_ratios / csa_compress_rotary_base into the PR that consumes them (#6405 per the scope section), or add an explicit TODO naming that PR so the dead fields are traceable:
# TODO(#6405): consumed by HybridModel layer construction, which selects the
# per-layer compress ratio and builds the compressed-KV rotary embedding.
# No runtime reader in this PR.
csa_compress_ratios: Optional[List[int]] = NoneThere was a problem hiding this comment.
Addressed in d02991d. I added an explicit TODO above these two fields. After the follow-on slices were recut, #6402—not #6405—is the actual consumer: it selects the per-layer compression ratio and constructs the compressed-KV rotary embedding. The comment also records that neither field has a production reader in this primitive-only PR.
| """Sparse core attention for CompressedSparseAttention. | ||
|
|
||
| Combines sliding window attention with compressed KV attention. The spec always | ||
| provides compressor and indexer submodule specs; this ``__init__`` inspects | ||
| ``config.csa_compress_ratios[layer_idx]`` and conditionally builds them: | ||
|
|
||
| * ``ratio == 0``: window-only (compressor and indexer NOT built) | ||
| * ``ratio == 4``: window + 4x compressed + learned Indexer (both built) | ||
| * ``ratio == 128``: window + 128x compressed, attend to all (compressor built only) | ||
| """ |
There was a problem hiding this comment.
[SUGGESTION Other] This docstring describes behavior the class does not have — a wrong docstring here is worse than none, because it points the reader at the wrong wiring contract.
Three inaccuracies:
- "this
__init__inspectsconfig.csa_compress_ratios[layer_idx]" — it does not.compress_ratioarrives as an explicit__init__argument (line 698) andcsa_compress_ratiosis never read in production code. - "
ratio == 128: window + 128x compressed, attend to all (compressor built only)" — the compressor is built for anyratio > 1(line 731), not specifically 128, and the indexer is built for exactlyratio == 4(line 747). The three enumerated cases read as an exhaustive dispatch table when the actual predicates areratio > 1andratio == 4. - "
ratio == 0: window-only" —ratio == 1also lands in window-only, and the default iscompress_ratio: int = 0. Theratio == 4case additionally depends onconfig.csa_dense_modebeing False, which isn't mentioned.
Suggested rewrite:
"""Sparse core attention for CompressedSparseAttention.
Combines sliding-window attention with compressed KV attention. The spec always
provides compressor and indexer submodule specs; which are actually built depends on
the ``compress_ratio`` passed by the caller:
* ``ratio <= 1``: window-only (neither compressor nor indexer is built).
* ``ratio > 1``: window + compressed KV via ``Compressor``.
* ``ratio == 4`` and not ``config.csa_dense_mode``: additionally builds ``CSAIndexer``
for learned top-k retrieval over compressed positions. Otherwise all causally
valid compressed positions are attended (dense over the compressed axis).
"""
There was a problem hiding this comment.
Addressed in d02991d. I rewrote the docstring to match the actual constructor contract: compress_ratio is caller-provided; ratios <= 1 are window-only; ratios > 1 build the compressor; and ratio 4 builds the indexer only when dense mode is disabled. It no longer claims to read config.csa_compress_ratios or presents 0/4/128 as an exhaustive dispatch table.
| self.n_local_heads = config.num_attention_heads | ||
|
|
||
| if softmax_scale is None: | ||
| softmax_scale = config.v_head_dim**-0.5 | ||
| self.softmax_scale = softmax_scale | ||
|
|
||
| # Learnable attention sink per head, kept in reference-checkpoint FP32. | ||
| self.attn_sink = mark_keep_in_fp32( | ||
| nn.Parameter(torch.zeros(self.n_local_heads, dtype=torch.float32)) | ||
| ) |
There was a problem hiding this comment.
[IMPORTANT Correctness] n_local_heads holds the global head count, and attn_sink is sized from it — so TP > 1 fails at an unrelated line with an opaque message.
config.num_attention_heads is the model-wide head count, not the TP-local one. Meanwhile unfused_compressed_sparse_attn does attn_sink.view(1, np_, 1, 1) (line 225), where np_ comes from query.size(2) — the TP-local head count. With tp_size == 2 this raises RuntimeError: shape '[1, 8, 1, 1]' is invalid for input of size 16 from inside the attention kernel, with nothing pointing at the head-count mismatch.
Two separate problems:
- Naming.
n_local_headsasserts the opposite of what the value is. Per the review guidance on parallel/distributed naming, this should be eithernum_attention_heads(what it is) or actually made local. - No guard. The PR scope is TP1, but nothing enforces it. A silent-until-cryptic failure at TP>1 is the worst outcome; either shard the parameter or reject the config up front.
Preferred (parameter is genuinely per-local-head, matching query):
tp_size = self.pg_collection.tp.size()
assert config.num_attention_heads % tp_size == 0, (
f"num_attention_heads ({config.num_attention_heads}) must be divisible by "
f"tensor-parallel size ({tp_size})"
)
self.num_local_attention_heads = config.num_attention_heads // tp_size
# Learnable attention sink per local head, kept in reference-checkpoint FP32.
self.attn_sink = mark_keep_in_fp32(
nn.Parameter(torch.zeros(self.num_local_attention_heads, dtype=torch.float32))
)Note that sharding attn_sink also needs tensor_model_parallel/partition_dim marking so it is checkpointed and all-reduced correctly. If that's out of scope for the unfused slice, the minimal alternative is an explicit rejection here:
if self.pg_collection.tp.size() > 1:
raise ValueError(
"CompressedSparseAttention does not support tensor parallelism yet "
f"(tp_size={self.pg_collection.tp.size()}); attn_sink is not sharded."
)Either way, a shape assertion at the query boundary in forward (assert np == self.attn_sink.numel()) would turn the remaining failure mode into a readable one.
There was a problem hiding this comment.
Addressed in d02991d. Since this native SBHD slice explicitly supports TP1 only, I chose the scoped rejection path rather than introducing partially sharded sink semantics. The constructor now rejects tp_size != 1 with a clear error, n_local_heads was renamed to num_attention_heads, and both the module boundary and direct attention helper validate the query/sink head count. Focused tests cover TP rejection, module-level head mismatch, and direct sink mismatch.
| def _normalize_indexer_teacher_target( | ||
| target: torch.Tensor, non_compressed_lse: torch.Tensor | None | ||
| ) -> torch.Tensor: | ||
| """L1-normalize teacher mass without changing the legacy DSA path.""" | ||
| if non_compressed_lse is None: | ||
| return dsa_indexer_loss.normalize_indexer_target(target) | ||
| return target / target.sum(dim=-1, keepdim=True).clamp_min(torch.finfo(torch.float32).tiny) | ||
|
|
||
|
|
There was a problem hiding this comment.
[IMPORTANT Implementation] The new branch silently changes the normalization floor from 1e-10 to 1.2e-38, and the zero-mass row it admits still produces a gradient.
normalize_indexer_target clamps with INDEXER_LOSS_EPS = 1e-10; this branch clamps with torch.finfo(torch.float32).tiny ≈ 1.18e-38. The looser floor is arguably more correct for CSA (with the full denominator, legitimate row masses can be far below 1e-10 without being degenerate, and 1e-10 would wrongly shrink them) — but the divergence is undocumented and the failure mode it opens is not handled.
In _compute_indexer_teacher_probabilities, each head's compressed mass is m_h = C_h / (W_h + C_h) where W_h = exp(non_compressed_lse). When the window+sink logits dominate — the expected regime early in training, before the compressor has learned anything — exp(masked_scores - full_lse) underflows: any compressed logit more than ~88 below full_lse becomes exactly 0.0 in FP32. If that holds for every head and every compressed key in a row, target.sum(dim=-1) is exactly 0, so 0 / tiny = 0 and the normalized target is all-zero.
That row is valid (it is not masked, query_valid_rows does not cover it), so bwd_fused_indexer_loss_naive line 872 computes grad = index_scores_softmax - 0 = predict, feeding a nonzero gradient toward a degenerate all-zero target. The KL term itself is 0, so the loss value gives no warning — this only shows up as unexplained indexer drift.
The underflow is avoidable entirely by staying in log space. The final target is mathematically sum_h m_h · softmax_h(compressed) / sum_h m_h, and m_h = sigmoid(compressed_lse_h - non_compressed_lse_h), which is stable for any finite input:
masked_scores = attention_scores.float().masked_fill(~expanded_valid_mask, float("-inf"))
compressed_lse = torch.logsumexp(masked_scores, dim=-1)
# m_h = C_h / (W_h + C_h) == sigmoid(log C_h - log W_h); no exp() underflow.
head_mass = torch.sigmoid(compressed_lse - non_compressed_lse.float())
probabilities = torch.softmax(masked_scores, dim=-1) * head_mass.unsqueeze(-1)
return torch.where(expanded_valid_mask, probabilities, torch.zeros_like(probabilities))At minimum, please add a comment here explaining why this path deliberately does not use INDEXER_LOSS_EPS, so the next reader doesn't "fix" it back.
There was a problem hiding this comment.
Thanks — fixed in d02991d. The external-mass path now separates the conditional distribution over compressed keys from each head’s compressed mass, keeps the latter in log space, and applies one common per-row shift across heads when the absolute mass would underflow; that shift cancels in the final L1 normalization. This path deliberately does not use INDEXER_LOSS_EPS: any positive compressed mass is normalized, while an exactly zero row stays zero. The manual KL gradient is now predict * target.sum(-1) - target, so zero-mass rows produce zero gradient. Added regressions for a 2-head/2-key LSE≈1000 float64 oracle, fully masked finite-zero target/autograd, valid zero-mass manual gradients, and unchanged moderate/legacy behavior. The exact-head full functional CI passed.
| non_compressed_lse = _compute_unfused_csa_non_compressed_lse( | ||
| query, kv, self.attn_sink, window_idxs, self.softmax_scale | ||
| ) |
There was a problem hiding this comment.
[IMPORTANT Performance] The sliding-window logits are computed twice per training step per CSA-4 layer.
_compute_unfused_csa_non_compressed_lse gathers the window KV and computes q · k over exactly the same (query, kv, window_idxs) triple that unfused_compressed_sparse_attn will recompute 50 lines later (line 904 → scores at line 218, whose first window_size columns are these same logits). The teacher only needs the row-wise LSE, which the main kernel already has all the inputs for.
At sq=4096, window=128, np=128, head_dim=128 that is an extra sq · window · np · head_dim ≈ 8.6 GFLOP per layer per step, in FP32, plus a second full gather of the window KV — on top of the redundant index_select. On a 60-layer DSv4 config with half the layers at ratio 4, that is a measurable fraction of attention time for a value that is a byproduct of work already being done.
Two ways out, either is fine:
- Have
unfused_compressed_sparse_attnoptionally return the per-row LSE restricted to the firstwindow_sizeindex columns (plus sink), and reorder so the sparse kernel runs before the indexer loss. This needs the topk concat to be split, so it's the more invasive option. - Cheaper: hoist the window-KV gather so both the teacher helper and the main kernel consume one
index_selectresult, and dropchunk_sizelooping in favor of the already-materialized gather.
If neither fits the unfused reference slice, please add a comment stating the duplication is accepted and that #6404 (fused kernels) is expected to eliminate it — right now nothing signals that this is a known cost.
There was a problem hiding this comment.
Thanks — this optimization is intentionally deferred in the native reference path. I added an explicit TODO documenting the duplicate window-logit/gather work, that #6404’s fused training backend avoids it, and that the native fallback should later share the gathered window KV/logits. This PR keeps the correctness-first implementation unchanged.
| causal_mask = ( | ||
| torch.arange(n_compressed, device=x.device).unsqueeze(0).expand(sq, -1) | ||
| ) | ||
| positions = torch.arange(1, sq + 1, device=x.device).unsqueeze(1) | ||
| causal_mask = ( | ||
| torch.where(causal_mask >= positions // self.compress_ratio, float("-inf"), 0.0) | ||
| .unsqueeze(0) | ||
| .expand(b, -1, -1) | ||
| ) # [b, sq, n_compressed] | ||
|
|
||
| if self.training and torch.is_grad_enabled(): | ||
| q_indexer, k_indexer, weights_indexer = self.indexer.forward_before_topk( | ||
| x_det, qr_det | ||
| ) | ||
| indexer_loss_coeff = self.config.dsa_indexer_loss_coeff or 0.0 | ||
| key_for_loss = compressed_kv.unsqueeze(2).expand(-1, -1, np, -1) | ||
| weights_for_unfused = weights_indexer.float() * self.indexer.softmax_scale | ||
| non_compressed_lse = _compute_unfused_csa_non_compressed_lse( | ||
| query, kv, self.attn_sink, window_idxs, self.softmax_scale | ||
| ) | ||
| topk_indices_compressed, indexer_loss = FusedDSAIndexerLoss.apply( | ||
| q_indexer, | ||
| weights_for_unfused, | ||
| k_indexer, | ||
| query.detach(), | ||
| key_for_loss.detach(), | ||
| self.softmax_scale, | ||
| min(self.indexer.index_topk, n_compressed), | ||
| indexer_loss_coeff, | ||
| causal_mask, | ||
| self.config.dsa_indexer_use_sparse_loss, | ||
| self.indexer.pg_collection, | ||
| None, | ||
| None, | ||
| None, | ||
| None, | ||
| self.config.calculate_per_token_loss, | ||
| True, | ||
| non_compressed_lse, | ||
| ) | ||
| if indexer_loss_coeff > 0: | ||
| DSAIndexerLossLoggingHelper.save_loss_to_tracker( | ||
| loss=indexer_loss, | ||
| layer_number=self.layer_number, | ||
| num_layers=self.config.num_layers + (self.config.mtp_num_layers or 0), | ||
| ) | ||
| else: | ||
| _, topk_indices_compressed = self.indexer(x_det, qr_det, mask=causal_mask) | ||
|
|
||
| n_valid_per_pos = positions // self.compress_ratio # [sq, 1] | ||
| valid = (topk_indices_compressed >= 0) & (topk_indices_compressed < n_valid_per_pos) | ||
| compress_topk_idxs = torch.where( | ||
| valid, topk_indices_compressed + offset, torch.tensor(-1, device=x.device) | ||
| ) |
There was a problem hiding this comment.
[SUGGESTION Simplification] Per-forward reconstruction of constants that only depend on (sq, n_compressed, ratio, b).
Three allocations here are invariant across every step of training for a fixed shape, yet are rebuilt on every layer's every forward:
causal_mask(lines 836-844): a[b, sq, n_compressed]FP32 tensor. Atsq=4096,n_compressed=1024that's a 16 MB allocation (stride-0 in batch, so 16 MB notb × 16 MB) plus four kernel launches, per layer per step.positions/n_valid_per_pos(lines 839, 885): tiny, but two more launches.torch.tensor(-1, device=x.device)(line 888): constructs a 0-dim tensor from a Python int, which is a host→device copy on the critical path.torch.whereaccepts a Python scalar for theotherargument, so this can just be-1.
The file already has the right pattern for this — _get_window_topk_idxs_cached / _get_compress_topk_idxs_cached. Adding a third @lru_cache'd helper would make the causal mask a lookup:
@lru_cache(maxsize=8)
def _get_compress_causal_mask_cached(
ratio: int, seqlen: int, n_compressed: int, device_str: str
) -> torch.Tensor:
"""Additive causal mask over compressed positions, [seqlen, n_compressed], -inf = invalid."""
compressed = torch.arange(n_compressed, device=device_str).unsqueeze(0)
positions = torch.arange(1, seqlen + 1, device=device_str).unsqueeze(1)
return torch.where(compressed >= positions // ratio, float("-inf"), 0.0)then in forward:
causal_mask = _get_compress_causal_mask_cached(
self.compress_ratio, sq, n_compressed, str(x.device)
).unsqueeze(0).expand(b, -1, -1) # [b, sq, n_compressed]and line 887-889 becomes:
compress_topk_idxs = torch.where(valid, topk_indices_compressed + offset, -1)Note that n_valid_per_pos would then need to come from a small cached arange too, rather than reusing the positions local — which is itself a readability improvement, since positions is currently defined for the mask and then reused 45 lines later for an unrelated validity check.
There was a problem hiding this comment.
Fixed in d02991d. The compressed causal mask and per-position valid counts now use bounded @lru_cache helpers keyed by shape/device inputs; the cached mask is batch-expanded as a view, and torch.where uses the Python scalar -1. Unit coverage verifies mask/count semantics and cache reuse.
|
@FDecaYed I think we can move the files under |
Signed-off-by: Deyu Fu <deyuf@nvidia.com>
|
/ok to test d02991d |
5 participants
Add unfused SBHD compressed sparse attention
Summary
Add the reference SBHD Compressed Sparse Attention implementation needed by DeepSeek-V4, while keeping ordinary DSA on the current-main implementation and API.
Related to #5795.
Scope
csa.py.ProcessGroupCollectionplumbing.Supported scope
BF16, eager, SBHD, fixed-length, TP1/PP1/CP1.
Non-goals
Review boundary and dependencies
#5929 and #5944 are merged in
main. This branch contains one residual commit:10ae87b793306371657c41ffc2ee4880c842907a.Please review the complete PR diff: 6 files, +2,460/-9.
Current-main ordinary DSA remains the source of truth. #6020 defines the main-first DSA versus dev-only CSA ownership boundary, while #5960 supplies the corrected full-denominator teacher semantics incorporated here. #5795 remains historical feature and test provenance, not the patch base.
Validation
Passed locally on the final changeset:
Signed-off-bytrailer;git diff --check;The macOS host has no usable Torch/CUDA runtime, so no local GPU/runtime result is claimed. Full functional CI passed on exact head
10ae87b793306371657c41ffc2ee4880c842907a: run 32114744957 (L1, 5 repeats, non-lightweight; 151 jobs succeeded, 8 expected skips, 0 failures, and the finalNemo_CICD_Testgate passed). No dependency or lock-file change is part of this commit.