diff --git a/tests/pytorch/test_dsv4_attention.py b/tests/pytorch/test_dsv4_attention.py new file mode 100644 index 00000000000..7a82dc42ec3 --- /dev/null +++ b/tests/pytorch/test_dsv4_attention.py @@ -0,0 +1,498 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""DSA primitives and DSv4 layer wiring, including packed sequence boundaries.""" + +import pytest +import torch + +from transformer_engine.pytorch.attention.sparse_attention import ( + Compressor, + Indexer, + compressor, + dsa_cudnn, + dsa_rope, + dsv4_attention, +) + + +@pytest.mark.parametrize("variant", ["hca", "csa"]) +@pytest.mark.parametrize("head_dim", [512, 576]) +def test_dsa_sparse_attention_forward_backward(variant, head_dim): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0): + pytest.skip("cuDNN sparse attention requires SM100") + cudnn = pytest.importorskip("cudnn") + dsa = getattr(cudnn, "DSA", None) + if not all( + hasattr(dsa, operation) + for operation in ( + "sparse_attention_forward_wrapper", + "sparse_attention_backward_wrapper", + ) + ): + pytest.skip("cuDNN sparse attention APIs are unavailable") + torch.manual_seed(17) + ratio = 128 if variant == "hca" else 4 + lengths = [256, 128] + comp_lengths = [n // ratio for n in lengths] + total, total_comp = sum(lengths), sum(comp_lengths) + cu = torch.tensor([0, 256, total], device="cuda", dtype=torch.int32) + cu_comp = torch.tensor([0, comp_lengths[0], total_comp], device="cuda", dtype=torch.int32) + query = ( + torch.randn(total, 64, head_dim, device="cuda", dtype=torch.bfloat16) * 0.1 + ).requires_grad_() + local = ( + torch.randn(total, head_dim, device="cuda", dtype=torch.bfloat16) * 0.1 + ).requires_grad_() + compressed = ( + torch.randn(total_comp, head_dim, device="cuda", dtype=torch.bfloat16) * 0.1 + ).requires_grad_() + sink = torch.randn(64, device="cuda", requires_grad=True) + leaves = (query, local, compressed, sink) + reference_leaves = tuple(x.detach().float().requires_grad_() for x in leaves) + indices = None + if variant == "csa": + # Select alternating compressed entries; distinct global IDs expose + # mistakes at the second sequence's compressed offset. + indices = torch.full((total, 32), -1, device="cuda", dtype=torch.int32) + start = comp_start = 0 + for length, count in zip(lengths, comp_lengths): + for position in range(length): + ids = torch.arange( + comp_start, comp_start + (position + 1) // ratio, 2, device="cuda" + ) + indices[start + position, : ids.numel()] = ids + start += length + comp_start += count + + output = dsa_cudnn.sparse_attention( + *leaves, + cu, + cu_comp, + window_size=32, + ratio=ratio, + indices=indices, + max_compressed_seqlen=max(comp_lengths) if indices is None else None, + ) + q, k, c, s = reference_leaves + reference = [] + start = comp_start = 0 + for length, count in zip(lengths, comp_lengths): + kv = torch.cat((k[start : start + length], c[comp_start : comp_start + count])) + pos = torch.arange(length, device="cuda") + local_mask = (pos[None, :] <= pos[:, None]) & (pos[None, :] > pos[:, None] - 32) + comp_ids = torch.arange(count, device="cuda") + comp_mask = (comp_ids[None, :] + 1) * ratio <= pos[:, None] + 1 + if indices is not None: + comp_mask &= (comp_ids % 2 == 0)[None, :] + mask = torch.cat((local_mask, comp_mask), dim=-1) + logits = torch.einsum("thd,kd->thk", q[start : start + length], kv) * (head_dim**-0.5) + logits = logits.masked_fill(~mask[:, None, :], float("-inf")) + probabilities = torch.cat((logits, s[None, :, None].expand(length, -1, -1)), -1).softmax(-1) + reference.append(torch.einsum("thk,kd->thd", probabilities[..., :-1], kv[..., :512])) + start += length + comp_start += count + expected = torch.cat(reference) + torch.testing.assert_close(output.float(), expected, atol=1e-3, rtol=1e-2) + gradient = torch.randn_like(output) + output.backward(gradient) + expected.backward(gradient.float()) + for actual, ref in zip(leaves, reference_leaves): + assert actual.grad is not None and torch.isfinite(actual.grad).all() + # Relative L2 error tolerates BF16 reduction rounding near zero entries. + relative_error = (actual.grad.float() - ref.grad).norm() / ref.grad.norm().clamp_min(1e-8) + assert relative_error < 0.02, relative_error.item() + + +def test_dsv4_rejects_unsupported_metadata(): + with pytest.raises(ValueError, match="positive"): + dsa_cudnn.sparse_attention( + torch.empty(1, 64, 512), + torch.empty(1, 512), + torch.empty(1, 512), + torch.empty(64), + None, + None, + window_size=0, + ratio=4, + ) + with pytest.raises(ValueError, match="query"): + dsa_cudnn.sparse_attention( + torch.empty(1, 4, 512), + None, + None, + None, + None, + None, + window_size=32, + ratio=4, + ) + + +def test_dsv4_attention_rejects_fp8(monkeypatch): + from transformer_engine.pytorch.quantization import FP8GlobalStateManager + + monkeypatch.setattr(FP8GlobalStateManager, "is_fp8_enabled", lambda: True) + layer = dsv4_attention.DSv4Attention.__new__(dsv4_attention.DSv4Attention) + torch.nn.Module.__init__(layer) + layer.qkv_format, layer.hidden_size = "sbhd", 8 + with pytest.raises(NotImplementedError, match="BF16 only"): + layer(torch.empty(4, 1, 8)) + + +@pytest.mark.parametrize("heads", [1, 64]) +def test_dsa_triton_rope_batched_forward_backward(monkeypatch, heads): + """The packed Triton kernels must reset positions at each BSHD sequence.""" + if not torch.cuda.is_available(): + pytest.skip("Triton DSA RoPE requires CUDA") + if dsa_rope._triton_rope_module() is None: + pytest.skip("Triton DSL is unavailable") + torch.manual_seed(23) + batch, seq, dim, width = 2, 17, 512, 64 + x = torch.randn(batch, seq, heads, dim, device="cuda", dtype=torch.bfloat16) + (cos, sin), _ = dsa_rope.rotary_embeddings(seq, 4, width, 160000.0, x.device) + cu = torch.arange(batch + 1, device=x.device, dtype=torch.int32) * seq + gradient = torch.randn_like(x) + + eager = dsa_rope._apply_rotary_eager + monkeypatch.setattr( + dsa_rope, + "_apply_rotary_eager", + lambda *_: pytest.fail("Triton dispatch fell back to eager RoPE"), + ) + candidate_input = x.detach().requires_grad_() + candidate = dsa_rope.apply_rotary(candidate_input, cos, sin, cu) + candidate_grad = torch.autograd.grad(candidate, candidate_input, gradient)[0] + + reference_input = x.detach().requires_grad_() + reference = eager(reference_input, cos, sin) + reference_grad = torch.autograd.grad(reference, reference_input, gradient)[0] + + assert candidate.data_ptr() != candidate_input.data_ptr() + torch.testing.assert_close(candidate, reference, atol=2e-2, rtol=2e-2) + grad_delta = candidate_grad.float() - reference_grad.float() + assert grad_delta.abs().max() <= 0.03125 + assert grad_delta.norm() / reference_grad.float().norm() < 0.002 + + query_source = x.detach().requires_grad_() + query_input = query_source.clone() + query_ptr = query_input.data_ptr() + query = dsa_rope._apply_rotary_query(query_input, cos, sin, cu) + query_grad = torch.autograd.grad(query, query_source, gradient)[0] + + assert query.data_ptr() == query_ptr + torch.testing.assert_close(query, reference, atol=2e-2, rtol=2e-2) + torch.testing.assert_close(query_grad, reference_grad, atol=0.03125, rtol=0.002) + + +def test_dsa_rope_non_power_of_two_forward_backward(monkeypatch): + """An accepted width of 96 must fall back for both rotary entry points.""" + if not torch.cuda.is_available(): + pytest.skip("DSA rotary dispatch requires CUDA") + if dsa_rope._triton_rope_module() is None: + pytest.skip("Triton DSL is unavailable") + torch.manual_seed(41) + x = torch.randn(17, 2, 128, device="cuda", dtype=torch.bfloat16) + (cos, sin), _ = dsa_rope.rotary_embeddings(17, 4, 96, 160000.0, x.device) + cu = torch.tensor([0, 17], device=x.device, dtype=torch.int32) + monkeypatch.setattr( + dsa_rope, "_launch_triton", lambda *_: pytest.fail("Width 96 reached Triton") + ) + assert not dsa_rope._can_use_triton(x, cos, sin, cu) + gradient = torch.randn_like(x) + reference_input = x.detach().requires_grad_() + reference = dsa_rope._apply_rotary_eager(reference_input, cos, sin, cu) + reference_grad = torch.autograd.grad(reference, reference_input, gradient)[0] + for apply in (dsa_rope.apply_rotary, dsa_rope._apply_rotary_query): + candidate_input = x.detach().requires_grad_() + candidate = apply(candidate_input, cos, sin, cu) + candidate_grad = torch.autograd.grad(candidate, candidate_input, gradient)[0] + torch.testing.assert_close(candidate, reference) + torch.testing.assert_close(candidate_grad, reference_grad) + + +@pytest.mark.parametrize("index_heads", [32, 64]) +def test_dsa_indexer_backward(index_heads): + """Train the standalone Indexer and its Compressor against a dense KL reference.""" + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0): + pytest.skip("cuDNN indexer training requires SM100") + cudnn = pytest.importorskip("cudnn") + if not all( + hasattr(getattr(cudnn, "DSA", None), op) + for op in ( + "dense_indexer_score_recompute_wrapper", + "dense_attn_score_recompute_wrapper", + "dense_indexer_backward_wrapper", + ) + ): + pytest.skip("cuDNN dense indexer training APIs are unavailable") + torch.manual_seed(43) + seq, ratio, window = 256, 4, 32 + layer = Indexer(32, 16, 128, index_heads, 8, ratio, 1e-6, "cuda", torch.bfloat16, fused=True) + x = torch.randn(1, seq, 32, device="cuda", dtype=torch.bfloat16) + q_residual = torch.randn(1, seq, 16, device="cuda", dtype=torch.bfloat16) + cu = torch.tensor([0, seq], device=x.device, dtype=torch.int32) + cu_comp = torch.tensor([0, seq // ratio], device=x.device, dtype=torch.int32) + token_rope, compressed_rope = dsa_rope.rotary_embeddings(seq, ratio, 64, 160000.0, x.device) + ids, query, key, weights = layer( + x, + q_residual, + cu, + cu_comp, + token_rope, + compressed_rope, + total_comp=seq // ratio, + max_seqlen=seq, + return_context=True, + ) + assert ids.shape == (seq, 8) + qi, ki, wi = (t.unsqueeze(0) for t in (query, key, weights)) + qa = (torch.randn(1, seq, 64, 512, device=x.device, dtype=torch.bfloat16) * 0.1).float() + kc = (torch.randn(1, seq // ratio, 512, device=x.device, dtype=torch.bfloat16) * 0.1).float() + local_kv = (torch.randn(seq, 512, device=x.device, dtype=torch.bfloat16) * 0.1).float() + sink = torch.randn(64, device=x.device) + pos = torch.arange(seq, device=x.device) + valid = (torch.arange(seq // ratio, device=x.device) + 1) * ratio <= pos[:, None] + 1 + local_mask = (pos[None, :] <= pos[:, None]) & (pos[None, :] > pos[:, None] - window) + scale_i, scale_a = (index_heads * 128) ** -0.5, 512**-0.5 + comp_logits = (torch.einsum("bshd,bkd->bshk", qa, kc) * scale_a).masked_fill( + ~valid[None, :, None, :], -torch.inf + ) + local_logits = (torch.einsum("bshd,kd->bshk", qa, local_kv) * scale_a).masked_fill( + ~local_mask[None, :, None, :], -torch.inf + ) + # The dense teacher denominator includes local keys, ALL compressed keys, and sink. + full_lse = torch.logaddexp( + torch.logaddexp(local_logits.logsumexp(-1), comp_logits.logsumexp(-1)), + sink[None, None, :], + ).contiguous() + target = (comp_logits - full_lse[..., None]).exp().sum(2) + target = target / target.sum(-1, keepdim=True).clamp_min(1e-30) + index_score, attn_score = dsa_cudnn.dense_indexer_loss_scores( + qi, + ki, + wi, + qa.to(torch.bfloat16), + kc.to(torch.bfloat16), + full_lse, + ratio=ratio, + index_scale=scale_i, + attention_scale=scale_a, + ) + actual_target = attn_score["out"] / attn_score["denom"].unsqueeze(-1).clamp_min(1e-30) + # Masked score-buffer entries are not valid teacher probabilities. + actual_target = actual_target.masked_fill(~valid[None], 0) + torch.testing.assert_close(actual_target, target, atol=2e-4, rtol=2e-2) + + qr, kr, wr = (t.detach().float().requires_grad_() for t in (qi, ki, wi)) + logits = (torch.einsum("bshd,bkd->bshk", qr, kr).relu() * wr[..., None]).sum(2) * scale_i + logits = logits.masked_fill(~valid[None], -torch.inf) + logits = torch.where(valid.any(-1)[None, :, None], logits, 0) + log_prob = logits.log_softmax(-1) + actual_log_prob = index_score["out"] - index_score["denom"].unsqueeze(-1) + torch.testing.assert_close(actual_log_prob[0][valid], log_prob[0][valid], atol=2e-3, rtol=2e-2) + coefficient, upstream = 0.01, 0.7 + loss = ( + coefficient * upstream * (target * (target.clamp_min(1e-30).log() - log_prob)).sum() / seq + ) + reference_grads = torch.autograd.grad(loss, (qr, kr, wr)) + actual_grads = dsa_cudnn.dense_indexer_loss_backward( + qi, + ki, + wi, + index_score, + attn_score, + torch.tensor(upstream, device=x.device), + ratio=ratio, + index_scale=scale_i, + loss_coeff=coefficient, + ) + for actual, expected in zip(actual_grads, reference_grads): + assert actual.shape == expected.shape and torch.isfinite(actual).all() + error = (actual.float() - expected).norm() / expected.norm().clamp_min(1e-8) + assert error < 0.03, error.item() + parameters = tuple(layer.parameters()) + candidate = torch.autograd.grad((qi, ki, wi), parameters, actual_grads) + # TE Linear releases its backward workspace; compare using a fresh forward. + _, ref_q, ref_k, ref_w = layer( + x, + q_residual, + cu, + cu_comp, + token_rope, + compressed_rope, + total_comp=seq // ratio, + max_seqlen=seq, + return_context=True, + ) + reference = torch.autograd.grad( + (ref_q, ref_k, ref_w), + parameters, + tuple(g[0].to(t.dtype) for g, t in zip(reference_grads, (ref_q, ref_k, ref_w))), + ) + for actual, expected in zip(candidate, reference): + assert torch.isfinite(actual).all() and actual.float().norm() > 0 + error = (actual.float() - expected.float()).norm() / expected.float().norm().clamp_min(1e-8) + assert error < 0.03, error.item() + + +def test_dsv4_thd_compressed_prefixes(): + """Floor each sequence independently, retaining a safe allocation capacity.""" + for lengths, capacity in (([129, 257], 96), ([131, 257], 97)): + x = torch.empty(sum(lengths), 8) + cu = torch.tensor([0, lengths[0], sum(lengths)], dtype=torch.int32) + cu_comp, total_comp, maximum = dsv4_attention._packed_prefixes(x, cu, 257, 4) + assert cu_comp.tolist() == [0, 32, 96] + assert total_comp == capacity + assert maximum == 257 + + +def test_dsa_compressor_clears_unused_capacity(monkeypatch): + """Unused cuDNN output must not contaminate normalization weight gradients.""" + layer = Compressor.__new__(Compressor) + torch.nn.Module.__init__(layer) + layer.kv_proj = layer.gate_proj = torch.nn.Identity() + layer.kv_norm = torch.nn.RMSNorm(4) + layer.position_bias = torch.nn.Parameter(torch.zeros(4, 4)) + layer.projected, layer.ratio, layer.overlap, layer.weight_width = 4, 4, False, 0 + monkeypatch.setattr( + compressor, + "compress", + lambda *args, **kwargs: torch.cat((torch.ones(1, 4), torch.full((1, 4), float("nan")))), + ) + cu = torch.tensor([0, 3, 8], dtype=torch.int32) + cu_comp = torch.tensor([0, 0, 1], dtype=torch.int32) + output = layer(torch.empty(8, 4), cu, cu_comp, total_comp=2) + output.sum().backward() + assert torch.isfinite(output).all() + assert torch.isfinite(layer.kv_norm.weight.grad).all() + assert torch.count_nonzero(output[1]) == 0 + + +@pytest.mark.parametrize( + "prefix,maximum", + [ + (None, 8), + ([0, 8], 8), + (torch.tensor([0, 8], dtype=torch.int64), 8), + (torch.tensor([[0, 8]], dtype=torch.int32), 8), + (torch.tensor([0], dtype=torch.int32), 8), + (torch.tensor([0, 0, 8, 8], dtype=torch.int32)[::2], 8), + (torch.tensor([0, 8], dtype=torch.int32), None), + (torch.tensor([0, 8], dtype=torch.int32), 0), + ], +) +def test_dsv4_thd_rejects_malformed_metadata(prefix, maximum): + with pytest.raises(ValueError, match="THD requires"): + dsv4_attention._packed_prefixes(torch.empty(8, 4), prefix, maximum, 4) + + +def test_dsv4_thd_rejects_hidden_batch_axis_and_multi_sequence_context(): + # These checks precede projections and do not require GPU-backed TE layers. + layer = dsv4_attention.DSv4Attention.__new__(dsv4_attention.DSv4Attention) + torch.nn.Module.__init__(layer) + layer.qkv_format, layer.hidden_size = "thd", 8 + layer.is_csa, layer.compression_ratio = True, 4 + with pytest.raises(ValueError, match="selected format"): + layer(torch.empty(16, 1, 8)) + with pytest.raises(ValueError, match="THD requires"): + layer(torch.empty(16, 8)) + with pytest.raises(ValueError, match="batch=1"): + layer( + torch.empty(16, 8), + cu_seqlens_q=torch.tensor([0, 8, 16], dtype=torch.int32), + max_seqlen_q=8, + return_indexer_context=True, + ) + + +def test_dsa_packed_rope_eager_forward_backward(): + """Reset per-sequence positions and safely rotate unused compressed capacity.""" + torch.manual_seed(31) + cu = torch.tensor([0, 3, 8], dtype=torch.int32) + (cos, sin), _ = dsa_rope.rotary_embeddings(5, 1, 4, 12345.0, "cpu") + x = torch.randn(9, 2, 8, requires_grad=True) + reference_input = x.detach().clone().requires_grad_() + actual = dsa_rope.apply_rotary(x, cos, sin, cu) + expected = torch.cat( + [ + dsa_rope._apply_rotary_eager( + reference_input[start:end].unsqueeze(0), + cos[:, : end - start], + sin[:, : end - start], + ).squeeze(0) + for start, end in ((0, 3), (3, 8), (8, 9)) + ] + ) + gradient = torch.randn_like(actual) + actual.backward(gradient) + expected.backward(gradient) + torch.testing.assert_close(actual, expected) + torch.testing.assert_close(x.grad, reference_input.grad) + + +@pytest.mark.parametrize("variant,ratio", [("hca", 128), ("csa", 4)]) +@pytest.mark.parametrize("lengths", [[129, 257], [255, 383]]) +def test_dsv4_hybrid_thd_forward_backward(variant, ratio, lengths): + """Packed full-layer output and gradients match independent dense sequences.""" + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0): + pytest.skip("DSv4 requires SM100") + cudnn = pytest.importorskip("cudnn") + if not hasattr(getattr(cudnn, "DSA", None), "sparse_attention_forward_wrapper"): + pytest.skip("DSv4 requires cuDNN Frontend 1.29 sparse attention APIs") + torch.manual_seed(37) + layer = dsv4_attention.DSv4Attention( + hidden_size=32, + q_lora_rank=16, + layer_type=( + "compressed_sparse_attention" if variant == "csa" else "heavily_compressed_attention" + ), + head_dim=512, + rope_head_dim=64, + sliding_window=32, + compression_ratio=ratio, + o_groups=8, + o_lora_rank=8, + index_topk=8, + rope=( + dsa_rope._DSv4RotaryEmbedding(ratio, 64, 12345.0, "cuda") if variant == "csa" else None + ), + params_dtype=torch.bfloat16, + qkv_format="thd", + ) + x = torch.randn(sum(lengths), 32, device="cuda", dtype=torch.bfloat16, requires_grad=True) + cu = torch.tensor([0, lengths[0], sum(lengths)], device=x.device, dtype=torch.int32) + actual = layer(x, cu_seqlens_q=cu, max_seqlen_q=max(lengths)) + assert actual.shape == x.shape + gradient = torch.randn_like(actual) + actual.backward(gradient) + packed_grads = { + name: p.grad.clone() for name, p in layer.named_parameters() if p.grad is not None + } + layer.zero_grad(set_to_none=True) + layer.qkv_format = "sbhd" if variant == "hca" else "bshd" + batch_axis = 1 if layer.qkv_format == "sbhd" else 0 + reference_input = x.detach().clone().requires_grad_() + parts = [] + for part in reference_input.split(lengths): + if variant == "csa": + output, context = layer(part.unsqueeze(batch_axis), return_indexer_context=True) + assert all(context[name].requires_grad for name in ("index_q", "index_k", "index_w")) + assert not any( + context[name].requires_grad + for name in ("attn_q", "local_kv", "compressed_kv", "sink") + ) + else: + output = layer(part.unsqueeze(batch_axis)) + parts.append(output.squeeze(batch_axis)) + expected = torch.cat(parts) + expected.backward(gradient) + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + for packed_grad, reference_grad in [(x.grad, reference_input.grad)] + [ + (packed_grads[name], p.grad) for name, p in layer.named_parameters() if name in packed_grads + ]: + assert reference_grad is not None + assert torch.isfinite(packed_grad).all() and torch.isfinite(reference_grad).all() + error = (packed_grad.float() - reference_grad.float()).norm() + assert error / reference_grad.float().norm().clamp_min(1e-8) < 0.03 diff --git a/transformer_engine/common/triton/dsa_rope.py b/transformer_engine/common/triton/dsa_rope.py new file mode 100644 index 00000000000..b7876e8220e --- /dev/null +++ b/transformer_engine/common/triton/dsa_rope.py @@ -0,0 +1,253 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""DSA RoPE kernels copied from NVIDIA/Megatron-LM. + +Source: megatron/core/fusions/fused_mla_yarn_rope_apply.py at +c4129afff4867ca89b1fb7b4b911128424cdfbe8. Kernel bodies and autotuning are +retained; TE owns tensor preparation and autograd. +""" + +import triton +import triton.language as tl + + +@triton.jit +def _get_thd_token_idx(cu_seqlens, pid_m, seq_num, cp_rank, cp_size): + # Cast ``pid_m`` and ``cu_seqlens`` loads to a single shared dtype so + # the loop-body reassignments don't surface as + # "initial value is int32 but redefined as int64" in newer Triton + # versions (which promote ``// Python_int`` to int64). + pid_m = pid_m.to(tl.int64) + token_idx = tl.full((), -1, dtype=tl.int64) + this_seq_len = tl.full((), 0, dtype=tl.int64) + seq_idx = 0 + last_cum_seqlen = tl.load(cu_seqlens).to(tl.int64) // cp_size + while seq_idx < seq_num: + cur_cum_seqlen = tl.load(cu_seqlens + seq_idx + 1).to(tl.int64) // cp_size + if token_idx == -1 and cur_cum_seqlen > pid_m: + token_idx = pid_m - last_cum_seqlen + this_seq_len = cur_cum_seqlen - last_cum_seqlen + last_cum_seqlen = cur_cum_seqlen + seq_idx += 1 + # Padding tokens beyond cu_seqlens[-1] (from THD CUDA-graph padding) + # never match any sequence, leaving token_idx == -1. Clamp to 0 so + # the cos/sin table loads stay in-bounds; the wrong RoPE result is + # harmless because padding positions are excluded by loss_mask. + if token_idx == -1: + token_idx = tl.full((), 0, dtype=tl.int64) + if cp_size > 1: + if token_idx < this_seq_len // 2: + token_idx = token_idx + cp_rank * this_seq_len // 2 + else: + token_idx = (token_idx - this_seq_len // 2) + ( + 2 * cp_size - cp_rank - 1 + ) * this_seq_len // 2 + return token_idx + + +@triton.jit +def _mla_rope_fwd_inplace_kernel( + Q, + COS, + SIN, + nope_dim, + emb_dim: tl.constexpr, + head_num: tl.constexpr, + batch_size, + seq_num, + cu_seqlens_q, + position_ids, + stride_x_seq, + stride_x_nheads, + stride_cos_seq, + stride_sin_seq, + cp_rank, + cp_size, + INVERSE: tl.constexpr, + REMOVE_INTERLEAVING: tl.constexpr, + BLOCK_H: tl.constexpr, +): + """ + Forward pass: apply RoPE inplace to the trailing emb_dim elements. + Reads from interleaved layout, writes back to interleaved layout. + + Input: + Q: [seq_len, batch_size, head_num, nope_dim + emb_dim] + or [total_seq_len, head_num, nope_dim + emb_dim] + COS/SIN: [max_seq_len, emb_dim] + + batch_size: batch size for sbhd format, not used for thd format + seq_num: number of sequences for thd format, not used for sbhd format + cu_seqlens_q: [seq_num + 1] accumulated sequence lengths for thd format + """ + pid_m = tl.program_id(axis=0) + pid_head = tl.program_id(axis=1) + + if position_ids is not None: + token_idx = tl.load(position_ids + pid_m) + elif cu_seqlens_q is None: + token_idx = pid_m // batch_size + else: + token_idx = _get_thd_token_idx(cu_seqlens_q, pid_m, seq_num, cp_rank, cp_size) + + cos_left = tl.load(COS + token_idx * stride_cos_seq + tl.arange(0, emb_dim // 2)) + sin_left = tl.load(SIN + token_idx * stride_sin_seq + tl.arange(0, emb_dim // 2)) + cos_right = tl.load( + COS + token_idx * stride_cos_seq + emb_dim // 2 + tl.arange(0, emb_dim // 2) + ) + sin_right = tl.load( + SIN + token_idx * stride_sin_seq + emb_dim // 2 + tl.arange(0, emb_dim // 2) + ) + if INVERSE: + sin_left = -sin_left + sin_right = -sin_right + cos_left = cos_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + sin_left = sin_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + cos_right = cos_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + sin_right = sin_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + + Q = Q + pid_m * stride_x_seq + pid_head * BLOCK_H * stride_x_nheads + + x_off = tl.arange(0, BLOCK_H)[:, None] * stride_x_nheads + nope_dim + mask = (pid_head * BLOCK_H + tl.arange(0, BLOCK_H))[:, None] < head_num + # x1 = t[..., 0::2], x2 = t[..., 1::2] + x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 + x_2_off = x_1_off + 1 + x_1 = tl.load(Q + x_1_off, mask=mask) + x_2 = tl.load(Q + x_2_off, mask=mask) + + x_left = x_1 * cos_left - x_2 * sin_left + x_right = x_2 * cos_right + x_1 * sin_right + + if REMOVE_INTERLEAVING: + tl.store(Q + x_1_off, x_left, mask=mask) + tl.store(Q + x_2_off, x_right, mask=mask) + else: + # The interleaved input and split output layouts alias. Finish all loads + # before any warp stores to the overlapping destination addresses. + tl.debug_barrier() + x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] + x_right_off = x_left_off + emb_dim // 2 + tl.store(Q + x_left_off, x_left, mask=mask) + tl.store(Q + x_right_off, x_right, mask=mask) + + +_autotuned_mla_rope_fwd_inplace_kernel = triton.autotune( + configs=[ + triton.Config({"BLOCK_H": 1}), + triton.Config({"BLOCK_H": 2}), + triton.Config({"BLOCK_H": 4}), + triton.Config({"BLOCK_H": 8}), + triton.Config({"BLOCK_H": 16}), + triton.Config({"BLOCK_H": 32}), + triton.Config({"BLOCK_H": 64}), + triton.Config({"BLOCK_H": 128}), + ], + key=["emb_dim", "head_num"], + restore_value=["Q"], +)(_mla_rope_fwd_inplace_kernel) + + +@triton.jit +def _mla_rope_bwd_inplace_kernel( + DO, + COS, + SIN, + nope_dim, + emb_dim: tl.constexpr, + head_num: tl.constexpr, + batch_size, + seq_num, + cu_seqlens_q, + position_ids, + stride_x_seq, + stride_x_nheads, + stride_cos_seq, + stride_sin_seq, + cp_rank, + cp_size, + INVERSE: tl.constexpr, + REMOVE_INTERLEAVING: tl.constexpr, + BLOCK_H: tl.constexpr, +): + """ + Backward pass: inverse RoPE inplace on the trailing emb_dim elements. + Reads from interleaved layout, writes to interleaved layout. + + Input: + DO: [seq_len, batch_size, head_num, nope_dim + emb_dim] + or [total_seq_len, head_num, nope_dim + emb_dim] + COS/SIN: [max_seq_len, emb_dim] + + batch_size, seq_num, and cu_seqlens_q are the same as in the forward pass + """ + pid_m = tl.program_id(axis=0) + pid_head = tl.program_id(axis=1) + + if position_ids is not None: + token_idx = tl.load(position_ids + pid_m) + elif cu_seqlens_q is None: + token_idx = pid_m // batch_size + else: + token_idx = _get_thd_token_idx(cu_seqlens_q, pid_m, seq_num, cp_rank, cp_size) + + cos_left = tl.load(COS + token_idx * stride_cos_seq + tl.arange(0, emb_dim // 2)) + sin_left = tl.load(SIN + token_idx * stride_sin_seq + tl.arange(0, emb_dim // 2)) + cos_right = tl.load( + COS + token_idx * stride_cos_seq + emb_dim // 2 + tl.arange(0, emb_dim // 2) + ) + sin_right = tl.load( + SIN + token_idx * stride_sin_seq + emb_dim // 2 + tl.arange(0, emb_dim // 2) + ) + if INVERSE: + sin_left = -sin_left + sin_right = -sin_right + cos_left = cos_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + sin_left = sin_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + cos_right = cos_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + sin_right = sin_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + + DO = DO + pid_m * stride_x_seq + pid_head * BLOCK_H * stride_x_nheads + + x_off = tl.arange(0, BLOCK_H)[:, None] * stride_x_nheads + nope_dim + mask = (pid_head * BLOCK_H + tl.arange(0, BLOCK_H))[:, None] < head_num + if REMOVE_INTERLEAVING: + x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 + x_2_off = x_1_off + 1 + x_left = tl.load(DO + x_1_off, mask=mask) + x_right = tl.load(DO + x_2_off, mask=mask) + else: + x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] + x_right_off = x_left_off + emb_dim // 2 + x_left = tl.load(DO + x_left_off, mask=mask) + x_right = tl.load(DO + x_right_off, mask=mask) + x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 + x_2_off = x_1_off + 1 + + x_1 = x_left * cos_left + x_right * sin_right + x_2 = -x_left * sin_left + x_right * cos_right + + if not REMOVE_INTERLEAVING: + # The split input and interleaved output layouts alias. Finish all loads + # before any warp stores to the overlapping destination addresses. + tl.debug_barrier() + tl.store(DO + x_1_off, x_1, mask=mask) + tl.store(DO + x_2_off, x_2, mask=mask) + + +_autotuned_mla_rope_bwd_inplace_kernel = triton.autotune( + configs=[ + triton.Config({"BLOCK_H": 1}), + triton.Config({"BLOCK_H": 2}), + triton.Config({"BLOCK_H": 4}), + triton.Config({"BLOCK_H": 8}), + triton.Config({"BLOCK_H": 16}), + triton.Config({"BLOCK_H": 32}), + triton.Config({"BLOCK_H": 64}), + triton.Config({"BLOCK_H": 128}), + ], + key=["emb_dim", "head_num"], + restore_value=["DO"], +)(_mla_rope_bwd_inplace_kernel) diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index 09f2a068332..2ea8384c2df 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -29,6 +29,12 @@ from transformer_engine.pytorch.module import destroy_ub from transformer_engine.pytorch.module import UserBufferQuantizationMode from transformer_engine.pytorch.attention import DotProductAttention +from transformer_engine.pytorch.attention import ( + Compressor, + DSv4Attention, + Indexer, + sparse_attention, +) from transformer_engine.pytorch.attention import GatedDeltaNetAttention from transformer_engine.pytorch.attention import GatedDeltaNet2Attention from transformer_engine.pytorch.attention import GatedDeltaProductAttention diff --git a/transformer_engine/pytorch/attention/__init__.py b/transformer_engine/pytorch/attention/__init__.py index f8c86dddede..4b145b28578 100644 --- a/transformer_engine/pytorch/attention/__init__.py +++ b/transformer_engine/pytorch/attention/__init__.py @@ -5,6 +5,8 @@ """Python interface for attention""" from .dot_product_attention import DotProductAttention +from . import sparse_attention +from .sparse_attention import Compressor, DSv4Attention, Indexer from .linear_attention import ( GatedDeltaNetAttention, GatedDeltaNet2Attention, @@ -17,6 +19,10 @@ __all__ = [ "DotProductAttention", + "DSv4Attention", + "Compressor", + "Indexer", + "sparse_attention", "GatedDeltaNetAttention", "GatedDeltaNet2Attention", "GatedDeltaProductAttention", diff --git a/transformer_engine/pytorch/attention/sparse_attention/README.md b/transformer_engine/pytorch/attention/sparse_attention/README.md new file mode 100644 index 00000000000..39532b9da66 --- /dev/null +++ b/transformer_engine/pytorch/attention/sparse_attention/README.md @@ -0,0 +1,131 @@ +# DSv4 attention API + +`DSv4Attention` is an experimental API in TE that composes cuDNN compression, +top-k indexing, and sparse attention kernels with Megatron's DSA RoPE kernels +through reusable building blocks such as `Compressor` and `Indexer`. cuDNN +supplies compression and sparse attention backward kernels +and separate indexer-loss score/backward primitives; top-k selection itself is +discrete. The layer maps hidden states to hidden states with the same shape. +`Compressor` and `Indexer` can also be composed into other attention layers; +`dsa_cudnn` exposes the underlying operations for prepared tensors. + +## Module hierarchy + +```text +DSv4Attention +├── TE query/local-KV projections + norms +├── Compressor +│ ├── TE KV/gate projections + learned position bias +│ ├── dsa_cudnn.compress -> cuDNN compressor +│ └── TE RMSNorm +├── Indexer (CSA only) +│ ├── TE index-query/head-weight projections +│ ├── own Compressor (index keys) +│ ├── dsa_rope.apply_rotary (query + index keys) +│ └── dsa_cudnn.select_indices -> cuDNN score + top-k +├── dsa_rope embeddings + rotary application +├── dsa_cudnn.sparse_attention -> cuDNN sparse attention +└── TE grouped output projection +``` + +## Inside DSv4Attention + +The following sketch expands the cuDNN calls underneath Compressor, Indexer, and attention. +Query and local KV are already projected, normalized, and rotated; tensors are contiguous CUDA BF16, with FP32 position bias/sink and CUDA INT32 sequence prefixes. + +```python +from transformer_engine.pytorch.attention.sparse_attention import ( + dsa_cudnn, +) + +pooled = dsa_cudnn.compress( + kv, gates, position_bias.float(), cu, cu_comp, + ratio=ratio, overlap=is_csa, total_comp=total_comp, +) +compressed_kv = ... # Apply RMSNorm and RoPE to compressed KV. + +indices = None +if is_csa: + index_pooled = dsa_cudnn.compress( + index_kv, index_gates, index_position_bias.float(), cu, cu_comp, + ratio=ratio, overlap=True, total_comp=total_comp, + ) + index_key = model.finish_index_key(index_pooled) + indices = dsa_cudnn.select_indices( + index_query, index_key, index_weights, cu, cu_comp, + top_k=top_k, ratio=ratio, max_seqlen=max_seqlen, + max_compressed_seqlen=max_compressed_seqlen, + scale=(index_head_dim * index_n_heads) ** -0.5, + ) + +output = dsa_cudnn.sparse_attention( + query, local_kv, compressed_kv, sink.float(), cu, cu_comp, + window_size=window_size, ratio=ratio, indices=indices, + max_compressed_seqlen=max_compressed_seqlen if indices is None else None, +) +output = ... # Apply inverse RoPE and grouped output projection. +``` + +`dsa_rope` supplies token/compressed-position embeddings and rotary application. +A caller may override the default embeddings with a `rope` module returning +token and compressed `(cos, sin)` pairs. + +## Usage + +CSA uses overlapping compression (see `Compressor`) and an Indexer to select +compressed KV entries. + +```python +import torch +from transformer_engine.pytorch.attention.sparse_attention import DSv4Attention + +csa = DSv4Attention( + hidden_size=4096, q_lora_rank=1536, + layer_type="compressed_sparse_attention", head_dim=512, rope_head_dim=64, + sliding_window=128, compression_ratio=4, o_groups=8, o_lora_rank=1024, + index_n_heads=32, index_topk=32, max_seqlen=2048, + params_dtype=torch.bfloat16, +) +output = csa(hidden_states) # [S, B, 4096] -> [S, B, 4096] +``` + +HCA uses non-overlapping compression and omits the Indexer: + +```python +hca = DSv4Attention( + hidden_size=4096, q_lora_rank=1536, + layer_type="heavily_compressed_attention", head_dim=512, rope_head_dim=64, + sliding_window=128, compression_ratio=128, o_groups=8, o_lora_rank=1024, + max_seqlen=2048, params_dtype=torch.bfloat16, +) +output = hca(hidden_states) # [S, B, 4096] -> [S, B, 4096] +``` + +Supported input tensor shapes are "batch-first" (shape `BSD`), "sequence-first" (shape `SBD`), +and packed (shape `TD`). Packed calls require contiguous CUDA INT32 +`cu_seqlens_q [B+1]` (from zero to `T`) and `max_seqlen_q`: + +```python +packed_csa = DSv4Attention( + hidden_size=4096, q_lora_rank=1536, + layer_type="compressed_sparse_attention", head_dim=512, rope_head_dim=64, + sliding_window=128, compression_ratio=4, o_groups=8, o_lora_rank=1024, + index_n_heads=32, index_topk=32, max_seqlen=2048, + params_dtype=torch.bfloat16, qkv_format="thd", +) +output = packed_csa(hidden_states, cu_seqlens_q=cu, max_seqlen_q=max_length) +``` + +## Note on support + +The current layer supports full-sequence BF16 forward/backward on SM100 only, +with 64 heads of width 512 and at least one complete compression window +per sequence. Sequences start at position zero; cache/decode, padding, +TP/SP/CP, and FP8 are unsupported. CSA and HCA include sliding-window attention; +standalone sliding-window-only layers are unsupported. + +Compression and attention support autograd; top-k selection is discrete. +CSA's `return_indexer_context=True` returns `(output, context)` for a +single-sequence, caller-owned indexer loss. Without that objective, the +indexer is not trained by the attention output alone. Pre-attention +normalization, residuals, and loss calculation remain model-owned. diff --git a/transformer_engine/pytorch/attention/sparse_attention/__init__.py b/transformer_engine/pytorch/attention/sparse_attention/__init__.py new file mode 100644 index 00000000000..f6cc5415a12 --- /dev/null +++ b/transformer_engine/pytorch/attention/sparse_attention/__init__.py @@ -0,0 +1,18 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Sparse attention. This feature is experimental and subject to change.""" + +from . import dsa_cudnn, dsv4_attention +from .compressor import Compressor +from .dsv4_attention import DSv4Attention +from .indexer import Indexer + +__all__ = [ + "dsa_cudnn", + "dsv4_attention", + "DSv4Attention", + "Compressor", + "Indexer", +] diff --git a/transformer_engine/pytorch/attention/sparse_attention/compressor.py b/transformer_engine/pytorch/attention/sparse_attention/compressor.py new file mode 100644 index 00000000000..edf2b73638d --- /dev/null +++ b/transformer_engine/pytorch/attention/sparse_attention/compressor.py @@ -0,0 +1,101 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Learned projections and normalization around cuDNN compression.""" + +import torch + +from transformer_engine.pytorch.module import Linear, RMSNorm + +from .dsa_cudnn import compress + + +class Compressor(torch.nn.Module): + """Project KV/gates, compress token windows with cuDNN, and apply RMSNorm. + + Non-overlapping compression combines each block of ``ratio`` tokens into + one compressed KV entry. Overlapping compression mixes the current block + with the entire preceding block to create one compressed KV entry, using + separate learned KV/gate projections for the current and preceding blocks. + Neighboring compressed KV entries therefore share an entire source block, + but use different projections of its tokens. + Weights are computed per KV channel as ``softmax(gates + position_bias)``. + + For ``ratio=2`` and ``[a b c d e f g h i j k]``: + + - Non-overlapping windows: ``[a b]``, ``[c d]``, ``[e f]``, ``[g h]``, ``[i j]``. + - Overlapping windows: ``[a b]``, ``[a b c d]``, ``[c d e f]``, + ``[e f g h]``, ``[g h i j]``. + + The first overlapping entry masks the missing previous block. Both modes emit + ``floor(length / ratio)`` entries per sequence; incomplete trailing blocks + (``k`` here) are excluded from compression. In DSv4Attention, these tokens + remain available through local sliding-window attention. + + Accept batch-major ``[B,S,D]`` or packed ``[T,D]`` hidden states with + token/compressed prefixes. Return ``[total_comp,width]`` rows before RoPE, + plus projected head weights when ``weight_width`` is nonzero in fused mode. + """ + + def __init__( + self, + hidden_size, + width, + ratio, + overlap, + eps, + device, + params_dtype, + *, + fused=False, + weight_width=0, + ): + super().__init__() + projected = width * (2 if overlap else 1) + kw = {"bias": False, "device": device, "params_dtype": params_dtype} + if fused: + # Fuse the KV, gate, and optional index-weight projections of x. + # Keep one contiguous weight for the shared projection call. + self.fused_proj = Linear(hidden_size, 2 * projected + weight_width, **kw) + else: + self.kv_proj = Linear(hidden_size, projected, **kw) + self.gate_proj = Linear(hidden_size, projected, **kw) + self.position_bias = torch.nn.Parameter( + torch.zeros(ratio, projected, device=device, dtype=params_dtype) + ) + self.kv_norm = RMSNorm(width, eps=eps, device=device, dtype=params_dtype) + self.ratio, self.overlap = ratio, overlap + self.projected, self.weight_width = projected, weight_width + + def forward(self, x, cu_seqlens, cu_seqlens_comp, *, total_comp): + """Project and compress KV rows, returning optional fused index weights.""" + if hasattr(self, "fused_proj"): + projected = self.fused_proj(x) + widths = ( + (self.projected, self.projected, self.weight_width) + if self.weight_width + else (self.projected, self.projected) + ) + parts = projected.split(widths, dim=-1) + kv, gate = parts[:2] + else: + kv, gate = self.kv_proj(x), self.gate_proj(x) + pooled = compress( + kv.reshape(-1, self.projected).contiguous(), + gate.reshape(-1, self.projected).contiguous(), + # cuDNN requires FP32 bias; the cast preserves gradients to the parameter. + self.position_bias.float(), + cu_seqlens, + cu_seqlens_comp, + ratio=self.ratio, + overlap=self.overlap, + total_comp=total_comp, + ) + if x.ndim == 2: + # cuDNN only guarantees valid rows below the final prefix. Clear + # unused capacity before RMSNorm so it cannot poison weight gradients. + valid = torch.arange(total_comp, device=x.device) < cu_seqlens_comp[-1] + pooled = pooled.masked_fill(~valid[:, None], 0) + compressed = self.kv_norm(pooled) + return (compressed, parts[2]) if self.weight_width else compressed diff --git a/transformer_engine/pytorch/attention/sparse_attention/dsa_cudnn.py b/transformer_engine/pytorch/attention/sparse_attention/dsa_cudnn.py new file mode 100644 index 00000000000..22256dac1cd --- /dev/null +++ b/transformer_engine/pytorch/attention/sparse_attention/dsa_cudnn.py @@ -0,0 +1,471 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""PyTorch wrappers for cuDNN DSA/CSA primitives.""" + +from importlib import import_module +from typing import Optional + +import torch +from torch.autograd.function import once_differentiable + +__all__ = [ + "compress", + "select_indices", + "sparse_attention", + "dense_indexer_loss_scores", + "dense_indexer_loss_backward", +] + + +def _namespace(name): + try: + return getattr(import_module("cudnn"), name) + except (ImportError, AttributeError) as exc: + raise ImportError( + f"Sparse attention requires nvidia-cudnn-frontend>=1.29.0 with the cudnn.{name}" + " namespace." + ) from exc + + +def _dsa_op(name): + try: + return getattr(_namespace("DSA"), name) + except AttributeError as exc: + raise ImportError( + f"Sparse attention requires cudnn.DSA.{name} from nvidia-cudnn-frontend>=1.29.0." + ) from exc + + +def _validate_tensors(reference, bf16, fp32, cu_seqlens, cu_seqlens_comp): + """Check metadata without copying packed sequence offsets to the CPU. + + Callers own prefix contents: start at zero, monotonic, final offsets equal + valid row counts, and compressed lengths floor(length / ratio). + """ + if not reference.is_cuda or torch.cuda.get_device_capability(reference.device) != ( + 10, + 0, + ): + raise ValueError("These cuDNN sparse attention kernels require an SM100 GPU.") + for tensor in (reference, *bf16, *fp32, cu_seqlens, cu_seqlens_comp): + if tensor.device != reference.device or not tensor.is_contiguous(): + raise ValueError("Inputs must be contiguous on the same CUDA device.") + if any(t.dtype != torch.bfloat16 for t in (reference, *bf16)): + raise TypeError("Activations must be BF16.") + if any(t.dtype != torch.float32 for t in fp32): + raise TypeError("Position biases and sinks must be FP32.") + for offsets in (cu_seqlens, cu_seqlens_comp): + if offsets.dtype != torch.int32 or offsets.ndim != 1 or offsets.numel() < 2: + raise ValueError("Sequence prefixes must be INT32 [batch+1].") + if cu_seqlens.shape != cu_seqlens_comp.shape: + raise ValueError("Local and compressed prefixes must describe the same batch.") + + +class _Compress(torch.autograd.Function): + @staticmethod + def forward(ctx, kv, score, ape, cu_seqlens, cu_seqlens_comp, ratio, overlap, total_comp): + """Run cuDNN compression and save its backward inputs.""" + coff = 2 if overlap else 1 + if kv.ndim != 2 or kv.shape[-1] % coff: + raise ValueError("kv must have shape [tokens, coff * head_dim].") + head_dim = kv.shape[-1] // coff + result = _namespace("CSA").csa_compressor_forward_wrapper( + kv, + score, + ape, + cu_seqlens, + cu_seqlens_comp, + ratio=ratio, + head_dim=head_dim, + coff=coff, + total_comp=total_comp, + ) + ctx.save_for_backward(kv, score, ape, cu_seqlens, cu_seqlens_comp) + ctx.config = ratio, head_dim, coff + return result["out"] + + @staticmethod + @once_differentiable + def backward(ctx, grad_out): + """Propagate gradients through cuDNN compression.""" + ratio, head_dim, coff = ctx.config + # cuDNN owns scratch allocation and dAPE zeroing. Autograd may supply + # a strided/expanded gradient, while this kernel requires contiguous data. + result = _namespace("CSA").csa_compressor_backward_wrapper( + *ctx.saved_tensors, + grad_out.contiguous(), + ratio=ratio, + head_dim=head_dim, + coff=coff, + ) + return ( + result["grad_kv"], + result["grad_score"], + result["grad_ape"], + None, + None, + None, + None, + None, + ) + + +def compress( + kv: torch.Tensor, + score: torch.Tensor, + ape: torch.Tensor, + cu_seqlens: torch.Tensor, + cu_seqlens_comp: torch.Tensor, + *, + ratio: int, + overlap: bool, + total_comp: int, +) -> torch.Tensor: + """Compress token windows with cuDNN forward and backward. + + kv/score: contiguous BF16 [T, coff*D]; ape: FP32 [ratio, coff*D]. + coff is 2 for overlapping compression, 1 otherwise. Sequence prefixes are + CUDA INT32 [B+1]; compressed segment lengths are floor(length / ratio). + Returns BF16 [total_comp, D]. Per-segment tails are dropped by cuDNN. + + total_comp specifies output capacity; only entries before + cu_seqlens_comp[-1] are valid. + """ + if ratio <= 0 or total_comp < 0: + raise ValueError("ratio must be positive and total_comp nonnegative.") + if kv.ndim != 2 or score.shape != kv.shape or ape.shape != (ratio, kv.shape[-1]): + raise ValueError("Expected kv/score [T,coff*D] and ape [ratio,coff*D].") + _validate_tensors(kv, (score,), (ape,), cu_seqlens, cu_seqlens_comp) + return _Compress.apply( + kv, + score, + ape, + cu_seqlens, + cu_seqlens_comp, + ratio, + overlap, + total_comp, + ) + + +@torch.no_grad() +def select_indices( + query: torch.Tensor, + compressed_key: torch.Tensor, + weights: torch.Tensor, + cu_seqlens: torch.Tensor, + cu_seqlens_comp: torch.Tensor, + *, + top_k: int, + ratio: int, + max_seqlen: int, + max_compressed_seqlen: int, + scale: Optional[float] = None, +) -> torch.Tensor: + """CSA selection using cuDNN's SM100 combined score + Top-K operation. + + Inputs are the indexer's projected Q/K/head weights. + BF16 query [T,H,128], compressed_key [Tc,128], weights [T,H]; H=32 or 64. + Full packed-sequence prefill only: causal timeline offsets start at zero. + Returns INT32 [T,top_k] indices into compressed_key's global packed rows; + -1 denotes an invalid slot. These are NOT IDs into local+compressed KV. + + Selection is discrete. Scale defaults to 1/sqrt(index head dim). If the + model also scales by 1/sqrt(index heads), fold that positive factor into + scale and pass raw BF16 projected weights. Scaling in FP32 then rounding + weights to BF16 can change close top-k boundaries. This helper does not + implement the separate indexer auxiliary loss or autograd through IDs. + max_compressed_seqlen is a compile bound and must cover the partial final + ratio; use ceil(max_seqlen / ratio) even though prefixes count only complete + compressed blocks. + """ + if query.ndim != 3 or query.shape[1] not in (32, 64) or query.shape[2] != 128: + raise ValueError("CSA index query must have shape [T,32 or 64,128].") + if compressed_key.ndim != 2 or compressed_key.shape[1] != 128: + raise ValueError("CSA index key must have shape [Tc,128].") + if weights.shape != query.shape[:2]: + raise ValueError("CSA index weights must have shape [T,H].") + if ratio <= 0 or top_k <= 0 or max_seqlen <= 0 or max_compressed_seqlen <= 0: + raise ValueError("CSA ratio, top_k and maximum sequence lengths must be positive.") + _validate_tensors(query, (compressed_key, weights), (), cu_seqlens, cu_seqlens_comp) + if scale is None: + scale = query.shape[-1] ** -0.5 + result = _dsa_op("indexer_forward_top_k_wrapper")( + query, + compressed_key.unsqueeze(1), + weights, + top_k=top_k, + ratio=ratio, + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens_comp, + max_seqlen_q=max_seqlen, + max_seqlen_k=max_compressed_seqlen, + precision="bf16", + topk_indices_global=True, + return_softmax=False, + sm_scale=scale, + ) + return result["indices"] + + +def _pad_index_heads(query, weights): + if query.ndim != 4 or query.shape[0] != 1 or query.shape[2] not in (32, 64): + raise ValueError("Dense indexer loss requires B1 [B,S,32 or 64,128] queries.") + if query.shape[-1] != 128 or weights.shape != query.shape[:3]: + raise ValueError("Dense indexer weights must match B1 index queries.") + if query.dtype != torch.bfloat16 or weights.dtype != torch.bfloat16: + raise TypeError("Dense indexer Q and weights must be BF16.") + pad = 64 - query.shape[2] + if pad: + return torch.nn.functional.pad(query, (0, 0, 0, pad)), torch.nn.functional.pad( + weights, (0, pad) + ) + return query, weights + + +def dense_indexer_loss_scores( + index_q, + index_k, + index_w, + attn_q, + compressed_kv, + full_lse, + *, + ratio: int, + index_scale: float, + attention_scale: float, +): + """Return cuDNN dense index/teacher score dictionaries for B1 BF16 prefill. + + full_lse is FP32 [1,S,64] and must include the local window, every eligible + compressed key, and the sink. The sparse-attention saved LSE covers only + selected keys and excludes the sink, so it cannot serve this dense teacher. + """ + if any(t.dtype != torch.bfloat16 for t in (index_k, attn_q, compressed_kv)): + raise TypeError("Dense indexer keys and attention Q/K must be BF16.") + if full_lse.dtype != torch.float32 or full_lse.shape != attn_q.shape[:3]: + raise ValueError("full_lse must be FP32 [B,S,attention_heads].") + padded_q, padded_w = _pad_index_heads(index_q, index_w) + index_score = _dsa_op("dense_indexer_score_recompute_wrapper")( + padded_q, + index_k.unsqueeze(2), + padded_w, + sm_scale=index_scale, + ratio=ratio, + ) + attn_score = _dsa_op("dense_attn_score_recompute_wrapper")( + attn_q, + compressed_kv.unsqueeze(2), + full_lse, + attention_scale, + ratio=ratio, + ) + return index_score, attn_score + + +def dense_indexer_loss_backward( + index_q, + index_k, + index_w, + index_score, + attn_score, + grad_loss, + *, + ratio: int, + index_scale: float, + loss_coeff: float, +): + """Return dQ/dK/dW for B1 BF16 mean-reduced KL, padding H32 at cuDNN. + + The tested scalar grad_loss/loss_coeff contract is Megatron's default + calculate_per_token_loss=False reduction; per-token reduction is untested. + """ + padded_q, padded_w = _pad_index_heads(index_q, index_w) + result = _dsa_op("dense_indexer_backward_wrapper")( + padded_q, + padded_w, + index_k, + attn_score["out"].clone(), + attn_score["denom"], + index_score["out"].clone(), + index_score["denom"], + grad_loss=grad_loss.contiguous(), + sm_scale=index_scale, + loss_coeff=loss_coeff, + ratio=ratio, + ) + heads = index_q.shape[2] + return ( + result["d_index_q"][:, :, :heads].contiguous(), + result["d_index_k"], + result["d_weights"][:, :, :heads].contiguous(), + ) + + +def _prepare_sparse_indices( + query: torch.Tensor, + cu_seqlens: torch.Tensor, + cu_seqlens_comp: torch.Tensor, + window_size: int, + ratio: int, + indices: Optional[torch.Tensor], + max_compressed_seqlen: Optional[int], +) -> torch.Tensor: + """Map local and compressed entries into one concatenated KV address space.""" + total_q = query.shape[0] + device = query.device + lengths = cu_seqlens[1:] - cu_seqlens[:-1] + seq_ids = torch.repeat_interleave( + torch.arange(lengths.numel(), device=device), + lengths, + output_size=total_q, + ) + token = torch.arange(total_q, device=device) + seq_start = cu_seqlens[seq_ids] + position = token - seq_start + + local_position = position[:, None] - torch.arange( + window_size - 1, + -1, + -1, + device=device, + ) + local = seq_start[:, None] + local_position + local = torch.where(local_position >= 0, local, -1) + + comp_start = cu_seqlens_comp[seq_ids] + comp_end = cu_seqlens_comp[seq_ids + 1] + if indices is None: + if max_compressed_seqlen is None: + raise ValueError("max_compressed_seqlen is required for HCA attention.") + comp_position = torch.arange(max_compressed_seqlen, device=device)[None, :] + compressed = comp_start[:, None] + comp_position + valid = (compressed < comp_end[:, None]) & ( + (comp_position + 1) * ratio <= position[:, None] + 1 + ) + else: + if indices.ndim != 2 or indices.shape[0] != total_q: + raise ValueError("CSA indices must have shape [total_tokens, top_k].") + compressed = indices + comp_position = compressed - comp_start[:, None] + valid = ( + (compressed >= comp_start[:, None]) + & (compressed < comp_end[:, None]) + & ((comp_position + 1) * ratio <= position[:, None] + 1) + ) + compressed = torch.where(valid, compressed + query.shape[0], -1) + # cuDNN accepts an arbitrary logical width and pads its internal kernel + # layout, so preserve only the model-visible local + compressed slots. + return torch.cat((local, compressed), dim=1).to(torch.int32).contiguous() + + +class _SparseAttention(torch.autograd.Function): + @staticmethod + def forward(ctx, query, kv, indices, sink, scale): + """Run cuDNN sparse prefill and save its backward inputs.""" + result = _dsa_op("sparse_attention_forward_wrapper")( + query, + kv, + indices, + attn_sink=sink, + softmax_scale=scale, + # Prefix LSE is only needed by the optional indexer-loss path. + indexer_topk=0, + ) + out, lse = result["out"], result["lse"] + ctx.save_for_backward(query, kv, out, lse, sink, indices) + ctx.scale = scale + return out + + @staticmethod + @once_differentiable + def backward(ctx, grad_out): + """Propagate output gradients through cuDNN sparse attention.""" + query, kv, out, lse, sink, indices = ctx.saved_tensors + result = _dsa_op("sparse_attention_backward_wrapper")( + query, + kv, + out, + grad_out.contiguous(), + lse, + sink, + indices, + softmax_scale=ctx.scale, + ) + return result["dq"], result["dkv"], None, result["d_sink"], None + + +def sparse_attention( + query: torch.Tensor, + local_kv: torch.Tensor, + compressed_kv: torch.Tensor, + sink: torch.Tensor, + cu_seqlens: torch.Tensor, + cu_seqlens_comp: torch.Tensor, + *, + window_size: int, + ratio: int, + indices: Optional[torch.Tensor] = None, + scale: Optional[float] = None, + max_compressed_seqlen: Optional[int] = None, +) -> torch.Tensor: + """Joint local + compressed sparse attention with autograd. + + Full packed-sequence prefill: query [T,H,D], local_kv [T,D], compressed_kv + [Tc,D], sink [H]. Q/KV have already had the model's norm/RoPE applied. + cu_seqlens describes both query and local KV; cu_seqlens_comp describes + valid compressed rows, excluding capacity padding. Values are the first + 512 channels of each shared KV row (Q/K width may also be 576). + + Local positions satisfy max(0,q-window_size+1) <= k <= q per sequence. + Compressed block j is eligible when (j+1)*ratio <= q+1. With indices (CSA), + use its selected, valid, causal compressed rows; without indices (HCA), + use every eligible compressed row. IDs refer to compressed_kv only. + + One cuDNN sparse-attention call normalizes the combined visible entries with + the sink and provides the KV-only LSE consumed by cuDNN backward. Output is + [T,H,512], before output unrotation/projection. Scale defaults to D**-0.5. + No dense attention mask is part of this interface. HCA must provide + max_compressed_seqlen to bound its dense compressed-history metadata without + reading a CUDA prefix on host. + """ + if query.ndim != 3 or query.shape[1] != 64: + raise ValueError("cuDNN sparse attention requires query [T,64,D].") + if query.shape[-1] not in (512, 576): + raise ValueError("attention head dimension must be 512 or 576.") + if local_kv.ndim != 2 or compressed_kv.ndim != 2: + raise ValueError("local_kv and compressed_kv must be 2D packed tensors.") + if local_kv.shape[-1] != query.shape[-1] or compressed_kv.shape[-1] != query.shape[-1]: + raise ValueError("Q and both KV tensors must have the same QK dimension.") + if local_kv.shape[0] != query.shape[0]: + raise ValueError("Full-sequence attention requires one local KV row per query.") + if sink.shape != (query.shape[1],): + raise ValueError("sink must have shape [64].") + if window_size <= 0 or ratio <= 0: + raise ValueError("window_size and ratio must be positive.") + if indices is None and (max_compressed_seqlen is None or max_compressed_seqlen < 0): + raise ValueError("HCA requires a nonnegative max_compressed_seqlen.") + _validate_tensors(query, (local_kv, compressed_kv), (sink,), cu_seqlens, cu_seqlens_comp) + if indices is not None: + if indices.dtype != torch.int32 or indices.device != query.device: + raise ValueError("CSA indices must be CUDA INT32 on the query device.") + if scale is None: + scale = query.shape[-1] ** -0.5 + selected = _prepare_sparse_indices( + query, + cu_seqlens, + cu_seqlens_comp, + window_size, + ratio, + indices, + max_compressed_seqlen, + ) + return _SparseAttention.apply( + query.contiguous(), + torch.cat((local_kv, compressed_kv), dim=0), + selected, + sink.contiguous(), + scale, + ) diff --git a/transformer_engine/pytorch/attention/sparse_attention/dsa_rope.py b/transformer_engine/pytorch/attention/sparse_attention/dsa_rope.py new file mode 100644 index 00000000000..26cf10b1fc6 --- /dev/null +++ b/transformer_engine/pytorch/attention/sparse_attention/dsa_rope.py @@ -0,0 +1,214 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Trailing-channel, interleaved DSA RoPE for full-sequence attention.""" + +from functools import lru_cache + +import torch + + +def rotary_embeddings(seq, ratio, width, theta, device): + """Return token and compressed-window (cos, sin) pairs in FP32.""" + inv_freq = theta ** (-torch.arange(0, width, 2, device=device, dtype=torch.float32) / width) + angles = torch.outer(torch.arange(seq, device=device, dtype=torch.float32), inv_freq) + cos = angles.cos().repeat_interleave(2, dim=-1)[None, :, None, :] + sin = angles.sin().repeat_interleave(2, dim=-1)[None, :, None, :] + window_positions = torch.arange(seq // ratio, device=device) * ratio + return (cos, sin), (cos[:, window_positions], sin[:, window_positions]) + + +class _DSv4RotaryEmbedding(torch.nn.Module): + """Cache shared FP32 token frequencies; compressed positions are window starts.""" + + def __init__(self, ratio, width, theta, device, max_seqlen=None): + super().__init__() + self.ratio, self.width, self.theta = ratio, width, theta + cos, sin = (None, None) + if max_seqlen is not None: + (cos, sin), _ = rotary_embeddings(max_seqlen, ratio, width, theta, device) + self.register_buffer("cos", cos, persistent=False) + self.register_buffer("sin", sin, persistent=False) + + def _apply(self, fn): + cos, sin = self.cos, self.sin + super()._apply(fn) + if cos is not None: + # Module.to(bfloat16) must not round the reusable frequencies. + self.cos = cos.to(device=self.cos.device) + self.sin = sin.to(device=self.sin.device) + return self + + def forward(self, seq, device): + """Return cached token and compressed-window frequency pairs.""" + if self.cos is None or self.cos.shape[1] < seq: + (self.cos, self.sin), _ = rotary_embeddings( + seq, self.ratio, self.width, self.theta, device + ) + n_comp = seq // self.ratio + return (self.cos[:, :seq], self.sin[:, :seq]), ( + self.cos[:, : n_comp * self.ratio : self.ratio], + self.sin[:, : n_comp * self.ratio : self.ratio], + ) + + +def _apply_rotary_eager(x, cos, sin, cu_seqlens=None): + """Rotate trailing interleaved pairs in FP32, then restore input dtype.""" + if x.ndim == 3: + if cu_seqlens is None: + raise ValueError("Packed RoPE requires cu_seqlens.") + rows = torch.arange(x.shape[0], device=x.device) + sequences = torch.bucketize(rows, cu_seqlens[1:], right=True) + positions = rows - cu_seqlens[sequences] + # Compressed tensors may have unused capacity past the final prefix. + positions = torch.where(rows < cu_seqlens[-1], positions, 0) + cos, sin = cos[0, positions], sin[0, positions] + width = cos.shape[-1] + tail = x[..., -width:] + pair = torch.stack((-tail[..., 1::2], tail[..., 0::2]), -1).flatten(-2) + rotated = (tail.float() * cos + pair.float() * sin).to(x.dtype) + return torch.cat((x[..., :-width], rotated), -1) + + +@lru_cache(maxsize=1) +def _triton_rope_module(): + """Load the optional Triton DSA RoPE implementation.""" + try: + from transformer_engine.common.triton import dsa_rope + except ImportError: + return None + return dsa_rope + + +def _can_use_triton(x, cos, sin, cu_seqlens): + """Keep unsupported inputs on the behaviorally identical eager path.""" + return ( + cu_seqlens is not None + and x.is_cuda + and x.ndim in (3, 4) + and x.stride(-1) == 1 + and cos.shape == sin.shape + and cos.shape[-1] >= 4 + # The Triton kernels use tl.arange over half the rotary width. + and cos.shape[-1] & (cos.shape[-1] - 1) == 0 + and not cos.requires_grad + and not sin.requires_grad + and _triton_rope_module() is not None + ) + + +def _rope_tables(table): + """Convert TE's repeated pairs to the Triton kernel's duplicated-half tables.""" + pairs = table[0, :, 0, ::2] + return torch.cat((pairs, pairs), dim=-1) + + +def _apply_rotary_triton(x, cos, sin, cu_seqlens): + """Apply Triton RoPE out of place so callers retain the eager aliasing contract.""" + module = _triton_rope_module() + if module is None: + raise RuntimeError("Triton RoPE is unavailable.") + return _TritonRotary.apply(x, _rope_tables(cos), _rope_tables(sin), cu_seqlens, module) + + +def _launch_triton(packed, cos, sin, cu_seqlens, module, backward): + kernel = ( + module._autotuned_mla_rope_bwd_inplace_kernel + if backward + else module._autotuned_mla_rope_fwd_inplace_kernel + ) + heads, dim = packed.shape[-2:] + + def grid(meta): + return packed.shape[0], module.triton.cdiv(heads, meta["BLOCK_H"]) + + kernel[grid]( + packed, + cos, + sin, + dim - cos.shape[-1], + cos.shape[-1], + heads, + 1, + cu_seqlens.numel() - 1, + cu_seqlens, + None, + packed.stride(0), + packed.stride(1), + cos.stride(0), + sin.stride(0), + 0, + 1, + INVERSE=False, + REMOVE_INTERLEAVING=True, + ) + + +class _TritonRotary(torch.autograd.Function): + """Pair the retained in-place Triton kernels behind an out-of-place autograd API.""" + + @staticmethod + def forward(ctx, x, cos, sin, cu_seqlens, module): + """Rotate a copy and save the frequency tables for backward.""" + # cuDNN attention saves its output for backward. Cloning also preserves + # apply_rotary's established out-of-place contract for every call site. + output = x.clone(memory_format=torch.contiguous_format) + packed = output.reshape(-1, x.shape[-2], x.shape[-1]) + _launch_triton(packed, cos, sin, cu_seqlens, module, False) + ctx.save_for_backward(cos, sin, cu_seqlens) + ctx.module = module + ctx.shape = x.shape + return output + + @staticmethod + def backward(ctx, grad_output): + """Apply the backward rotation to a copy of the output gradient.""" + cos, sin, cu_seqlens = ctx.saved_tensors + grad_input = grad_output.clone(memory_format=torch.contiguous_format) + packed = grad_input.reshape(-1, ctx.shape[-2], ctx.shape[-1]) + _launch_triton(packed, cos, sin, cu_seqlens, ctx.module, True) + return grad_input, None, None, None, None + + +class _TritonQueryRotary(torch.autograd.Function): + """Rotate the fresh normalized query in place; keep backward out of place.""" + + @staticmethod + def forward(ctx, query, cos, sin, cu_seqlens, module): + """Rotate the fresh query and record its in-place mutation.""" + packed = query.reshape(-1, query.shape[-2], query.shape[-1]) + _launch_triton(packed, cos, sin, cu_seqlens, module, False) + ctx.save_for_backward(cos, sin, cu_seqlens) + ctx.module = module + ctx.shape = query.shape + ctx.mark_dirty(query) + return query + + @staticmethod + def backward(ctx, grad_output): + """Apply the backward rotation without mutating the output gradient.""" + cos, sin, cu_seqlens = ctx.saved_tensors + grad_input = grad_output.clone(memory_format=torch.contiguous_format) + packed = grad_input.reshape(-1, ctx.shape[-2], ctx.shape[-1]) + _launch_triton(packed, cos, sin, cu_seqlens, ctx.module, True) + return grad_input, None, None, None, None + + +def _apply_rotary_query(query, cos, sin, cu_seqlens): + """Rotate fresh normalized queries in place with Triton, avoiding a copy. + + Unlike apply_rotary, this private path requires an input safe to mutate. + Its backward and eager fallback remain out of place. + """ + if not _can_use_triton(query, cos, sin, cu_seqlens): + return _apply_rotary_eager(query, cos, sin, cu_seqlens) + module = _triton_rope_module() + return _TritonQueryRotary.apply(query, _rope_tables(cos), _rope_tables(sin), cu_seqlens, module) + + +def apply_rotary(x, cos, sin, cu_seqlens=None): + """Rotate trailing interleaved pairs, using Triton when its contract is satisfied.""" + if _can_use_triton(x, cos, sin, cu_seqlens): + return _apply_rotary_triton(x, cos, sin, cu_seqlens) + return _apply_rotary_eager(x, cos, sin, cu_seqlens) diff --git a/transformer_engine/pytorch/attention/sparse_attention/dsv4_attention.py b/transformer_engine/pytorch/attention/sparse_attention/dsv4_attention.py new file mode 100644 index 00000000000..0b7b1dd4848 --- /dev/null +++ b/transformer_engine/pytorch/attention/sparse_attention/dsv4_attention.py @@ -0,0 +1,322 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""DSv4 CSA/HCA hidden-state attention layer.""" + +from typing import Optional + +import torch + +from transformer_engine.pytorch.module import Linear, RMSNorm + +from .compressor import Compressor +from .dsa_cudnn import sparse_attention as _sparse_attention +from .dsa_rope import _apply_rotary_query, _DSv4RotaryEmbedding, apply_rotary +from .indexer import Indexer + +__all__ = ["DSv4Attention"] + + +@torch.compile +def _compiled_q_rms_norm(q: torch.Tensor, eps: float) -> torch.Tensor: + """Compile DSv4's weightless per-head query normalization.""" + return torch.nn.functional.rms_norm(q, (q.shape[-1],), eps=eps) + + +def _packed_prefixes(hidden_states, cu_seqlens, max_seqlen, ratio): + """Build compressed sequence prefixes and output capacity.""" + if not isinstance(cu_seqlens, torch.Tensor): + raise ValueError("THD requires contiguous INT32 cu_seqlens_q [B+1] on the input device.") + if ( + cu_seqlens.dtype != torch.int32 + or cu_seqlens.ndim != 1 + or cu_seqlens.numel() < 2 + or not cu_seqlens.is_contiguous() + or cu_seqlens.device != hidden_states.device + ): + raise ValueError("THD requires contiguous INT32 cu_seqlens_q [B+1] on the input device.") + if not isinstance(max_seqlen, int) or max_seqlen < ratio: + raise ValueError("THD requires integer max_seqlen_q >= compression_ratio.") + # The sum of per-sequence floors can be smaller than floor(T / ratio). + # Allocate that safe capacity without reading CUDA prefix contents on host; + # cuDNN uses the true device prefix to exclude unused trailing rows. + lengths = cu_seqlens[1:] - cu_seqlens[:-1] + cu_comp = torch.cat((cu_seqlens.new_zeros(1), (lengths // ratio).cumsum(0, dtype=torch.int32))) + total_comp = hidden_states.shape[0] // ratio + if not total_comp: + raise ValueError("Sequences must contain at least one complete compression window.") + return cu_comp, total_comp, max_seqlen + + +class DSv4Attention(torch.nn.Module): + """DSv4 CSA/HCA attention in ``sbhd``, ``bshd``, or packed ``thd`` format. + + Composes TE projections/norms, Compressor, Indexer, and RoPE with cuDNN + sparse attention. Supports full-sequence BF16 forward/backward on SM100 only, + with 64 heads of width 512 and CSA index heads of width 128. Shared-input + projections are fused by default. Each sequence starts at position zero. + CSA and HCA include sliding-window attention; standalone sliding-window-only + layers are unsupported. Cache/decode, physical inter-sequence padding, + TP/SP/CP, and FP8 are unsupported. The model owns the indexer auxiliary loss; + discrete top-k selection does not train the indexer from attention output alone. + """ + + def __init__( + self, + hidden_size: int, + q_lora_rank: int, + *, + layer_type: str, + head_dim: int, + rope_head_dim: int, + sliding_window: int, + compression_ratio: int, + o_groups: int, + o_lora_rank: int, + index_n_heads: int = 64, + index_topk: int = 512, + rope_theta: float = 160000.0, + rope: Optional[torch.nn.Module] = None, + rms_norm_eps: float = 1e-6, + max_seqlen: Optional[int] = None, + device: str = "cuda", + params_dtype: torch.dtype = torch.float32, + qkv_format: str = "sbhd", + _fuse_projections: bool = True, + ): + super().__init__() + # Keep the current kernel geometry explicit without exposing untested sizes. + self.num_heads = 64 + self.index_head_dim = 128 + if layer_type not in ("compressed_sparse_attention", "heavily_compressed_attention"): + raise ValueError("Only DSv4 CSA and HCA are supported.") + if qkv_format not in ("sbhd", "bshd", "thd"): + raise ValueError("qkv_format must be 'sbhd', 'bshd', or 'thd'.") + if head_dim != 512 or rope_head_dim < 2 or rope_head_dim % 2 or rope_head_dim > head_dim: + raise ValueError("head_dim must be 512; rope_head_dim must be even and fit.") + if any( + v <= 0 + for v in ( + hidden_size, + q_lora_rank, + sliding_window, + compression_ratio, + o_groups, + o_lora_rank, + ) + ): + raise ValueError("DSv4 dimensions, window size and ratio must be positive.") + if self.num_heads * head_dim % o_groups: + raise ValueError("o_groups must divide the attention output channels.") + is_csa = layer_type == "compressed_sparse_attention" + if is_csa and ( + index_n_heads not in (32, 64) or index_topk <= 0 or rope_head_dim > self.index_head_dim + ): + raise ValueError( + "CSA requires 32 or 64 index heads, positive top-k, and RoPE width <= 128." + ) + + kw = {"bias": False, "device": device, "params_dtype": params_dtype} + if _fuse_projections: + # Query-down and local-KV projections read the same hidden states. + # Keep slices explicit below while one contiguous weight feeds the GEMM. + self.q_a_kv_proj = Linear( + hidden_size, + q_lora_rank + head_dim, + **kw, + ) + else: + self.q_a_proj = Linear(hidden_size, q_lora_rank, **kw) + self.kv_proj = Linear(hidden_size, head_dim, **kw) + # LayerNormLinear would normalize before projection; DSv4 normalizes after it. + self.q_a_norm = RMSNorm(q_lora_rank, eps=rms_norm_eps, device=device, dtype=params_dtype) + self.q_b_proj = Linear(q_lora_rank, self.num_heads * head_dim, **kw) + self.kv_norm = RMSNorm(head_dim, eps=rms_norm_eps, device=device, dtype=params_dtype) + self.compressor = Compressor( + hidden_size, + head_dim, + compression_ratio, + is_csa, + rms_norm_eps, + device, + params_dtype, + fused=_fuse_projections, + ) + if is_csa: + self.indexer = Indexer( + hidden_size, + q_lora_rank, + self.index_head_dim, + index_n_heads, + index_topk, + compression_ratio, + rms_norm_eps, + device, + params_dtype, + fused=_fuse_projections, + ) + # HF/Megatron use one independent block per output group. A normal TE + # Linear forward mixes groups, so use its weight in the block multiply. + self.o_a_proj = Linear(self.num_heads * head_dim // o_groups, o_groups * o_lora_rank, **kw) + self.o_b_proj = Linear(o_groups * o_lora_rank, hidden_size, **kw) + self.sinks = torch.nn.Parameter( + torch.zeros(self.num_heads, device=device, dtype=params_dtype) + ) + self.sliding_window = sliding_window + # A model can supply its RoPE frequencies (for example YaRN) without + # duplicating that policy here. It must return the same token/compressed + # (cos, sin) pairs as _DSv4RotaryEmbedding.forward. + self.rope = ( + rope + if rope is not None + else _DSv4RotaryEmbedding( + compression_ratio, rope_head_dim, rope_theta, device, max_seqlen + ) + ) + self.hidden_size, self.q_lora_rank = hidden_size, q_lora_rank + self.fused_projections = _fuse_projections + self.head_dim, self.rope_head_dim = head_dim, rope_head_dim + self.rope_theta = rope_theta if rope is None else None + self.rms_norm_eps = rms_norm_eps + self.compression_ratio = compression_ratio + self.o_groups, self.o_lora_rank = o_groups, o_lora_rank + self.is_csa = is_csa + self.qkv_format = qkv_format + + def forward( + self, + hidden_states: torch.Tensor, + *, + cu_seqlens_q: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + return_indexer_context: bool = False, + ): + """Run full self-attention and optionally expose single-sequence CSA context. + + Packed calls require contiguous INT32 ``cu_seqlens_q`` on the input + device and ``max_seqlen_q``. Prefixes have no padding/global offsets; + their contents are caller-owned and are not checked on the host. + These metadata arguments are accepted only for ``thd``. + """ + packed = self.qkv_format == "thd" + if ( + hidden_states.ndim != (2 if packed else 3) + or hidden_states.shape[-1] != self.hidden_size + ): + raise ValueError("hidden_states must have the selected format and hidden_size.") + from transformer_engine.pytorch.quantization import FP8GlobalStateManager + + if FP8GlobalStateManager.is_fp8_enabled(): + raise NotImplementedError("DSv4Attention supports BF16 only; disable TE FP8 autocast.") + if return_indexer_context and not self.is_csa: + raise ValueError("Only CSA has an indexer context.") + if self.qkv_format == "sbhd": + # cuDNN consumes batch-major packed rows. B=1 needs only a view; + # larger batches require a physical repack with the current core. + hidden_states = hidden_states.transpose(0, 1).contiguous() + if packed: + cu = cu_seqlens_q + cu_comp, total_comp, seq = _packed_prefixes( + hidden_states, cu, max_seqlen_q, self.compression_ratio + ) + batch = cu.numel() - 1 + comp_shape = (total_comp,) + else: + if cu_seqlens_q is not None or max_seqlen_q is not None: + raise ValueError("cu_seqlens_q and max_seqlen_q are only supported with THD input.") + batch, seq, _ = hidden_states.shape + n_comp = seq // self.compression_ratio + if not n_comp: + raise ValueError("Sequence must contain at least one complete compression window.") + cu = torch.arange(batch + 1, device=hidden_states.device, dtype=torch.int32) * seq + cu_comp = ( + torch.arange(batch + 1, device=hidden_states.device, dtype=torch.int32) * n_comp + ) + total_comp = batch * n_comp + comp_shape = (batch, n_comp) + if return_indexer_context and batch != 1: + raise ValueError("The CSA indexer context is currently supported for batch=1.") + token_shape = hidden_states.shape[:-1] + + token_rope, compressed_rope = self.rope(seq, hidden_states.device) + + if self.fused_projections: + q_a, local_kv_projected = self.q_a_kv_proj(hidden_states).split( + (self.q_lora_rank, self.head_dim), dim=-1 + ) + else: + q_a = self.q_a_proj(hidden_states) + local_kv_projected = self.kv_proj(hidden_states) + q_residual = self.q_a_norm(q_a) + q = self.q_b_proj(q_residual).reshape(*token_shape, self.num_heads, self.head_dim) + # Query-up norm is unweighted; TE RMSNorm would add a learned scale. + q = _compiled_q_rms_norm(q, self.rms_norm_eps) + q = ( + _apply_rotary_query(q, *token_rope, cu) + .reshape(-1, self.num_heads, self.head_dim) + .contiguous() + ) + local_kv = apply_rotary(self.kv_norm(local_kv_projected).unsqueeze(-2), *token_rope, cu) + local_kv = local_kv.reshape(-1, self.head_dim).contiguous() + compressed_kv = apply_rotary( + self.compressor(hidden_states, cu, cu_comp, total_comp=total_comp).reshape( + *comp_shape, 1, self.head_dim + ), + *compressed_rope, + cu_comp, + ) + compressed_kv = compressed_kv.reshape(total_comp, self.head_dim).contiguous() + + indices = None + index_q = index_k = index_w = None + if self.is_csa: + # Megatron trains this tower with a separate loss, without sending + # its gradients into the shared hidden/query projections. + index_result = self.indexer( + hidden_states.detach(), + q_residual.detach(), + cu, + cu_comp, + token_rope, + compressed_rope, + total_comp=total_comp, + max_seqlen=seq, + return_context=return_indexer_context, + ) + if return_indexer_context: + indices, index_q, index_k, index_w = index_result + else: + indices = index_result + + output = _sparse_attention( + q, + local_kv, + compressed_kv, + self.sinks.float(), + cu, + cu_comp, + window_size=self.sliding_window, + ratio=self.compression_ratio, + indices=indices, + max_compressed_seqlen=None if self.is_csa else seq // self.compression_ratio, + ).reshape(*token_shape, self.num_heads, self.head_dim) + cos, sin = token_rope + output = apply_rotary(output, cos, -sin, cu).reshape(*token_shape, self.o_groups, -1) + weight = self.o_a_proj.weight.reshape(self.o_groups, self.o_lora_rank, -1) + grouped = torch.einsum("...gd,grd->...gr", output, weight).flatten(-2) + result = self.o_b_proj(grouped) + if self.qkv_format == "sbhd": + result = result.transpose(0, 1).contiguous() + if return_indexer_context: + return result, { + "index_q": index_q, + "index_k": index_k, + "index_w": index_w, + "indices": indices, + "attn_q": q.detach(), + "local_kv": local_kv.detach(), + "compressed_kv": compressed_kv.detach(), + "sink": self.sinks.float().detach(), + } + return result diff --git a/transformer_engine/pytorch/attention/sparse_attention/indexer.py b/transformer_engine/pytorch/attention/sparse_attention/indexer.py new file mode 100644 index 00000000000..6b163ad4908 --- /dev/null +++ b/transformer_engine/pytorch/attention/sparse_attention/indexer.py @@ -0,0 +1,112 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Index projections, compression, and cuDNN block selection.""" + +import torch + +from transformer_engine.pytorch.module import Linear + +from .compressor import Compressor +from .dsa_cudnn import select_indices +from .dsa_rope import apply_rotary + + +class Indexer(torch.nn.Module): + """Select compressed rows using the indexer's own Q/K/weight projections. + + Accept batch-major or packed hidden states and query residuals, sequence + prefixes, and caller-provided RoPE tables. Return global compressed-row IDs; + ``return_context=True`` also returns live Q/K/head weights for a separate loss. + """ + + def __init__( + self, + hidden_size, + q_lora_rank, + head_dim, + n_heads, + top_k, + ratio, + eps, + device, + params_dtype, + *, + fused, + ): + super().__init__() + self.compressor = Compressor( + hidden_size, + head_dim, + ratio, + True, + eps, + device, + params_dtype, + fused=fused, + weight_width=n_heads if fused else 0, + ) + kw = {"bias": False, "device": device, "params_dtype": params_dtype} + self.q_proj = Linear(q_lora_rank, n_heads * head_dim, **kw) + if not fused: + self.weights_proj = Linear(hidden_size, n_heads, **kw) + self.head_dim, self.n_heads = head_dim, n_heads + self.top_k, self.ratio, self.fused = top_k, ratio, fused + + def forward( + self, + hidden_states, + q_residual, + cu, + cu_comp, + token_rope, + compressed_rope, + *, + total_comp, + max_seqlen, + return_context=False, + ): + """Select compressed rows and optionally expose live tensors for the indexer loss.""" + n_comp = max_seqlen // self.ratio + comp_shape = (total_comp,) if hidden_states.ndim == 2 else (hidden_states.shape[0], n_comp) + index_compressed = self.compressor(hidden_states, cu, cu_comp, total_comp=total_comp) + if self.fused: + index_compressed, weights = index_compressed + else: + weights = self.weights_proj(hidden_states) + # The singleton axis is the shared index-key head, not a tunable head count. + key = ( + apply_rotary( + index_compressed.reshape(*comp_shape, 1, self.head_dim), + *compressed_rope, + cu_comp, + ) + .reshape(total_comp, self.head_dim) + .contiguous() + ) + query = self.q_proj(q_residual).reshape( + *hidden_states.shape[:-1], self.n_heads, self.head_dim + ) + query = ( + apply_rotary(query, *token_rope, cu) + .reshape(-1, self.n_heads, self.head_dim) + .contiguous() + ) + # Top-k IDs are discrete; training index projections needs a separate loss. + weights = weights.reshape(-1, self.n_heads).contiguous() + indices = select_indices( + query, + key, + weights, + cu, + cu_comp, + top_k=min(self.top_k, n_comp), + ratio=self.ratio, + max_seqlen=max_seqlen, + # cuDNN's compile bound covers the partial trailing ratio even + # though cu_comp contains only complete compressed blocks. + max_compressed_seqlen=(max_seqlen + self.ratio - 1) // self.ratio, + scale=(self.head_dim * self.n_heads) ** -0.5, + ) + return (indices, query, key, weights) if return_context else indices