From b83d57e9502e72b9eca73ee34835e58c2f968a3b Mon Sep 17 00:00:00 2001 From: kunlunl Date: Fri, 26 Jun 2026 05:33:01 -0700 Subject: [PATCH 01/23] Add DSv4 THD context parallel support Signed-off-by: kunlunl --- megatron/core/datasets/data_schedule.py | 41 +- megatron/core/datasets/data_schedule_utils.py | 47 +- .../experimental_attention_variant/csa.py | 503 ++- .../csa_cp_layout_kernels.py | 3253 +++++++++++++++++ .../csa_cp_utils.py | 965 +++++ .../deepseek_v4_hybrid_attention.py | 219 +- .../experimental_attention_variant/dsa.py | 12 +- .../dsa_kernels.py | 251 +- .../core/transformer/transformer_config.py | 5 + megatron/training/training.py | 1 + tests/unit_tests/test_sequence_packing.py | 43 + .../test_csa_cp_layout_kernels.py | 1214 ++++++ .../test_csa_cp_utils.py | 682 ++++ .../test_dsv4_hybrid_attention.py | 4 +- .../test_dsv4_hybrid_attention_cp.py | 981 +++++ 15 files changed, 8105 insertions(+), 116 deletions(-) create mode 100644 megatron/core/transformer/experimental_attention_variant/csa_cp_layout_kernels.py create mode 100644 megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py create mode 100644 tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_layout_kernels.py create mode 100644 tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py create mode 100644 tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py diff --git a/megatron/core/datasets/data_schedule.py b/megatron/core/datasets/data_schedule.py index ab2d292ef60..6471d7465ff 100644 --- a/megatron/core/datasets/data_schedule.py +++ b/megatron/core/datasets/data_schedule.py @@ -613,6 +613,15 @@ def get_batch_on_this_rank_for_sequence_packing( is_first_or_last_stage = is_first_stage or is_last_stage dev = torch.cuda.current_device() + static_cp_target_len = None + if ( + config is not None + and getattr(config, 'experimental_attention_variant', None) == "dsv4_hybrid" + and cp_group.size() > 1 + and getattr(config, 'pad_packed_seq_alignment', None) == "max" + and getattr(config, 'max_seqlen_per_dp_cp_rank', None) is not None + ): + static_cp_target_len = int(config.max_seqlen_per_dp_cp_rank) # data_iterator should return a batch including the following keys. batch_keys = ['cu_seqlens', 'cu_seqlens_padded', 'max_seqlen'] @@ -658,7 +667,22 @@ def get_batch_on_this_rank_for_sequence_packing( cp_slice_keys = ['padding_mask'] if is_first_or_last_stage or mtp_on_this_rank: cp_slice_keys.extend(['tokens', 'position_ids', 'labels', 'loss_mask']) - get_cp_slice_for_thd(batch, cp_group, keys=cp_slice_keys) + csa_cp_partition_mode = None + if ( + config is not None + and getattr(config, 'experimental_attention_variant', None) == "dsv4_hybrid" + ): + csa_cp_partition_mode = config.csa_cp_partition_mode + partition_total_tokens = ( + static_cp_target_len * cp_group.size() if static_cp_target_len is not None else None + ) + get_cp_slice_for_thd( + batch, + cp_group, + keys=cp_slice_keys, + csa_cp_partition_mode=csa_cp_partition_mode, + partition_total_tokens=partition_total_tokens, + ) # Broadcast cu_seqlens_size because we need it to create placeholder for cu_seqlens and # cu_seqlens_padded for non TP 0 ranks. @@ -672,12 +696,15 @@ def get_batch_on_this_rank_for_sequence_packing( # Broadcast total_tokens because padding_mask is prepared on every PP stage. # Tokens/labels/loss_mask/position_ids use the same length on stages that own them. if is_tp_rank_0: - # Under VPP, the last PP stage has labels but no tokens, so derive - # total_tokens from cu_seqlens_padded, which is present on every - # stage. cu_seqlens_padded keeps the pre-CP packed length; divide - # by cp_size to match the already CP-sliced sequence tensors. - cp_world = cp_group.size() - total_tokens = (batch['cu_seqlens_padded'][-1].to(torch.int32) // cp_world).reshape(1) + if static_cp_target_len is not None: + total_tokens = torch.tensor([static_cp_target_len], dtype=torch.int32, device=dev) + else: + # Under VPP, the last PP stage has labels but no tokens, so derive + # total_tokens from cu_seqlens_padded, which is present on every + # stage. cu_seqlens_padded keeps the pre-CP packed length; divide + # by cp_size to match the already CP-sliced sequence tensors. + cp_world = cp_group.size() + total_tokens = (batch['cu_seqlens_padded'][-1].to(torch.int32) // cp_world).reshape(1) else: total_tokens = torch.empty(1, dtype=torch.int32, device=dev) broadcast_tensor(total_tokens, tp_src_rank, tp_group) diff --git a/megatron/core/datasets/data_schedule_utils.py b/megatron/core/datasets/data_schedule_utils.py index 5f8bb6b3e01..0f30a28d73c 100644 --- a/megatron/core/datasets/data_schedule_utils.py +++ b/megatron/core/datasets/data_schedule_utils.py @@ -11,16 +11,27 @@ from megatron.core.rerun_state_machine import RerunDataIterator -def get_cp_slice_for_thd(batch, cp_group, keys: Optional[Sequence[str]] = None): +def get_cp_slice_for_thd( + batch, + cp_group, + keys: Optional[Sequence[str]] = None, + csa_cp_partition_mode: Optional[str] = None, + partition_total_tokens: Optional[int] = None, +): """Partition sequence data for context parallelism in THD format. - Uses TE's THD partitioned indices to split the packed sequence across CP ranks. - Only keys present in the batch are sliced. + Uses TE's THD partitioned indices by default. When ``csa_cp_partition_mode`` + is provided, uses the CSA partition helper so data row order matches the + DSv4 attention layer's CP row order. Only keys present in the batch are sliced. Args: batch: Dict with packed sequence data. cp_group: Context parallel process group. keys: Sequence data keys to slice. Defaults to the original THD data tensors. + csa_cp_partition_mode: Optional CSA CP partition mode for DSv4 hybrid attention. + partition_total_tokens: Optional padded total used only for choosing CP row indices. + When set, tensors selected by ``keys`` are tail-padded to this length before slicing. Existing + cu_seqlens metadata is left unchanged. """ cp_size = cp_group.size() if cp_size <= 1: @@ -32,10 +43,36 @@ def get_cp_slice_for_thd(batch, cp_group, keys: Optional[Sequence[str]] = None): # Use cu_seqlens_padded[-1] for total_tokens instead of batch['tokens'].size(0): # under VPP, the last PP stage has labels/loss_mask but no tokens, so # batch['tokens'] is None on that stage. cu_seqlens_padded is always populated. - total_tokens = int(cu_seqlens[-1].item()) - index = get_thd_partitioned_indices(cu_seqlens, total_tokens, cp_size, cp_rank) + total_tokens = ( + int(cu_seqlens[-1].item()) + if partition_total_tokens is None + else int(partition_total_tokens) + ) if keys is None: keys = ('tokens', 'position_ids', 'labels', 'loss_mask') + if partition_total_tokens is not None: + for key in keys: + if key not in batch or batch[key] is None: + continue + pad_len = total_tokens - batch[key].numel() + if pad_len < 0: + raise RuntimeError( + f"partition_total_tokens={total_tokens} is smaller than {key} length " + f"{batch[key].numel()}." + ) + if pad_len > 0: + pad_value = True if key == 'padding_mask' else 0 + batch[key] = torch.cat([batch[key], batch[key].new_full((pad_len,), pad_value)]) + if csa_cp_partition_mode is None: + index = get_thd_partitioned_indices(cu_seqlens, total_tokens, cp_size, cp_rank) + else: + from megatron.core.transformer.experimental_attention_variant.csa_cp_utils import ( + thd_cp_local_row_indices, + ) + + index = thd_cp_local_row_indices( + csa_cp_partition_mode, total_tokens, cp_size, cp_rank, cu_seqlens.device + ) for key in keys: if key in batch and batch[key] is not None: batch[key] = batch[key].index_select(0, index) diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index 2e57ec0b9f8..535880ff7c9 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -13,6 +13,7 @@ from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.experimental_attention_variant import csa_cp_utils as cp_utils from megatron.core.transformer.experimental_attention_variant.dsa import ( DSAIndexerLossAutoScaler, DSAIndexerLossLoggingHelper, @@ -22,6 +23,7 @@ rotate_activation, ) from megatron.core.transformer.experimental_attention_variant.dsa_kernels import ( + FusedIndexerSparseAttnFromTopkFunc, batch_of_row, build_flat_topk_idxs, dsa_sparse_attn, @@ -617,6 +619,99 @@ def unfused_compressed_sparse_attn( return output.reshape(b, sq, np_ * hn).permute(1, 0, 2).contiguous() +def _unfused_indexer_sparse_attn_from_topk( + query: torch.Tensor, + kv_full: torch.Tensor, + attn_sink: torch.Tensor, + topk_indices: torch.Tensor, + q_indexer: torch.Tensor, + k_indexer: torch.Tensor, + weights: torch.Tensor, + indexer_topk_indices: torch.Tensor, + compressed_kv: torch.Tensor, + softmax_scale: float, + indexer_softmax_scale: float, + loss_coeff: float, + calculate_per_token_loss: bool, + global_query_rows: int, + q_padding_mask: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """PyTorch sparse attention plus caller-supplied top-k indexer loss for THD CP. + + This mirrors the fused CP top-k path at the tensor-contract level: + ``topk_indices`` indexes ``kv_full`` for sparse attention, and + ``indexer_topk_indices`` indexes rank-major compressed K for indexer loss. + """ + output = unfused_compressed_sparse_attn(query, kv_full, attn_sink, topk_indices, softmax_scale) + + total_q, np_, hn = query.shape + indexer_topk = indexer_topk_indices.shape[-1] + if indexer_topk == 0: + return output, query.new_zeros((), dtype=torch.float32) + + if q_padding_mask is not None: + if q_padding_mask.shape != (total_q,): + raise RuntimeError( + "q_padding_mask must have shape " + f"({total_q},), got {tuple(q_padding_mask.shape)}." + ) + if q_padding_mask.dtype != torch.bool: + raise RuntimeError(f"q_padding_mask must be bool, got {q_padding_mask.dtype}.") + if q_padding_mask.device != query.device: + raise RuntimeError( + f"q_padding_mask must be on {query.device}, got {q_padding_mask.device}." + ) + indexer_topk_indices = indexer_topk_indices.masked_fill(q_padding_mask.unsqueeze(-1), -1) + + valid = indexer_topk_indices >= 0 + row_valid = valid.any(dim=-1, keepdim=True) + safe_indices = indexer_topk_indices.clamp(min=0).long() + + index_scores_full = torch.einsum("rhd,kd->rhk", q_indexer.float(), k_indexer.float()) + index_scores_full = torch.relu(index_scores_full) + weights_scaled = weights.float() * float(indexer_softmax_scale) + index_scores_full = index_scores_full * weights_scaled.unsqueeze(-1) + index_scores_full = index_scores_full.sum(dim=1) + + predict_logits = torch.gather(index_scores_full, dim=-1, index=safe_indices) + predict_logits = predict_logits.masked_fill(~valid, float("-inf")) + predict_logits = predict_logits.masked_fill(~row_valid, 0.0) + predict = torch.softmax(predict_logits, dim=-1, dtype=torch.float32) + predict = predict * row_valid.float() + + selected_kv = compressed_kv.detach().index_select(0, safe_indices.reshape(-1)) + selected_kv = selected_kv.reshape(total_q, indexer_topk, hn) + + # This fallback does not run FlashMLA, so it cannot reuse the fused path's + # lse_indexer-based target helper. Recompute the same sink-aware target and + # KL terms explicitly to keep the no-fusion path numerically aligned. + attn_scores = torch.einsum("rhd,rkd->rhk", query.detach().float(), selected_kv.float()) + attn_scores = attn_scores * softmax_scale + attn_scores = attn_scores.masked_fill(~valid.unsqueeze(1), float("-inf")) + sink = attn_sink.view(1, np_, 1).float() + scores_max = torch.maximum(attn_scores.max(dim=-1, keepdim=True).values, sink) + exp_scores = torch.exp(attn_scores - scores_max) + exp_sink = torch.exp(sink - scores_max) + attn_probs = exp_scores / (exp_scores.sum(dim=-1, keepdim=True) + exp_sink) + attn_probs = attn_probs * row_valid.unsqueeze(1).float() + target = attn_probs.sum(dim=1) + + eps = torch.finfo(torch.float32).tiny + target = target.clamp(min=eps) + predict = predict.clamp(min=eps) + kl_per_row = (target * (torch.log(target) - torch.log(predict))).sum(dim=-1) + kl_per_row = torch.where(row_valid.squeeze(-1), kl_per_row, torch.zeros_like(kl_per_row)) + raw_local_loss = kl_per_row.sum() + + if calculate_per_token_loss: + indexer_loss = raw_local_loss * loss_coeff + else: + if global_query_rows <= 0: + raise RuntimeError(f"global_query_rows must be positive, got {global_query_rows}.") + indexer_loss = raw_local_loss * (float(loss_coeff) / float(global_query_rows)) + return output, indexer_loss + + # --------------------------------------------------------------------------- # Compressor # --------------------------------------------------------------------------- @@ -812,6 +907,7 @@ def _forward_thd( x: torch.Tensor, cu_seqlens: torch.Tensor, max_seqlen_q: Optional[int] = None, + compressed_group_ids: Optional[torch.Tensor] = None, fixed_total_comp: Optional[int] = None, ) -> Tuple[Optional[torch.Tensor], torch.Tensor]: """THD per-segment compression — fully vectorized. @@ -831,6 +927,8 @@ def _forward_thd( (matches ``packed_seq_params.cu_seqlens_q``). max_seqlen_q: max original sequence length (avoids a GPU→CPU sync when building the rotary table for ``ratio > 1``). + compressed_group_ids: CP compressor-prep group ids. When supplied, + ``x`` is already packed into ``ratio``-sized groups. fixed_total_comp: when set, overrides ``cu_seqlens_compressed[-1]`` as the output row count. Must be >= the true compressed count. @@ -845,6 +943,7 @@ def _forward_thd( ratio = self.compress_ratio device = x.device dtype = x.dtype + pre_grouped = compressed_group_ids is not None # Per-segment compressed lengths (vectorized). seq_lens = cu_seqlens[1:] - cu_seqlens[:-1] @@ -868,29 +967,39 @@ def _forward_thd( kv, _ = self.linear_wkv(x) # (total, 1, coff * head_dim) score, _ = self.linear_wgate(x) # (total, 1, coff * head_dim) - # Build gather index: (total_comp, ratio). ``total_comp`` can be a - # static capacity for CUDA graph capture, so rows beyond the true - # ``cu_seqlens_compressed[-1]`` are mapped to a safe source row and - # left as tail padding by downstream index lowering. - row_idx = torch.arange(total_comp, device=device, dtype=cu_seqlens_compressed.dtype) - batch_ids = batch_of_row(cu_seqlens_compressed, total_q=total_comp) - valid_comp = row_idx < cu_seqlens_compressed[-1] - local_pos = row_idx - cu_seqlens_compressed[batch_ids] - local_pos = torch.where(valid_comp, local_pos, torch.zeros_like(local_pos)) - # (total_comp, 1) + (1, ratio) → (total_comp, ratio) - base = cu_seqlens[batch_ids].unsqueeze(1) + local_pos.unsqueeze(1) * ratio - base = torch.where(valid_comp.unsqueeze(1), base, torch.zeros_like(base)) - offsets = torch.arange(ratio, device=device, dtype=base.dtype).unsqueeze(0) - gather_idx = base + offsets # (total_comp, ratio) - - kv_grouped = kv[gather_idx] # (total_comp, ratio, 1, coff * d) - score_grouped = score[gather_idx] + if pre_grouped: + # Compressor-prep already groups rows as ``[g * ratio, (g + 1) * ratio)``. + kv_grouped = kv.reshape(total_comp, ratio, 1, -1) + score_grouped = score.reshape(total_comp, ratio, 1, -1) + local_pos = None + else: + # Build gather index: (total_comp, ratio). ``total_comp`` can be a + # static capacity for CUDA graph capture, so rows beyond the true + # ``cu_seqlens_compressed[-1]`` are mapped to a safe source row and left + # as tail padding by downstream index lowering. + row_idx = torch.arange(total_comp, device=device, dtype=cu_seqlens_compressed.dtype) + batch_ids = batch_of_row(cu_seqlens_compressed, total_q=total_comp) + valid_comp = row_idx < cu_seqlens_compressed[-1] + local_pos = row_idx - cu_seqlens_compressed[batch_ids] + local_pos = torch.where(valid_comp, local_pos, torch.zeros_like(local_pos)) + # (total_comp, 1) + (1, ratio) → (total_comp, ratio) + base = cu_seqlens[batch_ids].unsqueeze(1) + local_pos.unsqueeze(1) * ratio + base = torch.where(valid_comp.unsqueeze(1), base, torch.zeros_like(base)) + offsets = torch.arange(ratio, device=device, dtype=base.dtype).unsqueeze(0) + gather_idx = base + offsets # (total_comp, ratio) + + # Index flat kv/score into ``(total_comp, ratio, 1, coff * d)`` groups. + kv_grouped = kv[gather_idx] # (total_comp, ratio, 1, coff * d) + score_grouped = score[gather_idx] # same # APE: (ratio, coff * d) → broadcast (1, ratio, 1, coff * d). score_grouped = score_grouped + self.ape.view(1, ratio, 1, -1) if self.overlap: - is_first = local_pos == 0 # (total_comp,) + if pre_grouped: + is_first = compressed_group_ids[:total_comp] == 0 + else: + is_first = local_pos == 0 # (total_comp,) kv_grouped = self._overlap_transform_thd(kv_grouped, is_first, fill_value=0) score_grouped = self._overlap_transform_thd( score_grouped, is_first, fill_value=float("-inf") @@ -903,20 +1012,34 @@ def _forward_thd( compressed_thd = self.norm(compressed_thd.to(dtype)) - # RoPE: applied in a single vectorized THD call. - max_seqlen_rope = (max_seqlen_q // ratio) * ratio if max_seqlen_q is not None else None - compressed_thd = _apply_rope( - compressed_thd, - self.head_dim - self.qk_pos_emb_head_dim, - self.qk_pos_emb_head_dim, - self.rotary_pos_emb, - self.config, - rotary_seq_len=0, - ratio=ratio, - cp_group=self.pg_collection.cp, - cu_seqlens=cu_seqlens_compressed, - max_seqlen_rope=max_seqlen_rope, - ) + if not pre_grouped: + # RoPE: applied in a single vectorized THD call. + max_seqlen_rope = (max_seqlen_q // ratio) * ratio if max_seqlen_q is not None else None + compressed_thd = _apply_rope( + compressed_thd, + self.head_dim - self.qk_pos_emb_head_dim, + self.qk_pos_emb_head_dim, + self.rotary_pos_emb, + self.config, + rotary_seq_len=0, + ratio=ratio, + cp_group=self.pg_collection.cp, + cu_seqlens=cu_seqlens_compressed, + max_seqlen_rope=max_seqlen_rope, + ) + else: + rotary_pos_cos, rotary_pos_sin = self.rotary_pos_emb.get_cached_cos_sin( + int(max_seqlen_q), dtype=compressed_thd.dtype, packed_seq=True, mscale=1.0 + ) + compressed_thd = cp_utils.apply_thd_cp_compressed_rope_fused( + compressed_thd, + rotary_pos_cos, + rotary_pos_sin, + compressed_group_ids, + ratio, + self.head_dim - self.qk_pos_emb_head_dim, + self.qk_pos_emb_head_dim, + ) if self.rotate: compressed_thd = rotate_activation(compressed_thd) @@ -1054,6 +1177,39 @@ def __init__( pg_collection=pg_collection, ) + def _forward_thd_query_weights_cp( + self, + x: torch.Tensor, + qr: torch.Tensor, + cu_seqlens_padded: torch.Tensor, + max_position: int, + chunk_ranges: Tuple[Tuple[int, int], ...], + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Compute local THD indexer Q/weights for the current CP row order.""" + local_rows, bsz, _ = x.size() + if bsz != 1: + raise RuntimeError(f"DSv4 THD CP indexer expects bsz=1, got {bsz}.") + + q, _ = self.linear_wq_b(qr) + q = q.reshape(local_rows, bsz, self.index_n_heads, self.index_head_dim).squeeze(1) + rotary_pos_cos, rotary_pos_sin = self.rotary_pos_emb.get_cached_cos_sin( + int(max_position), dtype=q.dtype, packed_seq=True, mscale=1.0 + ) + q = cp_utils.apply_thd_cp_local_rope_fused( + q, + rotary_pos_cos, + rotary_pos_sin, + self.index_head_dim - self.qk_pos_emb_head_dim, + self.qk_pos_emb_head_dim, + cu_seqlens_padded, + chunk_ranges, + ) + q = rotate_activation(q) + + weights, _ = self.linear_weights_proj(x) + weights = weights.squeeze(1) * (self.index_n_heads**-0.5) + return q, weights + def forward_before_topk( self, x: torch.Tensor, qr: torch.Tensor, packed_seq_params: Optional[PackedSeqParams] = None ) -> Union[ @@ -1569,6 +1725,8 @@ def forward( attention_mask: torch.Tensor, x: torch.Tensor = None, qr: torch.Tensor = None, + boundary_hidden: Optional[torch.Tensor] = None, + boundary_kv: Optional[torch.Tensor] = None, attn_mask_type: AttnMaskType = None, attention_bias: torch.Tensor = None, packed_seq_params: PackedSeqParams = None, @@ -1589,7 +1747,12 @@ def forward( nvtx_range_push("compressed_sparse_attn") if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': - output = self._forward_thd(query, key, x, qr, packed_seq_params) + if self.pg_collection.cp is not None and self.pg_collection.cp.size() > 1: + output = self._forward_thd_cp( + query, key, x, qr, boundary_hidden, boundary_kv, packed_seq_params + ) + else: + output = self._forward_thd(query, key, x, qr, packed_seq_params) nvtx_range_pop("compressed_sparse_attn") return output @@ -2026,6 +2189,280 @@ def _forward_fused_indexer_training_thd( output = output.unsqueeze(1) return output, indexer_loss + def _forward_thd_cp( + self, + query: torch.Tensor, + key: torch.Tensor, + x: torch.Tensor, + qr: torch.Tensor, + boundary_hidden: torch.Tensor, + boundary_kv: torch.Tensor, + packed_seq_params: PackedSeqParams, + ) -> torch.Tensor: + """THD-packed context-parallel branch. + + Build this rank's local KV context from boundary rows and fixed-capacity + compressed KV, then run sparse attention with optional indexer loss. + """ + # ---- Step 1: CP metadata and THD shape contract ---------------------- + cp_group = self.pg_collection.cp + cp_size = cp_group.size() + cp_rank = cp_group.rank() + + l_local = query.shape[0] + if l_local != key.shape[0]: + raise RuntimeError("DSv4 THD CP path currently supports self-attention only.") + cu_seqlens_q = ( + packed_seq_params.cu_seqlens_q_padded + if packed_seq_params.cu_seqlens_q_padded is not None + else packed_seq_params.cu_seqlens_q + ) + cu_seqlens_kv = ( + packed_seq_params.cu_seqlens_kv_padded + if packed_seq_params.cu_seqlens_kv_padded is not None + else packed_seq_params.cu_seqlens_kv + ) + max_seqlen_q = int(packed_seq_params.max_seqlen_q) + + # ---- Step 2: local CP rows, local KV, and boundary tensors ------------ + chunk_ranges = cp_utils.local_q_cp_chunk_ranges( + self.config.csa_cp_partition_mode, l_local, cp_size, cp_rank + ) + kv_local = key.squeeze(-2).squeeze(1) + if boundary_hidden is None or boundary_kv is None: + raise RuntimeError( + "DSv4 THD CP path requires boundary_hidden and boundary_kv from " + "the hidden-only boundary exchange and boundary KV projection path." + ) + boundary_kv = boundary_kv.squeeze(-2).squeeze(1) + d_window = boundary_hidden.shape[0] // len(chunk_ranges) + # Window-only defaults; compression fills rank-major buffers/metadata. + compressed_kv_rank_major = kv_local.new_empty((0, kv_local.shape[-1])) + cu_seqlens_compressed = None + seq_ids_rank_major = torch.empty((0,), dtype=torch.int32, device=query.device) + comp_ids_rank_major = torch.empty((0,), dtype=torch.int32, device=query.device) + valid_rank_major = torch.empty((0,), dtype=torch.bool, device=query.device) + + # Logical top-k uses per-sequence comp ids; the map lowers them to + # rank-major gathered rows when needed. + indexer_topk_logical = rank_major_by_seq_major = None + ratio = self.compress_ratio + indexer = self.indexer + indexer_loss_coeff = self.config.dsa_indexer_loss_coeff or 0.0 + training_with_grad = self.training and torch.is_grad_enabled() + + if self.compressor is not None and ratio > 1: + # ---- Step 3: build fixed-capacity compressor input ---------------- + + # ratio=4 overlaps one previous group; ratio=128 is non-overlap. + d_comp = 8 if ratio == 4 else ratio + + # One compressed row per full ratio-sized group; tails are dropped. + cu_seqlens_compressed = cp_utils.build_global_compressed_cu_seqlens( + cu_seqlens_kv, ratio + ) + seq_major_compressed_capacity = (l_local * cp_size) // ratio + + # Compact local+boundary tokens for Compressor. ``cu_compact`` is + # local token prefixes; ``comp_ids_local`` drives compressed RoPE. + hidden_compact, cu_compact, comp_ids_local = ( + cp_utils.build_cp_compressor_prep_compact_fused( + x, boundary_hidden, cu_seqlens_kv, chunk_ranges, ratio, d_comp, d_window + ) + ) + + # Rank-major metadata: sequence id, comp id, and canonical validity. + seq_ids_rank_major, comp_ids_rank_major, valid_rank_major = ( + cp_utils.build_cp_rank_major_compressed_metadata_fused( + cu_seqlens_kv, chunk_ranges, cp_size, ratio, d_comp + ) + ) + + if indexer is not None: + # ---- Step 4: optional indexer compressed path ----------------- + if ( + training_with_grad + and indexer_loss_coeff > 0 + and not self.config.dsa_indexer_use_sparse_loss + ): + raise RuntimeError( + "DSv4 THD CP path currently supports sparse indexer loss only. " + "Dense CP-aware indexer loss needs a CP-aware dense score kernel." + ) + + q_indexer_cp, weights_indexer_cp = indexer._forward_thd_query_weights_cp( + x.detach(), qr.detach(), cu_seqlens_q, max_seqlen_q, chunk_ranges=chunk_ranges + ) + + indexer_compressed_local, _ = indexer.compressor._forward_thd( + hidden_compact.detach(), + cu_compact, + max_seqlen_q=max_seqlen_q, + compressed_group_ids=comp_ids_local, + fixed_total_comp=hidden_compact.shape[0] // ratio, + ) + # Gather indexer compressed K in rank-major order. + k_indexer_rank_major = cp_utils.all_gather_fixed_cp_tensor( + indexer_compressed_local.squeeze(1), cp_group + ) + + # Repack valid compressed K to seq-major order for top-k. + k_indexer_seq_major, rank_major_by_seq_major = ( + cp_utils.repack_rank_major_compressed_to_seq_major_fused( + k_indexer_rank_major, + seq_ids_rank_major, + comp_ids_rank_major, + valid_rank_major, + cu_seqlens_compressed, + seq_major_rows=seq_major_compressed_capacity, + ) + ) + # Top-k is still in logical compressed coordinates here. + indexer_topk_logical = cp_utils.compute_cp_indexer_topk_logical_fused( + q_indexer_cp, + weights_indexer_cp, + k_indexer_seq_major, + cu_seqlens_q, + cu_seqlens_compressed, + chunk_ranges, + ratio, + indexer.index_topk, + indexer.softmax_scale, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_q // ratio, + ) + + # ---- Step 5: attention compressed KV path ------------------------- + compressed_kv_local, _ = self.compressor._forward_thd( + hidden_compact, + cu_compact, + max_seqlen_q=max_seqlen_q, + compressed_group_ids=comp_ids_local, + fixed_total_comp=hidden_compact.shape[0] // ratio, + ) + # Gather attention compressed KV in the same rank-major layout. + compressed_kv_rank_major = cp_utils.all_gather_fixed_cp_tensor( + compressed_kv_local.squeeze(1), cp_group + ) + # Two-chunk lowering needs the seq-major -> rank-major map. + if len(chunk_ranges) == 2 and rank_major_by_seq_major is None: + _, rank_major_by_seq_major = ( + cp_utils.repack_rank_major_compressed_to_seq_major_fused( + compressed_kv_rank_major, + seq_ids_rank_major, + comp_ids_rank_major, + valid_rank_major, + cu_seqlens_compressed, + seq_major_rows=seq_major_compressed_capacity, + ) + ) + + # ---- Step 6: pack the local full-KV layout --------------------------- + # Pack the local KV buffer addressed by the final CP attention indices. + kv_full_thd, shared_compressed_base = cp_utils.pack_cp_kv_full_fused( + kv_local, + boundary_kv, + compressed_kv_rank_major, + seq_ids_rank_major, + comp_ids_rank_major, + valid_rank_major, + cu_seqlens_kv, + chunk_ranges, + d_window, + ratio, + rank_major_by_seq_major=rank_major_by_seq_major, + cu_seqlens_compressed=cu_seqlens_compressed, + ) + if training_with_grad and indexer_topk_logical is not None: + # ---- Step 7a: indexer-loss path ---------------------------------- + # Build attention indices plus rank-major indices for indexer loss. + topk_idxs, indexer_topk_rank_major = cp_utils.build_cp_indexer_loss_indices_fused( + cu_seqlens_q, + cu_seqlens_compressed, + chunk_ranges, + d_window, + self.window_size, + ratio, + indexer_topk_logical, + rank_major_by_seq_major, + shared_compressed_base=shared_compressed_base, + ) + cu_seqlens_q_unpadded = None + if ( + packed_seq_params.cu_seqlens_q is not None + and packed_seq_params.cu_seqlens_q_padded is not None + and packed_seq_params.cu_seqlens_q.data_ptr() + != packed_seq_params.cu_seqlens_q_padded.data_ptr() + ): + cu_seqlens_q_unpadded = packed_seq_params.cu_seqlens_q + q_padding_mask = cp_utils.build_thd_cp_local_padding_row_mask( + cu_seqlens_q, + cu_seqlens_q_unpadded, + chunk_ranges, + total_padded_rows=l_local * cp_size, + ) + output, indexer_loss = ( + FusedIndexerSparseAttnFromTopkFunc.apply + if self.apply_dsa_kernel_fusion + else _unfused_indexer_sparse_attn_from_topk + )( + query, + kv_full_thd, + self.attn_sink.float(), + topk_idxs.int(), + q_indexer_cp, + k_indexer_rank_major, + weights_indexer_cp, + indexer_topk_rank_major.int(), + compressed_kv_rank_major, + self.softmax_scale, + indexer.softmax_scale, + indexer_loss_coeff, + self.config.calculate_per_token_loss, + l_local * cp_size, + q_padding_mask, + ) + 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), + reduce_group=cp_group, + ) + output = DSAIndexerLossAutoScaler.apply(output, indexer_loss) + return output.unsqueeze(1) + + # ---- Step 7b: sparse attention path ---------------------------------- + # No indexer-loss backward: lower final attention indices and run attention only. + topk_idxs, topk_length = cp_utils.build_cp_attention_indices_fused( + cu_seqlens_q, + chunk_ranges, + d_window, + self.window_size, + ratio, + indexer_topk_logical, + max_n_compressed=max_seqlen_q // ratio if ratio > 1 else 0, + rank_major_by_seq_major=rank_major_by_seq_major, + cu_seqlens_compressed=cu_seqlens_compressed, + shared_compressed_base=shared_compressed_base, + ) + + if self.apply_dsa_kernel_fusion: + output = dsa_sparse_attn( + query, + kv_full_thd, + self.attn_sink.float(), + topk_idxs.int(), + self.softmax_scale, + topk_length=topk_length, + is_thd=True, + ) + else: + output = unfused_compressed_sparse_attn( + query, kv_full_thd, self.attn_sink.float(), topk_idxs.int(), self.softmax_scale + ) + return output.unsqueeze(1) + def _forward_thd( self, query: torch.Tensor, # (total_q, np, hn) TE THD convention diff --git a/megatron/core/transformer/experimental_attention_variant/csa_cp_layout_kernels.py b/megatron/core/transformer/experimental_attention_variant/csa_cp_layout_kernels.py new file mode 100644 index 00000000000..ebf0df3ac6f --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/csa_cp_layout_kernels.py @@ -0,0 +1,3253 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +"""CuTeDSL kernels for DSv4 THD context parallel layout work. + +This module is intentionally below the MCore-facing CP utility layer. Callers +outside ``csa_cp_utils.py`` should not import this file directly. The functions +here allocate output tensors, compile/cache CuTeDSL launches, and execute small +layout kernels for the DSv4 CP path. +""" + +import math +import os +from functools import lru_cache +from typing import Optional, Tuple + +import torch + +try: + import cuda.bindings.driver as cuda + import cutlass + import cutlass.cute as cute + from cutlass._mlir.dialects import llvm + from cutlass.cute.runtime import from_dlpack, make_fake_stream + from cutlass.cutlass_dsl import T, dsl_user_op + + _CUTE_AVAILABLE = True +except ImportError: + cuda = None + cutlass = None + cute = None + llvm = None + T = None + dsl_user_op = None + from_dlpack = None + make_fake_stream = None + _CUTE_AVAILABLE = False + +# ============================================================================= +# CuTeDSL Kernel Definitions +# ============================================================================= +# Device kernels and CuTe launch functions. These blocks describe the GPU work; +# tensor allocation, validation, and autograd integration stay in the wrappers. +# ============================================================================= + +if _CUTE_AVAILABLE: + + @dsl_user_op + def _ptr_as_i64( + tensor: cute.Tensor, offset: cutlass.Int32, *, loc=None, ip=None + ) -> cutlass.Int64: + elem_ptr = tensor.iterator + cutlass.Int32(offset) + return cutlass.Int64(llvm.ptrtoint(T.i64(), elem_ptr.llvm_ptr, loc=loc, ip=ip)) + + @dsl_user_op + def _ld_global_v4_u64( + base_ptr: cutlass.Int64, *, loc=None, ip=None + ) -> Tuple[cutlass.Uint64, cutlass.Uint64, cutlass.Uint64, cutlass.Uint64]: + result = llvm.inline_asm( + llvm.StructType.get_literal([T.i64(), T.i64(), T.i64(), T.i64()]), + [cutlass.Int64(base_ptr).ir_value(loc=loc, ip=ip)], + "ld.global.L1::no_allocate.v4.u64 {$0, $1, $2, $3}, [$4];", + "=l,=l,=l,=l,l", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + v0 = llvm.extractvalue(T.i64(), result, [0], loc=loc, ip=ip) + v1 = llvm.extractvalue(T.i64(), result, [1], loc=loc, ip=ip) + v2 = llvm.extractvalue(T.i64(), result, [2], loc=loc, ip=ip) + v3 = llvm.extractvalue(T.i64(), result, [3], loc=loc, ip=ip) + return cutlass.Uint64(v0), cutlass.Uint64(v1), cutlass.Uint64(v2), cutlass.Uint64(v3) + + @dsl_user_op + def _st_global_v4_u64( + base_ptr: cutlass.Int64, + v0: cutlass.Uint64, + v1: cutlass.Uint64, + v2: cutlass.Uint64, + v3: cutlass.Uint64, + *, + loc=None, + ip=None, + ): + llvm.inline_asm( + None, + [ + cutlass.Int64(base_ptr).ir_value(loc=loc, ip=ip), + cutlass.Uint64(v0).ir_value(loc=loc, ip=ip), + cutlass.Uint64(v1).ir_value(loc=loc, ip=ip), + cutlass.Uint64(v2).ir_value(loc=loc, ip=ip), + cutlass.Uint64(v3).ir_value(loc=loc, ip=ip), + ], + "st.global.L1::evict_first.v4.u64 [$0], {$1, $2, $3, $4};", + "l,l,l,l,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + def _launch_named( + kernel, name: str, args: Tuple, grid: Tuple, block: Tuple, stream: cuda.CUstream + ): + """Launch one CSA CP CuTe kernel and keep its timeline prefix local.""" + kernel.set_name_prefix(name) + launcher = kernel(*args) + try: + launcher.launch(grid=grid, block=block, stream=stream) + finally: + if hasattr(kernel, "_name_prefix"): + kernel._name_prefix = None + dsl = getattr(launcher, "dsl", None) + if dsl is not None and hasattr(dsl, "_name_prefix"): + dsl._name_prefix = None + + def _launch_1d(kernel, name: str, args: Tuple, work: cutlass.Int32, stream: cuda.CUstream): + """Launch a 1-D CuTe kernel with 128 threads per block.""" + _launch_named( + kernel, + name, + args, + grid=(cute.ceil_div(work, 128), 1, 1), + block=(128, 1, 1), + stream=stream, + ) + + # Kernel contract: + # x/out: row-major tensor views, shape (local_rows, row_width), bf16/fp16/fp32. + # cos/sin: rotary tables, shape (max_seq_len, pos_dim), floating dtype. + # cu_seqlens: int32 global sequence prefix sums, shape (n_seq + 1,). + # chunk0/1_global_start and chunk_len map local rows back to two global + # packed-token chunks. Applies non-interleaved RoPE to each head's position + # dims; if adjoint is set, applies the backward/adjoint transform. + @cute.kernel + def _thd_local_rope_kernel( + x: cute.Tensor, + out: cute.Tensor, + cos: cute.Tensor, + sin: cute.Tensor, + cu_seqlens: cute.Tensor, + n_seq: cutlass.Constexpr, + chunk0_global_start: cutlass.Int32, + chunk1_global_start: cutlass.Int32, + chunk_len: cutlass.Int32, + clamp_to_valid_token: cutlass.Constexpr, + row_width: cutlass.Constexpr, + head_dim: cutlass.Constexpr, + nope_dim: cutlass.Constexpr, + pos_dim: cutlass.Constexpr, + inverse: cutlass.Constexpr, + adjoint: cutlass.Constexpr, + head_group: cutlass.Constexpr, + total_work: cutlass.Int32, + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + row = bidx + if row < total_work: + n_heads = row_width // head_dim + if cutlass.const_expr( + x.element_type.width == 16 and nope_dim % 16 == 0 and head_dim % 16 == 0 + ): + copy_vec_elems = cutlass.Int32(16) + copy_vecs_per_head = nope_dim // 16 + copy_use_u256 = True + elif cutlass.const_expr( + x.element_type.width == 32 and nope_dim % 8 == 0 and head_dim % 8 == 0 + ): + copy_vec_elems = cutlass.Int32(8) + copy_vecs_per_head = nope_dim // 8 + copy_use_u256 = True + elif cutlass.const_expr( + ( + x.element_type.width == 16 + and nope_dim % 4 == 0 + and head_dim % 4 == 0 + ) + or ( + x.element_type.width == 32 + and nope_dim % 2 == 0 + and head_dim % 2 == 0 + ) + ): + copy_use_u256 = False + if cutlass.const_expr(x.element_type.width == 16): + copy_vec_elems = cutlass.Int32(4) + copy_vecs_per_head = nope_dim // 4 + else: + copy_vec_elems = cutlass.Int32(2) + copy_vecs_per_head = nope_dim // 2 + else: + copy_vec_elems = cutlass.Int32(1) + copy_vecs_per_head = nope_dim + copy_use_u256 = False + + copy_work = tidx + while copy_work < n_heads * copy_vecs_per_head: + unit = copy_work % copy_vecs_per_head + head = copy_work // copy_vecs_per_head + offset = row * row_width + head * head_dim + unit * copy_vec_elems + if cutlass.const_expr( + ( + x.element_type.width == 16 + and nope_dim % 4 == 0 + and head_dim % 4 == 0 + ) + or ( + x.element_type.width == 32 + and nope_dim % 2 == 0 + and head_dim % 2 == 0 + ) + ): + if cutlass.const_expr(copy_use_u256): + src_ptr = _ptr_as_i64(x, offset) + dst_ptr = _ptr_as_i64(out, offset) + v0, v1, v2, v3 = _ld_global_v4_u64(src_ptr) + _st_global_v4_u64(dst_ptr, v0, v1, v2, v3) + else: + src_ptr = cute.recast_ptr(x.iterator + offset, dtype=cutlass.Int64) + dst_ptr = cute.recast_ptr(out.iterator + offset, dtype=cutlass.Int64) + value = cute.arch.load( + src_ptr.llvm_ptr, + cutlass.Int64, + level1_eviction_priority="evict_no_allocate", + ) + cute.arch.store( + dst_ptr.llvm_ptr, + value, + level1_eviction_priority="evict_first", + ) + else: + col = head * head_dim + unit + out[row, col] = x[row, col] + copy_work = copy_work + 128 + + rope_pos = 0 + if tidx == 0 or tidx == 32 or tidx == 64 or tidx == 96: + local_row = row + global_token = chunk0_global_start + local_row + if row >= chunk_len: + local_row = row - chunk_len + global_token = chunk1_global_start + local_row + + if clamp_to_valid_token != 0: + if global_token < 0: + global_token = 0 + last_token = cu_seqlens[n_seq] - 1 + if global_token > last_token: + global_token = last_token + + for seq in range(n_seq): + seq_start = cu_seqlens[seq] + seq_end = cu_seqlens[seq + 1] + if global_token >= seq_start and global_token < seq_end: + rope_pos = global_token - seq_start + + rope_pos = cute.arch.shuffle_sync(rope_pos, 0, mask=-1, mask_and_clamp=31) + + pos_pairs = pos_dim // 2 + grouped_heads = n_heads // head_group + work = tidx + while work < grouped_heads * pos_pairs: + pair = work % pos_pairs + head_block = work // pos_pairs + cos_left = cos[rope_pos, pair].to(cutlass.Float32) + sin_left = sin[rope_pos, pair].to(cutlass.Float32) + cos_right = cos[rope_pos, pos_pairs + pair].to(cutlass.Float32) + sin_right = sin[rope_pos, pos_pairs + pair].to(cutlass.Float32) + if inverse != 0: + sin_left = cutlass.Float32(0.0) - sin_left + sin_right = cutlass.Float32(0.0) - sin_right + + for slot in cutlass.range_constexpr(16): + if slot < head_group: + head = head_block * head_group + slot + col = head * head_dim + nope_dim + pair * 2 + x1 = x[row, col].to(cutlass.Float32) + x2 = x[row, col + 1].to(cutlass.Float32) + + x1_cos_left = (x1 * cos_left).to(out.element_type).to(cutlass.Float32) + x2_sin_left = (x2 * sin_left).to(out.element_type).to(cutlass.Float32) + x2_cos_right = (x2 * cos_right).to(out.element_type).to(cutlass.Float32) + x1_sin_right = (x1 * sin_right).to(out.element_type).to(cutlass.Float32) + value_left = x1_cos_left - x2_sin_left + value_right = x2_cos_right + x1_sin_right + if adjoint != 0: + x2_sin_right = (x2 * sin_right).to(out.element_type).to( + cutlass.Float32 + ) + x1_sin_left = (x1 * sin_left).to(out.element_type).to( + cutlass.Float32 + ) + value_left = x1_cos_left + x2_sin_right + value_right = cutlass.Float32(0.0) - x1_sin_left + x2_cos_right + out[row, col] = value_left.to(out.element_type) + out[row, col + 1] = value_right.to(out.element_type) + work = work + 128 + + # Launch contract for _thd_local_rope_kernel. + # row_width is static; total_work is the local row count. + @cute.jit + def _thd_local_rope_launch( + x: cute.Tensor, + out: cute.Tensor, + cos: cute.Tensor, + sin: cute.Tensor, + cu_seqlens: cute.Tensor, + n_seq: cutlass.Constexpr, + chunk0_global_start: cutlass.Int32, + chunk1_global_start: cutlass.Int32, + chunk_len: cutlass.Int32, + clamp_to_valid_token: cutlass.Constexpr, + row_width: cutlass.Constexpr, + head_dim: cutlass.Constexpr, + nope_dim: cutlass.Constexpr, + pos_dim: cutlass.Constexpr, + inverse: cutlass.Constexpr, + adjoint: cutlass.Constexpr, + head_group: cutlass.Constexpr, + total_work: cutlass.Int32, + stream: cuda.CUstream, + ): + _launch_named( + _thd_local_rope_kernel, + "dsv4_cp_thd_local_rope", + ( + x, + out, + cos, + sin, + cu_seqlens, + n_seq, + chunk0_global_start, + chunk1_global_start, + chunk_len, + clamp_to_valid_token, + row_width, + head_dim, + nope_dim, + pos_dim, + inverse, + adjoint, + head_group, + total_work, + ), + grid=(total_work, 1, 1), + block=(128, 1, 1), + stream=stream, + ) + + # Kernel contract: + # x/out: row-major compressed tensor views, shape (rows, row_width). + # cos/sin: rotary tables, shape (max_seq_len, pos_dim). + # comp_ids: int32 compressed group id per row, shape (rows,); -1 is treated as 0. + # ratio maps comp_id to original token position comp_id * ratio. + # Copies no-RoPE dims and rotates each non-interleaved pair in the position dims. + @cute.kernel + def _thd_compressed_rope_kernel( + x: cute.Tensor, + out: cute.Tensor, + cos: cute.Tensor, + sin: cute.Tensor, + comp_ids: cute.Tensor, + ratio: cutlass.Int32, + row_width: cutlass.Constexpr, + head_dim: cutlass.Int32, + nope_dim: cutlass.Int32, + pos_dim: cutlass.Int32, + inverse: cutlass.Int32, + adjoint: cutlass.Int32, + total_work: cutlass.Int32, + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + linear = bidx * 128 + tidx + if linear < total_work: + row = linear // row_width + col = linear - row * row_width + head_col = col - (col // head_dim) * head_dim + + if head_col < nope_dim: + out[row, col] = x[row, col] + else: + comp_id = comp_ids[row] + if comp_id < 0: + comp_id = 0 + rope_pos = comp_id * ratio + + pos_col = head_col - nope_dim + pair = pos_col // 2 + pair_base = col - head_col + nope_dim + pair * 2 + x1 = x[row, pair_base].to(cutlass.Float32) + x2 = x[row, pair_base + 1].to(cutlass.Float32) + + cos_left = cos[rope_pos, pair].to(cutlass.Float32) + sin_left = sin[rope_pos, pair].to(cutlass.Float32) + cos_right = cos[rope_pos, pos_dim // 2 + pair].to(cutlass.Float32) + sin_right = sin[rope_pos, pos_dim // 2 + pair].to(cutlass.Float32) + if inverse != 0: + sin_left = cutlass.Float32(0.0) - sin_left + sin_right = cutlass.Float32(0.0) - sin_right + + x1_cos_left = (x1 * cos_left).to(out.element_type).to(cutlass.Float32) + x2_sin_left = (x2 * sin_left).to(out.element_type).to(cutlass.Float32) + x2_cos_right = (x2 * cos_right).to(out.element_type).to(cutlass.Float32) + x1_sin_right = (x1 * sin_right).to(out.element_type).to(cutlass.Float32) + value = x1_cos_left - x2_sin_left + if pos_col - pair * 2 != 0: + value = x2_cos_right + x1_sin_right + if adjoint != 0: + x2_sin_right = (x2 * sin_right).to(out.element_type).to(cutlass.Float32) + x1_sin_left = (x1 * sin_left).to(out.element_type).to(cutlass.Float32) + value = x1_cos_left + x2_sin_right + if pos_col - pair * 2 != 0: + value = cutlass.Float32(0.0) - x1_sin_left + x2_cos_right + out[row, col] = value.to(out.element_type) + + # Launch contract for _thd_compressed_rope_kernel. + # row_width is static; total_work is rows * row_width. + @cute.jit + def _thd_compressed_rope_launch( + x: cute.Tensor, + out: cute.Tensor, + cos: cute.Tensor, + sin: cute.Tensor, + comp_ids: cute.Tensor, + ratio: cutlass.Int32, + row_width: cutlass.Constexpr, + head_dim: cutlass.Int32, + nope_dim: cutlass.Int32, + pos_dim: cutlass.Int32, + inverse: cutlass.Int32, + adjoint: cutlass.Int32, + total_work: cutlass.Int32, + stream: cuda.CUstream, + ): + _launch_1d( + _thd_compressed_rope_kernel, + "dsv4_cp_thd_compressed_rope", + ( + x, + out, + cos, + sin, + comp_ids, + ratio, + row_width, + head_dim, + nope_dim, + pos_dim, + inverse, + adjoint, + total_work, + ), + total_work, + stream, + ) + + # Kernel contract: + # hidden_local: local hidden rows, shape (l_local, row_width), bf16/fp16/fp32. + # boundary_hidden: left boundary rows, shape (d_window, row_width). + # hidden_compact: output compact rows, shape (compact_len, row_width). + # cu_compact: int32 output prefix, shape (n_seq + 1,). + # seq_ids/comp_ids/valid: compressed metadata outputs, shape (c_cap,). + # Enumerates visible full compression groups in [global_start-d_comp, + # global_start+l_local), copies their ratio tokens from boundary/local + # hidden into hidden_compact, and emits compact sequence metadata. + @cute.kernel + def _compressor_input_compact_fwd_kernel( + hidden_local: cute.Tensor, + boundary_hidden: cute.Tensor, + hidden_compact: cute.Tensor, + cu_seqlens: cute.Tensor, + cu_compact: cute.Tensor, + seq_ids: cute.Tensor, + comp_ids: cute.Tensor, + valid: cute.Tensor, + n_seq: cutlass.Int32, + global_start: cutlass.Int32, + l_local: cutlass.Int32, + ratio: cutlass.Constexpr, + d_comp: cutlass.Constexpr, + d_window: cutlass.Constexpr, + c_cap: cutlass.Int32, + compact_len: cutlass.Int32, + row_width: cutlass.Constexpr, + total_work: cutlass.Int32, + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + row = bidx + + if cutlass.const_expr(hidden_compact.element_type.width == 16 and row_width % 4 == 0): + vec_elems = cutlass.Int32(4) + vec_cols = row_width // 4 + else: + vec_elems = cutlass.Int32(1) + vec_cols = row_width + + range_start = global_start + range_end = global_start + l_local + first_range_group_start = range_start - d_comp + + src_global = cutlass.Int32(-1) + if row < compact_len: + if tidx == 0 or tidx == 32: + running_tokens = 0 + for seq in range(n_seq): + seq_start = cu_seqlens[seq] + seq_end = cu_seqlens[seq + 1] + local_seq_start = seq_start + if local_seq_start < range_start: + local_seq_start = range_start + local_seq_end = seq_end + if local_seq_end > range_end: + local_seq_end = range_end + + if local_seq_start < local_seq_end: + seq_full_group_count = (seq_end - seq_start) // ratio + first_visible_numer = first_range_group_start - seq_start + if first_visible_numer < 0: + first_visible_numer = 0 + first_visible_group = 0 + if first_visible_numer > 0: + first_visible_group = (first_visible_numer + ratio - 1) // ratio + stop_visible_group = (local_seq_end - seq_start) // ratio + if stop_visible_group > seq_full_group_count: + stop_visible_group = seq_full_group_count + visible_group_count = stop_visible_group - first_visible_group + if visible_group_count < 0: + visible_group_count = 0 + visible_token_count = visible_group_count * ratio + + if row >= running_tokens and row < running_tokens + visible_token_count: + local_visible_token = row - running_tokens + comp_id = first_visible_group + local_visible_token // ratio + token_in_group = ( + local_visible_token - (local_visible_token // ratio) * ratio + ) + src_global = seq_start + comp_id * ratio + token_in_group + running_tokens = running_tokens + visible_token_count + + src_global = cute.arch.shuffle_sync(src_global, 0, mask=-1, mask_and_clamp=31) + + vec_col = tidx + while vec_col < vec_cols: + dst_offset = row * row_width + vec_col * vec_elems + if cutlass.const_expr( + hidden_compact.element_type.width == 16 and row_width % 4 == 0 + ): + dst_ptr = cute.recast_ptr( + hidden_compact.iterator + dst_offset, dtype=cutlass.Int64 + ) + value = cutlass.Int64(0) + if src_global >= 0: + if src_global < range_start: + src_row = src_global - (range_start - d_window) + src_offset = src_row * row_width + vec_col * vec_elems + src_ptr = cute.recast_ptr( + boundary_hidden.iterator + src_offset, dtype=cutlass.Int64 + ) + value = cute.arch.load( + src_ptr.llvm_ptr, + cutlass.Int64, + level1_eviction_priority="evict_no_allocate", + ) + else: + src_row = src_global - range_start + src_offset = src_row * row_width + vec_col * vec_elems + src_ptr = cute.recast_ptr( + hidden_local.iterator + src_offset, dtype=cutlass.Int64 + ) + value = cute.arch.load( + src_ptr.llvm_ptr, + cutlass.Int64, + level1_eviction_priority="evict_no_allocate", + ) + cute.arch.store( + dst_ptr.llvm_ptr, + value, + level1_eviction_priority="evict_first", + ) + else: + value = cutlass.Float32(0.0).to(hidden_compact.element_type) + if src_global >= 0: + if src_global < range_start: + src_row = src_global - (range_start - d_window) + value = boundary_hidden[src_row, vec_col] + else: + src_row = src_global - range_start + value = hidden_local[src_row, vec_col] + hidden_compact[row, vec_col] = value + vec_col = vec_col + 64 + + if row < c_cap and tidx == 0: + slot = row + seq_ids[slot] = -1 + comp_ids[slot] = -1 + valid[slot] = False + + running = 0 + for seq in range(n_seq): + seq_start = cu_seqlens[seq] + seq_end = cu_seqlens[seq + 1] + local_seq_start = seq_start + if local_seq_start < range_start: + local_seq_start = range_start + local_seq_end = seq_end + if local_seq_end > range_end: + local_seq_end = range_end + + if local_seq_start < local_seq_end: + seq_full_group_count = (seq_end - seq_start) // ratio + first_visible_numer = first_range_group_start - seq_start + if first_visible_numer < 0: + first_visible_numer = 0 + first_visible_group = 0 + if first_visible_numer > 0: + first_visible_group = (first_visible_numer + ratio - 1) // ratio + stop_visible_group = (local_seq_end - seq_start) // ratio + if stop_visible_group > seq_full_group_count: + stop_visible_group = seq_full_group_count + visible_group_count = stop_visible_group - first_visible_group + if visible_group_count < 0: + visible_group_count = 0 + + if slot >= running and slot < running + visible_group_count: + comp_id = first_visible_group + slot - running + group_end = seq_start + (comp_id + 1) * ratio + seq_ids[slot] = seq + comp_ids[slot] = comp_id + valid[slot] = group_end - 1 >= range_start and group_end - 1 < range_end + running = running + visible_group_count + + if row < n_seq + 1 and tidx == 1: + upto = row + running = 0 + for seq in range(n_seq): + if seq < upto: + seq_start = cu_seqlens[seq] + seq_end = cu_seqlens[seq + 1] + local_seq_start = seq_start + if local_seq_start < range_start: + local_seq_start = range_start + local_seq_end = seq_end + if local_seq_end > range_end: + local_seq_end = range_end + + if local_seq_start < local_seq_end: + seq_full_group_count = (seq_end - seq_start) // ratio + first_visible_numer = first_range_group_start - seq_start + if first_visible_numer < 0: + first_visible_numer = 0 + first_visible_group = 0 + if first_visible_numer > 0: + first_visible_group = (first_visible_numer + ratio - 1) // ratio + stop_visible_group = (local_seq_end - seq_start) // ratio + if stop_visible_group > seq_full_group_count: + stop_visible_group = seq_full_group_count + visible_group_count = stop_visible_group - first_visible_group + if visible_group_count < 0: + visible_group_count = 0 + running = running + visible_group_count * ratio + cu_compact[row] = running + + # Launch contract for _compressor_input_compact_fwd_kernel. + # total_work covers compact rows plus metadata rows. + @cute.jit + def _compressor_input_compact_fwd_launch( + hidden_local: cute.Tensor, + boundary_hidden: cute.Tensor, + hidden_compact: cute.Tensor, + cu_seqlens: cute.Tensor, + cu_compact: cute.Tensor, + seq_ids: cute.Tensor, + comp_ids: cute.Tensor, + valid: cute.Tensor, + n_seq: cutlass.Int32, + global_start: cutlass.Int32, + l_local: cutlass.Int32, + ratio: cutlass.Constexpr, + d_comp: cutlass.Constexpr, + d_window: cutlass.Constexpr, + c_cap: cutlass.Int32, + compact_len: cutlass.Int32, + row_width: cutlass.Constexpr, + total_work: cutlass.Int32, + stream: cuda.CUstream, + ): + _launch_named( + _compressor_input_compact_fwd_kernel, + "dsv4_cp_compressor_input_compact_fwd", + ( + hidden_local, + boundary_hidden, + hidden_compact, + cu_seqlens, + cu_compact, + seq_ids, + comp_ids, + valid, + n_seq, + global_start, + l_local, + ratio, + d_comp, + d_window, + c_cap, + compact_len, + row_width, + total_work, + ), + grid=(total_work, 1, 1), + block=(64, 1, 1), + stream=stream, + ) + + # Kernel contract: + # grad_hidden_compact: gradient of compact rows, shape (compact_len, row_width). + # grad_hidden_local: output gradient, shape (l_local, row_width). + # grad_boundary_hidden: output gradient, shape (d_window, row_width). + # Reconstructs the same compact source mapping as forward and scatters each + # compact-row gradient back to either boundary or local hidden. + @cute.kernel + def _compressor_input_compact_bwd_kernel( + grad_hidden_compact: cute.Tensor, + grad_hidden_local: cute.Tensor, + grad_boundary_hidden: cute.Tensor, + cu_seqlens: cute.Tensor, + n_seq: cutlass.Int32, + global_start: cutlass.Int32, + l_local: cutlass.Int32, + ratio: cutlass.Constexpr, + d_comp: cutlass.Constexpr, + d_window: cutlass.Constexpr, + compact_len: cutlass.Int32, + row_width: cutlass.Constexpr, + total_work: cutlass.Int32, + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + out_row = bidx + + if cutlass.const_expr( + grad_hidden_compact.element_type.width == 16 and row_width % 4 == 0 + ): + vec_elems = cutlass.Int32(4) + vec_cols = row_width // 4 + else: + vec_elems = cutlass.Int32(1) + vec_cols = row_width + + compact_row = cutlass.Int32(-1) + if out_row < total_work: + range_start = global_start + range_end = global_start + l_local + first_range_group_start = range_start - d_comp + src_global = range_start + (out_row - d_window) + dst_row = out_row - d_window + if out_row < d_window: + src_global = range_start - d_window + out_row + dst_row = out_row + + if tidx == 0 or tidx == 32: + running_tokens = 0 + for seq in range(n_seq): + seq_start = cu_seqlens[seq] + seq_end = cu_seqlens[seq + 1] + local_seq_start = seq_start + if local_seq_start < range_start: + local_seq_start = range_start + local_seq_end = seq_end + if local_seq_end > range_end: + local_seq_end = range_end + + if local_seq_start < local_seq_end: + seq_full_group_count = (seq_end - seq_start) // ratio + first_visible_numer = first_range_group_start - seq_start + if first_visible_numer < 0: + first_visible_numer = 0 + first_visible_group = 0 + if first_visible_numer > 0: + first_visible_group = (first_visible_numer + ratio - 1) // ratio + stop_visible_group = (local_seq_end - seq_start) // ratio + if stop_visible_group > seq_full_group_count: + stop_visible_group = seq_full_group_count + visible_group_count = stop_visible_group - first_visible_group + if visible_group_count < 0: + visible_group_count = 0 + visible_token_count = visible_group_count * ratio + + if src_global >= seq_start and src_global < seq_end: + seq_offset = src_global - seq_start + comp_id = seq_offset // ratio + token_in_group = seq_offset - (seq_offset // ratio) * ratio + if comp_id >= first_visible_group and comp_id < stop_visible_group: + compact_row = ( + running_tokens + + (comp_id - first_visible_group) * ratio + + token_in_group + ) + running_tokens = running_tokens + visible_token_count + + compact_row = cute.arch.shuffle_sync(compact_row, 0, mask=-1, mask_and_clamp=31) + + vec_col = tidx + while vec_col < vec_cols: + dst_offset = dst_row * row_width + vec_col * vec_elems + if cutlass.const_expr( + grad_hidden_compact.element_type.width == 16 and row_width % 4 == 0 + ): + value = cutlass.Int64(0) + if compact_row >= 0 and compact_row < compact_len: + src_offset = compact_row * row_width + vec_col * vec_elems + src_ptr = cute.recast_ptr( + grad_hidden_compact.iterator + src_offset, dtype=cutlass.Int64 + ) + value = cute.arch.load( + src_ptr.llvm_ptr, + cutlass.Int64, + level1_eviction_priority="evict_no_allocate", + ) + if out_row < d_window: + dst_ptr = cute.recast_ptr( + grad_boundary_hidden.iterator + dst_offset, + dtype=cutlass.Int64, + ) + cute.arch.store( + dst_ptr.llvm_ptr, + value, + level1_eviction_priority="evict_first", + ) + else: + dst_ptr = cute.recast_ptr( + grad_hidden_local.iterator + dst_offset, dtype=cutlass.Int64 + ) + cute.arch.store( + dst_ptr.llvm_ptr, + value, + level1_eviction_priority="evict_first", + ) + else: + value = cutlass.Float32(0.0).to(grad_hidden_compact.element_type) + if compact_row >= 0 and compact_row < compact_len: + value = grad_hidden_compact[compact_row, vec_col] + if out_row < d_window: + grad_boundary_hidden[dst_row, vec_col] = value + else: + grad_hidden_local[dst_row, vec_col] = value + vec_col = vec_col + 64 + + # Launch contract for _compressor_input_compact_bwd_kernel. + # total_work is the number of boundary + local gradient rows. + @cute.jit + def _compressor_input_compact_bwd_launch( + grad_hidden_compact: cute.Tensor, + grad_hidden_local: cute.Tensor, + grad_boundary_hidden: cute.Tensor, + cu_seqlens: cute.Tensor, + n_seq: cutlass.Int32, + global_start: cutlass.Int32, + l_local: cutlass.Int32, + ratio: cutlass.Constexpr, + d_comp: cutlass.Constexpr, + d_window: cutlass.Constexpr, + compact_len: cutlass.Int32, + row_width: cutlass.Constexpr, + total_work: cutlass.Int32, + stream: cuda.CUstream, + ): + _launch_named( + _compressor_input_compact_bwd_kernel, + "dsv4_cp_compressor_input_compact_bwd", + ( + grad_hidden_compact, + grad_hidden_local, + grad_boundary_hidden, + cu_seqlens, + n_seq, + global_start, + l_local, + ratio, + d_comp, + d_window, + compact_len, + row_width, + total_work, + ), + grid=(total_work, 1, 1), + block=(64, 1, 1), + stream=stream, + ) + + # Kernel contract: + # kv_local/boundary_kv/compressed_rank_major: source tensors for forward, + # or output gradient tensors for backward, each flattened to + # ``(rows, row_width)``. + # kv_full: output tensor for forward, or grad_kv_full input for backward, + # shape ``(kv_full_capacity, row_width)``. + # rank-major metadata inputs describe compressed rows after all-gather. + # direction: 0 packs source tensors into kv_full; 1 scatters kv_full + # gradients back to the source tensors. + # Computes the physical KV-full row mapping inline. No source-map tensor is + # materialized or saved for backward; backward recomputes the same mapping. + @cute.kernel + def _thd_full_kv_pack_kernel( + kv_local: cute.Tensor, + boundary_kv: cute.Tensor, + compressed_rank_major: cute.Tensor, + kv_full: cute.Tensor, + seq_ids_rank_major: cute.Tensor, + comp_ids_rank_major: cute.Tensor, + valid_rank_major: cute.Tensor, + cu_seqlens: cute.Tensor, + cu_seqlens_compressed: cute.Tensor, + rank_major_by_seq_major: cute.Tensor, + n_seq: cutlass.Int32, + global_start: cutlass.Int32, + l_local: cutlass.Int32, + chunk0_start: cutlass.Int32, + chunk1_start: cutlass.Int32, + chunk_count: cutlass.Constexpr, + chunk_len: cutlass.Int32, + window_capacity_per_chunk: cutlass.Int32, + d_window: cutlass.Constexpr, + ratio: cutlass.Constexpr, + compressed_rows: cutlass.Int32, + use_seq_major_map: cutlass.Constexpr, + seq_major_rows: cutlass.Int32, + kv_full_capacity: cutlass.Int32, + row_width: cutlass.Constexpr, + direction: cutlass.Constexpr, + total_work: cutlass.Int32, + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + out_row = bidx + + if cutlass.const_expr(kv_full.element_type.width == 16 and row_width % 4 == 0): + vec_elems = cutlass.Int32(4) + vec_cols = row_width // 4 + else: + vec_elems = cutlass.Int32(1) + vec_cols = row_width + + if out_row < total_work: + source_kind = cutlass.Int32(-1) + source_index = cutlass.Int32(0) + resolve_row = tidx == 0 or tidx == 32 + + if resolve_row and chunk_count != 0: + shared_compressed_base = chunk_count * window_capacity_per_chunk + if out_row >= shared_compressed_base: + compressed_idx = out_row - shared_compressed_base + if compressed_idx >= 0 and compressed_idx < compressed_rows: + source_kind = 2 + source_index = compressed_idx + else: + chunk_id = out_row // window_capacity_per_chunk + chunk_row = out_row - chunk_id * window_capacity_per_chunk + chunk_start = chunk0_start + if chunk_id == 1: + chunk_start = chunk1_start + chunk_end = chunk_start + chunk_len + + offset = 0 + for seq in range(n_seq): + seq_start = cu_seqlens[seq] + seq_end = cu_seqlens[seq + 1] + local_seq_start = seq_start + if local_seq_start < chunk_start: + local_seq_start = chunk_start + local_seq_end = seq_end + if local_seq_end > chunk_end: + local_seq_end = chunk_end + if local_seq_start < local_seq_end: + window_start = local_seq_start - d_window + if window_start < seq_start: + window_start = seq_start + window_len = local_seq_end - window_start + if chunk_row >= offset and chunk_row < offset + window_len: + src_global = window_start + chunk_row - offset + if src_global < chunk_start: + source_kind = 1 + source_index = ( + chunk_id * d_window + src_global - (chunk_start - d_window) + ) + else: + source_kind = 0 + source_index = chunk_id * chunk_len + src_global - chunk_start + offset = offset + window_len + elif resolve_row: + offset = 0 + for seq in range(n_seq): + seq_start = cu_seqlens[seq] + seq_end = cu_seqlens[seq + 1] + global_end = global_start + l_local + local_seq_start = seq_start + if local_seq_start < global_start: + local_seq_start = global_start + local_seq_end = seq_end + if local_seq_end > global_end: + local_seq_end = global_end + + if local_seq_start < local_seq_end: + window_start = local_seq_start - d_window + if window_start < seq_start: + window_start = seq_start + window_len = local_seq_end - window_start + comp_len = 0 + comp_start = 0 + comp_end = 0 + if ratio > 1 and use_seq_major_map != 0: + comp_start = cu_seqlens_compressed[seq] + comp_end = cu_seqlens_compressed[seq + 1] + comp_len = comp_end - comp_start + elif ratio > 1: + comp_len = (seq_end - seq_start) // ratio + seq_total = window_len + comp_len + + if out_row >= offset and out_row < offset + seq_total: + inner = out_row - offset + if inner < window_len: + src_global = window_start + inner + if src_global < global_start: + source_kind = 1 + source_index = src_global - (global_start - d_window) + else: + source_kind = 0 + source_index = src_global - global_start + else: + comp_id = inner - window_len + if use_seq_major_map != 0: + seq_major = comp_start + comp_id + if seq_major >= comp_start and seq_major < comp_end: + if seq_major >= 0 and seq_major < seq_major_rows: + rank_major = rank_major_by_seq_major[seq_major] + if rank_major >= 0: + source_kind = 2 + source_index = rank_major + else: + # Slow fallback for no-indexer paths: linear scan is + # intentional and not expected to be the hot training path. + for comp_row in range(compressed_rows): + if ( + valid_rank_major[comp_row] + and seq_ids_rank_major[comp_row] == seq + and comp_ids_rank_major[comp_row] == comp_id + ): + source_kind = 2 + source_index = comp_row + offset = offset + seq_total + + source_kind = cute.arch.shuffle_sync(source_kind, 0, mask=-1, mask_and_clamp=31) + source_index = cute.arch.shuffle_sync(source_index, 0, mask=-1, mask_and_clamp=31) + + vec_col = tidx + while vec_col < vec_cols: + col = vec_col + if direction == 0: + if cutlass.const_expr(kv_full.element_type.width == 16 and row_width % 4 == 0): + dst_offset = out_row * row_width + vec_col * vec_elems + value = cutlass.Int64(0) + if source_kind == 0: + src_offset = source_index * row_width + vec_col * vec_elems + src_ptr = cute.recast_ptr( + kv_local.iterator + src_offset, dtype=cutlass.Int64 + ) + value = cute.arch.load( + src_ptr.llvm_ptr, + cutlass.Int64, + level1_eviction_priority="evict_no_allocate", + ) + elif source_kind == 1: + src_offset = source_index * row_width + vec_col * vec_elems + src_ptr = cute.recast_ptr( + boundary_kv.iterator + src_offset, dtype=cutlass.Int64 + ) + value = cute.arch.load( + src_ptr.llvm_ptr, + cutlass.Int64, + level1_eviction_priority="evict_no_allocate", + ) + elif source_kind == 2: + src_offset = source_index * row_width + vec_col * vec_elems + src_ptr = cute.recast_ptr( + compressed_rank_major.iterator + src_offset, + dtype=cutlass.Int64, + ) + value = cute.arch.load( + src_ptr.llvm_ptr, + cutlass.Int64, + level1_eviction_priority="evict_no_allocate", + ) + dst_ptr = cute.recast_ptr(kv_full.iterator + dst_offset, dtype=cutlass.Int64) + cute.arch.store( + dst_ptr.llvm_ptr, + value, + level1_eviction_priority="evict_first", + ) + else: + value = cutlass.Float32(0.0).to(kv_full.element_type) + if source_kind == 0: + value = kv_local[source_index, col] + elif source_kind == 1: + value = boundary_kv[source_index, col] + elif source_kind == 2: + value = compressed_rank_major[source_index, col] + kv_full[out_row, col] = value + else: + if cutlass.const_expr(kv_full.element_type.width == 16 and row_width % 4 == 0): + src_offset = out_row * row_width + vec_col * vec_elems + src_ptr = cute.recast_ptr(kv_full.iterator + src_offset, dtype=cutlass.Int64) + value = cute.arch.load( + src_ptr.llvm_ptr, + cutlass.Int64, + level1_eviction_priority="evict_no_allocate", + ) + if source_kind == 0: + dst_offset = source_index * row_width + vec_col * vec_elems + dst_ptr = cute.recast_ptr( + kv_local.iterator + dst_offset, dtype=cutlass.Int64 + ) + cute.arch.store( + dst_ptr.llvm_ptr, + value, + level1_eviction_priority="evict_first", + ) + elif source_kind == 1: + dst_offset = source_index * row_width + vec_col * vec_elems + dst_ptr = cute.recast_ptr( + boundary_kv.iterator + dst_offset, dtype=cutlass.Int64 + ) + cute.arch.store( + dst_ptr.llvm_ptr, + value, + level1_eviction_priority="evict_first", + ) + elif source_kind == 2: + dst_offset = source_index * row_width + vec_col * vec_elems + dst_ptr = cute.recast_ptr( + compressed_rank_major.iterator + dst_offset, + dtype=cutlass.Int64, + ) + cute.arch.store( + dst_ptr.llvm_ptr, + value, + level1_eviction_priority="evict_first", + ) + else: + if source_kind == 0: + kv_local[source_index, col] = kv_full[out_row, col] + elif source_kind == 1: + boundary_kv[source_index, col] = kv_full[out_row, col] + elif source_kind == 2: + compressed_rank_major[source_index, col] = kv_full[out_row, col] + vec_col = vec_col + 64 + + # Launch contract for _thd_full_kv_pack_kernel. + # Supports contiguous CP and two-chunk CP. ``use_seq_major_map`` selects + # whether compressed ids must be translated through rank_major_by_seq_major. + # ``row_width`` is constexpr to preserve the existing compiled layout key. + @cute.jit + def _thd_full_kv_pack_launch( + kv_local: cute.Tensor, + boundary_kv: cute.Tensor, + compressed_rank_major: cute.Tensor, + kv_full: cute.Tensor, + seq_ids_rank_major: cute.Tensor, + comp_ids_rank_major: cute.Tensor, + valid_rank_major: cute.Tensor, + cu_seqlens: cute.Tensor, + cu_seqlens_compressed: cute.Tensor, + rank_major_by_seq_major: cute.Tensor, + n_seq: cutlass.Int32, + global_start: cutlass.Int32, + l_local: cutlass.Int32, + chunk0_start: cutlass.Int32, + chunk1_start: cutlass.Int32, + chunk_count: cutlass.Constexpr, + chunk_len: cutlass.Int32, + window_capacity_per_chunk: cutlass.Int32, + d_window: cutlass.Constexpr, + ratio: cutlass.Constexpr, + compressed_rows: cutlass.Int32, + use_seq_major_map: cutlass.Constexpr, + seq_major_rows: cutlass.Int32, + kv_full_capacity: cutlass.Int32, + row_width: cutlass.Constexpr, + direction: cutlass.Constexpr, + total_work: cutlass.Int32, + stream: cuda.CUstream, + ): + _launch_named( + _thd_full_kv_pack_kernel, + "dsv4_cp_thd_full_kv_pack", + ( + kv_local, + boundary_kv, + compressed_rank_major, + kv_full, + seq_ids_rank_major, + comp_ids_rank_major, + valid_rank_major, + cu_seqlens, + cu_seqlens_compressed, + rank_major_by_seq_major, + n_seq, + global_start, + l_local, + chunk0_start, + chunk1_start, + chunk_count, + chunk_len, + window_capacity_per_chunk, + d_window, + ratio, + compressed_rows, + use_seq_major_map, + seq_major_rows, + kv_full_capacity, + row_width, + direction, + total_work, + ), + grid=(total_work, 1, 1), + block=(64, 1, 1), + stream=stream, + ) + + # Kernel contract: + # rank_by_seq_major: int32 output reverse map, shape (total_seq_major_rows,). + # Initializes the reverse map before the deterministic rank-map build. + @cute.kernel + def _repack_compressed_kv_init_rank_map_kernel( + rank_major_by_seq_major: cute.Tensor, seq_major_rows: cutlass.Int32 + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + linear = bidx * 128 + tidx + if linear < seq_major_rows: + rank_major_by_seq_major[linear] = -1 + + # Kernel contract: + # seq_ids/comp_ids/valid: rank-major metadata, shape (rows,). + # cu_compressed: int32 compressed sequence prefixes, shape (n_seq + 1,). + # rank_by_seq_major: int32 output reverse map, shape (total_seq_major_rows,). + # Builds a deterministic seq-major -> rank-major map. Multiple rank-major + # candidates can describe the same compressed group near CP chunk + # boundaries, so the kernel selects the greatest rank-major row with + # atomic max instead of racing concurrent value writes. + @cute.kernel + def _repack_compressed_kv_rank_map_kernel( + seq_ids_rank_major: cute.Tensor, + comp_ids_rank_major: cute.Tensor, + valid_rank_major: cute.Tensor, + cu_seqlens_compressed: cute.Tensor, + rank_major_by_seq_major: cute.Tensor, + rank_major_rows: cutlass.Int32, + seq_major_rows: cutlass.Int32, + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + rank_row = bidx * 128 + tidx + if rank_row < rank_major_rows and valid_rank_major[rank_row] == 1: + seq = seq_ids_rank_major[rank_row] + comp = comp_ids_rank_major[rank_row] + seq_major = cu_seqlens_compressed[seq] + comp + if seq_major >= 0 and seq_major < seq_major_rows: + rank_map_ptr = rank_major_by_seq_major.iterator + seq_major + cute.arch.atomic_max(rank_map_ptr.llvm_ptr, rank_row, sem="relaxed", scope="gpu") + + # Kernel contract: + # rank_major: compressed rows in rank-major order, shape (rows, row_width). + # rank_by_seq_major: deterministic reverse map built by + # _repack_compressed_kv_rank_map_kernel. + # out: seq-major compressed rows, shape (total_seq_major_rows, row_width). + # Gathers selected rank-major rows into seq-major order and zero-fills + # rows with no valid compressed group. + @cute.kernel + def _repack_compressed_kv_to_seq_major_kernel( + compressed_rank_major: cute.Tensor, + compressed_seq_major: cute.Tensor, + rank_major_by_seq_major: cute.Tensor, + rank_major_rows: cutlass.Int32, + seq_major_rows: cutlass.Int32, + row_width: cutlass.Constexpr, + total_work: cutlass.Int32, + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + linear = bidx * 128 + tidx + + if linear < seq_major_rows * row_width: + row = linear // row_width + col = linear - row * row_width + rank_row = rank_major_by_seq_major[row] + value = cutlass.Float32(0.0).to(compressed_seq_major.element_type) + if rank_row >= 0 and rank_row < rank_major_rows: + value = compressed_rank_major[rank_row, col] + compressed_seq_major[row, col] = value + + # Launch contract for _repack_compressed_kv_init_rank_map_kernel. + # launch_work is seq-major rows. + @cute.jit + def _repack_compressed_kv_init_rank_map_launch( + rank_major_by_seq_major: cute.Tensor, + seq_major_rows: cutlass.Int32, + launch_work: cutlass.Int32, + stream: cuda.CUstream, + ): + _launch_1d( + _repack_compressed_kv_init_rank_map_kernel, + "dsv4_cp_repack_compressed_kv_init_rank_map", + (rank_major_by_seq_major, seq_major_rows), + launch_work, + stream, + ) + + # Launch contract for _repack_compressed_kv_rank_map_kernel. + # launch_work is rank-major rows. + @cute.jit + def _repack_compressed_kv_rank_map_launch( + seq_ids_rank_major: cute.Tensor, + comp_ids_rank_major: cute.Tensor, + valid_rank_major: cute.Tensor, + cu_seqlens_compressed: cute.Tensor, + rank_major_by_seq_major: cute.Tensor, + rank_major_rows: cutlass.Int32, + seq_major_rows: cutlass.Int32, + launch_work: cutlass.Int32, + stream: cuda.CUstream, + ): + _launch_1d( + _repack_compressed_kv_rank_map_kernel, + "dsv4_cp_repack_compressed_kv_rank_map", + ( + seq_ids_rank_major, + comp_ids_rank_major, + valid_rank_major, + cu_seqlens_compressed, + rank_major_by_seq_major, + rank_major_rows, + seq_major_rows, + ), + launch_work, + stream, + ) + + # Launch contract for _repack_compressed_kv_to_seq_major_kernel. + # total_work is seq-major rows * row_width. + @cute.jit + def _repack_compressed_kv_to_seq_major_launch( + compressed_rank_major: cute.Tensor, + compressed_seq_major: cute.Tensor, + rank_major_by_seq_major: cute.Tensor, + rank_major_rows: cutlass.Int32, + seq_major_rows: cutlass.Int32, + row_width: cutlass.Constexpr, + total_work: cutlass.Int32, + stream: cuda.CUstream, + ): + _launch_1d( + _repack_compressed_kv_to_seq_major_kernel, + "dsv4_cp_repack_compressed_kv_to_seq_major", + ( + compressed_rank_major, + compressed_seq_major, + rank_major_by_seq_major, + rank_major_rows, + seq_major_rows, + row_width, + total_work, + ), + total_work, + stream, + ) + + # Kernel contract: + # cu_seqlens: int32, shape (n_seq + 1,), global packed sequence prefixes. + # seq_ids/comp_ids: int32 outputs, shape (total_rows,). + # valid: bool output, shape (total_rows,). + # total_rows is cp_size * c_cap_per_rank. For each rank-major compressed + # row, computes (seq_id, compressed group id) and marks whether that group + # ends inside the owning rank/chunk. + @cute.kernel + def _build_compressed_row_metadata_kernel( + cu_seqlens: cute.Tensor, + seq_ids: cute.Tensor, + comp_ids: cute.Tensor, + valid: cute.Tensor, + n_seq: cutlass.Int32, + cp_size: cutlass.Int32, + chunk_len: cutlass.Int32, + ratio: cutlass.Int32, + d_comp: cutlass.Int32, + c_cap_per_chunk: cutlass.Int32, + c_cap_per_rank: cutlass.Int32, + two_chunk: cutlass.Int32, + total_rows: cutlass.Int32, + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + row = bidx * 128 + tidx + if row < total_rows: + rank = row // c_cap_per_rank + rank_slot = row - rank * c_cap_per_rank + slot = rank_slot + rank_start = rank * chunk_len + row_is_padding = False + if two_chunk != 0: + local_chunk = rank_slot // c_cap_per_chunk + if local_chunk >= 2: + row_is_padding = True + else: + slot = rank_slot - local_chunk * c_cap_per_chunk + chunk_id = rank + if local_chunk != 0: + chunk_id = cp_size * 2 - 1 - rank + rank_start = chunk_id * chunk_len + rank_end = rank_start + chunk_len + + seq_ids[row] = -1 + comp_ids[row] = -1 + valid[row] = False + + if not row_is_padding: + running = 0 + for seq in range(n_seq): + seq_start = cu_seqlens[seq] + seq_end = cu_seqlens[seq + 1] + local_seq_start = seq_start + if local_seq_start < rank_start: + local_seq_start = rank_start + local_seq_end = seq_end + if local_seq_end > rank_end: + local_seq_end = rank_end + + if local_seq_start < local_seq_end: + n_full_groups = (seq_end - seq_start) // ratio + first_numer = rank_start - d_comp - seq_start + if first_numer < 0: + first_numer = 0 + first_group = 0 + if first_numer > 0: + first_group = (first_numer + ratio - 1) // ratio + stop_group = (local_seq_end - seq_start) // ratio + if stop_group > n_full_groups: + stop_group = n_full_groups + group_count = stop_group - first_group + if group_count < 0: + group_count = 0 + + if slot >= running and slot < running + group_count: + comp_id = first_group + slot - running + group_end = seq_start + (comp_id + 1) * ratio + seq_ids[row] = seq + comp_ids[row] = comp_id + valid[row] = group_end - 1 >= rank_start and group_end - 1 < rank_end + running = running + group_count + + # Launch contract for _build_compressed_row_metadata_kernel. + # two_chunk selects one contiguous chunk per rank or two chunks + # per rank; c_cap_per_rank is the fixed local compressed capacity. + @cute.jit + def _build_compressed_row_metadata_launch( + cu_seqlens: cute.Tensor, + seq_ids: cute.Tensor, + comp_ids: cute.Tensor, + valid: cute.Tensor, + n_seq: cutlass.Int32, + cp_size: cutlass.Int32, + chunk_len: cutlass.Int32, + ratio: cutlass.Int32, + d_comp: cutlass.Int32, + c_cap_per_chunk: cutlass.Int32, + c_cap_per_rank: cutlass.Int32, + two_chunk: cutlass.Int32, + total_rows: cutlass.Int32, + stream: cuda.CUstream, + ): + _launch_1d( + _build_compressed_row_metadata_kernel, + "dsv4_cp_build_compressed_row_metadata", + ( + cu_seqlens, + seq_ids, + comp_ids, + valid, + n_seq, + cp_size, + chunk_len, + ratio, + d_comp, + c_cap_per_chunk, + c_cap_per_rank, + two_chunk, + total_rows, + ), + total_rows, + stream, + ) + + # Kernel contract: + # k_indexer_seq_major: compressed indexer K, shape (seq_major_rows, row_width). + # cu_q: int32 query prefixes, shape (n_seq + 1,). + # cu_compressed: int32 compressed K prefixes, shape (n_seq + 1,). + # k_topk: output packed K for local trapezoid top-k, shape (k_rows, row_width). + # cu_q_topk/cu_k_topk/seq_lens: int32 outputs, shape (n_seq + 1)/(n_seq + 1)/(n_seq). + # Copies the compressed K rows visible to local query rows and emits local + # q/k prefix sums for the indexer kernel. + @cute.kernel + def _build_indexer_topk_metadata_kernel( + k_seq_major: cute.Tensor, + k_topk: cute.Tensor, + cu_seqlens_q: cute.Tensor, + cu_seqlens_compressed: cute.Tensor, + cu_q_topk: cute.Tensor, + cu_k_topk: cute.Tensor, + seq_lens: cute.Tensor, + n_seq: cutlass.Int32, + global_start: cutlass.Int32, + l_local: cutlass.Int32, + ratio: cutlass.Int32, + row_width: cutlass.Constexpr, + seq_major_rows: cutlass.Int32, + total_work: cutlass.Int32, + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + if bidx == 0 and tidx == 0: + global_end = global_start + l_local + q_prefix = 0 + k_prefix = 0 + cu_q_topk[0] = cutlass.Int32(0) + cu_k_topk[0] = cutlass.Int32(0) + for seq in range(n_seq): + seq_start = cu_seqlens_q[seq] + seq_end = cu_seqlens_q[seq + 1] + local_start = seq_start + if local_start < global_start: + local_start = global_start + local_end = seq_end + if local_end > global_end: + local_end = global_end + + q_len = 0 + k_len = 0 + if local_start < local_end: + q_len = local_end - local_start + seq_comp_start = cu_seqlens_compressed[seq] + seq_comp_end = cu_seqlens_compressed[seq + 1] + seq_comp_len = seq_comp_end - seq_comp_start + k_len = (local_end - seq_start) // ratio + if k_len > seq_comp_len: + k_len = seq_comp_len + + q_prefix = q_prefix + q_len + k_prefix = k_prefix + k_len + cu_q_topk[seq + 1] = q_prefix + cu_k_topk[seq + 1] = k_prefix + for row in range(q_len): + seq_lens[q_prefix - q_len + row] = k_len + + actual_total = cu_seqlens_q[n_seq] + padding_q = 0 + if global_end > actual_total: + padding_start = actual_total + if padding_start < global_start: + padding_start = global_start + if padding_start < global_end: + padding_q = global_end - padding_start + cu_q_topk[n_seq + 1] = q_prefix + padding_q + cu_k_topk[n_seq + 1] = k_prefix + for row in range(padding_q): + seq_lens[q_prefix + row] = cutlass.Int32(0) + + linear = bidx * 128 + tidx + if linear < total_work: + dst_row = linear // row_width + col = linear - dst_row * row_width + k_topk[dst_row, col] = cutlass.Float32(0.0).to(k_topk.element_type) + + global_end = global_start + l_local + k_prefix = 0 + src_row = -1 + for seq in range(n_seq): + seq_start = cu_seqlens_q[seq] + seq_end = cu_seqlens_q[seq + 1] + local_start = seq_start + if local_start < global_start: + local_start = global_start + local_end = seq_end + if local_end > global_end: + local_end = global_end + + k_len = 0 + if local_start < local_end: + seq_comp_start = cu_seqlens_compressed[seq] + seq_comp_end = cu_seqlens_compressed[seq + 1] + seq_comp_len = seq_comp_end - seq_comp_start + k_len = (local_end - seq_start) // ratio + if k_len > seq_comp_len: + k_len = seq_comp_len + if dst_row >= k_prefix and dst_row < k_prefix + k_len: + src_row = seq_comp_start + dst_row - k_prefix + k_prefix = k_prefix + k_len + + if src_row >= 0 and src_row < seq_major_rows: + k_topk[dst_row, col] = k_seq_major[src_row, col] + + # Launch contract for _build_indexer_topk_metadata_kernel. + # total_work covers max(k_rows * row_width, n_seq + 1). + @cute.jit + def _build_indexer_topk_metadata_launch( + k_seq_major: cute.Tensor, + k_topk: cute.Tensor, + cu_seqlens_q: cute.Tensor, + cu_seqlens_compressed: cute.Tensor, + cu_q_topk: cute.Tensor, + cu_k_topk: cute.Tensor, + seq_lens: cute.Tensor, + n_seq: cutlass.Int32, + global_start: cutlass.Int32, + l_local: cutlass.Int32, + ratio: cutlass.Int32, + row_width: cutlass.Constexpr, + seq_major_rows: cutlass.Int32, + total_work: cutlass.Int32, + launch_work: cutlass.Int32, + stream: cuda.CUstream, + ): + _launch_1d( + _build_indexer_topk_metadata_kernel, + "dsv4_cp_build_indexer_topk_metadata", + ( + k_seq_major, + k_topk, + cu_seqlens_q, + cu_seqlens_compressed, + cu_q_topk, + cu_k_topk, + seq_lens, + n_seq, + global_start, + l_local, + ratio, + row_width, + seq_major_rows, + total_work, + ), + launch_work, + stream, + ) + + # Kernel contract: + # cu_seqlens: int32 global query prefixes, shape (n_seq + 1,). + # optional indexer_topk: int32 compressed logical ids, shape (l_local, compressed_width). + # rank_major_by_seq_major: int32 reverse compressed map when compression is active. + # topk_idxs: int32 output, shape (l_local, window_size + compressed_width). + # topk_length: int32 output valid count, shape (l_local,). + # Lowers each local query row to physical KV-full ids: window ids first, + # then compressed ids from either dense causal range or indexer top-k. + @cute.kernel + def _build_attention_indices_kernel( + cu_seqlens: cute.Tensor, + cu_seqlens_compressed: cute.Tensor, + indexer_topk: cute.Tensor, + rank_major_by_seq_major: cute.Tensor, + topk_idxs: cute.Tensor, + topk_length: cute.Tensor, + n_seq: cutlass.Int32, + global_start: cutlass.Int32, + l_local: cutlass.Int32, + chunk0_start: cutlass.Int32, + chunk1_start: cutlass.Int32, + chunk_count: cutlass.Int32, + chunk_len: cutlass.Int32, + window_capacity_per_chunk: cutlass.Int32, + shared_compressed_base: cutlass.Int32, + d_window: cutlass.Int32, + window_size: cutlass.Int32, + ratio: cutlass.Int32, + compressed_width: cutlass.Int32, + seq_major_rows: cutlass.Int32, + has_indexer_topk: cutlass.Int32, + has_compressed: cutlass.Int32, + total_width: cutlass.Int32, + parallel_dense: cutlass.Int32, + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + linear = bidx * 128 + tidx + row_count = l_local + if chunk_count != 0: + row_count = chunk_count * chunk_len + + row = 0 + if parallel_dense != 0: + if linear < row_count * total_width: + row = linear // total_width + col = linear - row * total_width + topk_value = -1 + length = 0 + + global_q = global_start + row + seq_id = -1 + seq_start_found = 0 + seq_end_found = 0 + for seq in range(n_seq): + seq_start = cu_seqlens[seq] + seq_end = cu_seqlens[seq + 1] + if global_q >= seq_start and global_q < seq_end: + seq_id = seq + seq_start_found = seq_start + seq_end_found = seq_end + + if seq_id >= 0: + seq_offset = 0 + seq_window_start = 0 + seq_window_len = 0 + seq_comp_len = 0 + global_end = global_start + row_count + for seq in range(n_seq): + seq_start = cu_seqlens[seq] + seq_end = cu_seqlens[seq + 1] + local_seq_start = seq_start + if local_seq_start < global_start: + local_seq_start = global_start + local_seq_end = seq_end + if local_seq_end > global_end: + local_seq_end = global_end + if local_seq_start < local_seq_end: + window_start = local_seq_start - d_window + if window_start < seq_start: + window_start = seq_start + window_len = local_seq_end - window_start + comp_len = 0 + if has_compressed != 0 and ratio > 1: + comp_len = (seq_end - seq_start) // ratio + if seq == seq_id: + seq_window_start = window_start + seq_window_len = window_len + seq_comp_len = comp_len + if seq < seq_id: + seq_offset = seq_offset + window_len + comp_len + + window_start_for_q = global_q - window_size + 1 + if window_start_for_q < seq_start_found: + window_start_for_q = seq_start_found + if window_start_for_q < seq_window_start: + window_start_for_q = seq_window_start + window_count = global_q - window_start_for_q + 1 + if window_count < 0: + window_count = 0 + if window_count > window_size: + window_count = window_size + + comp_count = 0 + if has_compressed != 0 and compressed_width > 0 and ratio > 1: + comp_count = (global_q - seq_start_found + 1) // ratio + if comp_count > compressed_width: + comp_count = compressed_width + if comp_count > seq_comp_len: + comp_count = seq_comp_len + length = window_count + comp_count + if col < window_count: + topk_value = seq_offset + window_start_for_q - seq_window_start + col + elif col < length: + topk_value = seq_offset + seq_window_len + col - window_count + elif total_width > 0: + # Fused DSA sparse attention expects at least one KV id per + # query row. Padded THD rows carry zero loss gradient, so a + # dummy in-range id is enough to keep the kernel contract. + length = 1 + if col == 0: + topk_value = 0 + + topk_idxs[row, col] = topk_value + if col == 0: + topk_length[row] = length + else: + row = linear + if row >= row_count: + row = row_count + if parallel_dense == 0 and row < row_count: + for col in range(total_width): + topk_idxs[row, col] = -1 + topk_length[row] = 0 + + local_start = global_start + row_in_local = row + chunk_id = 0 + if chunk_count != 0: + chunk_id = row // chunk_len + row_in_local = row - chunk_id * chunk_len + local_start = chunk0_start + if chunk_id == 1: + local_start = chunk1_start + global_q = local_start + row_in_local + seq_id = -1 + seq_start_found = 0 + seq_end_found = 0 + for seq in range(n_seq): + seq_start = cu_seqlens[seq] + seq_end = cu_seqlens[seq + 1] + if global_q >= seq_start and global_q < seq_end: + seq_id = seq + seq_start_found = seq_start + seq_end_found = seq_end + + if seq_id >= 0: + seq_offset = 0 + seq_window_start = 0 + seq_window_len = 0 + seq_comp_len = 0 + seq_comp_start = 0 + seq_comp_end = 0 + global_end = local_start + row_count + if chunk_count != 0: + global_end = local_start + chunk_len + for seq in range(n_seq): + seq_start = cu_seqlens[seq] + seq_end = cu_seqlens[seq + 1] + local_seq_start = seq_start + if local_seq_start < local_start: + local_seq_start = local_start + local_seq_end = seq_end + if local_seq_end > global_end: + local_seq_end = global_end + if local_seq_start < local_seq_end: + window_start = local_seq_start - d_window + if window_start < seq_start: + window_start = seq_start + window_len = local_seq_end - window_start + comp_len = 0 + if has_compressed != 0 and ratio > 1: + if chunk_count != 0: + comp_len = ( + cu_seqlens_compressed[seq + 1] - cu_seqlens_compressed[seq] + ) + else: + comp_len = (seq_end - seq_start) // ratio + if seq == seq_id: + seq_window_start = window_start + seq_window_len = window_len + seq_comp_len = comp_len + if chunk_count != 0 and has_compressed != 0 and ratio > 1: + seq_comp_start = cu_seqlens_compressed[seq] + seq_comp_end = cu_seqlens_compressed[seq + 1] + if seq < seq_id: + seq_offset = seq_offset + window_len + if chunk_count == 0: + seq_offset = seq_offset + comp_len + + write_col = 0 + window_start_for_q = global_q - window_size + 1 + if window_start_for_q < seq_start_found: + window_start_for_q = seq_start_found + window_count = global_q - window_start_for_q + 1 + for w in range(window_size): + pos = window_start_for_q + w + if ( + w < window_count + and pos >= seq_window_start + and pos < seq_window_start + seq_window_len + ): + window_value = seq_offset + pos - seq_window_start + if chunk_count != 0: + window_value = chunk_id * window_capacity_per_chunk + window_value + topk_idxs[row, write_col] = window_value + write_col = write_col + 1 + + if has_compressed != 0 and compressed_width > 0 and ratio > 1: + pos_in_seq = global_q - seq_start_found + for j in range(compressed_width): + comp_id = -1 + if has_indexer_topk != 0: + comp_id = indexer_topk[row, j] + else: + n_visible = (pos_in_seq + 1) // ratio + if n_visible > compressed_width: + n_visible = compressed_width + if j < n_visible: + comp_id = j + if comp_id >= 0 and comp_id < seq_comp_len: + if chunk_count != 0: + seq_major_id = seq_comp_start + comp_id + if seq_major_id >= 0 and seq_major_id < seq_major_rows: + rank_major_id = rank_major_by_seq_major[seq_major_id] + if rank_major_id >= 0: + topk_idxs[row, write_col] = ( + shared_compressed_base + rank_major_id + ) + write_col = write_col + 1 + else: + topk_idxs[row, write_col] = seq_offset + seq_window_len + comp_id + write_col = write_col + 1 + topk_length[row] = write_col + elif total_width > 0: + # Same padding-row contract as the parallel dense path above. + topk_idxs[row, 0] = 0 + topk_length[row] = 1 + + # Launch contract for _build_attention_indices_kernel. + # Supports contiguous and two-chunk CP. total_work is + # l_local * (window_size + compressed_width). + @cute.jit + def _build_attention_indices_launch( + cu_seqlens: cute.Tensor, + cu_seqlens_compressed: cute.Tensor, + indexer_topk: cute.Tensor, + rank_major_by_seq_major: cute.Tensor, + topk_idxs: cute.Tensor, + topk_length: cute.Tensor, + n_seq: cutlass.Int32, + global_start: cutlass.Int32, + l_local: cutlass.Int32, + chunk0_start: cutlass.Int32, + chunk1_start: cutlass.Int32, + chunk_count: cutlass.Int32, + chunk_len: cutlass.Int32, + window_capacity_per_chunk: cutlass.Int32, + shared_compressed_base: cutlass.Int32, + d_window: cutlass.Int32, + window_size: cutlass.Int32, + ratio: cutlass.Int32, + compressed_width: cutlass.Int32, + seq_major_rows: cutlass.Int32, + has_indexer_topk: cutlass.Int32, + has_compressed: cutlass.Int32, + total_width: cutlass.Int32, + parallel_dense: cutlass.Int32, + launch_work: cutlass.Int32, + stream: cuda.CUstream, + ): + _launch_1d( + _build_attention_indices_kernel, + "dsv4_cp_build_attention_indices", + ( + cu_seqlens, + cu_seqlens_compressed, + indexer_topk, + rank_major_by_seq_major, + topk_idxs, + topk_length, + n_seq, + global_start, + l_local, + chunk0_start, + chunk1_start, + chunk_count, + chunk_len, + window_capacity_per_chunk, + shared_compressed_base, + d_window, + window_size, + ratio, + compressed_width, + seq_major_rows, + has_indexer_topk, + has_compressed, + total_width, + parallel_dense, + ), + launch_work, + stream, + ) + + # Kernel contract: + # Same inputs as _build_attention_indices_kernel, but emits compressed-first + # indexer-loss order. + # topk_idxs: int32 output, shape (l_local, compressed_width + window_size), + # compressed ids first then window ids. + # indexer_rank_major: int32 output, shape (l_local, compressed_width), mapping + # each compressed logical top-k id to rank-major indexer-K row. + @cute.kernel + def _build_indexer_loss_indices_kernel( + cu_seqlens: cute.Tensor, + cu_seqlens_compressed: cute.Tensor, + indexer_topk_logical: cute.Tensor, + indexer_rank_by_seq_major: cute.Tensor, + topk_idxs: cute.Tensor, + indexer_rank_major: cute.Tensor, + n_seq: cutlass.Int32, + global_start: cutlass.Int32, + l_local: cutlass.Int32, + chunk0_start: cutlass.Int32, + chunk1_start: cutlass.Int32, + chunk_count: cutlass.Int32, + chunk_len: cutlass.Int32, + window_capacity_per_chunk: cutlass.Int32, + shared_compressed_base: cutlass.Int32, + d_window: cutlass.Int32, + window_size: cutlass.Int32, + ratio: cutlass.Int32, + compressed_width: cutlass.Int32, + seq_major_rows: cutlass.Int32, + total_width: cutlass.Int32, + total_work: cutlass.Int32, + ): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + linear = bidx * 128 + tidx + if linear < total_work: + row = linear // total_width + col = linear - row * total_width + topk_value = -1 + rank_major_value = -1 + + local_start = global_start + row_in_local = row + chunk_id = 0 + if chunk_count != 0: + chunk_id = row // chunk_len + row_in_local = row - chunk_id * chunk_len + local_start = chunk0_start + if chunk_id == 1: + local_start = chunk1_start + global_q = local_start + row_in_local + seq_id = -1 + seq_start_found = 0 + for seq in range(n_seq): + seq_start = cu_seqlens[seq] + seq_end = cu_seqlens[seq + 1] + if global_q >= seq_start and global_q < seq_end: + seq_id = seq + seq_start_found = seq_start + + if seq_id >= 0: + seq_offset = 0 + seq_window_start = 0 + seq_window_len = 0 + seq_comp_len = 0 + global_end = local_start + l_local + if chunk_count != 0: + global_end = local_start + chunk_len + for seq in range(n_seq): + seq_start = cu_seqlens[seq] + seq_end = cu_seqlens[seq + 1] + local_seq_start = seq_start + if local_seq_start < local_start: + local_seq_start = local_start + local_seq_end = seq_end + if local_seq_end > global_end: + local_seq_end = global_end + if local_seq_start < local_seq_end: + window_start = local_seq_start - d_window + if window_start < seq_start: + window_start = seq_start + window_len = local_seq_end - window_start + comp_len = 0 + if ratio > 1: + comp_len = cu_seqlens_compressed[seq + 1] - cu_seqlens_compressed[seq] + if seq == seq_id: + seq_window_start = window_start + seq_window_len = window_len + seq_comp_len = comp_len + if seq < seq_id: + seq_offset = seq_offset + window_len + if chunk_count == 0: + seq_offset = seq_offset + comp_len + + if col < compressed_width: + comp_id = indexer_topk_logical[row, col] + if comp_id >= 0 and comp_id < seq_comp_len: + seq_major_id = cu_seqlens_compressed[seq_id] + comp_id + if seq_major_id >= 0 and seq_major_id < seq_major_rows: + rank_major_id = indexer_rank_by_seq_major[seq_major_id] + if rank_major_id >= 0: + if chunk_count != 0: + topk_value = shared_compressed_base + rank_major_id + else: + topk_value = seq_offset + seq_window_len + comp_id + rank_major_value = rank_major_id + else: + window_col = col - compressed_width + window_start_for_q = global_q - window_size + 1 + if window_start_for_q < seq_start_found: + window_start_for_q = seq_start_found + window_count = global_q - window_start_for_q + 1 + pos = window_start_for_q + window_col + if ( + window_col < window_count + and pos >= seq_window_start + and pos < seq_window_start + seq_window_len + ): + topk_value = seq_offset + pos - seq_window_start + if chunk_count != 0: + topk_value = chunk_id * window_capacity_per_chunk + topk_value + + topk_idxs[row, col] = topk_value + if col < compressed_width: + indexer_rank_major[row, col] = rank_major_value + + # Launch contract for _build_indexer_loss_indices_kernel. + # Supports contiguous and two-chunk CP. total_work is + # l_local * (compressed_width + window_size). + @cute.jit + def _build_indexer_loss_indices_launch( + cu_seqlens: cute.Tensor, + cu_seqlens_compressed: cute.Tensor, + indexer_topk_logical: cute.Tensor, + indexer_rank_by_seq_major: cute.Tensor, + topk_idxs: cute.Tensor, + indexer_rank_major: cute.Tensor, + n_seq: cutlass.Int32, + global_start: cutlass.Int32, + l_local: cutlass.Int32, + chunk0_start: cutlass.Int32, + chunk1_start: cutlass.Int32, + chunk_count: cutlass.Int32, + chunk_len: cutlass.Int32, + window_capacity_per_chunk: cutlass.Int32, + shared_compressed_base: cutlass.Int32, + d_window: cutlass.Int32, + window_size: cutlass.Int32, + ratio: cutlass.Int32, + compressed_width: cutlass.Int32, + seq_major_rows: cutlass.Int32, + total_width: cutlass.Int32, + total_work: cutlass.Int32, + stream: cuda.CUstream, + ): + _launch_1d( + _build_indexer_loss_indices_kernel, + "dsv4_cp_build_indexer_loss_indices", + ( + cu_seqlens, + cu_seqlens_compressed, + indexer_topk_logical, + indexer_rank_by_seq_major, + topk_idxs, + indexer_rank_major, + n_seq, + global_start, + l_local, + chunk0_start, + chunk1_start, + chunk_count, + chunk_len, + window_capacity_per_chunk, + shared_compressed_base, + d_window, + window_size, + ratio, + compressed_width, + seq_major_rows, + total_width, + total_work, + ), + total_work, + stream, + ) + + +# ============================================================================= +# Torch Wrapper Functions +# ============================================================================= +# Torch-facing entry points called by csa_cp_utils.py. They validate inputs, +# allocate outputs, and dispatch the CuTeDSL kernels defined above. +# ============================================================================= + + +def _require_cute(message: str, *tensors: Optional[torch.Tensor]) -> None: + """Raise ``RuntimeError`` when a wrapper cannot use CuTeDSL kernels.""" + if ( + not _CUTE_AVAILABLE + or os.environ.get("DSV4_CP_DISABLE_CUTE_KERNELS") + or not all(tensor is None or tensor.is_cuda for tensor in tensors) + ): + raise RuntimeError(message) + + +def _make_compiled_launch_runner(): + """Build the CuTe launch helper with private helpers and persistent caches.""" + compiled_launch_cache = {} + + def _cute_scalar(value): + if isinstance(value, bool): + return cutlass.Int32(1 if value else 0) + if isinstance(value, int): + return cutlass.Int32(value) + if isinstance(value, float): + return cutlass.Float32(value) + return value + + @lru_cache(maxsize=None) + def _gpu_arch_flag() -> str: + cap = torch.cuda.get_device_capability() + arch_map = {(9, 0): "sm_90a", (10, 0): "sm_100a", (10, 3): "sm_103a"} + arch = arch_map.get(cap) + if arch is None: + raise RuntimeError(f"Unsupported GPU compute capability {cap} for CSA CP CuTe kernels.") + return arch + + def _run_compiled_launch( + launch_fn, + tensor_args: Tuple[torch.Tensor, ...], + scalar_args: Tuple, + key_arg_indices: Optional[Tuple[int, ...]] = None, + static_arg_indices: Optional[Tuple[int, ...]] = None, + ) -> None: + """Compile/cache a CuTe launch function and invoke it on the current stream. + + Inputs: + launch_fn: ``@cute.jit`` launcher. + tensor_args: Torch CUDA tensors passed through DLPack with dynamic leading + layouts. + scalar_args: Python scalar launch parameters. + key_arg_indices: Scalar indices that participate in the compile cache key. + static_arg_indices: Scalar indices passed as compile-time constexpr values. + + Output: + None. The launch writes into output tensors passed in ``tensor_args``. + """ + key_arg_indices = ( + tuple(range(len(scalar_args))) if key_arg_indices is None else tuple(key_arg_indices) + ) + static_arg_indices = () if static_arg_indices is None else tuple(static_arg_indices) + static_arg_set = set(static_arg_indices) + key = ( + launch_fn.__name__, + tuple( + (tensor.dtype, tuple(tensor.shape), tuple(tensor.stride())) + for tensor in tensor_args + ), + key_arg_indices, + tuple(scalar_args[i] for i in key_arg_indices), + static_arg_indices, + ) + compiled = compiled_launch_cache.get(key) + if compiled is None: + fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) + cute_tensor_args = [] + for tensor in tensor_args: + cute_tensor = from_dlpack(tensor.detach(), assumed_align=16, enable_tvm_ffi=True) + if tensor.ndim != 0: + cute_tensor = cute_tensor.mark_layout_dynamic(leading_dim=tensor.ndim - 1) + cute_tensor_args.append(cute_tensor) + compiled = cute.compile( + launch_fn, + *cute_tensor_args, + *( + arg if i in static_arg_set else _cute_scalar(arg) + for i, arg in enumerate(scalar_args) + ), + fake_stream, + options=f"--enable-tvm-ffi --gpu-arch {_gpu_arch_flag()}", + ) + compiled_launch_cache[key] = compiled + compiled( + *tensor_args, + *(_cute_scalar(arg) for i, arg in enumerate(scalar_args) if i not in static_arg_set), + cuda.CUstream(torch.cuda.current_stream(tensor_args[0].device).cuda_stream), + ) + + return _run_compiled_launch + + +_run_compiled_launch = _make_compiled_launch_runner() + + +def _validate_rope_inputs( + x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, nope_dim: int, pos_dim: int +) -> Tuple[int, int, int]: + """Validate common RoPE tensor dimensions and return flattened row metadata. + + Inputs: + x: Tensor with first dimension as rows and remaining dims flattened per row. + cos/sin: Rotary tables with first dimension as sequence position and at + least ``pos_dim`` values per row. + nope_dim/pos_dim: Per-head no-RoPE and RoPE dimensions. + + Output: + ``(rows, row_width, head_dim)`` where ``head_dim = nope_dim + pos_dim``. + """ + rows = x.shape[0] + row_width = math.prod(x.shape[1:]) + head_dim = nope_dim + pos_dim + if pos_dim <= 0 or pos_dim % 2 != 0: + raise RuntimeError(f"DSv4 CP RoPE expects an even positive pos_dim, got {pos_dim}.") + if head_dim <= 0 or row_width % head_dim != 0: + raise RuntimeError( + "DSv4 CP RoPE expects the flattened row width to be a multiple of " + f"head_dim={head_dim}, got row_width={row_width}." + ) + cos_width = math.prod(cos.shape[1:]) + sin_width = math.prod(sin.shape[1:]) + if cos_width < pos_dim or sin_width < pos_dim: + raise RuntimeError( + "DSv4 CP RoPE cos/sin tables are narrower than pos_dim: " + f"cos={cos_width}, sin={sin_width}, pos_dim={pos_dim}." + ) + return rows, row_width, head_dim + + +class ThdLocalRope(torch.autograd.Function): + """Apply non-interleaved RoPE for THD rows owned by one or two CP chunks. + + Inputs: + x: CUDA tensor, shape ``(local_rows, ...)``; each row flattens to + ``num_heads * (nope_dim + pos_dim)`` and dtype is preserved. + cos/sin: CUDA tensors, shape ``(max_seq_len, >= pos_dim)``. + cu_seqlens_padded: int32 CUDA tensor, shape ``(n_seq + 1,)``. + chunk0_global_start/chunk1_global_start: global packed-token row starts. + chunk_len: positive row count per chunk. ``local_rows <= 2 * chunk_len``. + inverse/adjoint: select inverse RoPE or backward adjoint math. + + Output: + Tensor with same shape and dtype as ``x``. The wrapper maps each local + row back to a global THD token position, derives position within its + sequence from ``cu_seqlens_padded``, and launches the fused RoPE kernel. + """ + + @staticmethod + def _run( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + cu_seqlens_padded: torch.Tensor, + chunk0_global_start: int, + chunk_len: int, + nope_dim: int, + pos_dim: int, + chunk1_global_start: int = 0, + inverse: bool = False, + clamp_to_valid_token: bool = False, + adjoint: bool = False, + ) -> torch.Tensor: + _require_cute( + "DSv4 THD CP RoPE requires CUDA tensors and available CuTe kernels.", + x, + cos, + sin, + cu_seqlens_padded, + ) + rows, row_width, head_dim = _validate_rope_inputs(x, cos, sin, nope_dim, pos_dim) + chunk_len = int(chunk_len) + if chunk_len <= 0: + raise RuntimeError(f"DSv4 THD CP RoPE expects positive chunk_len, got {chunk_len}.") + if rows > 2 * chunk_len: + raise RuntimeError( + "DSv4 THD CP RoPE supports at most two chunks: " + f"rows={rows}, chunk_len={chunk_len}." + ) + if rows == 0: + return x + out = torch.empty_like(x) + n_heads = row_width // head_dim + head_group = 1 + for candidate in (16, 8, 4, 2): + if n_heads % candidate == 0: + head_group = candidate + break + total_work = rows + _run_compiled_launch( + _thd_local_rope_launch, + ( + x.reshape(rows, row_width), + out.reshape(rows, row_width), + cos.reshape(cos.shape[0], math.prod(cos.shape[1:])), + sin.reshape(sin.shape[0], math.prod(sin.shape[1:])), + cu_seqlens_padded, + ), + ( + cu_seqlens_padded.shape[0] - 1, + int(chunk0_global_start), + int(chunk1_global_start), + chunk_len, + 1 if clamp_to_valid_token else 0, + row_width, + head_dim, + int(nope_dim), + int(pos_dim), + 1 if inverse else 0, + 1 if adjoint else 0, + head_group, + total_work, + ), + key_arg_indices=(0, 4, 5, 6, 7, 8, 9, 10, 11), + static_arg_indices=(0, 4, 5, 6, 7, 8, 9, 10, 11), + ) + return out + + @staticmethod + def forward( + ctx, + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + cu_seqlens_padded: torch.Tensor, + chunk0_global_start: int, + chunk_len: int, + nope_dim: int, + pos_dim: int, + chunk1_global_start: int, + inverse: bool, + clamp_to_valid_token: bool, + ) -> torch.Tensor: + """Apply local THD CP RoPE and save metadata for backward.""" + ctx.rope_args = ( + int(chunk0_global_start), + int(chunk_len), + int(nope_dim), + int(pos_dim), + int(chunk1_global_start), + bool(inverse), + bool(clamp_to_valid_token), + ) + ctx.save_for_backward(cos, sin, cu_seqlens_padded) + return ThdLocalRope._run(x, cos, sin, cu_seqlens_padded, *ctx.rope_args) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + """Apply the adjoint local THD CP RoPE transform to gradients.""" + cos, sin, cu_seqlens_padded = ctx.saved_tensors + grad_x = ThdLocalRope._run( + grad_output.contiguous(), cos, sin, cu_seqlens_padded, *ctx.rope_args, adjoint=True + ) + return (grad_x, *([None] * 10)) + + +class ThdCompressedRope(torch.autograd.Function): + """Apply RoPE to compressed rows using per-row compressed group ids. + + Inputs: + x: CUDA tensor, shape ``(rows, ...)``; flattened row width must be a + multiple of ``nope_dim + pos_dim``. + cos/sin: CUDA tensors, shape ``(max_seq_len, >= pos_dim)``. + comp_ids_local: int32 CUDA tensor, shape at least ``(rows,)``; -1 maps to 0. + ratio: compression ratio, used as ``rope_position = comp_id * ratio``. + + Output: + Tensor with same shape and dtype as ``x``. No-RoPE dims are copied; RoPE + dims are rotated pairwise, with inverse/adjoint variants controlled by + flags. + """ + + @staticmethod + def _run( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + comp_ids_local: torch.Tensor, + ratio: int, + nope_dim: int, + pos_dim: int, + inverse: bool = False, + adjoint: bool = False, + ) -> torch.Tensor: + _require_cute( + "DSv4 CP compressed RoPE requires CUDA tensors and available CuTe kernels.", + x, + cos, + sin, + comp_ids_local, + ) + rows, row_width, head_dim = _validate_rope_inputs(x, cos, sin, nope_dim, pos_dim) + if comp_ids_local.shape[0] < rows: + raise RuntimeError( + "DSv4 CP compressed RoPE received too few comp ids: " + f"comp_ids={comp_ids_local.shape[0]}, rows={rows}." + ) + if rows == 0: + return x + out = torch.empty_like(x) + total_work = rows * row_width + _run_compiled_launch( + _thd_compressed_rope_launch, + ( + x.reshape(rows, row_width), + out.reshape(rows, row_width), + cos.reshape(cos.shape[0], math.prod(cos.shape[1:])), + sin.reshape(sin.shape[0], math.prod(sin.shape[1:])), + comp_ids_local[:rows], + ), + ( + int(ratio), + row_width, + head_dim, + int(nope_dim), + int(pos_dim), + 1 if inverse else 0, + 1 if adjoint else 0, + total_work, + ), + key_arg_indices=(1,), + static_arg_indices=(1,), + ) + return out + + @staticmethod + def forward( + ctx, + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + comp_ids_local: torch.Tensor, + ratio: int, + nope_dim: int, + pos_dim: int, + inverse: bool, + ) -> torch.Tensor: + """Apply compressed-row RoPE and save metadata for backward.""" + ctx.rope_args = (int(ratio), int(nope_dim), int(pos_dim), bool(inverse)) + ctx.save_for_backward(cos, sin, comp_ids_local) + return ThdCompressedRope._run(x, cos, sin, comp_ids_local, *ctx.rope_args) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + """Apply the adjoint compressed-row RoPE transform to gradients.""" + cos, sin, comp_ids_local = ctx.saved_tensors + grad_x = ThdCompressedRope._run( + grad_output.contiguous(), cos, sin, comp_ids_local, *ctx.rope_args, adjoint=True + ) + return (grad_x, *([None] * 7)) + + +class CompressorInputCompact(torch.autograd.Function): + """Compact local and boundary hidden rows for compressor input. + + Inputs: + hidden_local: CUDA tensor, shape ``(l_local, ...)``. + boundary_hidden: CUDA tensor, shape ``(d_window, ...)``. + cu_seqlens: int32 CUDA tensor, shape ``(n_seq + 1,)``. + global_start/l_local: local contiguous global row range. + ratio/d_comp/d_window/c_cap: fixed compressor window and capacity. + + Outputs: + ``hidden_compact`` shape ``(c_cap * ratio, ...)``, same dtype as hidden; + ``cu_compact`` int32 shape ``(n_seq + 1,)``; ``seq_ids`` int32, + ``comp_ids`` int32, and ``valid`` bool each shape ``(c_cap,)``. + The kernel copies the ratio source tokens for each visible compressed + group and emits metadata for the compressed rows. + """ + + @staticmethod + def _forward_kernel( + hidden_local: torch.Tensor, + boundary_hidden: torch.Tensor, + cu_seqlens: torch.Tensor, + global_start: int, + l_local: int, + ratio: int, + d_comp: int, + d_window: int, + c_cap: int, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + compact_len = c_cap * ratio + output_shape = (compact_len,) + tuple(hidden_local.shape[1:]) + hidden_compact = hidden_local.new_empty(output_shape) + cu_compact = torch.empty_like(cu_seqlens) + seq_ids = torch.empty((c_cap,), dtype=torch.int32, device=hidden_local.device) + comp_ids = torch.empty((c_cap,), dtype=torch.int32, device=hidden_local.device) + valid = torch.empty((c_cap,), dtype=torch.bool, device=hidden_local.device) + + row_width = math.prod(hidden_local.shape[1:]) + hidden_local_flat = hidden_local.reshape(l_local, row_width) + boundary_flat = boundary_hidden.reshape(d_window, row_width) + compact_flat = hidden_compact.reshape(compact_len, row_width) + total_work = max(compact_len, c_cap, cu_seqlens.shape[0]) + _run_compiled_launch( + _compressor_input_compact_fwd_launch, + ( + hidden_local_flat, + boundary_flat, + compact_flat, + cu_seqlens, + cu_compact, + seq_ids, + comp_ids, + valid, + ), + ( + cu_seqlens.shape[0] - 1, + global_start, + l_local, + ratio, + d_comp, + d_window, + c_cap, + compact_len, + row_width, + total_work, + ), + key_arg_indices=(3, 4, 5, 8), + static_arg_indices=(3, 4, 5, 8), + ) + return hidden_compact, cu_compact, seq_ids, comp_ids, valid + + @staticmethod + def _backward_kernel( + grad_hidden_compact: torch.Tensor, + hidden_shape: Tuple[int, ...], + boundary_shape: Tuple[int, ...], + cu_seqlens: torch.Tensor, + global_start: int, + l_local: int, + ratio: int, + d_comp: int, + d_window: int, + ) -> Tuple[torch.Tensor, torch.Tensor]: + grad_hidden = grad_hidden_compact.new_empty(hidden_shape) + grad_boundary = grad_hidden_compact.new_empty(boundary_shape) + compact_len = grad_hidden_compact.shape[0] + row_width = math.prod(hidden_shape[1:]) + + total_work = l_local + d_window + if total_work > 0: + _run_compiled_launch( + _compressor_input_compact_bwd_launch, + ( + grad_hidden_compact.reshape(compact_len, row_width), + grad_hidden.reshape(l_local, row_width), + grad_boundary.reshape(d_window, row_width), + cu_seqlens, + ), + ( + cu_seqlens.shape[0] - 1, + global_start, + l_local, + ratio, + d_comp, + d_window, + compact_len, + row_width, + total_work, + ), + key_arg_indices=(3, 4, 5, 7), + static_arg_indices=(3, 4, 5, 7), + ) + return grad_hidden, grad_boundary + + @staticmethod + def forward( + ctx, + hidden_local: torch.Tensor, + boundary_hidden: torch.Tensor, + cu_seqlens: torch.Tensor, + global_start: int, + l_local: int, + ratio: int, + d_comp: int, + d_window: int, + c_cap: int, + ): + """Compact local and boundary hidden rows into compressor input rows.""" + ctx.hidden_shape = tuple(hidden_local.shape) + ctx.boundary_shape = tuple(boundary_hidden.shape) + ctx.compact_args = (int(global_start), int(l_local), int(ratio), int(d_comp), int(d_window)) + ctx.save_for_backward(cu_seqlens) + return CompressorInputCompact._forward_kernel( + hidden_local, + boundary_hidden, + cu_seqlens, + int(global_start), + int(l_local), + int(ratio), + int(d_comp), + int(d_window), + int(c_cap), + ) + + @staticmethod + def backward( + ctx, + grad_hidden_compact: torch.Tensor, + _grad_cu_compact: torch.Tensor, + _grad_seq_ids: torch.Tensor, + _grad_comp_ids: torch.Tensor, + _grad_valid: torch.Tensor, + ): + """Scatter compacted compressor gradients to local and boundary rows.""" + (cu_seqlens,) = ctx.saved_tensors + grad_hidden, grad_boundary = CompressorInputCompact._backward_kernel( + grad_hidden_compact.contiguous(), + ctx.hidden_shape, + ctx.boundary_shape, + cu_seqlens, + *ctx.compact_args, + ) + return (grad_hidden, grad_boundary, *([None] * 7)) + + +class ThdFullKvPack(torch.autograd.Function): + """Autograd wrapper for static-shape KV-full packing. + + Forward builds a physical ``kv_full`` row space from local window rows, + boundary rows, and all-gathered compressed rows. Backward recomputes the + same row mapping and scatters gradients back to those three sources. + """ + + @staticmethod + def forward( + ctx, + kv_local: torch.Tensor, + boundary_kv: torch.Tensor, + compressed_rank_major: torch.Tensor, + seq_ids_rank_major: torch.Tensor, + comp_ids_rank_major: torch.Tensor, + valid_rank_major: torch.Tensor, + cu_seqlens: torch.Tensor, + rank_major_by_seq_major: torch.Tensor, + cu_seqlens_compressed: torch.Tensor, + global_start: int, + l_local: int, + d_window: int, + ratio: int, + kv_full_capacity: int, + chunk0_start: int, + chunk1_start: int, + chunk_count: int, + chunk_len: int, + window_capacity_per_chunk: int, + ) -> torch.Tensor: + """Pack KV rows into the fixed physical sparse-attention row order. + + Inputs: + kv_local: CUDA tensor, shape ``(local_rows, ...)``. + boundary_kv: CUDA tensor, shape ``(d_window or 2*d_window, ...)``. + compressed_rank_major: CUDA tensor, shape ``(compressed_rows, ...)``. + seq_ids/comp_ids/valid: rank-major compressed metadata. + cu_seqlens: int32 CUDA tensor, shape ``(n_seq + 1,)``. + optional compressed reverse-map metadata for contiguous CP. + chunk fields select contiguous or two-chunk layout. + + Output: + ``kv_full`` CUDA tensor, shape ``(kv_full_capacity, ...)``, same dtype + as ``kv_local``. Invalid tail rows are zero-filled. + """ + ctx.pack_shapes = ( + tuple(kv_local.shape), + tuple(boundary_kv.shape), + tuple(compressed_rank_major.shape), + ) + ctx.kv_full_capacity = int(kv_full_capacity) + + output_shape = (int(kv_full_capacity),) + tuple(kv_local.shape[1:]) + kv_full = kv_local.new_empty(output_shape) + row_width = math.prod(kv_local.shape[1:]) + local_rows = kv_local.shape[0] + boundary_rows = boundary_kv.shape[0] + kv_local_flat = kv_local.reshape(local_rows, row_width) + boundary_flat = boundary_kv.reshape(boundary_rows, row_width) + compressed_rows = compressed_rank_major.shape[0] + compressed_flat = compressed_rank_major.reshape(compressed_rows, row_width) + kv_full_flat = kv_full.reshape(int(kv_full_capacity), row_width) + total_work = int(kv_full_capacity) + use_seq_major_map = ( + int(chunk_count) == 0 + and ratio > 1 + and compressed_rows > 0 + and rank_major_by_seq_major.shape[0] > 1 + ) + _run_compiled_launch( + _thd_full_kv_pack_launch, + ( + kv_local_flat, + boundary_flat, + compressed_flat, + kv_full_flat, + seq_ids_rank_major, + comp_ids_rank_major, + valid_rank_major, + cu_seqlens, + cu_seqlens_compressed, + rank_major_by_seq_major, + ), + ( + cu_seqlens.shape[0] - 1, + int(global_start), + int(l_local), + int(chunk0_start), + int(chunk1_start), + int(chunk_count), + int(chunk_len), + int(window_capacity_per_chunk), + int(d_window), + int(ratio), + compressed_rows, + 1 if use_seq_major_map else 0, + rank_major_by_seq_major.shape[0] if use_seq_major_map else 0, + int(kv_full_capacity), + row_width, + 0, + total_work, + ), + key_arg_indices=(5, 8, 9, 11, 14, 15), + static_arg_indices=(5, 8, 9, 11, 14, 15), + ) + ctx.save_for_backward( + seq_ids_rank_major, + comp_ids_rank_major, + valid_rank_major, + cu_seqlens, + cu_seqlens_compressed, + rank_major_by_seq_major, + ) + ctx.pack_args = ( + int(global_start), + int(l_local), + int(chunk0_start), + int(chunk1_start), + int(chunk_count), + int(chunk_len), + int(window_capacity_per_chunk), + int(d_window), + int(ratio), + compressed_rows, + 1 if use_seq_major_map else 0, + rank_major_by_seq_major.shape[0] if use_seq_major_map else 0, + ) + return kv_full + + @staticmethod + def backward(ctx, grad_kv_full: torch.Tensor): + """Scatter ``kv_full`` gradients back to local, boundary, and compressed KV.""" + local_shape, boundary_shape, compressed_shape = ctx.pack_shapes + grad_kv_local = grad_kv_full.new_empty(local_shape) + grad_boundary_kv = grad_kv_full.new_empty(boundary_shape) + grad_compressed = grad_kv_full.new_empty(compressed_shape) + + compressed_rows = compressed_shape[0] + kv_full_capacity = ctx.kv_full_capacity + row_width = math.prod(local_shape[1:]) + ( + global_start, + l_local, + chunk0_start, + chunk1_start, + chunk_count, + chunk_len, + window_capacity_per_chunk, + d_window, + ratio, + compressed_rows, + use_seq_major_map, + seq_major_rows, + ) = ctx.pack_args + grad_kv_local.zero_() + grad_boundary_kv.zero_() + grad_compressed.zero_() + + total_work = kv_full_capacity + if total_work > 0: + ( + seq_ids_rank_major, + comp_ids_rank_major, + valid_rank_major, + cu_seqlens, + cu_seqlens_compressed, + rank_major_by_seq_major, + ) = ctx.saved_tensors + _run_compiled_launch( + _thd_full_kv_pack_launch, + ( + grad_kv_local.reshape(local_shape[0], row_width), + grad_boundary_kv.reshape(boundary_shape[0], row_width), + grad_compressed.reshape(compressed_rows, row_width), + grad_kv_full.reshape(kv_full_capacity, row_width), + seq_ids_rank_major, + comp_ids_rank_major, + valid_rank_major, + cu_seqlens, + cu_seqlens_compressed, + rank_major_by_seq_major, + ), + ( + cu_seqlens.shape[0] - 1, + global_start, + l_local, + chunk0_start, + chunk1_start, + chunk_count, + chunk_len, + window_capacity_per_chunk, + d_window, + ratio, + compressed_rows, + use_seq_major_map, + seq_major_rows, + kv_full_capacity, + row_width, + 1, + total_work, + ), + key_arg_indices=(5, 8, 9, 11, 14, 15), + static_arg_indices=(5, 8, 9, 11, 14, 15), + ) + return (grad_kv_local, grad_boundary_kv, grad_compressed, *([None] * 16)) + + +def repack_compressed_kv_to_seq_major( + compressed_rank_major: torch.Tensor, + seq_ids_rank_major: torch.Tensor, + comp_ids_rank_major: torch.Tensor, + valid_rank_major: torch.Tensor, + cu_seqlens_compressed: torch.Tensor, + seq_major_rows: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Repack all-gathered compressed rows from rank-major to seq-major order. + + Inputs: + compressed_rank_major: CUDA tensor, shape ``(rank_major_rows, ...)``. + seq_ids_rank_major/comp_ids_rank_major: int32 metadata, shape + ``(rank_major_rows,)``. + valid_rank_major: bool metadata, shape ``(rank_major_rows,)``. + cu_seqlens_compressed: int32 prefixes, shape ``(n_seq + 1,)``. + seq_major_rows: total compressed seq-major rows. + + Outputs: + ``compressed_seq_major`` shape ``(seq_major_rows, ...)`` and + ``rank_major_by_seq_major`` int32 shape ``(seq_major_rows,)``. + """ + output_shape = (seq_major_rows,) + tuple(compressed_rank_major.shape[1:]) + compressed_seq_major = compressed_rank_major.new_empty(output_shape) + rank_major_by_seq_major = torch.empty( + (seq_major_rows,), dtype=torch.int32, device=compressed_rank_major.device + ) + rank_major_rows = compressed_rank_major.shape[0] + row_width = math.prod(compressed_rank_major.shape[1:]) + map_work = int(rank_major_rows) + gather_work = int(seq_major_rows) * int(row_width) + _run_compiled_launch( + _repack_compressed_kv_init_rank_map_launch, + (rank_major_by_seq_major,), + (seq_major_rows, max(1, seq_major_rows)), + ) + _run_compiled_launch( + _repack_compressed_kv_rank_map_launch, + ( + seq_ids_rank_major, + comp_ids_rank_major, + valid_rank_major, + cu_seqlens_compressed, + rank_major_by_seq_major, + ), + (rank_major_rows, seq_major_rows, max(1, map_work)), + ) + _run_compiled_launch( + _repack_compressed_kv_to_seq_major_launch, + ( + compressed_rank_major.reshape(rank_major_rows, row_width), + compressed_seq_major.reshape(seq_major_rows, row_width), + rank_major_by_seq_major, + ), + (rank_major_rows, seq_major_rows, row_width, max(1, gather_work)), + key_arg_indices=(2,), + static_arg_indices=(2,), + ) + return compressed_seq_major, rank_major_by_seq_major + + +def build_compressed_row_metadata( + cu_seqlens: torch.Tensor, + cp_size: int, + l_local: int, + ratio: int, + d_comp: int, + c_cap: int, + *, + c_cap_per_rank: Optional[int] = None, + use_two_chunk: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build rank-major compressed metadata for all CP ranks. + + Inputs: + cu_seqlens: int32 CUDA tensor, shape ``(n_seq + 1,)``. + cp_size: CP world size. + l_local: rows per contiguous rank, or chunk length when ``use_two_chunk``. + ratio/d_comp: compression parameters. + c_cap: fixed compressed capacity per chunk. + c_cap_per_rank: optional fixed compressed capacity per rank. + use_two_chunk: whether each rank owns two chunks. + + Outputs: + ``seq_ids`` int32, ``comp_ids`` int32, and ``valid`` bool, each shape + ``(cp_size * c_cap_per_rank,)``. Rows are ordered rank-major and describe + which compressed group each all-gathered row represents. + """ + c_cap_per_rank = int(c_cap if c_cap_per_rank is None else c_cap_per_rank) + total_rows = int(cp_size) * c_cap_per_rank + seq_ids = torch.empty((total_rows,), dtype=torch.int32, device=cu_seqlens.device) + comp_ids = torch.empty((total_rows,), dtype=torch.int32, device=cu_seqlens.device) + valid = torch.empty((total_rows,), dtype=torch.bool, device=cu_seqlens.device) + _run_compiled_launch( + _build_compressed_row_metadata_launch, + (cu_seqlens, seq_ids, comp_ids, valid), + ( + cu_seqlens.shape[0] - 1, + int(cp_size), + int(l_local), + int(ratio), + int(d_comp), + int(c_cap), + c_cap_per_rank, + 1 if use_two_chunk else 0, + total_rows, + ), + key_arg_indices=(), + ) + return seq_ids, comp_ids, valid + + +def build_indexer_topk_metadata( + k_indexer_seq_major: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_compressed: torch.Tensor, + global_start: int, + l_local: int, + ratio: int, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Build local trapezoid inputs for the indexer top-k kernel. + + Inputs: + k_indexer_seq_major: CUDA tensor, shape ``(seq_major_rows, ...)``. + cu_seqlens_q/cu_seqlens_compressed: int32 prefixes, shape ``(n_seq + 1,)``. + global_start/l_local: local query range. + ratio: compression ratio. + + Outputs: + ``k_topk`` tensor with same shape as ``k_indexer_seq_major``; + ``cu_q_topk`` and ``cu_k_topk`` int32 tensors, shape ``(n_seq + 2,)``; + ``seq_lens`` int32 tensor, shape ``(l_local,)``. The metadata represents + the local Q/K trapezoid that the indexer kernel consumes. + """ + seq_major_rows = k_indexer_seq_major.shape[0] + row_width = math.prod(k_indexer_seq_major.shape[1:]) + k_topk = k_indexer_seq_major.new_empty(k_indexer_seq_major.shape) + cu_q_topk = torch.empty( + (cu_seqlens_q.shape[0] + 1,), dtype=cu_seqlens_q.dtype, device=cu_seqlens_q.device + ) + cu_k_topk = torch.empty( + (cu_seqlens_q.shape[0] + 1,), dtype=cu_seqlens_compressed.dtype, device=cu_seqlens_q.device + ) + seq_lens = torch.empty((l_local,), dtype=torch.int32, device=cu_seqlens_q.device) + n_seq = cu_seqlens_q.shape[0] - 1 + total_work = seq_major_rows * row_width + _run_compiled_launch( + _build_indexer_topk_metadata_launch, + ( + k_indexer_seq_major.reshape(seq_major_rows, row_width), + k_topk.reshape(seq_major_rows, row_width), + cu_seqlens_q, + cu_seqlens_compressed, + cu_q_topk, + cu_k_topk, + seq_lens, + ), + ( + n_seq, + global_start, + l_local, + ratio, + row_width, + seq_major_rows, + total_work, + max(1, total_work), + ), + key_arg_indices=(4,), + static_arg_indices=(4,), + ) + return k_topk, cu_q_topk, cu_k_topk, seq_lens + + +def build_attention_indices( + cu_seqlens: torch.Tensor, + global_start: int, + l_local: int, + d_window: int, + window_size: int, + ratio: int, + compressed_width: int, + indexer_topk_compressed_logical_ids: Optional[torch.Tensor] = None, + cu_seqlens_compressed: Optional[torch.Tensor] = None, + rank_major_by_seq_major: Optional[torch.Tensor] = None, + chunk_starts: Optional[Tuple[int, ...]] = None, + chunk_len: int = 0, + window_capacity_per_chunk: int = 0, + shared_compressed_base: int = 0, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Build final sparse-attention indices in physical ``kv_full`` row space. + + Inputs: + cu_seqlens: int32 CUDA tensor, shape ``(n_seq + 1,)``. + global_start/l_local: contiguous local query range, ignored for two-chunk mode. + d_window/window_size: physical left-window capacity and per-query window width. + ratio/compressed_width: compression mode and number of compressed columns. + indexer_topk_compressed_logical_ids: optional int32 tensor, shape + ``(l_local, compressed_width)``. + cu_seqlens_compressed/rank_major_by_seq_major/chunk metadata: required for + two-chunk compressed lowering. + + Outputs: + ``topk_idxs`` int32, shape ``(l_local, window_size + compressed_width)``; + ``topk_length`` int32, shape ``(l_local,)``. Window ids appear before + compressed ids and are already lowered into ``kv_full`` rows. + """ + chunk_count = 0 + chunk1_start = 0 + if chunk_starts is not None: + _require_cute( + "DSv4 two-chunk shared final idx requires CUDA tensors and CuTeDSL.", + cu_seqlens_compressed, + rank_major_by_seq_major, + ) + if not chunk_starts or len(chunk_starts) > 2: + raise RuntimeError( + "DSv4 two-chunk shared final idx supports one or two chunks, " + f"got {len(chunk_starts)}." + ) + chunk_count = len(chunk_starts) + l_local = int(chunk_count) * int(chunk_len) + global_start = 0 + chunk1_start = int(chunk_starts[1]) if chunk_count > 1 else 0 + else: + cu_seqlens_compressed = cu_seqlens + rank_major_by_seq_major = torch.empty((1,), dtype=torch.int32, device=cu_seqlens.device) + + total_width = window_size + compressed_width + topk_idxs = torch.empty((l_local, total_width), dtype=torch.int32, device=cu_seqlens.device) + topk_length = torch.empty((l_local,), dtype=torch.int32, device=cu_seqlens.device) + if indexer_topk_compressed_logical_ids is None: + dummy = torch.empty((1, 1), dtype=torch.int32, device=cu_seqlens.device) + has_indexer = 0 + else: + dummy = indexer_topk_compressed_logical_ids + has_indexer = 1 + parallel_dense = 1 if has_indexer == 0 and chunk_count == 0 else 0 + launch_work = l_local * total_width if parallel_dense else l_local + _run_compiled_launch( + _build_attention_indices_launch, + (cu_seqlens, cu_seqlens_compressed, dummy, rank_major_by_seq_major, topk_idxs, topk_length), + ( + cu_seqlens.shape[0] - 1, + global_start, + l_local, + int(chunk_starts[0]) if chunk_starts is not None else 0, + chunk1_start, + chunk_count, + int(chunk_len), + int(window_capacity_per_chunk), + int(shared_compressed_base), + d_window, + window_size, + ratio, + compressed_width, + rank_major_by_seq_major.shape[0], + has_indexer, + 1 if ratio > 1 and compressed_width > 0 else 0, + total_width, + parallel_dense, + launch_work, + ), + key_arg_indices=(), + ) + return topk_idxs, topk_length + + +def build_indexer_loss_indices( + cu_seqlens: torch.Tensor, + cu_seqlens_compressed: torch.Tensor, + global_start: int, + l_local: int, + d_window: int, + window_size: int, + ratio: int, + indexer_topk_compressed_logical_ids: torch.Tensor, + indexer_rank_by_seq_major: torch.Tensor, + chunk_starts: Optional[Tuple[int, ...]] = None, + chunk_len: int = 0, + window_capacity_per_chunk: int = 0, + shared_compressed_base: int = 0, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Build compressed-first sparse-attention ids and indexer-loss rank-major ids. + + Inputs: + cu_seqlens/cu_seqlens_compressed: int32 CUDA prefixes, shape ``(n_seq + 1,)``. + global_start/l_local: contiguous local query range, ignored for two-chunk mode. + d_window/window_size/ratio: final-id lowering parameters. + indexer_topk_compressed_logical_ids: int32 tensor, shape + ``(l_local, compressed_width)``. + indexer_rank_by_seq_major: int32 reverse map, shape ``(seq_major_rows,)``. + optional chunk metadata selects two-chunk lowering. + + Outputs: + ``topk_idxs`` int32, shape ``(l_local, compressed_width + window_size)``, + with compressed ids first; ``indexer_rank_major`` int32, shape + ``(l_local, compressed_width)`` for indexer-loss K rows. + """ + _require_cute( + "DSv4 CP indexer-loss final idx generation requires CUDA tensors and CuTeDSL.", + cu_seqlens, + cu_seqlens_compressed, + indexer_topk_compressed_logical_ids, + indexer_rank_by_seq_major, + ) + chunk_count = 0 + chunk1_start = 0 + if chunk_starts is not None: + if not chunk_starts or len(chunk_starts) > 2: + raise RuntimeError( + "DSv4 two-chunk shared indexer-loss final idx supports one or two chunks, " + f"got {len(chunk_starts)}." + ) + chunk_count = len(chunk_starts) + l_local = int(chunk_count) * int(chunk_len) + global_start = 0 + chunk1_start = int(chunk_starts[1]) if chunk_count > 1 else 0 + + compressed_width = indexer_topk_compressed_logical_ids.shape[-1] + total_width = compressed_width + window_size + topk_idxs = torch.empty((l_local, total_width), dtype=torch.int32, device=cu_seqlens.device) + indexer_rank_major = torch.empty( + (l_local, compressed_width), dtype=torch.int32, device=cu_seqlens.device + ) + total_work = l_local * total_width + if total_work == 0: + return topk_idxs, indexer_rank_major + _run_compiled_launch( + _build_indexer_loss_indices_launch, + ( + cu_seqlens, + cu_seqlens_compressed, + indexer_topk_compressed_logical_ids, + indexer_rank_by_seq_major, + topk_idxs, + indexer_rank_major, + ), + ( + cu_seqlens.shape[0] - 1, + global_start, + l_local, + int(chunk_starts[0]) if chunk_starts is not None else 0, + chunk1_start, + chunk_count, + int(chunk_len), + int(window_capacity_per_chunk), + int(shared_compressed_base), + d_window, + window_size, + ratio, + compressed_width, + indexer_rank_by_seq_major.shape[0], + total_width, + total_work, + ), + key_arg_indices=(), + ) + return topk_idxs, indexer_rank_major diff --git a/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py b/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py new file mode 100644 index 00000000000..c6b24b12a9f --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py @@ -0,0 +1,965 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +"""MCore-facing utilities for the DSv4 THD context-parallel path. + +This module is the boundary between attention code and the small CP layout +kernels in ``csa_cp_layout_kernels.py``. It owns CP partition metadata, boundary +communication, fixed-shape collectives, and the typed wrappers used by CSA/DSv4 +attention modules. +""" + +import math +from typing import Optional, Sequence, Tuple + +import torch +import torch.distributed as dist + +from megatron.core.transformer.experimental_attention_variant import csa_cp_layout_kernels +from megatron.core.transformer.experimental_attention_variant.dsa_kernels import ( + batch_of_row, + indexer_topk, +) + +# ============================================================================= +# CP Partition Modes And Layout +# ============================================================================= + +DSV4_CP_PARTITION_CONTIGUOUS = "contiguous" +DSV4_CP_PARTITION_TWO_CHUNK = "two_chunk" + + +def local_q_cp_chunk_ranges( + partition_mode: Optional[str], local_rows: int, cp_size: int, cp_rank: int +) -> Tuple[Tuple[int, int], ...]: + """Return this rank's local rows as global packed-token ranges.""" + mode = DSV4_CP_PARTITION_CONTIGUOUS if partition_mode is None else partition_mode + if mode not in (DSV4_CP_PARTITION_CONTIGUOUS, DSV4_CP_PARTITION_TWO_CHUNK): + raise RuntimeError( + "Unsupported CSA CP partition mode: " f"{mode!r}. Expected contiguous or two_chunk." + ) + local_rows, cp_size, cp_rank = int(local_rows), int(cp_size), int(cp_rank) + if local_rows <= 0: + raise RuntimeError(f"local_rows must be positive, got {local_rows}.") + if cp_size < 1 or cp_rank < 0 or cp_rank >= cp_size: + raise RuntimeError(f"Invalid CP rank/size: cp_rank={cp_rank}, cp_size={cp_size}.") + + if mode == DSV4_CP_PARTITION_CONTIGUOUS: + start = cp_rank * local_rows + return ((start, start + local_rows),) + + if cp_size == 1: + return ((0, local_rows),) + if local_rows % 2 != 0: + raise RuntimeError( + "DSv4 two-chunk CP partition expects even local_rows: " + f"local_rows={local_rows}, cp_size={cp_size}." + ) + chunk_len = local_rows // 2 + total_chunks = 2 * cp_size + chunk_ids = (cp_rank, total_chunks - 1 - cp_rank) + return tuple((chunk_id * chunk_len, (chunk_id + 1) * chunk_len) for chunk_id in chunk_ids) + + +def local_kv_cp_chunk_ranges( + partition_mode: Optional[str], local_rows: int, boundary_rows: int, cp_size: int, cp_rank: int +) -> Tuple[Tuple[int, int], ...]: + """Return boundary rows followed by local rows as global packed-token ranges.""" + local_ranges = local_q_cp_chunk_ranges(partition_mode, local_rows, cp_size, cp_rank) + boundary_rows = int(boundary_rows) + if boundary_rows <= 0: + raise RuntimeError(f"boundary_rows must be positive, got {boundary_rows}.") + if boundary_rows % len(local_ranges) != 0: + raise RuntimeError( + "CP boundary rows must be divisible by local chunk count: " + f"boundary_rows={boundary_rows}, chunks={len(local_ranges)}." + ) + d_window = boundary_rows // len(local_ranges) + boundary_ranges = tuple((int(start) - d_window, int(start)) for start, _ in local_ranges) + return boundary_ranges + local_ranges + + +def _compressed_group_capacity(rows: int, ratio: int, d_comp: int) -> int: + """Return compressed group capacity with row alignment for TE FP8 linear inputs.""" + c_cap = max(1, (int(rows) + int(d_comp)) // int(ratio)) + group_alignment = 32 // math.gcd(32, int(ratio)) + return ((c_cap + group_alignment - 1) // group_alignment) * group_alignment + + +def thd_cp_local_row_indices( + partition_mode: Optional[str], + padded_total_rows: int, + cp_size: int, + cp_rank: int, + device: torch.device, +) -> torch.Tensor: + """Return global THD row indices selected by this CP rank. + + ``padded_total_rows`` is the final padded packed-token length. The output is + an int64 tensor of shape ``(padded_total_rows / cp_size,)``. Its row order is + exactly the order expected by ``local_q_cp_chunk_ranges`` for the selected + CSA CP partition mode. + """ + padded_total_rows = int(padded_total_rows) + cp_size, cp_rank = int(cp_size), int(cp_rank) + if padded_total_rows <= 0: + raise RuntimeError(f"padded_total_rows must be positive, got {padded_total_rows}.") + if cp_size < 1 or cp_rank < 0 or cp_rank >= cp_size: + raise RuntimeError(f"Invalid CP rank/size: cp_rank={cp_rank}, cp_size={cp_size}.") + if padded_total_rows % cp_size != 0: + raise RuntimeError( + "DSv4 THD CP partition expects padded_total_rows divisible by cp_size: " + f"padded_total_rows={padded_total_rows}, cp_size={cp_size}." + ) + + local_rows = padded_total_rows // cp_size + return torch.cat( + [ + torch.arange(start, end, dtype=torch.long, device=device) + for start, end in local_q_cp_chunk_ranges( + partition_mode, local_rows=local_rows, cp_size=cp_size, cp_rank=cp_rank + ) + ], + dim=0, + ) + + +def _normalize_row_ranges( + chunk_ranges: Sequence[Tuple[int, int]], op_name: str +) -> Tuple[Tuple[Tuple[int, int], ...], Tuple[int, ...], int]: + """Validate row ranges and return normalized ranges, lengths, and total rows.""" + if not chunk_ranges: + raise RuntimeError(f"{op_name} expects at least one chunk range.") + normalized = [] + lengths = [] + for start, end in chunk_ranges: + start = int(start) + end = int(end) + length = end - start + if length <= 0: + raise RuntimeError(f"{op_name} expects positive chunk lengths, got {length}.") + normalized.append((start, end)) + lengths.append(length) + return tuple(normalized), tuple(lengths), sum(lengths) + + +def _two_chunk_layout( + chunk_ranges: Sequence[Tuple[int, int]], op_name: str +) -> Tuple[Tuple[int, int], int, int]: + ranges, lengths, l_local = _normalize_row_ranges(chunk_ranges, op_name) + if len(ranges) != 2 or lengths[0] != lengths[1]: + raise RuntimeError(f"{op_name} expects exactly two equal-length chunks.") + return (int(ranges[0][0]), int(ranges[1][0])), int(lengths[0]), int(l_local) + + +def build_thd_cp_local_padding_row_mask( + cu_seqlens_padded: torch.Tensor, + cu_seqlens_unpadded: Optional[torch.Tensor], + chunk_ranges: Sequence[Tuple[int, int]], + total_padded_rows: int, +) -> Optional[torch.Tensor]: + """Return this CP rank's local query rows that are CUDA-graph padding rows.""" + if cu_seqlens_unpadded is None: + return None + if cu_seqlens_padded.shape != cu_seqlens_unpadded.shape: + raise RuntimeError( + "padded and unpadded THD cu_seqlens must have the same shape, got " + f"{tuple(cu_seqlens_padded.shape)} and {tuple(cu_seqlens_unpadded.shape)}." + ) + + chunk_ranges, _, _ = _normalize_row_ranges(chunk_ranges, "DSv4 CP padding-row mask") + device = cu_seqlens_padded.device + global_rows = torch.cat( + [ + torch.arange(start, end, device=device, dtype=torch.int64) + for start, end in chunk_ranges + ], + dim=0, + ) + + row_batch_ids_full = batch_of_row(cu_seqlens_padded, total_q=int(total_padded_rows)) + row_batch_ids = row_batch_ids_full.index_select(0, global_rows) + + real_seg_lens = cu_seqlens_unpadded[1:] - cu_seqlens_unpadded[:-1] + pos_in_seg = global_rows.to(cu_seqlens_padded.dtype) - cu_seqlens_padded[row_batch_ids] + return pos_in_seg >= real_seg_lens[row_batch_ids] + + +# ============================================================================= +# RoPE Wrappers +# ============================================================================= + + +def apply_thd_cp_local_rope_fused( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + nope_dim: int, + pos_dim: int, + cu_seqlens_padded: torch.Tensor, + chunk_ranges: Sequence[Tuple[int, int]], + inverse: bool = False, + clamp_to_valid_token: bool = False, +) -> torch.Tensor: + """Apply fused non-interleaved RoPE to local THD CP rows.""" + chunk_ranges, lengths, l_local = _normalize_row_ranges(chunk_ranges, "DSv4 THD CP local RoPE") + if x.shape[0] != l_local: + raise RuntimeError( + f"DSv4 THD CP local RoPE expects x rows to be {l_local}, got {x.shape[0]}." + ) + if len(chunk_ranges) == 1 or (len(chunk_ranges) == 2 and lengths[0] == lengths[1]): + return csa_cp_layout_kernels.ThdLocalRope.apply( + x, + cos, + sin, + cu_seqlens_padded, + int(chunk_ranges[0][0]), + int(lengths[0]), + nope_dim, + pos_dim, + int(chunk_ranges[1][0]) if len(chunk_ranges) == 2 else 0, + bool(inverse), + bool(clamp_to_valid_token), + ) + + parts = [] + row_start = 0 + for chunk_range, length in zip(chunk_ranges, lengths): + parts.append( + csa_cp_layout_kernels.ThdLocalRope.apply( + x.narrow(0, row_start, int(length)), + cos, + sin, + cu_seqlens_padded, + int(chunk_range[0]), + int(length), + nope_dim, + pos_dim, + 0, + bool(inverse), + bool(clamp_to_valid_token), + ) + ) + row_start += int(length) + return torch.cat(parts, dim=0) + + +def apply_thd_cp_compressed_rope_fused( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + comp_ids_local: torch.Tensor, + ratio: int, + nope_dim: int, + pos_dim: int, + inverse: bool = False, +) -> torch.Tensor: + """Apply fused RoPE to compressed rows using their original compression ids.""" + return csa_cp_layout_kernels.ThdCompressedRope.apply( + x, cos, sin, comp_ids_local, ratio, nope_dim, pos_dim, bool(inverse) + ) + + +# ============================================================================= +# Boundary Hidden Exchange +# ============================================================================= + + +class _LeftBoundaryExchange(torch.autograd.Function): + """Exchange fixed left-boundary windows and scatter gradients back to senders.""" + + @staticmethod + def forward( + ctx, + tensor: torch.Tensor, + d_window: int, + cp_group: torch.distributed.ProcessGroup, + two_chunk: bool, + ): + """Receive fixed left-boundary hidden rows needed by this CP rank.""" + cp_size = 1 if cp_group is None else cp_group.size() + cp_rank = 0 if cp_group is None else cp_group.rank() + if cp_size <= 1: + total_chunks, local_chunk_ids = 1, (0,) + owner_ranks = (0,) + elif two_chunk: + total_chunks = 2 * cp_size + local_chunk_ids = (cp_rank, total_chunks - 1 - cp_rank) + owner_ranks = tuple( + min(chunk_id, total_chunks - 1 - chunk_id) for chunk_id in range(total_chunks) + ) + else: + total_chunks, local_chunk_ids = cp_size, (cp_rank,) + owner_ranks = tuple(range(cp_size)) + ctx.cp_group = cp_group + ctx.d_window = d_window + ctx.input_shape = tensor.shape + ctx.local_chunk_ids = local_chunk_ids + ctx.owner_ranks = owner_ranks + + local_chunks = len(local_chunk_ids) + if tensor.shape[0] % local_chunks != 0: + raise RuntimeError( + "DSv4 CP boundary exchange expects equal local chunks: " + f"local={tensor.shape[0]}, chunks={local_chunks}." + ) + chunk_len = tensor.shape[0] // local_chunks + ctx.chunk_len = chunk_len + if chunk_len < d_window: + raise RuntimeError( + "DSv4 CP boundary exchange requires chunk_len >= D_window: " + f"chunk_len={chunk_len}, D_window={d_window}." + ) + if cp_size <= 1: + return tensor.new_zeros((local_chunks * d_window,) + tuple(tensor.shape[1:])) + + recv = tensor.new_zeros((local_chunks, d_window) + tuple(tensor.shape[1:])) + ops = [] + send_tensors = [] + to_peer = dist.get_global_rank if hasattr(dist, "get_global_rank") else lambda _, rank: rank + + for local_idx, chunk_id in enumerate(local_chunk_ids): + if chunk_id > 0: + prev_chunk_id = chunk_id - 1 + source_rank = owner_ranks[prev_chunk_id] + if source_rank == cp_rank: + source_idx = local_chunk_ids.index(prev_chunk_id) + source_start = source_idx * chunk_len + recv[local_idx].copy_( + tensor[source_start + chunk_len - d_window : source_start + chunk_len] + ) + else: + ops.append( + dist.P2POp( + dist.irecv, recv[local_idx], to_peer(cp_group, source_rank), cp_group + ) + ) + + next_chunk_id = chunk_id + 1 + if next_chunk_id < total_chunks: + target_rank = owner_ranks[next_chunk_id] + if target_rank != cp_rank: + start = local_idx * chunk_len + send_tensors.append( + tensor[start + chunk_len - d_window : start + chunk_len].contiguous() + ) + ops.append( + dist.P2POp( + dist.isend, send_tensors[-1], to_peer(cp_group, target_rank), cp_group + ) + ) + + if ops: + for req in dist.batch_isend_irecv(ops): + req.wait() + return recv.reshape((local_chunks * d_window,) + tuple(tensor.shape[1:])) + + @staticmethod + def backward(ctx, grad_boundary: torch.Tensor): + """Send boundary gradients back to ranks that own those hidden rows.""" + cp_group = ctx.cp_group + cp_size = 1 if cp_group is None else cp_group.size() + cp_rank = 0 if cp_group is None else cp_group.rank() + d_window = ctx.d_window + chunk_len = ctx.chunk_len + grad_input = grad_boundary.new_zeros(ctx.input_shape) + if cp_size <= 1: + return grad_input, None, None, None + + local_chunk_ids = ctx.local_chunk_ids + owner_ranks = ctx.owner_ranks + total_chunks = len(owner_ranks) + local_chunks = len(local_chunk_ids) + grad_chunks = grad_boundary.reshape( + (local_chunks, d_window) + tuple(grad_boundary.shape[1:]) + ) + ops = [] + send_tensors = [] + recv_specs = [] + to_peer = dist.get_global_rank if hasattr(dist, "get_global_rank") else lambda _, rank: rank + + for local_idx, chunk_id in enumerate(local_chunk_ids): + if chunk_id > 0: + prev_chunk_id = chunk_id - 1 + source_rank = owner_ranks[prev_chunk_id] + if source_rank == cp_rank: + source_idx = local_chunk_ids.index(prev_chunk_id) + source_start = source_idx * chunk_len + grad_input[ + source_start + chunk_len - d_window : source_start + chunk_len + ] += grad_chunks[local_idx] + else: + send_tensors.append(grad_chunks[local_idx].contiguous()) + ops.append( + dist.P2POp( + dist.isend, send_tensors[-1], to_peer(cp_group, source_rank), cp_group + ) + ) + + next_chunk_id = chunk_id + 1 + if next_chunk_id < total_chunks: + target_rank = owner_ranks[next_chunk_id] + if target_rank != cp_rank: + recv = grad_chunks.new_zeros((d_window,) + tuple(grad_chunks.shape[2:])) + recv_specs.append((local_idx, recv)) + ops.append( + dist.P2POp(dist.irecv, recv, to_peer(cp_group, target_rank), cp_group) + ) + + if ops: + for req in dist.batch_isend_irecv(ops): + req.wait() + for local_idx, recv in recv_specs: + start = local_idx * chunk_len + grad_input[start + chunk_len - d_window : start + chunk_len] += recv + return grad_input, None, None, None + + +def exchange_cp_boundary_hidden( + hidden_states: torch.Tensor, + csa_compress_ratios: Sequence[int], + csa_window_size: int, + partition_mode: Optional[str], + cp_group: torch.distributed.ProcessGroup, +) -> torch.Tensor: + """Exchange hidden-state boundary rows for the selected CSA CP partition mode.""" + d_comp = max( + (8 if ratio == 4 else ratio for ratio in csa_compress_ratios if ratio and ratio > 1), + default=0, + ) + d_window = max(int(csa_window_size), d_comp) + hidden_flat = hidden_states.view(hidden_states.shape[0], -1) + cp_size = 1 if cp_group is None else cp_group.size() + cp_rank = 0 if cp_group is None else cp_group.rank() + local_chunks = len( + local_q_cp_chunk_ranges(partition_mode, hidden_flat.shape[0], cp_size, cp_rank) + ) + boundary_hidden = _LeftBoundaryExchange.apply( + hidden_flat, d_window, cp_group, local_chunks == 2 + ) + return boundary_hidden.reshape((local_chunks * d_window,) + tuple(hidden_states.shape[1:])) + + +# ============================================================================= +# Compressed Metadata And Compressor Inputs +# ============================================================================= + + +def build_global_compressed_cu_seqlens(cu_seqlens_padded: torch.Tensor, ratio: int) -> torch.Tensor: + """Build fixed-shape compressed sequence prefix sums. + + Inputs: + cu_seqlens_padded: int32 CUDA tensor, shape ``(n_seq + 1,)``. + ratio: positive compression ratio. + + Output: + int32 CUDA tensor, shape ``(n_seq + 1,)`` where each prefix increments by + ``floor(seq_len / ratio)``. This is intentionally a PyTorch utility, not + a CuTe kernel, because it is fixed-shape and does not need a host sync. + """ + if ratio <= 1: + raise RuntimeError(f"DSv4 CP compressed cu_seqlens expects ratio > 1, got {ratio}.") + out = torch.empty_like(cu_seqlens_padded) + out[:1].zero_() + compressed_lens = torch.div( + cu_seqlens_padded[1:] - cu_seqlens_padded[:-1], int(ratio), rounding_mode="floor" + ) + out[1:].copy_(torch.cumsum(compressed_lens, dim=0, dtype=torch.int32)) + return out + + +def build_cp_rank_major_compressed_metadata_fused( + cu_seqlens: torch.Tensor, + chunk_ranges: Sequence[Tuple[int, int]], + cp_size: int, + ratio: int, + d_comp: int, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build rank-major compressed metadata for the selected CP partition mode.""" + if len(chunk_ranges) == 2: + _, chunk_len, _ = _two_chunk_layout(chunk_ranges, "DSv4 two-chunk compressed metadata") + c_cap_per_chunk = (chunk_len + int(d_comp)) // int(ratio) + c_cap_per_rank = _compressed_group_capacity(2 * chunk_len, ratio, 2 * d_comp) + return csa_cp_layout_kernels.build_compressed_row_metadata( + cu_seqlens, + int(cp_size), + chunk_len, + ratio, + d_comp, + c_cap_per_chunk, + c_cap_per_rank=c_cap_per_rank, + use_two_chunk=True, + ) + _, _, l_local = _normalize_row_ranges(chunk_ranges, "DSv4 CP compressed metadata") + c_cap = _compressed_group_capacity(l_local, ratio, d_comp) + return csa_cp_layout_kernels.build_compressed_row_metadata( + cu_seqlens, int(cp_size), l_local, ratio, d_comp, c_cap + ) + + +def build_cp_compressor_prep_compact_fused( + hidden_local: torch.Tensor, + boundary_hidden: torch.Tensor, + cu_seqlens: torch.Tensor, + chunk_ranges: Sequence[Tuple[int, int]], + ratio: int, + d_comp: int, + d_window: int, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build fixed-capacity compressor input for the selected CP partition mode. + + Returns: + ``hidden_compact``: rank-local compressor input, shape + ``(compact_group_capacity * ratio, ...)``. + ``cu_compact``: rank-local token-row prefix sums over global sequence + ids, shape ``(n_seq + 1,)``. The compressor divides these deltas + by ``ratio`` to recover this rank's visible compact groups. + ``comp_ids_local``: original per-sequence compressed group id for each + compact group, shape ``(compact_group_capacity,)``. For example, + with ``ratio=4``, ``comp_id=3`` maps to RoPE position ``12``. + """ + d_window = int(d_window) + + two_chunk = len(chunk_ranges) == 2 + if two_chunk: + chunk_starts, chunk_len, l_local = _two_chunk_layout( + chunk_ranges, "DSv4 two-chunk compressor-prep" + ) + else: + start, end = chunk_ranges[0] + chunk_ranges = ((int(start), int(end)),) + l_local = chunk_ranges[0][1] - chunk_ranges[0][0] + if hidden_local.shape[0] != l_local: + raise RuntimeError( + "DSv4 CP compressor-prep expects hidden rows to be " + f"{l_local}, got {hidden_local.shape[0]}." + ) + expected_boundary_rows = len(chunk_ranges) * d_window + if boundary_hidden.shape[0] != expected_boundary_rows: + raise RuntimeError( + "DSv4 CP compressor-prep expects one boundary window per chunk: " + f"boundary={boundary_hidden.shape[0]}, expected={expected_boundary_rows}." + ) + + if two_chunk: + chunk_specs = [ + (0, chunk_len, 0, chunk_starts[0]), + (chunk_len, chunk_len, d_window, chunk_starts[1]), + ] + else: + chunk_specs = [(0, l_local, 0, chunk_ranges[0][0])] + + parts = [] + for row_start, rows, boundary_start, global_start in chunk_specs: + c_cap = ( + (int(rows) + int(d_comp)) // int(ratio) + if two_chunk + else _compressed_group_capacity(rows, ratio, d_comp) + ) + parts.append( + csa_cp_layout_kernels.CompressorInputCompact.apply( + hidden_local.narrow(0, row_start, rows), + boundary_hidden.narrow(0, boundary_start, d_window), + cu_seqlens, + global_start, + rows, + ratio, + d_comp, + d_window, + c_cap, + ) + ) + if not two_chunk: + hidden_compact, cu_compact, _seq_ids, comp_ids, _valid = parts[0] + return hidden_compact, cu_compact, comp_ids + hidden_parts, cu_parts, _seq_parts, comp_parts, _valid_parts = zip(*parts) + target_c_cap = _compressed_group_capacity(l_local, ratio, len(chunk_ranges) * d_comp) + current_c_cap = sum(part.shape[0] for part in comp_parts) + if current_c_cap > target_c_cap: + raise RuntimeError( + "DSv4 two-chunk compressor-prep produced more groups than its aligned capacity: " + f"current={current_c_cap}, capacity={target_c_cap}." + ) + if current_c_cap < target_c_cap: + pad_groups = target_c_cap - current_c_cap + hidden_pad_rows = pad_groups * int(ratio) + hidden_parts = hidden_parts + ( + hidden_local.new_zeros((hidden_pad_rows,) + tuple(hidden_local.shape[1:])), + ) + comp_parts = comp_parts + ( + torch.full((pad_groups,), -1, dtype=torch.int32, device=hidden_local.device), + ) + cu_deltas = cu_parts[0][1:] - cu_parts[0][:-1] + for cu in cu_parts[1:]: + cu_deltas = cu_deltas + (cu[1:] - cu[:-1]) + cu_compact = torch.cat( + ( + torch.zeros((1,), dtype=cu_seqlens.dtype, device=cu_seqlens.device), + torch.cumsum(cu_deltas, dim=0), + ) + ) + return ( + torch.cat(hidden_parts, dim=0), + cu_compact, + torch.cat(comp_parts, dim=0), + ) + + +# ============================================================================= +# Fixed-Shape CP Collectives +# ============================================================================= + + +class _AllGatherFixedCPTensor(torch.autograd.Function): + """All-gather a fixed local tensor and reduce-scatter gradients in backward.""" + + @staticmethod + def forward(ctx, tensor: torch.Tensor, cp_group) -> torch.Tensor: + """Gather fixed-size CP-local rows into rank-major order.""" + cp_size = 1 if cp_group is None else cp_group.size() + ctx.cp_group = cp_group + ctx.input_shape = tuple(tensor.shape) + ctx.cp_size = cp_size + + if cp_size == 1: + return tensor + + local = tensor.contiguous() + output_shape = (cp_size * local.shape[0],) + tuple(local.shape[1:]) + output = local.new_empty(output_shape) + if hasattr(dist, "all_gather_into_tensor"): + dist.all_gather_into_tensor(output, local, group=cp_group) + elif hasattr(dist, "_all_gather_base"): + dist._all_gather_base(output, local, group=cp_group) + else: + raise RuntimeError("DSv4 CP fixed all-gather requires all_gather_into_tensor support") + return output + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + """Reduce-scatter rank-major gradients back to the local CP shard.""" + cp_size = ctx.cp_size + if cp_size == 1: + return grad_output, None + + expected_shape = (cp_size * ctx.input_shape[0],) + tuple(ctx.input_shape[1:]) + if tuple(grad_output.shape) != expected_shape: + raise RuntimeError( + "DSv4 CP fixed all-gather backward received an unexpected gradient shape: " + f"grad={tuple(grad_output.shape)}, expected={expected_shape}" + ) + + grad_output = grad_output.contiguous() + grad_input = grad_output.new_empty(ctx.input_shape) + if hasattr(dist, "reduce_scatter_tensor"): + dist.reduce_scatter_tensor( + grad_input, grad_output, op=dist.ReduceOp.SUM, group=ctx.cp_group + ) + elif hasattr(dist, "_reduce_scatter_base"): + dist._reduce_scatter_base( + grad_input, grad_output, op=dist.ReduceOp.SUM, group=ctx.cp_group + ) + else: + raise RuntimeError( + "DSv4 CP fixed all-gather backward requires reduce_scatter_tensor support" + ) + return grad_input, None + + +def all_gather_fixed_cp_tensor( + tensor: torch.Tensor, cp_group: torch.distributed.ProcessGroup +) -> torch.Tensor: + """All-gather fixed-capacity CP tensors along dim 0 with autograd support.""" + return _AllGatherFixedCPTensor.apply(tensor, cp_group) + + +# ============================================================================= +# KV Packing And Compressed KV Repacking +# ============================================================================= + + +def pack_cp_kv_full_fused( + kv_local: torch.Tensor, + boundary_kv: torch.Tensor, + compressed_kv_rank_major: torch.Tensor, + seq_ids_rank_major: torch.Tensor, + comp_ids_rank_major: torch.Tensor, + valid_rank_major: torch.Tensor, + cu_seqlens: torch.Tensor, + chunk_ranges: Sequence[Tuple[int, int]], + d_window: int, + ratio: int, + rank_major_by_seq_major: Optional[torch.Tensor] = None, + cu_seqlens_compressed: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, Optional[int]]: + """Pack full KV for the selected CP partition mode.""" + d_window = int(d_window) + shared_compressed_base = None + if len(chunk_ranges) == 2: + chunk_starts, chunk_len, l_local = _two_chunk_layout(chunk_ranges, "DSv4 two-chunk KV pack") + window_capacity = max(1, int(chunk_len) + d_window * (cu_seqlens.shape[0] - 1)) + shared_compressed_base = 2 * window_capacity + total_capacity = max(1, shared_compressed_base + int(compressed_kv_rank_major.shape[0])) + dummy_ids = torch.empty((1,), dtype=torch.int32, device=kv_local.device) + dummy_valid = torch.empty((1,), dtype=torch.bool, device=kv_local.device) + seq_ids_rank_major = comp_ids_rank_major = rank_major_by_seq_major = dummy_ids + valid_rank_major = dummy_valid + cu_seqlens_compressed = cu_seqlens + global_start = local_rows = kernel_ratio = 0 + chunk0_start, chunk1_start, chunk_count = int(chunk_starts[0]), int(chunk_starts[1]), 2 + else: + chunk_ranges, _, l_local = _normalize_row_ranges(chunk_ranges, "DSv4 CP KV pack") + total_capacity = max( + 1, + int(l_local) + + d_window * (cu_seqlens.shape[0] - 1) + + int(compressed_kv_rank_major.shape[0]), + ) + if rank_major_by_seq_major is None: + rank_major_by_seq_major = torch.empty((1,), dtype=torch.int32, device=kv_local.device) + if cu_seqlens_compressed is None: + cu_seqlens_compressed = cu_seqlens + global_start = int(chunk_ranges[0][0]) + local_rows = int(l_local) + kernel_ratio = int(ratio) + chunk0_start = chunk1_start = chunk_count = chunk_len = window_capacity = 0 + if kv_local.shape[0] != l_local: + raise RuntimeError( + f"DSv4 CP KV pack expects local KV rows to be {l_local}, got {kv_local.shape[0]}." + ) + expected_boundary_rows = len(chunk_ranges) * d_window + if boundary_kv.shape[0] != expected_boundary_rows: + raise RuntimeError( + "DSv4 CP KV pack expects one boundary window per chunk: " + f"boundary={boundary_kv.shape[0]}, expected={expected_boundary_rows}." + ) + + kv_full = csa_cp_layout_kernels.ThdFullKvPack.apply( + kv_local, + boundary_kv, + compressed_kv_rank_major, + seq_ids_rank_major, + comp_ids_rank_major, + valid_rank_major, + cu_seqlens, + rank_major_by_seq_major, + cu_seqlens_compressed, + global_start, + local_rows, + d_window, + kernel_ratio, + int(total_capacity), + chunk0_start, + chunk1_start, + chunk_count, + int(chunk_len), + window_capacity, + ) + return kv_full, shared_compressed_base + + +def repack_rank_major_compressed_to_seq_major_fused( + rank_major: torch.Tensor, + seq_ids_rank_major: torch.Tensor, + comp_ids_rank_major: torch.Tensor, + valid_rank_major: torch.Tensor, + cu_seqlens_compressed: torch.Tensor, + seq_major_rows: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Repack gathered compressed KV from rank-major rows to sequence-major rows.""" + return csa_cp_layout_kernels.repack_compressed_kv_to_seq_major( + rank_major, + seq_ids_rank_major, + comp_ids_rank_major, + valid_rank_major, + cu_seqlens_compressed, + seq_major_rows, + ) + + +def compute_cp_indexer_topk_logical_fused( + q_indexer_local: torch.Tensor, + weights_indexer_local: torch.Tensor, + k_indexer_seq_major: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_compressed: torch.Tensor, + chunk_ranges: Sequence[Tuple[int, int]], + ratio: int, + topk_width: int, + indexer_softmax_scale: float, + max_seqlen_q: int, + max_seqlen_kv: int, +) -> Optional[torch.Tensor]: + """Run CP-aware indexer top-k for the selected CP partition mode. + + This is a workaround for the DSA fused indexer forward mask contract. A CP + chunk may need a mask whose first local query already sees a compressed K + prefix, for example: + + ``` + 1 1 1 0 0 0 0 0 + 1 1 1 1 0 0 0 0 + 1 1 1 1 1 0 0 0 + ``` + + The fused kernel assumes the last Q row can see all K rows, then derives + earlier rows' causal masks from that endpoint. A CP chunk's last Q row may + only see a prefix of the global K rows, so this helper copies that visible + prefix before calling the kernel. + """ + topk_width = int(topk_width) + if topk_width == 0 or k_indexer_seq_major.shape[0] == 0: + return None + + chunk_ranges, lengths, l_local = _normalize_row_ranges(chunk_ranges, "DSv4 CP indexer top-k") + if q_indexer_local.shape[0] != l_local: + raise RuntimeError( + "DSv4 CP indexer top-k expects q_indexer_local rows to be " + f"{l_local}, got {q_indexer_local.shape[0]}." + ) + if weights_indexer_local.shape[0] != l_local: + raise RuntimeError( + "DSv4 CP indexer top-k expects weights rows to be " + f"{l_local}, got {weights_indexer_local.shape[0]}." + ) + + outputs = [] + row_start = 0 + for (global_start, _), length in zip(chunk_ranges, lengths): + q_for_topk = q_indexer_local.narrow(0, row_start, int(length)) + weights_for_topk = weights_indexer_local.narrow(0, row_start, int(length)) + k_for_topk, cu_q_topk, cu_k_topk, seq_lens_topk = ( + csa_cp_layout_kernels.build_indexer_topk_metadata( + k_indexer_seq_major, + cu_seqlens_q, + cu_seqlens_compressed, + int(global_start), + int(length), + ratio, + ) + ) + topk_chunk, _ = indexer_topk( + q_for_topk, + k_for_topk, + weights_for_topk, + topk=topk_width, + ratio=ratio, + indexer_softmax_scale=indexer_softmax_scale, + cu_seqlens_q=cu_q_topk, + cu_seqlens_kv=cu_k_topk, + max_seqlen_q=int(max_seqlen_q), + max_seqlen_kv=int(max_seqlen_kv), + visible_k_lengths=seq_lens_topk, + ) + outputs.append(topk_chunk) + row_start += int(length) + return torch.cat(outputs, dim=0) + + +# ============================================================================= +# Final Attention And Indexer-Loss Indices +# ============================================================================= + + +def build_cp_attention_indices_fused( + cu_seqlens: torch.Tensor, + chunk_ranges: Sequence[Tuple[int, int]], + d_window: int, + window_size: int, + ratio: int, + indexer_topk_compressed_logical_ids: Optional[torch.Tensor] = None, + max_n_compressed: int = 0, + rank_major_by_seq_major: Optional[torch.Tensor] = None, + cu_seqlens_compressed: Optional[torch.Tensor] = None, + shared_compressed_base: Optional[int] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Build final sparse-attention indices for the selected CP partition mode.""" + compressed_width = max_n_compressed if ratio > 1 else 0 + if indexer_topk_compressed_logical_ids is not None: + compressed_width = indexer_topk_compressed_logical_ids.shape[-1] + + if len(chunk_ranges) == 2: + chunk_starts, chunk_len, l_local = _two_chunk_layout( + chunk_ranges, "DSv4 two-chunk final idx" + ) + if shared_compressed_base is None: + raise RuntimeError("DSv4 two-chunk final idx requires a shared compressed base.") + if compressed_width > 0 and ( + rank_major_by_seq_major is None or cu_seqlens_compressed is None + ): + raise RuntimeError("DSv4 two-chunk final idx with compression requires metadata.") + if compressed_width == 0: + cu_seqlens_compressed = cu_seqlens + rank_major_by_seq_major = torch.empty((1,), dtype=torch.int32, device=cu_seqlens.device) + global_start, window_capacity_per_chunk = 0, int(shared_compressed_base) // 2 + else: + chunk_ranges, _, l_local = _normalize_row_ranges(chunk_ranges, "DSv4 CP final idx") + global_start = chunk_ranges[0][0] + chunk_starts, chunk_len, window_capacity_per_chunk, shared_compressed_base = None, 0, 0, 0 + if ( + indexer_topk_compressed_logical_ids is not None + and indexer_topk_compressed_logical_ids.shape[0] != l_local + ): + raise RuntimeError( + "DSv4 CP final idx expects indexer rows to be " + f"{l_local}, got {indexer_topk_compressed_logical_ids.shape[0]}." + ) + return csa_cp_layout_kernels.build_attention_indices( + cu_seqlens, + global_start, + l_local, + d_window, + window_size, + ratio, + compressed_width, + indexer_topk_compressed_logical_ids, + cu_seqlens_compressed=cu_seqlens_compressed, + rank_major_by_seq_major=rank_major_by_seq_major, + chunk_starts=chunk_starts, + chunk_len=chunk_len, + window_capacity_per_chunk=window_capacity_per_chunk, + shared_compressed_base=shared_compressed_base, + ) + + +def build_cp_indexer_loss_indices_fused( + cu_seqlens: torch.Tensor, + cu_seqlens_compressed: torch.Tensor, + chunk_ranges: Sequence[Tuple[int, int]], + d_window: int, + window_size: int, + ratio: int, + indexer_topk_compressed_logical_ids: torch.Tensor, + rank_major_by_seq_major: torch.Tensor, + shared_compressed_base: Optional[int] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Build indexer-loss indices for the selected CP partition mode.""" + if len(chunk_ranges) == 2: + chunk_starts, chunk_len, l_local = _two_chunk_layout( + chunk_ranges, "DSv4 two-chunk indexer-loss idx" + ) + if shared_compressed_base is None: + raise RuntimeError("DSv4 two-chunk indexer-loss idx requires a shared compressed base.") + global_start, window_capacity_per_chunk = 0, int(shared_compressed_base) // 2 + else: + chunk_ranges, _, l_local = _normalize_row_ranges(chunk_ranges, "DSv4 CP indexer-loss idx") + global_start = chunk_ranges[0][0] + chunk_starts, chunk_len, window_capacity_per_chunk, shared_compressed_base = None, 0, 0, 0 + if indexer_topk_compressed_logical_ids.shape[0] != l_local: + raise RuntimeError( + "DSv4 CP indexer-loss idx expects indexer rows to be " + f"{l_local}, got {indexer_topk_compressed_logical_ids.shape[0]}." + ) + return csa_cp_layout_kernels.build_indexer_loss_indices( + cu_seqlens, + cu_seqlens_compressed, + global_start, + l_local, + d_window, + window_size, + ratio, + indexer_topk_compressed_logical_ids, + rank_major_by_seq_major, + chunk_starts=chunk_starts, + chunk_len=chunk_len, + window_capacity_per_chunk=window_capacity_per_chunk, + shared_compressed_base=shared_compressed_base, + ) diff --git a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py index 67dfc365128..0a281b732df 100644 --- a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py +++ b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py @@ -19,6 +19,7 @@ from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.transformer.attention import Attention from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.experimental_attention_variant import csa_cp_utils as cp_utils from megatron.core.transformer.spec_utils import ModuleSpec, build_module from megatron.core.transformer.torch_norm import LayerNormBuilder from megatron.core.transformer.transformer_config import MLATransformerConfig @@ -262,17 +263,36 @@ def forward( assert packed_seq_params.cp_group is not None, "cp_group must be set in dynamic-cp mode" self.pg_collection.cp = packed_seq_params.cp_group + cp_size = self.pg_collection.cp.size() + qkv_format = packed_seq_params.qkv_format if packed_seq_params is not None else None + if cp_size > 1 and qkv_format != 'thd': + raise RuntimeError( + "DSv4 context parallel support requires THD packed sequence input; " + f"got cp_size={cp_size}, qkv_format={qkv_format!r}." + ) + + boundary_hidden = None + if cp_size > 1: + boundary_hidden = cp_utils.exchange_cp_boundary_hidden( + hidden_states, + self.config.csa_compress_ratios, + self.config.csa_window_size, + self.config.csa_cp_partition_mode, + self.pg_collection.cp, + ) + # ===================== # Query, Key, and Value # ===================== # Get the query, key and value tensors based on the type of attention - # self or cross attn. - query, key, value, q_compressed, kv_compressed = self.get_query_key_value_tensors( + query, key, value, q_compressed, _, boundary_kv = self.get_query_key_value_tensors( hidden_states, key_value_states, position_ids, packed_seq_params, inference_context=inference_context, + boundary_hidden=boundary_hidden, ) # TODO: Currently, TE can only accept contiguous tensors for MLA @@ -296,9 +316,14 @@ def forward( packed_seq_params=packed_seq_params, x=hidden_states, qr=q_compressed, + boundary_hidden=boundary_hidden, + boundary_kv=boundary_kv, ) + forced_released_tensors = [query, key, value] + if boundary_kv is not None: + forced_released_tensors.append(boundary_kv) core_attn_out = core_attn_manager.group_offload( - core_attn_out, forced_released_tensors=[query, key, value] + core_attn_out, forced_released_tensors=forced_released_tensors ) if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': @@ -355,22 +380,42 @@ def forward( else: rotary_pos_emb = self.rotary_pos_emb(rope_seqlen, packed_seq=packed_seq) if self.config.apply_rope_fusion: - if packed_seq: - core_attn_out = core_attn_out.squeeze(1) - core_attn_out = fused_mla_rope_inplace( - core_attn_out, - rotary_pos_cos, - rotary_pos_sin, - nope_dim, - pos_dim, - cu_seqlens_kv, - self.pg_collection.cp.rank(), - self.pg_collection.cp.size(), - inverse=True, - remove_interleaving=True, - ) - if packed_seq: - core_attn_out = core_attn_out.unsqueeze(1) + if cp_size > 1: + if rope_seqlen is None: + raise RuntimeError("DSv4 THD CP inverse RoPE requires max_seqlen_kv.") + chunk_ranges = cp_utils.local_q_cp_chunk_ranges( + self.config.csa_cp_partition_mode, + core_attn_out.shape[0], + cp_size, + self.pg_collection.cp.rank(), + ) + core_attn_out = cp_utils.apply_thd_cp_local_rope_fused( + core_attn_out, + rotary_pos_cos, + rotary_pos_sin, + nope_dim, + pos_dim, + cu_seqlens_kv, + chunk_ranges, + inverse=True, + ) + else: + if packed_seq: + core_attn_out = core_attn_out.squeeze(1) + core_attn_out = fused_mla_rope_inplace( + core_attn_out, + rotary_pos_cos, + rotary_pos_sin, + nope_dim, + pos_dim, + cu_seqlens_kv, + self.pg_collection.cp.rank(), + self.pg_collection.cp.size(), + inverse=True, + remove_interleaving=True, + ) + if packed_seq: + core_attn_out = core_attn_out.unsqueeze(1) else: content_part, rot_part = torch.split( core_attn_out, [core_attn_out.size(-1) - pos_dim, pos_dim], dim=-1 @@ -532,9 +577,14 @@ def get_query_key_value_tensors( inference_context=None, *, inference_params=None, + boundary_hidden=None, ): """ Derives `query`, `key` and `value` tensors from `hidden_states`. + + Returns: + Tuple of ``(query, key, value, q_compressed, kv_compressed, boundary_kv)``. + ``boundary_kv`` carries CP boundary rows when the DSv4 CP path is active. """ # s = sequence length, b = batch size, h = hidden size, n = num attention heads # Attention heads [s, b, n*h] @@ -600,6 +650,7 @@ def get_query_key_value_tensors( q_compressed, _ = self.linear_q_down_proj(hidden_states) kv_compressed = hidden_states + boundary_kv_compressed = boundary_hidden k_pos_emb = None if packed_seq_params is not None: @@ -608,6 +659,8 @@ def get_query_key_value_tensors( # So we need to reshape qkv from [t, 1, h, d] to [t, h, d]. q_compressed = q_compressed.squeeze(1) kv_compressed = kv_compressed.squeeze(1) + if boundary_kv_compressed is not None: + boundary_kv_compressed = boundary_kv_compressed.squeeze(1) # ========================================= # Apply norm @@ -621,7 +674,9 @@ def get_query_key_value_tensors( # QKV up projection and RoPE apply # ========================================= - def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb): + def qkv_up_proj_and_rope_apply( + q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb, boundary_kv_compressed + ): """ Apply the up projection and RoPE to the query and key. When sequence packing enabled, the input tensors adopt a packed shape of [t, ...]; @@ -636,42 +691,97 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po q = q.view(*q.size()[:-1], self.num_attention_heads_per_partition, self.q_head_dim) q = _q_rms_norm(q, self.config.layernorm_epsilon) - kv, _ = self.linear_kv_proj(kv_compressed) + boundary_rows = 0 + if boundary_kv_compressed is not None: + boundary_rows = boundary_kv_compressed.shape[0] + kv_projection_input = torch.cat([boundary_kv_compressed, kv_compressed], dim=0) + else: + kv_projection_input = kv_compressed + + kv, _ = self.linear_kv_proj(kv_projection_input) kv = self.kv_layernorm(kv) + boundary_kv = None # [num_tokens, qk_pos_emb_head_dim] -> [num_tokens, 1, qk_pos_emb_head_dim] if k_pos_emb is not None: k_pos_emb = torch.unsqueeze(k_pos_emb, -2) + cp_size = self.pg_collection.cp.size() if self.config.apply_rope_fusion: - cp_rank = self.pg_collection.cp.rank() - cp_size = self.pg_collection.cp.size() - query = fused_mla_rope_inplace( - q, - rotary_pos_cos, - rotary_pos_sin, - self.config.qk_head_dim, - self.config.qk_pos_emb_head_dim, - cu_seqlens_q, - cp_rank, - cp_size, - remove_interleaving=True, - ) - kv = kv.unsqueeze(-2) - kv = fused_mla_rope_inplace( - kv, - rotary_pos_cos, - rotary_pos_sin, - self.config.qk_head_dim, - self.config.qk_pos_emb_head_dim, - cu_seqlens_q, - cp_rank, - cp_size, - remove_interleaving=True, - ) + if cp_size > 1: + cp_rank = self.pg_collection.cp.rank() + # CP local rows may start mid-sequence; use global THD ranges for RoPE. + assert packed_seq, "DSv4 CP fused RoPE expects THD packed sequence input." + chunk_ranges = cp_utils.local_q_cp_chunk_ranges( + self.config.csa_cp_partition_mode, q.shape[0], cp_size, cp_rank + ) + query = cp_utils.apply_thd_cp_local_rope_fused( + q, + rotary_pos_cos, + rotary_pos_sin, + self.config.qk_head_dim, + self.config.qk_pos_emb_head_dim, + cu_seqlens_q, + chunk_ranges, + ) + kv_chunk_ranges = ( + cp_utils.local_kv_cp_chunk_ranges( + self.config.csa_cp_partition_mode, + kv_compressed.shape[0], + boundary_rows, + cp_size, + cp_rank, + ) + if boundary_rows + else chunk_ranges + ) + kv = kv.unsqueeze(-2) + kv = cp_utils.apply_thd_cp_local_rope_fused( + kv, + rotary_pos_cos, + rotary_pos_sin, + self.config.qk_head_dim, + self.config.qk_pos_emb_head_dim, + cu_seqlens_q, + kv_chunk_ranges, + clamp_to_valid_token=bool(boundary_rows), + ) + if boundary_rows: + boundary_kv = kv[:boundary_rows] + kv = kv[boundary_rows:] + else: + query = fused_mla_rope_inplace( + q, + rotary_pos_cos, + rotary_pos_sin, + self.config.qk_head_dim, + self.config.qk_pos_emb_head_dim, + cu_seqlens_q, + self.pg_collection.cp.rank(), + self.pg_collection.cp.size(), + remove_interleaving=True, + ) + kv = kv.unsqueeze(-2) + kv = fused_mla_rope_inplace( + kv, + rotary_pos_cos, + rotary_pos_sin, + self.config.qk_head_dim, + self.config.qk_pos_emb_head_dim, + cu_seqlens_q, + self.pg_collection.cp.rank(), + self.pg_collection.cp.size(), + remove_interleaving=True, + ) key = kv value = kv else: + if packed_seq and cp_size > 1: + raise RuntimeError( + "DSv4 THD CP requires apply_rope_fusion=True; " + "the unfused THD CP RoPE path is not implemented." + ) + q_len = q.size()[0] if packed_seq_params is None or self.config.context_parallel_size == 1: # Shorten rotary_pos_emb to the sequence length when inference_params @@ -730,21 +840,28 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po query = query.contiguous() key = key.contiguous() value = value.contiguous() + if boundary_kv is not None: + boundary_kv = boundary_kv.contiguous() - return query, key, value + return query, key, value, boundary_kv if self.recompute_up_proj: quantization = self.config.fp8 or self.config.fp4 self.qkv_up_checkpoint = tensor_parallel.CheckpointWithoutOutput(fp8=quantization) - query, key, value = self.qkv_up_checkpoint.checkpoint( - qkv_up_proj_and_rope_apply, q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb + query, key, value, boundary_kv = self.qkv_up_checkpoint.checkpoint( + qkv_up_proj_and_rope_apply, + q_compressed, + kv_compressed, + k_pos_emb, + rotary_pos_emb, + boundary_kv_compressed, ) else: - query, key, value = qkv_up_proj_and_rope_apply( - q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb + query, key, value, boundary_kv = qkv_up_proj_and_rope_apply( + q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb, boundary_kv_compressed ) - return query, key, value, q_compressed, kv_compressed + return query, key, value, q_compressed, kv_compressed, boundary_kv def backward_dw(self) -> NoReturn: """Execute weight gradient computation""" diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py index 4e78030b6c8..0a6b8bf6079 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa.py +++ b/megatron/core/transformer/experimental_attention_variant/dsa.py @@ -91,13 +91,15 @@ def save_loss_to_tracker( tracker["avg_group"] = avg_group @staticmethod - def clean_loss_in_tracker(): + def clean_loss_in_tracker(preserve_groups: bool = False): """Clear the indexer losses.""" tracker = DSAIndexerLossLoggingHelper.tracker + reduce_group = tracker.get("reduce_group") if preserve_groups else None + avg_group = tracker.get("avg_group") if preserve_groups else None if "values" in tracker: tracker["values"].zero_() - tracker["reduce_group"] = None - tracker["avg_group"] = None + tracker["reduce_group"] = reduce_group + tracker["avg_group"] = avg_group @staticmethod def reduce_loss_in_tracker(num_layers: Optional[int] = None): @@ -172,6 +174,7 @@ def track_indexer_metrics( per_layer_logging: bool = False, num_layers: Optional[int] = None, csa_compress_ratios: Optional[List[int]] = None, + preserve_groups: bool = False, ): """Track the sparse attention indexer metrics for logging. @@ -191,6 +194,7 @@ def track_indexer_metrics( ``ratio == 4`` (the only ratio that owns an indexer) as the divisor. Otherwise (legacy DSA path) every layer is assumed to be an indexer layer and the divisor is the tracker tensor size. + preserve_groups: Keep tracker reduction groups after logging, for graph-captured runs. """ DSAIndexerLossLoggingHelper.reduce_loss_in_tracker(num_layers=num_layers) tracker = DSAIndexerLossLoggingHelper.tracker @@ -221,7 +225,7 @@ def track_indexer_metrics( if wandb_writer is not None: wandb_writer.log({"indexer loss": avg_indexer_loss}, iteration) - DSAIndexerLossLoggingHelper.clean_loss_in_tracker() + DSAIndexerLossLoggingHelper.clean_loss_in_tracker(preserve_groups=preserve_groups) def compute_dsa_indexer_loss( diff --git a/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py b/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py index 0d12954bb1e..60c5eea56bc 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py +++ b/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py @@ -479,6 +479,7 @@ def _indexer_topk_core( cu_seqlens_kv: Optional[Tensor] = None, max_seqlen_q: Optional[int] = None, max_seqlen_kv: Optional[int] = None, + visible_k_lengths: Optional[Tensor] = None, ) -> Tuple[Tensor, Tensor, Tensor]: """Layout-agnostic core for :func:`indexer_topk`. @@ -544,19 +545,36 @@ def _indexer_topk_core( sk = int(max_seqlen_kv) total_q = q.shape[0] - # Per-row valid KV length: for token ``i`` in batch ``b``, - # pos_in_seq = i - cu_seqlens_q[b] - # valid = min((pos_in_seq + 1) // ratio, seqlen_kv[b]) - row_idx = torch.arange(total_q, device=device, dtype=torch.int32) - row_batch_ids = batch_of_row(cu_seqlens_q, total_q=total_q) - row_valid = row_idx < cu_seqlens_q[-1] - pos_in_seq = row_idx - cu_seqlens_q[row_batch_ids] - pos_in_seq = torch.where(row_valid, pos_in_seq, torch.zeros_like(pos_in_seq)) - seqlen_kv_per_row = (cu_seqlens_kv[1:] - cu_seqlens_kv[:-1])[row_batch_ids] - seq_lens = ( - ((pos_in_seq + 1) // ratio).clamp(max=seqlen_kv_per_row).to(torch.int32).contiguous() - ) - seq_lens = torch.where(row_valid, seq_lens, torch.zeros_like(seq_lens)) + if visible_k_lengths is not None: + if visible_k_lengths.shape[0] != total_q: + raise ValueError( + "visible_k_lengths must have one entry per THD query row: " + f"got={visible_k_lengths.shape[0]}, expected={total_q}." + ) + if visible_k_lengths.dtype != torch.int32: + raise ValueError(f"visible_k_lengths must be int32, got {visible_k_lengths.dtype}.") + if visible_k_lengths.device != device: + raise ValueError("visible_k_lengths must be on the same device as q_indexer.") + seq_lens = visible_k_lengths + else: + # Per-row valid KV length: for token ``i`` in batch ``b``, + # pos_in_seq = i - cu_seqlens_q[b] + # valid = min((pos_in_seq + 1) // ratio, seqlen_kv[b]) + # CP callers with a trapezoid mask pass ``visible_k_lengths`` + # directly because q/k local ranges can differ from this formula. + row_idx = torch.arange(total_q, device=device, dtype=torch.int32) + row_batch_ids = batch_of_row(cu_seqlens_q, total_q=total_q) + row_valid = row_idx < cu_seqlens_q[-1] + pos_in_seq = row_idx - cu_seqlens_q[row_batch_ids] + pos_in_seq = torch.where(row_valid, pos_in_seq, torch.zeros_like(pos_in_seq)) + seqlen_kv_per_row = (cu_seqlens_kv[1:] - cu_seqlens_kv[:-1])[row_batch_ids] + seq_lens = ( + ((pos_in_seq + 1) // ratio) + .clamp(max=seqlen_kv_per_row) + .to(torch.int32) + .contiguous() + ) + seq_lens = torch.where(row_valid, seq_lens, torch.zeros_like(seq_lens)) else: # Kernel wants k as 4-D ``(b, sk, h_kv, idx_hd)``. scores = _DSA.indexer_forward_wrapper(q, k.unsqueeze(2), w, ratio=ratio)[ @@ -584,7 +602,17 @@ def _indexer_topk_core( pad = torch.full((total_q, topk - topk_k), -1, dtype=torch.int32, device=device) topk_indices = torch.cat([topk_indices, pad], dim=-1) - topk_length = (topk_indices >= 0).sum(dim=-1).int() # (total_q,) + row_valid = (topk_indices >= 0) & (topk_indices < seq_lens.unsqueeze(1)) + topk_indices = topk_indices.masked_fill(~row_valid, -1) + + if is_thd: + safe_topk = topk_indices.clamp(min=0, max=max(sk - 1, 0)).to(torch.long) + selected_scores = torch.gather(scores_flat, dim=-1, index=safe_topk) + selected_valid = (topk_indices >= 0) & (topk_indices < sk) & torch.isfinite(selected_scores) + topk_indices = topk_indices.masked_fill(~selected_valid, -1) + topk_length = (topk_indices >= 0).sum(dim=-1).int() + else: + topk_length = (topk_indices >= 0).sum(dim=-1).int() # (total_q,) # ---------------- Layout-specific output reshape -------------------- if is_thd: @@ -604,6 +632,7 @@ def indexer_topk( cu_seqlens_kv: Optional[Tensor] = None, max_seqlen_q: Optional[int] = None, max_seqlen_kv: Optional[int] = None, + visible_k_lengths: Optional[Tensor] = None, ) -> Tuple[Tensor, Tensor]: """Score + top-K selection for inference (no KL loss, no backward). @@ -625,6 +654,10 @@ def indexer_topk( cu_seqlens_kv: THD only — ``(B+1,)`` int32 CUDA cumulative KV lens. max_seqlen_q: THD only — per-batch max Q length. max_seqlen_kv: THD only — per-batch max KV length. + visible_k_lengths: THD only — optional ``(total_q,)`` int32 CUDA + tensor with the per-row visible K length for the radix top-K wrapper. + Callers pass this from a layout kernel to avoid rebuilding the + same metadata with generic PyTorch indexing during CUDA graph replay. Returns: SBHD: ``(topk_indices (b, sq, topk), topk_length (b, sq))`` int32 @@ -638,6 +671,8 @@ def indexer_topk( "indexer_topk THD mode requires cu_seqlens_q, cu_seqlens_kv, " "max_seqlen_q, and max_seqlen_kv to all be supplied." ) + if not is_thd and visible_k_lengths is not None: + raise ValueError("visible_k_lengths is only supported in THD mode.") # ``indexer_softmax_scale`` is applied via the # ``relu(c·x) = c·relu(x)`` trick (the cudnn kernel does the relu), @@ -665,6 +700,7 @@ def indexer_topk( cu_seqlens_kv=cu_seqlens_kv, max_seqlen_q=int(max_seqlen_q) if max_seqlen_q is not None else None, max_seqlen_kv=int(max_seqlen_kv) if max_seqlen_kv is not None else None, + visible_k_lengths=visible_k_lengths, ) return topk_indices, topk_length @@ -1459,6 +1495,193 @@ def backward(ctx, grad_output, grad_loss): ) +class FusedIndexerSparseAttnFromTopkFunc(torch.autograd.Function): + """Sparse attention with caller-supplied indexer top-k. + + The caller owns top-k selection. Sparse attention and indexer-loss + backward still use FlashMLA / cuDNN DSA wrappers. + """ + + @staticmethod + def forward( + ctx, + query: Tensor, + kv_full: Tensor, + attn_sink: Tensor, + topk_idxs: Tensor, + q_indexer: Tensor, + k_indexer: Tensor, + weights: Tensor, + indexer_topk_idxs: Tensor, + compressed_kv: Tensor, + softmax_scale: float, + indexer_softmax_scale: float, + loss_coeff: float, + calculate_per_token_loss: bool, + global_query_rows: int, + q_padding_mask: Optional[Tensor] = None, + ) -> Tuple[Tensor, Tensor]: + """Run fused sparse attention using caller-supplied top-k indices.""" + _ensure_dsa_namespace() + + total_q, np_ = query.shape[:2] + idx_nh, idx_hd = q_indexer.shape[1], q_indexer.shape[2] + total_comp = k_indexer.shape[0] + indexer_topk = indexer_topk_idxs.shape[-1] + + indexer_topk_idxs_for_loss = indexer_topk_idxs + if q_padding_mask is not None: + if q_padding_mask.shape != (total_q,): + raise RuntimeError( + "q_padding_mask must have shape " + f"({total_q},), got {tuple(q_padding_mask.shape)}." + ) + if q_padding_mask.dtype != torch.bool: + raise RuntimeError(f"q_padding_mask must be bool, got {q_padding_mask.dtype}.") + if q_padding_mask.device != query.device: + raise RuntimeError( + f"q_padding_mask must be on {query.device}, got {q_padding_mask.device}." + ) + indexer_topk_idxs_for_loss = indexer_topk_idxs.masked_fill( + q_padding_mask.unsqueeze(-1), -1 + ) + + out_flat, lse, lse_indexer = _dsa_fwd_flash_mla( + query, + kv_full, + topk_idxs, + softmax_scale, + attn_sink=attn_sink, + topk_length=None, + indexer_topk=indexer_topk, + ) + if lse_indexer is None: + raise RuntimeError("Indexer sparse attention from top-k requires lse_indexer.") + + if indexer_softmax_scale != 1.0: + weights_scaled = (weights.float() * indexer_softmax_scale).to(weights.dtype) + else: + weights_scaled = weights + predict = _compute_indexer_predict( + q_indexer, + k_indexer, + weights_scaled, + indexer_topk_idxs_for_loss, + qhead_per_kv_head=idx_nh, + topk_indices_global=True, + ) + target = _compute_attn_target( + query.detach(), + compressed_kv.detach(), + lse_indexer.detach(), + indexer_topk_idxs_for_loss, + softmax_scale, + qhead_per_kv_head=np_, + topk_indices_global=True, + ) + + raw_local_loss = _kl_loss_from_target_predict( + target, + predict, + indexer_topk_idxs_for_loss, + loss_coeff, + calculate_per_token_loss=True, + ) + if calculate_per_token_loss: + indexer_loss = raw_local_loss + bwd_loss_coeff = loss_coeff * total_q + else: + if global_query_rows <= 0: + raise RuntimeError(f"global_query_rows must be positive, got {global_query_rows}.") + indexer_loss = raw_local_loss / float(global_query_rows) + bwd_loss_coeff = loss_coeff * float(total_q) / float(global_query_rows) + + if loss_coeff > 0: + ig = _DSA.indexer_backward_wrapper( + q_indexer.view(1, total_q, idx_nh, idx_hd), + weights.view(1, total_q, idx_nh), + k_indexer.view(1, total_comp, idx_hd), + target.view(1, total_q, indexer_topk), + predict.view(1, total_q, indexer_topk), + indexer_topk_idxs_for_loss.view(1, total_q, indexer_topk), + sm_scale=indexer_softmax_scale, + loss_coeff=bwd_loss_coeff, + grad_loss=torch.ones((), device=query.device, dtype=torch.float32), + block_I=128, + ) + saved_grad_q_indexer = ig["d_index_q"].view(total_q, idx_nh, idx_hd) + saved_grad_k_indexer = ig["d_index_k"].view(total_comp, idx_hd) + saved_grad_weights = ig["d_weights"].view(total_q, idx_nh) + if q_padding_mask is not None: + saved_grad_q_indexer[q_padding_mask] = 0 + saved_grad_weights[q_padding_mask] = 0 + else: + saved_grad_q_indexer = torch.zeros_like(q_indexer) + saved_grad_k_indexer = torch.zeros_like(k_indexer) + saved_grad_weights = torch.zeros_like(weights) + + ctx.save_for_backward( + query, + kv_full, + attn_sink, + topk_idxs, + out_flat, + lse, + saved_grad_q_indexer, + saved_grad_k_indexer, + saved_grad_weights, + ) + ctx.softmax_scale = softmax_scale + + return out_flat.reshape(total_q, np_ * out_flat.shape[-1]), indexer_loss + + @staticmethod + def backward(ctx, grad_output, grad_loss): + """Run sparse-attention and indexer-loss backward kernels.""" + _ensure_dsa_namespace() + ( + query, + kv_full, + attn_sink, + topk_idxs, + out_flat, + lse, + saved_grad_q_indexer, + saved_grad_k_indexer, + saved_grad_weights, + ) = ctx.saved_tensors + + dO_flat = grad_output.reshape(query.shape[0], query.shape[1], out_flat.shape[-1]) + attn_bwd = _DSA.sparse_attention_backward_wrapper( + query, + kv_full, + out_flat, + dO_flat, + lse, + attn_sink, + topk_idxs, + softmax_scale=ctx.softmax_scale, + topk_length=None, + ) + return ( + attn_bwd["dq"], + attn_bwd["dkv"], + attn_bwd["d_sink"], + None, + saved_grad_q_indexer * grad_loss, + saved_grad_k_indexer * grad_loss, + saved_grad_weights * grad_loss, + None, + None, + None, + None, + None, + None, + None, + None, + ) + + def fused_indexer_sparse_attn( query: Tensor, kv_full: Tensor, diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 1a44adcb0d7..a8f88aafda8 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -343,6 +343,11 @@ class TransformerConfig(ModelParallelConfig): """Whether to use dense mode for compressed sparse attention. If True, the CSA indexer will be disabled.""" + csa_cp_partition_mode: str = "contiguous" + """CSA CP token partition mode. ``contiguous`` keeps the usual one-block-per-rank + split. ``two_chunk`` gives each rank two separated token ranges, which helps balance + attention work for a single long packed sequence.""" + apply_dsa_kernel_fusion: bool = False """If True, use fused DSA sparse-attention kernels (FlashMLA forward + cuDNN DSA backward, indexer scoring, and top-K selection). Requires ``flash_mla`` and ``nvidia-cudnn-frontend`` diff --git a/megatron/training/training.py b/megatron/training/training.py index ae6216c260b..ee450d919ad 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -2973,6 +2973,7 @@ def training_log( total_loss_dict=total_loss_dict, num_layers=args.num_layers + (args.mtp_num_layers or 0), csa_compress_ratios=args.csa_compress_ratios, + preserve_groups=args.cuda_graph_impl != "none", ) # Dump memory snapshot and print metrics to stdout. diff --git a/tests/unit_tests/test_sequence_packing.py b/tests/unit_tests/test_sequence_packing.py index f1fbcb53dea..643e1bf02f1 100644 --- a/tests/unit_tests/test_sequence_packing.py +++ b/tests/unit_tests/test_sequence_packing.py @@ -169,6 +169,49 @@ def __next__(self): return batch +class _MockCPGroup: + def __init__(self, size, rank): + self._size = size + self._rank = rank + + def size(self): + return self._size + + def rank(self): + return self._rank + + +def test_dsv4_thd_cp_slice_uses_static_partition_total(): + from megatron.core.datasets.data_schedule_utils import get_cp_slice_for_thd + from megatron.core.transformer.experimental_attention_variant.csa_cp_utils import ( + DSV4_CP_PARTITION_CONTIGUOUS, + ) + + batch = { + "tokens": torch.arange(8, dtype=torch.int64), + "position_ids": torch.arange(8, dtype=torch.int64), + "labels": torch.arange(100, 108, dtype=torch.int64), + "loss_mask": torch.ones(8, dtype=torch.float32), + "cu_seqlens": torch.tensor([0, 8], dtype=torch.int32), + "cu_seqlens_padded": torch.tensor([0, 8], dtype=torch.int32), + "max_seqlen": torch.tensor([8], dtype=torch.int32), + } + + get_cp_slice_for_thd( + batch, + _MockCPGroup(size=4, rank=2), + csa_cp_partition_mode=DSV4_CP_PARTITION_CONTIGUOUS, + partition_total_tokens=16, + ) + + assert torch.equal(batch["tokens"], torch.zeros(4, dtype=torch.int64)) + assert torch.equal(batch["position_ids"], torch.zeros(4, dtype=torch.int64)) + assert torch.equal(batch["labels"], torch.zeros(4, dtype=torch.int64)) + assert torch.equal(batch["loss_mask"], torch.zeros(4, dtype=torch.float32)) + assert torch.equal(batch["cu_seqlens"], torch.tensor([0, 8], dtype=torch.int32)) + assert torch.equal(batch["cu_seqlens_padded"], torch.tensor([0, 8], dtype=torch.int32)) + + def _gather_tensor_from_tp_group(tensor): """Gather tensors from all TP ranks for comparison.""" assert tensor is not None, "Tensor should not be None" diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_layout_kernels.py b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_layout_kernels.py new file mode 100644 index 00000000000..6518bf8c18f --- /dev/null +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_layout_kernels.py @@ -0,0 +1,1214 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import os +from typing import List, Tuple + +import pytest +import torch + +from megatron.core.transformer.experimental_attention_variant import csa_cp_layout_kernels + +# This file guards only DSv4 CP layout/metadata kernels. Layer-level CUDA graph +# tests guard graph capture/replay behavior. + +_E2E_RAGGED_SEG_LENS = (1, 127, 1000, 23, 129, 900, 55, 257, 800, 95, 509, 148) +_E2E_RAGGED_PADDED_SEG_LENS = (1, 127, 1000, 23, 129, 900, 55, 257, 800, 95, 509, 200) +_E2E_CP_SIZE = 4 +_E2E_CP_RANK = 1 + + +def _require_cute_cuda(): + if not torch.cuda.is_available(): + pytest.skip("DSv4 CP CuTe kernels require CUDA.") + if not csa_cp_layout_kernels._CUTE_AVAILABLE: + pytest.skip("DSv4 CP CuTe kernels are not available in this environment.") + + +def _make_e2e_like_cu_seqlens(device: str = "cuda", *, padded: bool = True) -> torch.Tensor: + """Return the ragged THD prefix pattern used by DSv4 CP e2e tests. + + The padded variant sums to 4096, so CP4 owns 1024 local rows. It contains + short, long, boundary-crossing, and padded-tail sequences without allocating + full e2e hidden sizes in these focused kernel tests. + """ + lengths = _E2E_RAGGED_PADDED_SEG_LENS if padded else _E2E_RAGGED_SEG_LENS + return torch.tensor([0] + list(torch.tensor(lengths).cumsum(0).tolist()), dtype=torch.int32, device=device) + + +def _compressed_cu_seqlens(cu_seqlens: torch.Tensor, ratio: int) -> torch.Tensor: + """Return floor-compressed sequence prefixes for the supplied THD layout.""" + return torch.cat( + ( + torch.zeros((1,), dtype=torch.int32, device=cu_seqlens.device), + torch.cumsum( + torch.div(cu_seqlens[1:] - cu_seqlens[:-1], int(ratio), rounding_mode="floor"), + dim=0, + dtype=torch.int32, + ), + ) + ) + + +def _e2e_like_contiguous_range(cp_rank: int = _E2E_CP_RANK) -> Tuple[int, int]: + """Return one CP4 contiguous local range for the e2e-like padded THD layout.""" + total = sum(_E2E_RAGGED_PADDED_SEG_LENS) + local = total // _E2E_CP_SIZE + start = int(cp_rank) * local + return start, local + + +def _e2e_like_two_chunk_ranges(cp_rank: int = _E2E_CP_RANK) -> Tuple[Tuple[int, int], ...]: + """Return one CP4 two-chunk local row order for the e2e-like THD layout.""" + total = sum(_E2E_RAGGED_PADDED_SEG_LENS) + local = total // _E2E_CP_SIZE + chunk_len = local // 2 + total_chunks = 2 * _E2E_CP_SIZE + chunk_ids = (int(cp_rank), total_chunks - 1 - int(cp_rank)) + return tuple((chunk_id * chunk_len, (chunk_id + 1) * chunk_len) for chunk_id in chunk_ids) + + +def _cosine_sim(actual: torch.Tensor, expected: torch.Tensor) -> float: + a = actual.detach().float().reshape(-1) + b = expected.detach().float().reshape(-1) + denom = a.norm() * b.norm() + if denom == 0: + return 1.0 if torch.equal(a, b) else 0.0 + return float(torch.dot(a, b) / denom) + + +def _tensor_sim(actual: torch.Tensor, expected: torch.Tensor) -> float: + a = actual.detach().float() + b = expected.detach().float() + denom = b.norm().clamp_min(1e-12) + return float(1.0 - (a - b).norm() / denom) + + +def _assert_rope_close(actual: torch.Tensor, expected: torch.Tensor, name: str) -> None: + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) + cos = _cosine_sim(actual, expected) + sim = _tensor_sim(actual, expected) + assert cos > 0.9999, f"{name} cosine similarity too low: {cos}" + assert sim > 0.9999, f"{name} tensor similarity too low: {sim}" + + +def _num_bytes(*tensors: torch.Tensor) -> int: + return sum(t.numel() * t.element_size() for t in tensors) + + +def _time_cuda(fn, *, warmup: int = 3, iters: int = 5) -> float: + iters = int(os.environ.get("DSV4_CP_KERNEL_BENCH_ITERS", iters)) + for _ in range(warmup): + fn() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return float(start.elapsed_time(end) / iters) + + +def _print_perf(name: str, cute_ms: float, torch_ms: float, effective_bytes: int) -> None: + cute_gbps = effective_bytes / max(cute_ms, 1e-9) / 1.0e6 + torch_gbps = effective_bytes / max(torch_ms, 1e-9) / 1.0e6 + print( + f"{name}: cute={cute_ms:.4f} ms ({cute_gbps:.2f} GB/s effective), " + f"torch={torch_ms:.4f} ms ({torch_gbps:.2f} GB/s effective), " + f"speedup={torch_ms / max(cute_ms, 1e-9):.2f}x" + ) + + +def _rope_reference( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + positions: torch.Tensor, + nope_dim: int, + pos_dim: int, + inverse: bool = False, +) -> torch.Tensor: + x_nope, x_pos = torch.split(x, [nope_dim, pos_dim], dim=-1) + cos_flat = cos.reshape(cos.shape[0], -1).index_select(0, positions.to(torch.long)) + sin_flat = sin.reshape(sin.shape[0], -1).index_select(0, positions.to(torch.long)) + view_shape = (x.shape[0],) + (1,) * (x.ndim - 2) + (pos_dim,) + cos_pos = cos_flat[:, :pos_dim].view(view_shape) + sin_pos = sin_flat[:, :pos_dim].view(view_shape) + if inverse: + sin_pos = -sin_pos + + half = pos_dim // 2 + x1 = x_pos[..., 0::2] + x2 = x_pos[..., 1::2] + left = x1 * cos_pos[..., :half] - x2 * sin_pos[..., :half] + right = x2 * cos_pos[..., half:] + x1 * sin_pos[..., half:] + x_rot = torch.stack((left, right), dim=-1).flatten(-2) + return torch.cat((x_nope, x_rot), dim=-1) + + +def _seq_positions_from_rows( + cu_seqlens: torch.Tensor, global_rows: torch.Tensor, clamp_to_valid_token: bool = False +) -> torch.Tensor: + cu = cu_seqlens.cpu().to(torch.long) + rows = global_rows.cpu().to(torch.long) + if clamp_to_valid_token: + rows = rows.clamp(min=0, max=int(cu[-1]) - 1) + seq_ids = torch.searchsorted(cu, rows, right=True) - 1 + seq_ids = seq_ids.clamp(min=0, max=cu.numel() - 2) + return (rows - cu[seq_ids]).to(global_rows.device) + + +def _compressed_groups( + cu_seqlens: torch.Tensor, global_start: int, l_local: int, ratio: int, d_comp: int +) -> List[Tuple[int, int]]: + cu = [int(x) for x in cu_seqlens.cpu().tolist()] + groups: List[Tuple[int, int]] = [] + global_end = int(global_start) + int(l_local) + for seq, (seq_start, seq_end) in enumerate(zip(cu[:-1], cu[1:])): + local_start = max(seq_start, int(global_start)) + local_end = min(seq_end, global_end) + if local_start >= local_end: + continue + n_full_groups = (seq_end - seq_start) // int(ratio) + first_numer = max(0, int(global_start) - int(d_comp) - seq_start) + first_group = (first_numer + int(ratio) - 1) // int(ratio) if first_numer > 0 else 0 + stop_group = min((local_end - seq_start) // int(ratio), n_full_groups) + for comp_id in range(first_group, max(first_group, stop_group)): + groups.append((seq, comp_id)) + return groups + + +def _native_compressor_input_compact( + hidden_local: torch.Tensor, + boundary_hidden: torch.Tensor, + cu_seqlens: torch.Tensor, + global_start: int, + l_local: int, + ratio: int, + d_comp: int, + d_window: int, + c_cap: int, +): + groups = _compressed_groups(cu_seqlens, global_start, l_local, ratio, d_comp) + cu = [int(x) for x in cu_seqlens.cpu().tolist()] + pieces = [] + for seq, comp_id in groups: + seq_start = cu[seq] + for token in range(ratio): + src_global = seq_start + comp_id * ratio + token + if src_global < global_start: + pieces.append( + boundary_hidden[ + src_global + - (global_start - d_window) : src_global + - (global_start - d_window) + + 1 + ] + ) + else: + pieces.append( + hidden_local[src_global - global_start : src_global - global_start + 1] + ) + compact_len = int(c_cap) * int(ratio) + if pieces: + hidden_compact = torch.cat(pieces, dim=0) + else: + hidden_compact = hidden_local.new_empty((0,) + tuple(hidden_local.shape[1:])) + if hidden_compact.shape[0] < compact_len: + hidden_compact = torch.cat( + ( + hidden_compact, + hidden_local.new_zeros( + (compact_len - hidden_compact.shape[0],) + tuple(hidden_local.shape[1:]) + ), + ), + dim=0, + ) + + cu_compact = torch.zeros_like(cu_seqlens) + running = 0 + for seq in range(len(cu) - 1): + running += sum(1 for group_seq, _ in groups if group_seq == seq) * int(ratio) + cu_compact[seq + 1] = running + + seq_ids = torch.full((c_cap,), -1, dtype=torch.int32, device=hidden_local.device) + comp_ids = torch.full((c_cap,), -1, dtype=torch.int32, device=hidden_local.device) + valid = torch.zeros((c_cap,), dtype=torch.bool, device=hidden_local.device) + for slot, (seq, comp_id) in enumerate(groups[:c_cap]): + seq_start = cu[seq] + group_end = seq_start + (comp_id + 1) * int(ratio) + seq_ids[slot] = seq + comp_ids[slot] = comp_id + valid[slot] = group_end - 1 >= global_start and group_end - 1 < global_start + l_local + return hidden_compact, cu_compact, seq_ids, comp_ids, valid + + +def _native_build_compressed_row_metadata( + cu_seqlens: torch.Tensor, + cp_size: int, + chunk_len: int, + ratio: int, + d_comp: int, + c_cap_per_chunk: int, + c_cap_per_rank: int, + use_two_chunk: bool, +): + total_rows = int(cp_size) * int(c_cap_per_rank) + seq_ids = torch.full((total_rows,), -1, dtype=torch.int32, device=cu_seqlens.device) + comp_ids = torch.full((total_rows,), -1, dtype=torch.int32, device=cu_seqlens.device) + valid = torch.zeros((total_rows,), dtype=torch.bool, device=cu_seqlens.device) + for row in range(total_rows): + rank = row // c_cap_per_rank + rank_slot = row - rank * c_cap_per_rank + local_chunk = 0 + slot = rank_slot + chunk_id = rank + if use_two_chunk: + local_chunk = rank_slot // c_cap_per_chunk + if local_chunk >= 2: + continue + slot = rank_slot - local_chunk * c_cap_per_chunk + chunk_id = rank if local_chunk == 0 else cp_size * 2 - 1 - rank + rank_start = chunk_id * chunk_len + groups = _compressed_groups(cu_seqlens, rank_start, chunk_len, ratio, d_comp) + if slot < len(groups): + seq, comp = groups[slot] + seq_start = int(cu_seqlens[seq].item()) + group_end = seq_start + (comp + 1) * ratio + seq_ids[row] = seq + comp_ids[row] = comp + valid[row] = group_end - 1 >= rank_start and group_end - 1 < rank_start + chunk_len + return seq_ids, comp_ids, valid + + +def _native_repack_compressed_kv_to_seq_major( + compressed_rank_major: torch.Tensor, + seq_ids: torch.Tensor, + comp_ids: torch.Tensor, + valid: torch.Tensor, + cu_seqlens_compressed: torch.Tensor, + seq_major_rows: int, +): + out = compressed_rank_major.new_zeros( + (seq_major_rows,) + tuple(compressed_rank_major.shape[1:]) + ) + rank_by_seq = torch.full( + (seq_major_rows,), -1, dtype=torch.int32, device=compressed_rank_major.device + ) + for rank_row in range(compressed_rank_major.shape[0]): + if bool(valid[rank_row]): + seq = int(seq_ids[rank_row]) + comp = int(comp_ids[rank_row]) + seq_major = int(cu_seqlens_compressed[seq]) + comp + if 0 <= seq_major < seq_major_rows: + out[seq_major] = compressed_rank_major[rank_row] + rank_by_seq[seq_major] = rank_row + return out, rank_by_seq + + +def _native_indexer_topk_metadata( + k_seq_major: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_compressed: torch.Tensor, + global_start: int, + l_local: int, + ratio: int, +): + n_seq = cu_seqlens_q.shape[0] - 1 + k_topk = torch.zeros_like(k_seq_major) + cu_q = torch.zeros((n_seq + 2,), dtype=torch.int32, device=k_seq_major.device) + cu_k = torch.zeros((n_seq + 2,), dtype=torch.int32, device=k_seq_major.device) + seq_lens = torch.empty((l_local,), dtype=torch.int32, device=k_seq_major.device) + q_prefix = 0 + k_prefix = 0 + for seq in range(n_seq): + seq_start = int(cu_seqlens_q[seq]) + seq_end = int(cu_seqlens_q[seq + 1]) + local_start = max(seq_start, int(global_start)) + local_end = min(seq_end, int(global_start) + int(l_local)) + q_len = 0 + k_len = 0 + if local_start < local_end: + q_len = local_end - local_start + seq_comp_start = int(cu_seqlens_compressed[seq]) + seq_comp_end = int(cu_seqlens_compressed[seq + 1]) + k_len = min((local_end - seq_start) // int(ratio), seq_comp_end - seq_comp_start) + if k_len: + k_topk[k_prefix : k_prefix + k_len] = k_seq_major[ + seq_comp_start : seq_comp_start + k_len + ] + if q_len: + seq_lens[q_prefix : q_prefix + q_len] = k_len + q_prefix += q_len + k_prefix += k_len + cu_q[seq + 1] = q_prefix + cu_k[seq + 1] = k_prefix + actual_total = int(cu_seqlens_q[-1]) + padding_start = max(actual_total, int(global_start)) + padding_q = max(0, int(global_start) + int(l_local) - padding_start) + if padding_q: + seq_lens[q_prefix : q_prefix + padding_q] = 0 + cu_q[n_seq + 1] = q_prefix + padding_q + cu_k[n_seq + 1] = k_prefix + return k_topk, cu_q, cu_k, seq_lens + + +def _native_thd_full_kv_pack( + kv_local: torch.Tensor, + boundary_kv: torch.Tensor, + compressed_rank_major: torch.Tensor, + seq_ids: torch.Tensor, + comp_ids: torch.Tensor, + valid: torch.Tensor, + cu_seqlens: torch.Tensor, + global_start: int, + l_local: int, + d_window: int, + ratio: int, + capacity: int, +): + rows = [] + n_seq = cu_seqlens.shape[0] - 1 + for seq in range(n_seq): + seq_start = int(cu_seqlens[seq]) + seq_end = int(cu_seqlens[seq + 1]) + local_start = max(seq_start, int(global_start)) + local_end = min(seq_end, int(global_start) + int(l_local)) + if local_start >= local_end: + continue + window_start = max(seq_start, local_start - int(d_window)) + for src_global in range(window_start, local_end): + if src_global < global_start: + rows.append( + boundary_kv[ + src_global + - (global_start - d_window) : src_global + - (global_start - d_window) + + 1 + ] + ) + else: + rows.append(kv_local[src_global - global_start : src_global - global_start + 1]) + if ratio > 1: + for comp_id in range((seq_end - seq_start) // int(ratio)): + for rank_row in range(compressed_rank_major.shape[0]): + if ( + bool(valid[rank_row]) + and int(seq_ids[rank_row]) == seq + and int(comp_ids[rank_row]) == comp_id + ): + rows.append(compressed_rank_major[rank_row : rank_row + 1]) + if rows: + out = torch.cat(rows, dim=0) + else: + out = kv_local.new_empty((0,) + tuple(kv_local.shape[1:])) + if out.shape[0] < capacity: + out = torch.cat( + (out, kv_local.new_zeros((capacity - out.shape[0],) + tuple(kv_local.shape[1:]))), dim=0 + ) + return out + + +def _native_attention_indices_one_seq( + l_local: int, window_size: int, ratio: int, compressed_width: int, device: torch.device +): + rows = torch.arange(l_local, dtype=torch.int32, device=device).unsqueeze(1) + window_cols = torch.arange(window_size, dtype=torch.int32, device=device).unsqueeze(0) + window_start = (rows - int(window_size) + 1).clamp_min(0) + window_count = rows - window_start + 1 + window_values = window_start + window_cols + window_values = torch.where( + window_cols < window_count, window_values, torch.full_like(window_values, -1) + ) + if ratio > 1 and compressed_width > 0: + comp_cols = torch.arange(compressed_width, dtype=torch.int32, device=device).unsqueeze(0) + visible = ((rows + 1) // int(ratio)).clamp(max=int(compressed_width)) + comp_values = int(l_local) + comp_cols + comp_values = torch.where( + comp_cols < visible, comp_values, torch.full_like(comp_values, -1) + ) + topk = torch.cat((window_values, comp_values.expand(l_local, -1)), dim=1) + topk_length = window_count.squeeze(1) + visible.squeeze(1) + else: + topk = window_values + topk_length = window_count.squeeze(1) + return topk, topk_length + + +def _native_attention_indices( + cu_seqlens: torch.Tensor, + global_start: int, + l_local: int, + d_window: int, + window_size: int, + ratio: int, + compressed_width: int, + indexer_topk: torch.Tensor = None, +): + """Reference contiguous final-index lowering for ragged THD CP rows.""" + device = cu_seqlens.device + total_width = int(window_size) + int(compressed_width) + topk = torch.full((int(l_local), total_width), -1, dtype=torch.int32, device=device) + lengths = torch.zeros((int(l_local),), dtype=torch.int32, device=device) + cu = [int(x) for x in cu_seqlens.cpu().tolist()] + seq_infos = [] + running = 0 + global_end = int(global_start) + int(l_local) + for seq, (seq_start, seq_end) in enumerate(zip(cu[:-1], cu[1:])): + local_start = max(seq_start, int(global_start)) + local_end = min(seq_end, global_end) + if local_start >= local_end: + continue + window_start = max(seq_start, local_start - int(d_window)) + window_len = local_end - window_start + comp_len = (seq_end - seq_start) // int(ratio) if int(ratio) > 1 else 0 + seq_infos.append((seq, seq_start, seq_end, window_start, window_len, comp_len, running)) + running += window_len + comp_len + + for row in range(int(l_local)): + global_q = int(global_start) + row + info = next((item for item in seq_infos if item[1] <= global_q < item[2]), None) + if info is None: + if total_width > 0: + topk[row, 0] = 0 + lengths[row] = 1 + continue + _seq, seq_start, _seq_end, seq_window_start, seq_window_len, seq_comp_len, seq_offset = info + write_col = 0 + window_start_for_q = max(global_q - int(window_size) + 1, seq_start, seq_window_start) + window_count = max(0, global_q - window_start_for_q + 1) + for w in range(int(window_size)): + pos = window_start_for_q + w + if w < window_count and seq_window_start <= pos < seq_window_start + seq_window_len: + topk[row, write_col] = seq_offset + pos - seq_window_start + write_col += 1 + if int(ratio) > 1 and int(compressed_width) > 0: + pos_in_seq = global_q - seq_start + for j in range(int(compressed_width)): + if indexer_topk is not None: + comp_id = int(indexer_topk[row, j]) + else: + n_visible = min((pos_in_seq + 1) // int(ratio), int(compressed_width)) + comp_id = j if j < n_visible else -1 + if 0 <= comp_id < seq_comp_len: + topk[row, write_col] = seq_offset + seq_window_len + comp_id + write_col += 1 + lengths[row] = write_col + return topk, lengths + + +def _native_indexer_loss_indices_one_seq( + l_local: int, + window_size: int, + ratio: int, + logical_ids: torch.Tensor, + rank_by_seq_major: torch.Tensor, +): + device = logical_ids.device + compressed_width = logical_ids.shape[1] + attention_topk, _ = _native_attention_indices_one_seq(l_local, window_size, ratio, 0, device) + compressed = torch.full((l_local, compressed_width), -1, dtype=torch.int32, device=device) + rank_major = torch.full_like(compressed, -1) + valid = (logical_ids >= 0) & (logical_ids < rank_by_seq_major.shape[0]) + if valid.any(): + lowered = int(l_local) + logical_ids.clamp(min=0) + rank_ids = rank_by_seq_major.index_select(0, logical_ids.clamp(min=0).reshape(-1)).view_as( + logical_ids + ) + compressed = torch.where(valid, lowered, compressed) + rank_major = torch.where(valid, rank_ids, rank_major) + return torch.cat((compressed, attention_topk), dim=1), rank_major + + +def _native_indexer_loss_indices( + cu_seqlens: torch.Tensor, + cu_seqlens_compressed: torch.Tensor, + global_start: int, + l_local: int, + d_window: int, + window_size: int, + ratio: int, + logical_ids: torch.Tensor, + rank_by_seq_major: torch.Tensor, +): + """Reference contiguous compressed-first indexer-loss lowering for ragged THD rows.""" + device = logical_ids.device + compressed_width = logical_ids.shape[1] + total_width = compressed_width + int(window_size) + topk = torch.full((int(l_local), total_width), -1, dtype=torch.int32, device=device) + rank_major = torch.full((int(l_local), compressed_width), -1, dtype=torch.int32, device=device) + cu = [int(x) for x in cu_seqlens.cpu().tolist()] + cu_comp = [int(x) for x in cu_seqlens_compressed.cpu().tolist()] + seq_infos = [] + running = 0 + global_end = int(global_start) + int(l_local) + for seq, (seq_start, seq_end) in enumerate(zip(cu[:-1], cu[1:])): + local_start = max(seq_start, int(global_start)) + local_end = min(seq_end, global_end) + if local_start >= local_end: + continue + window_start = max(seq_start, local_start - int(d_window)) + window_len = local_end - window_start + comp_len = cu_comp[seq + 1] - cu_comp[seq] if int(ratio) > 1 else 0 + seq_infos.append((seq, seq_start, seq_end, window_start, window_len, comp_len, running)) + running += window_len + comp_len + + for row in range(int(l_local)): + global_q = int(global_start) + row + info = next((item for item in seq_infos if item[1] <= global_q < item[2]), None) + if info is None: + continue + seq_id, seq_start, _seq_end, window_start, window_len, comp_len, offset = info + + for col in range(compressed_width): + comp_id = int(logical_ids[row, col]) + if 0 <= comp_id < comp_len: + seq_major = cu_comp[seq_id] + comp_id + rank_id = int(rank_by_seq_major[seq_major]) + if rank_id >= 0: + topk[row, col] = offset + window_len + comp_id + rank_major[row, col] = rank_id + + window_start_for_q = max(global_q - int(window_size) + 1, seq_start) + window_count = global_q - window_start_for_q + 1 + for window_col in range(int(window_size)): + pos = window_start_for_q + window_col + if ( + window_col < window_count + and pos >= window_start + and pos < window_start + window_len + ): + topk[row, compressed_width + window_col] = offset + pos - window_start + + return topk, rank_major + + +def test_thd_local_rope_matches_native_forward_backward_and_reports_bandwidth(): + _require_cute_cuda() + torch.manual_seed(11) + cu = _make_e2e_like_cu_seqlens() + chunk_ranges = _e2e_like_two_chunk_ranges() + rows = sum(end - start for start, end in chunk_ranges) + nope_dim = 4 + pos_dim = 4 + x = torch.randn(rows, 2, nope_dim + pos_dim, device="cuda", requires_grad=True) + cos = torch.randn(max(_E2E_RAGGED_PADDED_SEG_LENS), 1, 1, pos_dim, device="cuda") + sin = torch.randn(max(_E2E_RAGGED_PADDED_SEG_LENS), 1, 1, pos_dim, device="cuda") + global_rows = torch.cat( + [torch.arange(start, end, device="cuda") for start, end in chunk_ranges] + ) + positions = _seq_positions_from_rows(cu, global_rows) + ref_x = x.detach().clone().requires_grad_(True) + ref = _rope_reference(ref_x, cos, sin, positions, nope_dim, pos_dim) + grad = torch.randn_like(ref) + ref.backward(grad) + + fused_x = x.detach().clone().requires_grad_(True) + fused = csa_cp_layout_kernels.ThdLocalRope.apply( + fused_x, + cos, + sin, + cu, + chunk_ranges[0][0], + chunk_ranges[0][1] - chunk_ranges[0][0], + nope_dim, + pos_dim, + chunk_ranges[1][0], + False, + False, + ) + fused.backward(grad) + _assert_rope_close(fused, ref, "thd_local_rope.forward") + _assert_rope_close(fused_x.grad, ref_x.grad, "thd_local_rope.backward") + + bench_x = torch.randn(512, 8, nope_dim + pos_dim, device="cuda") + bench_grad = torch.randn_like(bench_x) + bench_cu = torch.tensor([0, 2048], dtype=torch.int32, device="cuda") + bench_cos = torch.randn(2048, 1, 1, pos_dim, device="cuda") + bench_sin = torch.randn(2048, 1, 1, pos_dim, device="cuda") + bench_pos = torch.arange(512, device="cuda") + + def run_cute(): + bx = bench_x.detach().clone().requires_grad_(True) + y = csa_cp_layout_kernels.ThdLocalRope.apply( + bx, bench_cos, bench_sin, bench_cu, 0, 512, nope_dim, pos_dim, 0, False, False + ) + y.backward(bench_grad) + + def run_torch(): + bx = bench_x.detach().clone().requires_grad_(True) + y = _rope_reference(bx, bench_cos, bench_sin, bench_pos, nope_dim, pos_dim) + y.backward(bench_grad) + + effective_bytes = 4 * _num_bytes(bench_x) + _print_perf( + "thd_local_rope_fwd_bwd", _time_cuda(run_cute), _time_cuda(run_torch), effective_bytes + ) + + +def test_thd_compressed_rope_matches_native_forward_backward_and_reports_bandwidth(): + _require_cute_cuda() + torch.manual_seed(22) + comp_ids = torch.tensor([0, 2, -1, 17, 31, 127], dtype=torch.int32, device="cuda") + ratio = 4 + nope_dim = 3 + pos_dim = 4 + x = torch.randn(comp_ids.numel(), 1, 2, nope_dim + pos_dim, device="cuda", requires_grad=True) + cos = torch.randn(max(_E2E_RAGGED_PADDED_SEG_LENS), 1, 1, pos_dim, device="cuda") + sin = torch.randn(max(_E2E_RAGGED_PADDED_SEG_LENS), 1, 1, pos_dim, device="cuda") + positions = comp_ids.clamp(min=0).to(torch.long) * ratio + ref_x = x.detach().clone().requires_grad_(True) + ref = _rope_reference(ref_x, cos, sin, positions, nope_dim, pos_dim) + grad = torch.randn_like(ref) + ref.backward(grad) + + fused_x = x.detach().clone().requires_grad_(True) + fused = csa_cp_layout_kernels.ThdCompressedRope.apply( + fused_x, cos, sin, comp_ids, ratio, nope_dim, pos_dim, False + ) + fused.backward(grad) + _assert_rope_close(fused, ref, "thd_compressed_rope.forward") + _assert_rope_close(fused_x.grad, ref_x.grad, "thd_compressed_rope.backward") + + bench_x = torch.randn(1024, 4, nope_dim + pos_dim, device="cuda") + bench_grad = torch.randn_like(bench_x) + bench_ids = torch.arange(1024, dtype=torch.int32, device="cuda") % 8 + bench_pos = bench_ids.to(torch.long) * ratio + + def run_cute(): + bx = bench_x.detach().clone().requires_grad_(True) + y = csa_cp_layout_kernels.ThdCompressedRope.apply( + bx, cos, sin, bench_ids, ratio, nope_dim, pos_dim, False + ) + y.backward(bench_grad) + + def run_torch(): + bx = bench_x.detach().clone().requires_grad_(True) + y = _rope_reference(bx, cos, sin, bench_pos, nope_dim, pos_dim) + y.backward(bench_grad) + + effective_bytes = 4 * _num_bytes(bench_x) + _print_perf( + "thd_compressed_rope_fwd_bwd", _time_cuda(run_cute), _time_cuda(run_torch), effective_bytes + ) + + +def test_compressor_input_compact_matches_native_forward_backward_and_reports_bandwidth(): + _require_cute_cuda() + cu = _make_e2e_like_cu_seqlens() + global_start, l_local = _e2e_like_contiguous_range() + ratio = 128 + d_comp = 128 + d_window = 128 + c_cap = (l_local + d_comp) // ratio + hidden = ( + torch.arange(global_start, global_start + l_local, dtype=torch.float32, device="cuda")[ + :, None + ] + .repeat(1, 3) + .to(torch.bfloat16) + .requires_grad_(True) + ) + boundary = ( + torch.arange(global_start - d_window, global_start, dtype=torch.float32, device="cuda")[ + :, None + ] + .repeat(1, 3) + .to(torch.bfloat16) + .requires_grad_(True) + ) + + ref = _native_compressor_input_compact( + hidden, boundary, cu, global_start, l_local, ratio, d_comp, d_window, c_cap + ) + grad = torch.randn_like(ref[0]) + ref[0].backward(grad) + ref_hidden_grad = hidden.grad.detach().clone() + ref_boundary_grad = boundary.grad.detach().clone() + hidden.grad.zero_() + boundary.grad.zero_() + + fused = csa_cp_layout_kernels.CompressorInputCompact.apply( + hidden, boundary, cu, global_start, l_local, ratio, d_comp, d_window, c_cap + ) + fused[0].backward(grad) + assert torch.equal(fused[0], ref[0]) + for actual, expected in zip(fused[1:], ref[1:]): + assert torch.equal(actual, expected) + assert torch.equal(hidden.grad, ref_hidden_grad) + assert torch.equal(boundary.grad, ref_boundary_grad) + + bench_cu = torch.tensor([0, 16384], dtype=torch.int32, device="cuda") + bench_global_start = 4096 + bench_l_local = 8192 + bench_ratio = 4 + bench_d_comp = 128 + bench_d_window = 128 + bench_c_cap = (bench_l_local + bench_d_comp) // bench_ratio + bench_hidden = torch.randn(bench_l_local, 64, dtype=torch.bfloat16, device="cuda") + bench_boundary = torch.randn(bench_d_window, 64, dtype=torch.bfloat16, device="cuda") + bench_grad = torch.randn(bench_c_cap * bench_ratio, 64, dtype=torch.bfloat16, device="cuda") + + def run_cute(): + h = bench_hidden.detach().clone().requires_grad_(True) + b = bench_boundary.detach().clone().requires_grad_(True) + y = csa_cp_layout_kernels.CompressorInputCompact.apply( + h, + b, + bench_cu, + bench_global_start, + bench_l_local, + bench_ratio, + bench_d_comp, + bench_d_window, + bench_c_cap, + )[0] + y.backward(bench_grad) + + def run_torch(): + h = bench_hidden.detach().clone().requires_grad_(True) + b = bench_boundary.detach().clone().requires_grad_(True) + y = _native_compressor_input_compact( + h, + b, + bench_cu, + bench_global_start, + bench_l_local, + bench_ratio, + bench_d_comp, + bench_d_window, + bench_c_cap, + )[0] + y.backward(bench_grad) + + effective_bytes = 2 * (_num_bytes(bench_hidden, bench_boundary, bench_grad)) + _print_perf( + "compressor_input_compact_fwd_bwd", + _time_cuda(run_cute), + _time_cuda(run_torch), + effective_bytes, + ) + + +def test_thd_full_kv_pack_matches_native_forward_backward_and_reports_bandwidth(): + _require_cute_cuda() + cu = _make_e2e_like_cu_seqlens() + global_start, l_local = _e2e_like_contiguous_range() + d_window = 128 + ratio = 128 + d_comp = 128 + c_cap = (l_local + d_comp) // ratio + kv_local = torch.randn(l_local, 4, dtype=torch.bfloat16, device="cuda").requires_grad_(True) + boundary = torch.randn(d_window, 4, dtype=torch.bfloat16, device="cuda").requires_grad_(True) + seq_ids, comp_ids, valid = _native_build_compressed_row_metadata( + cu, _E2E_CP_SIZE, l_local, ratio, d_comp, c_cap, c_cap, False + ) + compressed = torch.randn(seq_ids.shape[0], 4, dtype=torch.bfloat16, device="cuda").requires_grad_( + True + ) + dummy = torch.empty((1,), dtype=torch.int32, device="cuda") + capacity = l_local + d_window * (cu.shape[0] - 1) + compressed.shape[0] + + ref = _native_thd_full_kv_pack( + kv_local, + boundary, + compressed, + seq_ids, + comp_ids, + valid, + cu, + global_start, + l_local, + d_window, + ratio, + capacity, + ) + grad = torch.randn_like(ref) + ref.backward(grad) + ref_grads = ( + kv_local.grad.detach().clone(), + boundary.grad.detach().clone(), + compressed.grad.detach().clone(), + ) + kv_local.grad.zero_() + boundary.grad.zero_() + compressed.grad.zero_() + + fused = csa_cp_layout_kernels.ThdFullKvPack.apply( + kv_local, + boundary, + compressed, + seq_ids, + comp_ids, + valid, + cu, + dummy, + cu, + global_start, + l_local, + d_window, + ratio, + capacity, + 0, + 0, + 0, + 0, + 0, + ) + fused.backward(grad) + assert torch.equal(fused, ref) + assert torch.equal(kv_local.grad, ref_grads[0]) + assert torch.equal(boundary.grad, ref_grads[1]) + assert torch.equal(compressed.grad, ref_grads[2]) + + bench_cu = torch.tensor([0, 8192], dtype=torch.int32, device="cuda") + bench_global_start = 2048 + bench_l_local = 4096 + bench_d_window = 128 + bench_ratio = 128 + bench_compressed_rows = int(bench_cu[-1]) // bench_ratio + bench_capacity = bench_d_window + bench_l_local + bench_compressed_rows + bench_kv = torch.randn(bench_l_local, 64, dtype=torch.bfloat16, device="cuda") + bench_boundary = torch.randn(bench_d_window, 64, dtype=torch.bfloat16, device="cuda") + bench_compressed = torch.randn(bench_compressed_rows, 64, dtype=torch.bfloat16, device="cuda") + bench_seq_ids = torch.zeros(bench_compressed_rows, dtype=torch.int32, device="cuda") + bench_comp_ids = torch.arange(bench_compressed_rows, dtype=torch.int32, device="cuda") + bench_valid = torch.ones(bench_compressed_rows, dtype=torch.bool, device="cuda") + bench_dummy = torch.empty((1,), dtype=torch.int32, device="cuda") + bench_grad = torch.randn(bench_capacity, 64, dtype=torch.bfloat16, device="cuda") + + def run_cute(): + k = bench_kv.detach().clone().requires_grad_(True) + b = bench_boundary.detach().clone().requires_grad_(True) + c = bench_compressed.detach().clone().requires_grad_(True) + y = csa_cp_layout_kernels.ThdFullKvPack.apply( + k, + b, + c, + bench_seq_ids, + bench_comp_ids, + bench_valid, + bench_cu, + bench_dummy, + bench_cu, + bench_global_start, + bench_l_local, + bench_d_window, + bench_ratio, + bench_capacity, + 0, + 0, + 0, + 0, + 0, + ) + y.backward(bench_grad) + + def run_torch(): + k = bench_kv.detach().clone().requires_grad_(True) + b = bench_boundary.detach().clone().requires_grad_(True) + c = bench_compressed.detach().clone().requires_grad_(True) + y = _native_thd_full_kv_pack( + k, + b, + c, + bench_seq_ids, + bench_comp_ids, + bench_valid, + bench_cu, + bench_global_start, + bench_l_local, + bench_d_window, + bench_ratio, + bench_capacity, + ) + y.backward(bench_grad) + + effective_bytes = 2 * _num_bytes(bench_kv, bench_boundary, bench_compressed, bench_grad) + _print_perf( + "thd_full_kv_pack_fwd_bwd", _time_cuda(run_cute), _time_cuda(run_torch), effective_bytes + ) + + +def test_repack_compressed_kv_to_seq_major_matches_native_and_reports_bandwidth(): + _require_cute_cuda() + cu = _make_e2e_like_cu_seqlens() + ratio = 128 + d_comp = 128 + global_start, l_local = _e2e_like_contiguous_range() + c_cap = (l_local + d_comp) // ratio + seq_ids, comp_ids, valid = _native_build_compressed_row_metadata( + cu, _E2E_CP_SIZE, l_local, ratio, d_comp, c_cap, c_cap, False + ) + rank_major = torch.arange(seq_ids.numel() * 4, dtype=torch.bfloat16, device="cuda").reshape( + seq_ids.numel(), 4 + ) + cu_comp = _compressed_cu_seqlens(cu, ratio) + ref = _native_repack_compressed_kv_to_seq_major( + rank_major, seq_ids, comp_ids, valid, cu_comp, int(cu_comp[-1]) + ) + fused = csa_cp_layout_kernels.repack_compressed_kv_to_seq_major( + rank_major, seq_ids, comp_ids, valid, cu_comp, int(cu_comp[-1]) + ) + assert torch.equal(fused[0], ref[0]) + assert torch.equal(fused[1], ref[1]) + + bench_rows = 4096 + bench_width = 64 + bench_rank_major = torch.randn(bench_rows, bench_width, dtype=torch.bfloat16, device="cuda") + bench_seq_ids = torch.zeros(bench_rows, dtype=torch.int32, device="cuda") + bench_comp_ids = torch.arange(bench_rows, dtype=torch.int32, device="cuda") + bench_valid = torch.ones(bench_rows, dtype=torch.bool, device="cuda") + bench_cu_comp = torch.tensor([0, bench_rows], dtype=torch.int32, device="cuda") + + def run_cute(): + csa_cp_layout_kernels.repack_compressed_kv_to_seq_major( + bench_rank_major, bench_seq_ids, bench_comp_ids, bench_valid, bench_cu_comp, bench_rows + ) + + def run_torch(): + _native_repack_compressed_kv_to_seq_major( + bench_rank_major, bench_seq_ids, bench_comp_ids, bench_valid, bench_cu_comp, bench_rows + ) + + effective_bytes = _num_bytes(bench_rank_major) * 2 + bench_rows * 4 + _print_perf( + "repack_compressed_kv_to_seq_major", + _time_cuda(run_cute), + _time_cuda(run_torch), + effective_bytes, + ) + + +def test_build_compressed_row_metadata_matches_native_and_reports_bandwidth(): + _require_cute_cuda() + cu = _make_e2e_like_cu_seqlens() + cp_size = _E2E_CP_SIZE + chunk_len = (sum(_E2E_RAGGED_PADDED_SEG_LENS) // cp_size) // 2 + ratio = 4 + d_comp = 8 + c_cap_per_chunk = (chunk_len + d_comp) // ratio + c_cap_per_rank = 2 * c_cap_per_chunk + ref = _native_build_compressed_row_metadata( + cu, cp_size, chunk_len, ratio, d_comp, c_cap_per_chunk, c_cap_per_rank, True + ) + fused = csa_cp_layout_kernels.build_compressed_row_metadata( + cu, + cp_size, + chunk_len, + ratio, + d_comp, + c_cap_per_chunk, + c_cap_per_rank=c_cap_per_rank, + use_two_chunk=True, + ) + for actual, expected in zip(fused, ref): + assert torch.equal(actual, expected) + + bench_cu = torch.tensor([0, 65536], dtype=torch.int32, device="cuda") + bench_cp_size = 4 + bench_chunk_len = 8192 + bench_ratio = 128 + bench_d_comp = 128 + bench_c_cap = (bench_chunk_len + bench_d_comp) // bench_ratio + bench_c_cap_per_rank = bench_c_cap * 2 + + def run_cute(): + csa_cp_layout_kernels.build_compressed_row_metadata( + bench_cu, + bench_cp_size, + bench_chunk_len, + bench_ratio, + bench_d_comp, + bench_c_cap, + c_cap_per_rank=bench_c_cap_per_rank, + use_two_chunk=True, + ) + + def run_torch(): + _native_build_compressed_row_metadata( + bench_cu, + bench_cp_size, + bench_chunk_len, + bench_ratio, + bench_d_comp, + bench_c_cap, + bench_c_cap_per_rank, + True, + ) + + effective_bytes = bench_cp_size * bench_c_cap_per_rank * (4 + 4 + 1) + _print_perf( + "build_compressed_row_metadata", + _time_cuda(run_cute), + _time_cuda(run_torch), + effective_bytes, + ) + + +def test_build_indexer_topk_metadata_matches_native_and_reports_bandwidth(): + _require_cute_cuda() + ratio = 4 + cu_q = _make_e2e_like_cu_seqlens() + cu_comp = _compressed_cu_seqlens(cu_q, ratio) + k_seq = torch.arange(int(cu_comp[-1]) * 2, dtype=torch.bfloat16, device="cuda").reshape( + int(cu_comp[-1]), 2 + ) + global_start, l_local = _e2e_like_contiguous_range() + ref = _native_indexer_topk_metadata(k_seq, cu_q, cu_comp, global_start, l_local, ratio) + fused = csa_cp_layout_kernels.build_indexer_topk_metadata( + k_seq, cu_q, cu_comp, global_start, l_local, ratio + ) + for actual, expected in zip(fused, ref): + assert torch.equal(actual, expected) + + bench_ratio = 128 + bench_l_local = 8192 + bench_rows = 512 + bench_k_seq = torch.randn(bench_rows, 64, dtype=torch.bfloat16, device="cuda") + bench_cu_q = torch.tensor([0, bench_l_local], dtype=torch.int32, device="cuda") + bench_cu_comp = torch.tensor([0, bench_rows], dtype=torch.int32, device="cuda") + + def run_cute(): + csa_cp_layout_kernels.build_indexer_topk_metadata( + bench_k_seq, bench_cu_q, bench_cu_comp, 0, bench_l_local, bench_ratio + ) + + def run_torch(): + _native_indexer_topk_metadata( + bench_k_seq, bench_cu_q, bench_cu_comp, 0, bench_l_local, bench_ratio + ) + + effective_bytes = _num_bytes(bench_k_seq) * 2 + bench_l_local * 4 + _print_perf( + "build_indexer_topk_metadata", _time_cuda(run_cute), _time_cuda(run_torch), effective_bytes + ) + + +def test_build_attention_indices_matches_native_and_reports_bandwidth(): + _require_cute_cuda() + cu = _make_e2e_like_cu_seqlens() + global_start, l_local = _e2e_like_contiguous_range() + d_window = 128 + window = 16 + ratio = 4 + compressed_width = 16 + fused = csa_cp_layout_kernels.build_attention_indices( + cu, global_start, l_local, d_window, window, ratio, compressed_width + ) + expected = _native_attention_indices( + cu, global_start, l_local, d_window, window, ratio, compressed_width + ) + assert torch.equal(fused[0], expected[0]) + assert torch.equal(fused[1], expected[1]) + + padded = csa_cp_layout_kernels.build_attention_indices( + torch.tensor([0, 8], dtype=torch.int32, device="cuda"), 0, 10, 2, 2, 0, 0 + ) + # Padded THD query rows still need one in-range dummy KV id for fused DSA. + assert torch.equal( + padded[0][8:10], torch.tensor([[0, -1], [0, -1]], dtype=torch.int32, device="cuda") + ) + assert torch.equal(padded[1][8:10], torch.tensor([1, 1], dtype=torch.int32, device="cuda")) + + bench_l_local = 4096 + bench_window = 128 + bench_ratio = 4 + bench_compressed_width = bench_l_local // bench_ratio + bench_cu = torch.tensor([0, bench_l_local], dtype=torch.int32, device="cuda") + bench_ref = _native_attention_indices_one_seq( + bench_l_local, bench_window, bench_ratio, bench_compressed_width, torch.device("cuda") + ) + + def run_cute(): + csa_cp_layout_kernels.build_attention_indices( + bench_cu, + 0, + bench_l_local, + bench_window, + bench_window, + bench_ratio, + bench_compressed_width, + ) + + def run_torch(): + _native_attention_indices_one_seq( + bench_l_local, bench_window, bench_ratio, bench_compressed_width, torch.device("cuda") + ) + + effective_bytes = _num_bytes(*bench_ref) + _print_perf( + "build_attention_indices", _time_cuda(run_cute), _time_cuda(run_torch), effective_bytes + ) + + +def test_build_indexer_loss_indices_matches_native_and_reports_bandwidth(): + _require_cute_cuda() + cu = _make_e2e_like_cu_seqlens() + ratio = 4 + cu_comp = _compressed_cu_seqlens(cu, ratio) + global_start, l_local = _e2e_like_contiguous_range() + d_window = 128 + window = 16 + compressed_width = 8 + logical_ids = ( + torch.arange(compressed_width, dtype=torch.int32, device="cuda").unsqueeze(0).repeat(l_local, 1) + ) + logical_ids[1::5, -1] = -1 + rank_by_seq = torch.arange(int(cu_comp[-1]), dtype=torch.int32, device="cuda") + fused = csa_cp_layout_kernels.build_indexer_loss_indices( + cu, cu_comp, global_start, l_local, d_window, window, ratio, logical_ids, rank_by_seq + ) + expected = _native_indexer_loss_indices( + cu, + cu_comp, + global_start, + l_local, + d_window, + window, + ratio, + logical_ids, + rank_by_seq, + ) + assert torch.equal(fused[0], expected[0]) + assert torch.equal(fused[1], expected[1]) + + bench_l_local = 4096 + bench_window = 128 + bench_ratio = 4 + bench_compressed_width = bench_l_local // bench_ratio + bench_cu = torch.tensor([0, bench_l_local], dtype=torch.int32, device="cuda") + bench_cu_comp = torch.tensor([0, bench_compressed_width], dtype=torch.int32, device="cuda") + bench_logical = ( + torch.arange(bench_compressed_width, dtype=torch.int32, device="cuda") + .unsqueeze(0) + .expand(bench_l_local, -1) + ) + bench_rank_by_seq = torch.arange(bench_compressed_width, dtype=torch.int32, device="cuda") + bench_ref = _native_indexer_loss_indices_one_seq( + bench_l_local, bench_window, bench_ratio, bench_logical, bench_rank_by_seq + ) + + def run_cute(): + csa_cp_layout_kernels.build_indexer_loss_indices( + bench_cu, + bench_cu_comp, + 0, + bench_l_local, + bench_window, + bench_window, + bench_ratio, + bench_logical, + bench_rank_by_seq, + ) + + def run_torch(): + _native_indexer_loss_indices_one_seq( + bench_l_local, bench_window, bench_ratio, bench_logical, bench_rank_by_seq + ) + + effective_bytes = _num_bytes(bench_logical, bench_rank_by_seq, *bench_ref) + _print_perf( + "build_indexer_loss_indices", _time_cuda(run_cute), _time_cuda(run_torch), effective_bytes + ) diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py new file mode 100644 index 00000000000..b9a69a2ce77 --- /dev/null +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py @@ -0,0 +1,682 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import pytest +import torch + +import megatron.core.parallel_state as parallel_state +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer.experimental_attention_variant import csa_cp_utils +from megatron.core.transformer.experimental_attention_variant.csa_cp_utils import ( + DSV4_CP_PARTITION_CONTIGUOUS, + DSV4_CP_PARTITION_TWO_CHUNK, + all_gather_fixed_cp_tensor, + apply_thd_cp_compressed_rope_fused, + apply_thd_cp_local_rope_fused, + build_cp_attention_indices_fused, + build_cp_compressor_prep_compact_fused, + build_cp_indexer_loss_indices_fused, + build_cp_rank_major_compressed_metadata_fused, + build_global_compressed_cu_seqlens, + compute_cp_indexer_topk_logical_fused, + exchange_cp_boundary_hidden, + local_kv_cp_chunk_ranges, + local_q_cp_chunk_ranges, + pack_cp_kv_full_fused, + repack_rank_major_compressed_to_seq_major_fused, + thd_cp_local_row_indices, +) +from tests.unit_tests.test_utilities import Utils + + +def _require_cuda(): + if not torch.cuda.is_available(): + pytest.skip("DSv4 CP CUDA utility tests require CUDA.") + + +def test_csa_cp_partition_mode_contract(): + """Validate the public CSA CP partition-mode names. + + Expected: omitted mode defaults to the original contiguous partition, + supported strings round-trip exactly, and unknown values fail before any CP + layout helper can silently choose the wrong row order. + """ + assert local_q_cp_chunk_ranges(None, 4, 1, 0) == ((0, 4),) + assert local_q_cp_chunk_ranges(DSV4_CP_PARTITION_CONTIGUOUS, 4, 1, 0) == ((0, 4),) + assert local_q_cp_chunk_ranges(DSV4_CP_PARTITION_TWO_CHUNK, 4, 1, 0) == ((0, 4),) + + with pytest.raises(RuntimeError, match="Unsupported CSA CP partition mode"): + local_q_cp_chunk_ranges("invalid_mode", 4, 1, 0) + + +@pytest.mark.parametrize( + "partition_mode, uses_two_chunk", + [(DSV4_CP_PARTITION_CONTIGUOUS, False), (DSV4_CP_PARTITION_TWO_CHUNK, True)], + ids=["contiguous", "two_chunk"], +) +def test_thd_cp_partition_mode_selects_expected_row_order(partition_mode, uses_two_chunk): + """Validate the CP row order selected by each partition mode. + + Expected: ``contiguous`` returns this rank's single local range, while + ``two_chunk`` returns chunk ``rank`` followed by chunk + ``2*cp_size-1-rank``. A failure means CP data slicing and layer-side layout + helpers could disagree about local packed-token order. + """ + local_rows = 16 + cp_size = 4 + + for cp_rank in range(cp_size): + ranges = local_q_cp_chunk_ranges(partition_mode, local_rows, cp_size, cp_rank) + if uses_two_chunk: + chunk_len = local_rows // 2 + total_chunks = 2 * cp_size + chunk_ids = (cp_rank, total_chunks - 1 - cp_rank) + expected = tuple( + (chunk_id * chunk_len, (chunk_id + 1) * chunk_len) for chunk_id in chunk_ids + ) + else: + expected = ((cp_rank * local_rows, (cp_rank + 1) * local_rows),) + assert ranges == expected + + +def test_two_chunk_cp_ranges_match_expected_order(): + """Validate the two-chunk CP partition contract. + + Expected: rank r owns global chunk r followed by chunk 2*cp_size-1-r. + A failure here means DSv4 chunk-aware CP helpers would disagree with the + two-chunk local row order. + """ + assert local_q_cp_chunk_ranges(DSV4_CP_PARTITION_TWO_CHUNK, 16, 1, 0) == ((0, 16),) + expected_by_rank = { + 0: ((0, 2), (14, 16)), + 1: ((2, 4), (12, 14)), + 2: ((4, 6), (10, 12)), + 3: ((6, 8), (8, 10)), + } + for rank, expected in expected_by_rank.items(): + assert local_q_cp_chunk_ranges(DSV4_CP_PARTITION_TWO_CHUNK, 4, 4, rank) == expected + + with pytest.raises(RuntimeError, match="even local_rows"): + local_q_cp_chunk_ranges(DSV4_CP_PARTITION_TWO_CHUNK, 5, 4, 0) + + +def test_cp_chunk_ranges_match_partition_mode(): + """Validate local and left-boundary ranges for each CP partition mode. + + Expected: each mode returns the current rank's local rows and matching + left-boundary rows in global packed-token coordinates. + """ + assert local_q_cp_chunk_ranges(None, local_rows=4, cp_size=4, cp_rank=2) == ((8, 12),) + assert local_kv_cp_chunk_ranges( + DSV4_CP_PARTITION_CONTIGUOUS, local_rows=4, boundary_rows=2, cp_size=4, cp_rank=2 + ) == ((6, 8), (8, 12)) + + assert local_q_cp_chunk_ranges( + DSV4_CP_PARTITION_TWO_CHUNK, local_rows=4, cp_size=4, cp_rank=0 + ) == ((0, 2), (14, 16)) + assert local_kv_cp_chunk_ranges( + DSV4_CP_PARTITION_TWO_CHUNK, local_rows=4, boundary_rows=4, cp_size=4, cp_rank=0 + ) == ((-2, 0), (12, 14), (0, 2), (14, 16)) + + with pytest.raises(RuntimeError, match="boundary rows must be divisible"): + local_kv_cp_chunk_ranges( + DSV4_CP_PARTITION_TWO_CHUNK, local_rows=4, boundary_rows=3, cp_size=4, cp_rank=0 + ) + + +def test_thd_cp_local_row_indices_match_chunk_ranges(): + """Validate THD batch row selection for each CSA CP partition mode. + + Expected: data loading selects the same global packed-token rows that the + attention layer later uses for RoPE, boundary exchange, and index lowering. + """ + assert torch.equal( + thd_cp_local_row_indices(DSV4_CP_PARTITION_CONTIGUOUS, 16, 4, 2, torch.device("cpu")), + torch.tensor([8, 9, 10, 11]), + ) + assert torch.equal( + thd_cp_local_row_indices(DSV4_CP_PARTITION_TWO_CHUNK, 16, 4, 0, torch.device("cpu")), + torch.tensor([0, 1, 14, 15]), + ) + assert torch.equal( + thd_cp_local_row_indices(DSV4_CP_PARTITION_TWO_CHUNK, 16, 4, 3, torch.device("cpu")), + torch.tensor([6, 7, 8, 9]), + ) + + with pytest.raises(RuntimeError, match="divisible by cp_size"): + thd_cp_local_row_indices(DSV4_CP_PARTITION_CONTIGUOUS, 15, 4, 0, torch.device("cpu")) + + +def test_boundary_exchange_single_rank_returns_zero_boundary_and_zero_grad(): + """Validate the no-CP boundary exchange contract. + + Expected: with cp_group=None, the fixed left boundary is zero-filled and its + backward path contributes no gradient to the local tensor. A failure here + means the CP path could invent boundary tokens or bogus local gradients. + """ + local = torch.arange(12, dtype=torch.float32).reshape(4, 3).requires_grad_(True) + + boundary = exchange_cp_boundary_hidden(local, [], 2, DSV4_CP_PARTITION_CONTIGUOUS, None) + + assert torch.equal(boundary, torch.zeros(2, 3)) + boundary.sum().backward() + assert torch.equal(local.grad, torch.zeros_like(local)) + + +def test_two_chunk_boundary_exchange_single_rank_returns_zero_boundary_and_zero_grad(): + """Validate the no-CP two-chunk boundary exchange contract. + + Expected: with cp_group=None, the helper degenerates to one zero-filled + boundary and contributes no local gradient. A failure here means chunk-aware + CP boundary plumbing could perturb disabled-CP paths. + """ + local = torch.arange(12, dtype=torch.float32).reshape(4, 3).requires_grad_(True) + + boundary = exchange_cp_boundary_hidden(local, [], 2, DSV4_CP_PARTITION_TWO_CHUNK, None) + + assert torch.equal(boundary, torch.zeros(2, 3)) + boundary.sum().backward() + assert torch.equal(local.grad, torch.zeros_like(local)) + + +def test_thd_cp_left_boundary_exchange_forward_backward(): + """Validate distributed CP boundary exchange forward/backward. + + Expected: forward receives the previous rank's tail window, and backward + sends gradient to this rank's tail only when the next rank consumed it as a + left boundary. + """ + _require_cuda() + if Utils.world_size < 2: + pytest.skip("Distributed CP boundary exchange requires at least 2 ranks.") + + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + context_parallel_size=Utils.world_size, + ) + try: + cp_group = ProcessGroupCollection.use_mpu_process_groups().cp + cp_rank = parallel_state.get_context_parallel_rank() + cp_size = parallel_state.get_context_parallel_world_size() + d_window = 2 + local_len = 4 + width = 3 + local_numel = local_len * width + local_start = cp_rank * local_numel + values = torch.arange( + local_start, local_start + local_numel, device='cuda', dtype=torch.float32 + ).reshape(local_len, width) + local = values.detach().clone().requires_grad_(True) + + boundary = exchange_cp_boundary_hidden( + local, [], d_window, DSV4_CP_PARTITION_CONTIGUOUS, cp_group + ) + if cp_rank == 0: + expected_boundary = torch.zeros_like(boundary) + else: + left_rank_start = (cp_rank - 1) * local_numel + expected_boundary = torch.arange( + left_rank_start + (local_len - d_window) * width, + left_rank_start + local_numel, + device='cuda', + dtype=torch.float32, + ).reshape(d_window, width) + assert torch.equal(boundary, expected_boundary) + + boundary.sum().backward() + expected_grad = torch.zeros_like(local) + if cp_rank + 1 < cp_size: + expected_grad[-d_window:] = 1 + assert torch.equal(local.grad, expected_grad) + finally: + Utils.destroy_model_parallel() + + +def test_global_compressed_cu_seqlens_matches_reference(): + """Validate compressed-prefix metadata for ragged padded sequences. + + Expected: each sequence contributes ``padded_seq_len // ratio`` compressed + rows to the global seq-major prefix. A failure here means rank-major + compressed rows could be repacked with the wrong sequence offsets. + """ + _require_cuda() + cu_seqlens_cpu = torch.tensor([0, 7, 128, 381, 512], dtype=torch.int32) + ratio = 4 + + actual = build_global_compressed_cu_seqlens(cu_seqlens_cpu.cuda(), ratio) + + assert torch.equal(actual.cpu(), torch.tensor([0, 1, 31, 94, 126], dtype=torch.int32)) + + +def test_apply_thd_cp_local_rope_fused_dispatches_equal_and_ragged_chunks(monkeypatch): + """Validate the public local-RoPE wrapper dispatch contract. + + Expected: equal one/two-chunk layouts call the fused kernel once with both + chunk starts, while ragged multi-chunk layouts call once per chunk and + concatenate the chunk results in local row order. + """ + calls = [] + + def fake_apply( + x, + cos, + sin, + cu, + chunk0_start, + chunk_len, + nope_dim, + pos_dim, + chunk1_start, + inverse, + clamp_to_valid_token, + ): + calls.append( + ( + tuple(x.shape), + int(chunk0_start), + int(chunk_len), + int(nope_dim), + int(pos_dim), + int(chunk1_start), + bool(inverse), + bool(clamp_to_valid_token), + ) + ) + return x + len(calls) * 10 + + monkeypatch.setattr( + csa_cp_utils.csa_cp_layout_kernels.ThdLocalRope, "apply", staticmethod(fake_apply) + ) + cos = sin = torch.empty(8, 2) + cu = torch.tensor([0, 8], dtype=torch.int32) + + equal = torch.arange(8, dtype=torch.float32).reshape(4, 2) + out = apply_thd_cp_local_rope_fused( + equal, cos, sin, 1, 1, cu, ((2, 4), (10, 12)), inverse=True + ) + assert torch.equal(out, equal + 10) + assert calls[-1] == ((4, 2), 2, 2, 1, 1, 10, True, False) + + ragged = torch.arange(10, dtype=torch.float32).reshape(5, 2) + out = apply_thd_cp_local_rope_fused( + ragged, cos, sin, 1, 1, cu, ((2, 4), (10, 13)), clamp_to_valid_token=True + ) + assert torch.equal(out, torch.cat((ragged[:2] + 20, ragged[2:] + 30), dim=0)) + assert calls[-2:] == [ + ((2, 2), 2, 2, 1, 1, 0, False, True), + ((3, 2), 10, 3, 1, 1, 0, False, True), + ] + + with pytest.raises(RuntimeError, match="expects x rows"): + apply_thd_cp_local_rope_fused(equal[:3], cos, sin, 1, 1, cu, ((2, 4), (10, 12))) + + +def test_apply_thd_cp_compressed_rope_fused_forwards_kernel_arguments(monkeypatch): + """Validate the public compressed-RoPE wrapper is a stable typed facade. + + Expected: the wrapper forwards tensor arguments, dimensions, ratio, and + inverse flag unchanged to the fused kernel entry point. + """ + captured = {} + + def fake_apply(x, cos, sin, comp_ids, ratio, nope_dim, pos_dim, inverse): + captured.update( + ratio=int(ratio), + nope_dim=int(nope_dim), + pos_dim=int(pos_dim), + inverse=bool(inverse), + comp_ids=comp_ids, + ) + return x + 1 + + monkeypatch.setattr( + csa_cp_utils.csa_cp_layout_kernels.ThdCompressedRope, "apply", staticmethod(fake_apply) + ) + x = torch.arange(6, dtype=torch.float32).reshape(3, 2) + comp_ids = torch.tensor([0, 2, -1], dtype=torch.int32) + out = apply_thd_cp_compressed_rope_fused( + x, torch.empty(16, 1), torch.empty(16, 1), comp_ids, 4, 1, 1, inverse=True + ) + + assert torch.equal(out, x + 1) + assert captured["ratio"] == 4 + assert captured["nope_dim"] == 1 + assert captured["pos_dim"] == 1 + assert captured["inverse"] is True + assert captured["comp_ids"] is comp_ids + + +def test_rank_major_compressed_metadata_wrapper_computes_capacities(monkeypatch): + """Validate public compressed-metadata capacity and partition dispatch. + + Expected: contiguous and two-chunk modes compute aligned capacities before + delegating to the kernel wrapper, so all ranks use fixed-size all-gather + rows independent of the number of valid compressed groups. + """ + calls = [] + + def fake_build(cu, cp_size, l_local, ratio, d_comp, c_cap, **kwargs): + calls.append( + ( + int(cp_size), + int(l_local), + int(ratio), + int(d_comp), + int(c_cap), + kwargs.get("c_cap_per_rank"), + kwargs.get("use_two_chunk", False), + ) + ) + rows = int(cp_size) * int(kwargs.get("c_cap_per_rank", c_cap)) + return ( + torch.full((rows,), -1, dtype=torch.int32), + torch.full((rows,), -1, dtype=torch.int32), + torch.zeros((rows,), dtype=torch.bool), + ) + + monkeypatch.setattr(csa_cp_utils.csa_cp_layout_kernels, "build_compressed_row_metadata", fake_build) + cu = torch.tensor([0, 32, 64], dtype=torch.int32) + + build_cp_rank_major_compressed_metadata_fused(cu, ((16, 32),), 4, 4, 8) + build_cp_rank_major_compressed_metadata_fused(cu, ((0, 4), (28, 32)), 4, 4, 8) + + assert calls == [ + (4, 16, 4, 8, 8, None, False), + (4, 4, 4, 8, 3, 8, True), + ] + + +def test_compressor_prep_compact_wrapper_merges_two_chunk_outputs(monkeypatch): + """Validate public compressor-prep wrapper splitting, merge, and padding. + + Expected: two chunks are compacted independently, local compact group ids + are concatenated in local chunk order, and the wrapper pads to the shared + aligned capacity used by later fixed-shape collectives. + """ + calls = [] + + def fake_apply(hidden, boundary, cu, global_start, rows, ratio, d_comp, d_window, c_cap): + calls.append((int(global_start), int(rows), int(c_cap), tuple(boundary.shape))) + marker = len(calls) + hidden_compact = torch.full((int(c_cap) * int(ratio), 1), marker, dtype=hidden.dtype) + cu_compact = torch.tensor([0, int(c_cap) * int(ratio)], dtype=cu.dtype) + seq_ids = torch.full((int(c_cap),), marker, dtype=torch.int32) + comp_ids = torch.arange(int(c_cap), dtype=torch.int32) + marker * 10 + valid = torch.ones((int(c_cap),), dtype=torch.bool) + return hidden_compact, cu_compact, seq_ids, comp_ids, valid + + monkeypatch.setattr( + csa_cp_utils.csa_cp_layout_kernels.CompressorInputCompact, "apply", staticmethod(fake_apply) + ) + hidden = torch.arange(6, dtype=torch.float32).reshape(6, 1) + boundary = torch.arange(4, dtype=torch.float32).reshape(4, 1) + cu = torch.tensor([0, 6], dtype=torch.int32) + + hidden_compact, cu_compact, comp_ids = build_cp_compressor_prep_compact_fused( + hidden, boundary, cu, ((0, 3), (10, 13)), ratio=2, d_comp=2, d_window=2 + ) + + assert calls == [(0, 3, 2, (2, 1)), (10, 3, 2, (2, 1))] + assert hidden_compact.shape == (32, 1) + assert torch.equal(cu_compact, torch.tensor([0, 8], dtype=torch.int32)) + assert torch.equal(comp_ids[:4], torch.tensor([10, 11, 20, 21], dtype=torch.int32)) + assert torch.equal(comp_ids[4:], torch.full((12,), -1, dtype=torch.int32)) + + +def test_all_gather_fixed_cp_tensor_single_rank_preserves_autograd(): + """Validate the public fixed CP all-gather no-CP fallback. + + Expected: with cp_group=None, the helper returns the local tensor and + backward is the identity, matching the disabled-CP execution path. + """ + x = torch.arange(6, dtype=torch.float32).reshape(3, 2).requires_grad_(True) + + out = all_gather_fixed_cp_tensor(x, None) + + assert torch.equal(out, x) + out.square().sum().backward() + assert torch.equal(x.grad, 2 * x.detach()) + + +def test_pack_cp_kv_full_fused_wrapper_selects_contiguous_and_two_chunk_layouts(monkeypatch): + """Validate public KV-pack wrapper argument lowering for both layouts. + + Expected: contiguous CP passes the physical local range through directly, + while two-chunk CP switches to shared-window layout and returns the shared + compressed base used by final index lowering. + """ + calls = [] + + def fake_apply(*args): + calls.append(args) + kv_local = args[0] + capacity = int(args[13]) + return kv_local.new_full((capacity,) + tuple(kv_local.shape[1:]), len(calls)) + + monkeypatch.setattr(csa_cp_utils.csa_cp_layout_kernels.ThdFullKvPack, "apply", staticmethod(fake_apply)) + kv_local = torch.ones(4, 2) + boundary = torch.ones(2, 2) + compressed = torch.ones(6, 2) + ids = torch.zeros(6, dtype=torch.int32) + valid = torch.ones(6, dtype=torch.bool) + cu = torch.tensor([0, 8], dtype=torch.int32) + rank_by_seq = torch.arange(2, dtype=torch.int32) + cu_comp = torch.tensor([0, 2], dtype=torch.int32) + + out, shared = pack_cp_kv_full_fused( + kv_local, + boundary, + compressed, + ids, + ids, + valid, + cu, + ((4, 8),), + d_window=2, + ratio=4, + rank_major_by_seq_major=rank_by_seq, + cu_seqlens_compressed=cu_comp, + ) + assert shared is None + assert out.shape == (12, 2) + assert calls[-1][9:19] == (4, 4, 2, 4, 12, 0, 0, 0, 0, 0) + + out, shared = pack_cp_kv_full_fused( + kv_local, + torch.ones(4, 2), + compressed, + ids, + ids, + valid, + cu, + ((0, 2), (14, 16)), + d_window=2, + ratio=4, + ) + assert shared == 8 + assert out.shape == (14, 2) + assert calls[-1][9:19] == (0, 0, 2, 0, 14, 0, 14, 2, 2, 4) + + +def test_repack_rank_major_compressed_to_seq_major_fused_forwards_arguments(monkeypatch): + """Validate the public compressed-KV repack wrapper. + + Expected: the wrapper forwards rank-major tensors and requested seq-major + capacity unchanged to the fused repack kernel. + """ + captured = {} + + def fake_repack(rank_major, seq_ids, comp_ids, valid, cu_comp, rows): + captured.update(rows=int(rows), rank_major=rank_major, seq_ids=seq_ids) + return rank_major[:rows].clone(), torch.arange(rows, dtype=torch.int32) + + monkeypatch.setattr(csa_cp_utils.csa_cp_layout_kernels, "repack_compressed_kv_to_seq_major", fake_repack) + rank_major = torch.arange(12, dtype=torch.float32).reshape(6, 2) + seq_ids = torch.zeros(6, dtype=torch.int32) + comp_ids = torch.arange(6, dtype=torch.int32) + valid = torch.ones(6, dtype=torch.bool) + cu_comp = torch.tensor([0, 3], dtype=torch.int32) + + out, rank_by_seq = repack_rank_major_compressed_to_seq_major_fused( + rank_major, seq_ids, comp_ids, valid, cu_comp, 3 + ) + + assert captured["rows"] == 3 + assert captured["rank_major"] is rank_major + assert captured["seq_ids"] is seq_ids + assert torch.equal(out, rank_major[:3]) + assert torch.equal(rank_by_seq, torch.arange(3, dtype=torch.int32)) + + +def test_compute_cp_indexer_topk_logical_fused_splits_chunks_and_passes_visible_lengths( + monkeypatch, +): + """Validate public CP indexer-topk wrapper chunk splitting. + + Expected: each local chunk gets its own trapezoid metadata and indexer-topk + call, visible K lengths are passed through, and chunk outputs are + concatenated in local row order. + """ + metadata_calls = [] + topk_calls = [] + + def fake_metadata(k_seq, cu_q, cu_comp, global_start, length, ratio): + metadata_calls.append((int(global_start), int(length), int(ratio))) + marker = len(metadata_calls) + return ( + k_seq + marker, + torch.tensor([0, int(length)], dtype=torch.int32), + torch.tensor([0, 1], dtype=torch.int32), + torch.full((int(length),), marker, dtype=torch.int32), + ) + + def fake_indexer_topk( + q, + k, + weights, + *, + topk, + ratio, + indexer_softmax_scale, + cu_seqlens_q, + cu_seqlens_kv, + max_seqlen_q, + max_seqlen_kv, + visible_k_lengths, + ): + topk_calls.append( + ( + tuple(q.shape), + int(topk), + int(ratio), + float(indexer_softmax_scale), + int(max_seqlen_q), + int(max_seqlen_kv), + visible_k_lengths.clone(), + ) + ) + return torch.full((q.shape[0], int(topk)), len(topk_calls), dtype=torch.int32), None + + monkeypatch.setattr(csa_cp_utils.csa_cp_layout_kernels, "build_indexer_topk_metadata", fake_metadata) + monkeypatch.setattr(csa_cp_utils, "indexer_topk", fake_indexer_topk) + + q = torch.randn(5, 2) + weights = torch.randn(5, 1) + k_seq = torch.randn(4, 2) + cu_q = torch.tensor([0, 10], dtype=torch.int32) + cu_comp = torch.tensor([0, 3], dtype=torch.int32) + + out = compute_cp_indexer_topk_logical_fused( + q, + weights, + k_seq, + cu_q, + cu_comp, + ((2, 4), (8, 11)), + ratio=4, + topk_width=2, + indexer_softmax_scale=0.5, + max_seqlen_q=10, + max_seqlen_kv=3, + ) + + assert metadata_calls == [(2, 2, 4), (8, 3, 4)] + assert torch.equal(out, torch.tensor([[1, 1], [1, 1], [2, 2], [2, 2], [2, 2]])) + assert torch.equal(topk_calls[0][-1], torch.ones(2, dtype=torch.int32)) + assert torch.equal(topk_calls[1][-1], torch.full((3,), 2, dtype=torch.int32)) + assert compute_cp_indexer_topk_logical_fused( + q, weights, k_seq[:0], cu_q, cu_comp, ((2, 4),), 4, 2, 1.0, 10, 3 + ) is None + + +def test_final_index_wrappers_lower_two_chunk_arguments(monkeypatch): + """Validate public final attention and indexer-loss index wrappers. + + Expected: two-chunk wrappers require shared compressed metadata and pass + chunk starts, chunk length, shared base, and local row count to the fused + kernel entry points. + """ + captured = {} + + def fake_attention(*args, **kwargs): + captured["attention"] = (args, kwargs) + l_local = int(args[2]) + width = int(args[6]) + return torch.zeros(l_local, width, dtype=torch.int32), torch.zeros(l_local, dtype=torch.int32) + + def fake_loss(*args, **kwargs): + captured["loss"] = (args, kwargs) + l_local = int(args[3]) + width = args[7].shape[1] + int(args[5]) + return torch.zeros(l_local, width, dtype=torch.int32), torch.zeros_like(args[7]) + + monkeypatch.setattr(csa_cp_utils.csa_cp_layout_kernels, "build_attention_indices", fake_attention) + monkeypatch.setattr(csa_cp_utils.csa_cp_layout_kernels, "build_indexer_loss_indices", fake_loss) + + cu = torch.tensor([0, 16], dtype=torch.int32) + cu_comp = torch.tensor([0, 4], dtype=torch.int32) + rank_by_seq = torch.arange(4, dtype=torch.int32) + logical = torch.zeros(4, 2, dtype=torch.int32) + chunks = ((0, 2), (14, 16)) + + with pytest.raises(RuntimeError, match="requires a shared compressed base"): + build_cp_attention_indices_fused(cu, chunks, 2, 2, 4, logical, 2) + + build_cp_attention_indices_fused( + cu, + chunks, + d_window=2, + window_size=2, + ratio=4, + indexer_topk_compressed_logical_ids=logical, + max_n_compressed=2, + rank_major_by_seq_major=rank_by_seq, + cu_seqlens_compressed=cu_comp, + shared_compressed_base=8, + ) + args, kwargs = captured["attention"] + assert args[0] is cu + assert args[1:7] == (0, 4, 2, 2, 4, 2) + assert kwargs["chunk_starts"] == (0, 14) + assert kwargs["chunk_len"] == 2 + assert kwargs["shared_compressed_base"] == 8 + + build_cp_indexer_loss_indices_fused( + cu, + cu_comp, + chunks, + d_window=2, + window_size=2, + ratio=4, + indexer_topk_compressed_logical_ids=logical, + rank_major_by_seq_major=rank_by_seq, + shared_compressed_base=8, + ) + args, kwargs = captured["loss"] + assert args[0] is cu + assert args[1] is cu_comp + assert args[2:7] == (0, 4, 2, 2, 4) + assert args[7] is logical + assert kwargs["chunk_starts"] == (0, 14) + assert kwargs["chunk_len"] == 2 + assert kwargs["shared_compressed_base"] == 8 diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention.py index 6a8ec00f624..725488f6122 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention.py @@ -357,7 +357,7 @@ def test_qkv_shapes(self): seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 ).cuda() - q, k, v, q_compressed, kv_compressed = attn.get_query_key_value_tensors(hidden) + q, k, v, q_compressed, kv_compressed, _ = attn.get_query_key_value_tensors(hidden) n_heads = self.config.num_attention_heads v_dim = self.config.v_head_dim @@ -380,7 +380,7 @@ def test_key_equals_value(self): seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 ).cuda() - q, k, v, _, _ = attn.get_query_key_value_tensors(hidden) + q, k, v, _, _, _ = attn.get_query_key_value_tensors(hidden) assert torch.equal(k, v), "key and value should be identical in wkv path" diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py new file mode 100644 index 00000000000..040495b1bff --- /dev/null +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py @@ -0,0 +1,981 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import gc +import os +import statistics +from contextlib import contextmanager, nullcontext + +os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") + +import pytest +import torch +import torch.distributed as dist +import torch.nn.functional as F + +import megatron.core.parallel_state as parallel_state +from megatron.core.extensions.transformer_engine import HAVE_TE +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.experimental_attention_variant.csa_cp_utils import ( + DSV4_CP_PARTITION_CONTIGUOUS, + DSV4_CP_PARTITION_TWO_CHUNK, + thd_cp_local_row_indices, +) +from tests.unit_tests.test_utilities import Utils +from tests.unit_tests.transformer.experimental_attention_variant.test_dsv4_hybrid_attention import ( + _SEED, + _build_attention, + _make_config, +) +from tests.unit_tests.transformer.experimental_attention_variant.test_dsv4_hybrid_native_parity import ( + _DSV4_VARIANTS, +) + +# =========================================================================== +# THD CP parity tests +# =========================================================================== + + +_DSV4_CP_PARITY_EPS = 1e-3 + +# CUDA graph with fused kernels may not be deterministic against eager, so this +# path uses similarity plus rtol/atol gates. The unfused CUDA graph path is +# still checked for bitwise parity. +_DSV4_CP_GRAPH_FUSED_SIM_EPS = 1e-6 +_DSV4_CP_GRAPH_FUSED_RTOL = 1e-6 +_DSV4_CP_GRAPH_FUSED_BF16_ATOL = 1.0 +_DSV4_CP_GRAPH_FUSED_FP32_ATOL = 2e-2 + +# Actual measured peak allocated delta scale is below these limits. They leave +# room for allocator noise while still catching regressions that allocate +# full-size activation buffers on each CP rank. +_DSV4_CP_MEMORY_RATIO_LIMITS = {2: 0.65, 4: 0.35} + +# Actual measured graph replay scale is below these limits. They leave room for +# machine timing jitter while still catching regressions that reintroduce +# full-size work or slow dynamic operations into the CP layer path. +_DSV4_CP_GRAPH_TIME_RATIO_LIMITS = {2: 0.70, 4: 0.50} + +_DSV4_CP_MEMORY_WARMUP_ITERS = 3 +_DSV4_CP_GRAPH_TIMING_WARMUP_REPLAYS = 2 +_DSV4_CP_GRAPH_TIMING_MEASURE_REPLAYS = 5 +_DSV4_CP_TEST_VARIANT = "flash" + +# Padded total is 4096, divisible by CP2/CP4. Only the final segment has tail +# padding, so the intermediate sequence boundaries match the unpadded layout. +_DSV4_CP_RAGGED_SEG_LENS = (1, 127, 1000, 23, 129, 900, 55, 257, 800, 95, 509, 148) +_DSV4_CP_RAGGED_PADDED_SEG_LENS = (1, 127, 1000, 23, 129, 900, 55, 257, 800, 95, 509, 200) +# Same shape, padded total, and max padded sequence length as +# _DSV4_CP_RAGGED_PADDED_SEG_LENS, but the padding is distributed across many +# sequences instead of only the tail. CUDA graph replay must tolerate this value +# change because the CP path rebuilds compressed-row metadata from device +# cu_seqlens_padded rather than from capture-time host sizes. +_DSV4_CP_REPLAY_PADDED_SEG_LENS = (8, 128, 1000, 32, 132, 904, 64, 260, 804, 96, 512, 156) + + +class _ReferenceCPGroup: + def rank(self): + """Return the CP rank used by the CP1 reference path.""" + return 0 + + def size(self): + """Return the CP size used by the CP1 reference path.""" + return 1 + + +def _make_thd_packed_seq_params(seg_lens, padded_seg_lens=None, device='cuda'): + """Build THD PackedSeqParams from raw and padded segment lengths.""" + if padded_seg_lens is None: + padded_seg_lens = seg_lens + cu = torch.tensor( + [0] + list(torch.tensor(seg_lens).cumsum(0).tolist()), dtype=torch.int32, device=device + ) + cu_padded = torch.tensor( + [0] + list(torch.tensor(padded_seg_lens).cumsum(0).tolist()), + dtype=torch.int32, + device=device, + ) + max_len = int(max(padded_seg_lens)) if padded_seg_lens else 0 + return PackedSeqParams( + cu_seqlens_q=cu, + cu_seqlens_q_padded=cu_padded, + cu_seqlens_kv=cu, + cu_seqlens_kv_padded=cu_padded, + max_seqlen_q=max_len, + max_seqlen_kv=max_len, + qkv_format='thd', + ) + + +def _cosine_sim(a: torch.Tensor, b: torch.Tensor) -> float: + """Return cosine similarity between two tensors as a Python float.""" + return F.cosine_similarity( + a.flatten().double().unsqueeze(0), b.flatten().double().unsqueeze(0) + ).item() + + +def _tensor_sim(a: torch.Tensor, b: torch.Tensor) -> float: + """Return scale-invariant tensor similarity between two tensors.""" + a, b = a.double(), b.double() + denom = (a * a + b * b).sum() + return (2.0 * (a * b).sum() / denom).item() if denom else 1.0 + + +def _assert_cp_tensor_match(actual: torch.Tensor, expected: torch.Tensor, label: str): + """Assert CP output matches reference by similarity metrics.""" + assert ( + actual.shape == expected.shape + ), f"{label}: shape {tuple(actual.shape)} != {tuple(expected.shape)}" + assert torch.isfinite(actual).all(), f"{label}: actual has non-finite values" + assert torch.isfinite(expected).all(), f"{label}: expected has non-finite values" + + diff = (actual - expected).abs() + max_abs = diff.max().item() if diff.numel() else 0.0 + cosine_sim = _cosine_sim(actual, expected) + tensor_sim = _tensor_sim(actual, expected) + actual_norm = actual.double().norm().item() + expected_norm = expected.double().norm().item() + assert cosine_sim > 1 - _DSV4_CP_PARITY_EPS, ( + f"{label}: cosine_sim={cosine_sim:.10f}, " + f"tensor_sim={tensor_sim:.10f}, max_abs={max_abs:.6e}, " + f"actual_norm={actual_norm:.6e}, expected_norm={expected_norm:.6e}, " + f"eps={_DSV4_CP_PARITY_EPS}" + ) + assert tensor_sim > 1 - _DSV4_CP_PARITY_EPS, ( + f"{label}: tensor_sim={tensor_sim:.10f}, " + f"cosine_sim={cosine_sim:.10f}, max_abs={max_abs:.6e}, " + f"actual_norm={actual_norm:.6e}, expected_norm={expected_norm:.6e}, " + f"eps={_DSV4_CP_PARITY_EPS}" + ) + + +def _assert_cp_graph_bitwise_match(actual: torch.Tensor, expected: torch.Tensor, label: str): + """Assert graph replay output is bitwise equal to eager output.""" + assert ( + actual.shape == expected.shape + ), f"{label}: shape {tuple(actual.shape)} != {tuple(expected.shape)}" + if torch.equal(actual, expected): + return + diff = (actual - expected).abs() + max_abs = diff.max().item() if diff.numel() else 0.0 + cosine_sim = _cosine_sim(actual, expected) + tensor_sim = _tensor_sim(actual, expected) + rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else -1 + print( + f"[rank{rank}] {label}: graph/eager not bitwise; max_abs={max_abs:.6e}, " + f"cosine_sim={cosine_sim:.10f}, tensor_sim={tensor_sim:.10f}", + flush=True, + ) + raise AssertionError( + f"{label}: graph/eager must be bitwise equal; max_abs={max_abs:.6e}, " + f"cosine_sim={cosine_sim:.10f}, tensor_sim={tensor_sim:.10f}" + ) + + +def _assert_cp_graph_fused_grad_match(actual: torch.Tensor, expected: torch.Tensor, label: str): + """Assert fused graph gradients match eager within fused-kernel tolerances.""" + assert ( + actual.shape == expected.shape + ), f"{label}: shape {tuple(actual.shape)} != {tuple(expected.shape)}" + assert actual.dtype == expected.dtype, f"{label}: dtype {actual.dtype} != {expected.dtype}" + assert torch.isfinite(actual).all(), f"{label}: actual has non-finite values" + assert torch.isfinite(expected).all(), f"{label}: expected has non-finite values" + + diff = (actual.float() - expected.float()).abs() + max_abs = diff.max().item() if diff.numel() else 0.0 + cosine_sim = _cosine_sim(actual, expected) + tensor_sim = _tensor_sim(actual, expected) + assert cosine_sim > 1 - _DSV4_CP_GRAPH_FUSED_SIM_EPS, ( + f"{label}: cosine_sim={cosine_sim:.10f}, " + f"max_abs={max_abs:.6e}, eps={_DSV4_CP_GRAPH_FUSED_SIM_EPS}" + ) + assert tensor_sim > 1 - _DSV4_CP_GRAPH_FUSED_SIM_EPS, ( + f"{label}: tensor_sim={tensor_sim:.10f}, " + f"max_abs={max_abs:.6e}, eps={_DSV4_CP_GRAPH_FUSED_SIM_EPS}" + ) + + if actual.dtype == torch.bfloat16: + atol = _DSV4_CP_GRAPH_FUSED_BF16_ATOL + elif actual.dtype == torch.float32: + atol = _DSV4_CP_GRAPH_FUSED_FP32_ATOL + else: + raise AssertionError( + f"{label}: unsupported dtype for fused graph close check: {actual.dtype}" + ) + try: + torch.testing.assert_close( + actual, + expected, + rtol=_DSV4_CP_GRAPH_FUSED_RTOL, + atol=atol, + msg=( + f"{label}: fused graph/eager gradient mismatch; " + f"rtol={_DSV4_CP_GRAPH_FUSED_RTOL}, atol={atol}, " + f"cosine_sim={cosine_sim:.10f}, tensor_sim={tensor_sim:.10f}, " + f"max_abs={max_abs:.6e}" + ), + ) + except AssertionError: + rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else -1 + print( + f"[rank{rank}] {label}: fused gradient close failed; " + f"rtol={_DSV4_CP_GRAPH_FUSED_RTOL}, atol={atol}, " + f"cosine_sim={cosine_sim:.10f}, tensor_sim={tensor_sim:.10f}, " + f"max_abs={max_abs:.6e}", + flush=True, + ) + raise + + +@contextmanager +def _deterministic_torch_algorithms(): + """Temporarily enable deterministic PyTorch algorithms.""" + old_enabled = torch.are_deterministic_algorithms_enabled() + old_warn_only = torch.is_deterministic_algorithms_warn_only_enabled() + torch.use_deterministic_algorithms(True) + try: + yield + finally: + torch.use_deterministic_algorithms(old_enabled, warn_only=old_warn_only) + + +def _make_dsv4_cp_config( + *, + context_parallel_size, + dsa_indexer_loss_coeff=1.0, + dsa_indexer_use_sparse_loss=True, + apply_dsa_kernel_fusion=True, + apply_rope_fusion=True, + csa_cp_partition_mode=DSV4_CP_PARTITION_CONTIGUOUS, +): + """Build the DSv4 flash attention config used by CP tests.""" + shape = _DSV4_VARIANTS[_DSV4_CP_TEST_VARIANT] + return _make_config( + hidden_size=shape["hidden_size"], + num_attention_heads=shape["num_attention_heads"], + v_head_dim=shape["v_head_dim"], + qk_pos_emb_head_dim=shape["qk_pos_emb_head_dim"], + q_lora_rank=shape["q_lora_rank"], + o_groups=shape["o_groups"], + o_lora_rank=shape["o_lora_rank"], + csa_compress_ratios=[0, 4, 128, 4], + csa_window_size=128, + dsa_indexer_n_heads=64, + dsa_indexer_head_dim=128, + dsa_indexer_topk=shape["dsa_indexer_topk"], + dsa_indexer_loss_coeff=dsa_indexer_loss_coeff, + dsa_indexer_use_sparse_loss=dsa_indexer_use_sparse_loss, + context_parallel_size=context_parallel_size, + csa_dense_mode=False, + csa_compress_rotary_base=shape["csa_compress_rotary_base"], + layernorm_epsilon=1e-6, + normalization="RMSNorm", + qk_layernorm=True, + layernorm_zero_centered_gamma=False, + expert_model_parallel_size=1, + apply_dsa_kernel_fusion=apply_dsa_kernel_fusion, + apply_rope_fusion=apply_rope_fusion, + csa_cp_partition_mode=csa_cp_partition_mode, + ) + + +def _copy_module_parameters(src, dst): + """Copy named parameters from ``src`` into ``dst``.""" + src_params = dict(src.named_parameters()) + for name, param in dst.named_parameters(): + assert name in src_params + param.data.copy_(src_params[name].data) + + +def _make_cp_partition_indices(partition_mode, padded_total_tokens, cp_size, device='cuda'): + """Return local row indices for every CP rank.""" + return tuple( + thd_cp_local_row_indices(partition_mode, padded_total_tokens, cp_size, rank, device) + for rank in range(cp_size) + ) + + +def _make_ragged_cp_case(partition_mode, cp_size, cp_rank): + """Build the ragged THD packed params and local rows for one CP rank.""" + padded_total_tokens = sum(_DSV4_CP_RAGGED_PADDED_SEG_LENS) + packed = _make_thd_packed_seq_params(_DSV4_CP_RAGGED_SEG_LENS, _DSV4_CP_RAGGED_PADDED_SEG_LENS) + partition_indices = _make_cp_partition_indices(partition_mode, padded_total_tokens, cp_size) + return packed, padded_total_tokens, partition_indices[cp_rank] + + +def _make_hidden_and_grad(padded_total_tokens, hidden_size): + """Create matching hidden states and output gradients for a layer run.""" + hidden = torch.randn(padded_total_tokens, 1, hidden_size, dtype=torch.bfloat16, device='cuda') + return hidden, torch.randn_like(hidden) + + +def _run_dsv4_attention_forward_backward( + attn, hidden, grad, packed_seq_params, *, collect_result=True +): + """Run one eager DSv4 attention forward/backward pass.""" + hidden.grad = None + attn.zero_grad(set_to_none=True) + output, _ = attn(hidden_states=hidden, attention_mask=None, packed_seq_params=packed_seq_params) + output.backward(grad) + if not collect_result: + return None + param_grads = { + name: param.grad.detach().clone() + for name, param in attn.named_parameters() + if param.grad is not None + } + return output.detach().clone(), hidden.grad.detach().clone(), param_grads + + +def _capture_dsv4_attention_forward_backward(attn, static_hidden, static_grad, packed_seq_params): + """Capture a DSv4 attention forward/backward pass in a CUDA graph.""" + warmup_stream = torch.cuda.Stream() + warmup_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup_stream): + for _ in range(3): + _run_dsv4_attention_forward_backward( + attn, static_hidden, static_grad, packed_seq_params, collect_result=False + ) + torch.cuda.current_stream().wait_stream(warmup_stream) + torch.cuda.synchronize() + + static_hidden.grad = None + attn.zero_grad(set_to_none=True) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, capture_error_mode="thread_local"): + graph_output, _ = attn( + hidden_states=static_hidden, attention_mask=None, packed_seq_params=packed_seq_params + ) + graph_output.backward(static_grad) + torch.cuda.synchronize() + return graph, graph_output + + +def _zero_existing_grads(attn, hidden): + """Zero existing parameter and hidden gradients in place.""" + if hidden.grad is not None: + hidden.grad.zero_() + for param in attn.parameters(): + if param.grad is not None: + param.grad.zero_() + + +def _run_dsv4_attention_forward_backward_reuse_grads(attn, hidden, grad, packed_seq_params): + """Run eager forward/backward while reusing existing grad buffers.""" + _zero_existing_grads(attn, hidden) + output, _ = attn(hidden_states=hidden, attention_mask=None, packed_seq_params=packed_seq_params) + output.backward(grad) + return output + + +def _max_int_across_world(value: int) -> int: + """Return the max integer value across distributed ranks.""" + tensor = torch.tensor(value, dtype=torch.int64, device='cuda') + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(tensor, op=dist.ReduceOp.MAX) + return int(tensor.item()) + + +def _max_float_across_world(value: float) -> float: + """Return the max floating-point value across distributed ranks.""" + tensor = torch.tensor(value, dtype=torch.float64, device='cuda') + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(tensor, op=dist.ReduceOp.MAX) + return float(tensor.item()) + + +def _measure_peak_allocated_delta(attn, hidden, grad, packed_seq_params) -> int: + """Measure post-warmup peak allocated memory delta for one layer pass.""" + attn.train() + for _ in range(_DSV4_CP_MEMORY_WARMUP_ITERS): + _run_dsv4_attention_forward_backward_reuse_grads(attn, hidden, grad, packed_seq_params) + torch.cuda.synchronize() + _zero_existing_grads(attn, hidden) + torch.cuda.reset_peak_memory_stats() + memory_start = torch.cuda.memory_allocated() + _run_dsv4_attention_forward_backward_reuse_grads(attn, hidden, grad, packed_seq_params) + torch.cuda.synchronize() + return _max_int_across_world(torch.cuda.max_memory_allocated() - memory_start) + + +def _measure_cuda_graph_time(attn, static_hidden, static_grad, packed_seq_params) -> float: + """Measure slowest-rank median CUDA graph replay time.""" + attn.train() + graph, _graph_output = _capture_dsv4_attention_forward_backward( + attn, static_hidden, static_grad, packed_seq_params + ) + for _ in range(_DSV4_CP_GRAPH_TIMING_WARMUP_REPLAYS): + graph.replay() + torch.cuda.synchronize() + + timings = [] + for _ in range(_DSV4_CP_GRAPH_TIMING_MEASURE_REPLAYS): + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + start_event.record() + graph.replay() + end_event.record() + torch.cuda.synchronize() + timings.append(float(start_event.elapsed_time(end_event))) + return _max_float_across_world(statistics.median(timings)) + + +def _format_scale_report( + metric, layer_number, cp_size, partition_mode, baseline, measured, scale, limit +): + """Format a CP-vs-CP1 scale report line.""" + return ( + f"DSv4 THD CP {metric} scale layer={layer_number} cp{cp_size}: " + f"partition={partition_mode}, " + f"cp1={baseline:.6f}, cp{cp_size}={measured:.6f}, " + f"scale={scale:.6f}, limit={limit:.6f}" + ) + + +def _clear_cuda_test_state(): + """Synchronize and release cached CUDA memory between cases.""" + torch.cuda.synchronize() + gc.collect() + torch.cuda.empty_cache() + torch.cuda.synchronize() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available") +class TestDSv4HybridAttentionTHDCP: + """CP-sliced THD attention should match full-sequence THD reference output and gradients.""" + + @pytest.fixture(scope='class', autouse=True, params=(2, 4), ids=lambda cp: f"cp{cp}") + def setup_method(self, request): + """Initialize model-parallel groups for each CP size.""" + cp_size = request.param + if Utils.world_size < cp_size: + pytest.skip(f"THD CP path test requires at least {cp_size} distributed ranks") + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + context_parallel_size=cp_size, + ) + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + cls = request.cls + cls.cp_size = cp_size + cls.cp_rank = parallel_state.get_context_parallel_rank() + cls.pg = ProcessGroupCollection.use_mpu_process_groups() + # Tradeoff: reuse the current model-parallel groups and only disable CP for the + # full-reference path to avoid extra group initialization. This is valid for + # these CP-only tests; rebuild the reference groups if future tests add other + # parallel dimensions. + cls.ref_pg = ProcessGroupCollection.use_mpu_process_groups() + cls.ref_pg.cp = _ReferenceCPGroup() + cls._cp1_memory_delta_cache = {} + cls._cp1_graph_time_cache = {} + + yield + _clear_cuda_test_state() + Utils.destroy_model_parallel() + + @pytest.fixture(autouse=True) + def clear_cuda_test_case(self): + """Clear CUDA state before and after each test case.""" + _clear_cuda_test_state() + yield + _clear_cuda_test_state() + + def _measure_cp1_peak_allocated_delta(self, layer_number, packed, padded_tokens): + """Measure and cache CP1 peak memory for a layer.""" + cached = self._cp1_memory_delta_cache.get(layer_number) + if cached is not None: + return cached + + torch.manual_seed(_SEED + 900 + layer_number) + model_parallel_cuda_manual_seed(_SEED + 900 + layer_number) + config = _make_dsv4_cp_config( + context_parallel_size=1, + dsa_indexer_loss_coeff=1.0, + dsa_indexer_use_sparse_loss=True, + apply_rope_fusion=True, + ) + attn = _build_attention(config, layer_number=layer_number, pg_collection=self.ref_pg).cuda() + hidden_values, grad_values = _make_hidden_and_grad(padded_tokens, config.hidden_size) + hidden = hidden_values.detach().clone().requires_grad_(True) + grad = grad_values.detach().clone() + measured = _measure_peak_allocated_delta(attn, hidden, grad, packed) + + del attn, hidden_values, grad_values, hidden, grad + _clear_cuda_test_state() + self._cp1_memory_delta_cache[layer_number] = measured + return measured + + def _measure_cp1_cuda_graph_time(self, layer_number, packed, padded_tokens): + """Measure and cache CP1 CUDA graph replay time for a layer.""" + cached = self._cp1_graph_time_cache.get(layer_number) + if cached is not None: + return cached + + torch.manual_seed(_SEED + 1000 + layer_number) + model_parallel_cuda_manual_seed(_SEED + 1000 + layer_number) + config = _make_dsv4_cp_config( + context_parallel_size=1, + dsa_indexer_loss_coeff=1.0, + dsa_indexer_use_sparse_loss=True, + apply_rope_fusion=True, + ) + attn = _build_attention(config, layer_number=layer_number, pg_collection=self.ref_pg).cuda() + hidden_values, grad_values = _make_hidden_and_grad(padded_tokens, config.hidden_size) + hidden = hidden_values.detach().clone().requires_grad_(True) + grad = grad_values.detach().clone() + measured = _measure_cuda_graph_time(attn, hidden, grad, packed) + + del attn, hidden_values, grad_values, hidden, grad + _clear_cuda_test_state() + self._cp1_graph_time_cache[layer_number] = measured + return measured + + @pytest.mark.parametrize( + "layer_number", + [1, 2, 3], + ids=["ratio_0_window_only", "ratio_4_indexer", "ratio_128_compressor"], + ) + @pytest.mark.parametrize( + "partition_mode", + [DSV4_CP_PARTITION_CONTIGUOUS, DSV4_CP_PARTITION_TWO_CHUNK], + ids=["contiguous", "two_chunk"], + ) + def test_thd_cp_matches_full_reference_forward_backward(self, layer_number, partition_mode): + """CP path matches the full-sequence THD reference on ragged + packed inputs using the DSv4 layer configuration. + + Verifies local CP output, hidden grad, and reduced parameter grads + against the sliced full-reference tensors. The test body is shared by + contiguous and two-chunk CP; only the partition mode and + local row indices differ. + """ + packed, padded_tokens, local_idx = _make_ragged_cp_case( + partition_mode, self.cp_size, self.cp_rank + ) + + torch.manual_seed(_SEED + layer_number) + model_parallel_cuda_manual_seed(_SEED + layer_number) + config_cp = _make_dsv4_cp_config( + context_parallel_size=self.cp_size, + dsa_indexer_loss_coeff=1.0, + dsa_indexer_use_sparse_loss=True, + apply_rope_fusion=True, + csa_cp_partition_mode=partition_mode, + ) + config_ref = _make_dsv4_cp_config( + context_parallel_size=1, + dsa_indexer_loss_coeff=1.0, + dsa_indexer_use_sparse_loss=True, + # Numerical parity uses the unfused CP1 reference so failures point + # at the CP path instead of non-CP fused RoPE behavior. + apply_rope_fusion=False, + ) + cp_attn = _build_attention( + config_cp, layer_number=layer_number, pg_collection=self.pg + ).cuda() + ref_attn = _build_attention( + config_ref, layer_number=layer_number, pg_collection=self.ref_pg + ).cuda() + _copy_module_parameters(cp_attn, ref_attn) + + full_hidden = torch.randn( + padded_tokens, 1, config_cp.hidden_size, dtype=torch.bfloat16, device='cuda' + ) + local_hidden = full_hidden.index_select(0, local_idx).detach().clone().requires_grad_(True) + ref_hidden = full_hidden.detach().clone().requires_grad_(True) + + local_out, _ = cp_attn( + hidden_states=local_hidden, attention_mask=None, packed_seq_params=packed + ) + ref_out, _ = ref_attn( + hidden_states=ref_hidden, attention_mask=None, packed_seq_params=packed + ) + _assert_cp_tensor_match( + local_out.detach(), + ref_out.detach().index_select(0, local_idx), + f"layer={layer_number}:{partition_mode}:output", + ) + + grad = torch.randn_like(ref_out) + local_out.backward(grad.index_select(0, local_idx)) + ref_out.backward(grad) + _assert_cp_tensor_match( + local_hidden.grad.detach(), + ref_hidden.grad.index_select(0, local_idx), + f"layer={layer_number}:{partition_mode}:hidden_grad", + ) + + ref_params = dict(ref_attn.named_parameters()) + for name, param in cp_attn.named_parameters(): + ref_grad = ref_params[name].grad + assert param.grad is not None, f"Missing CP grad for {name}" + assert ref_grad is not None, f"Missing reference grad for {name}" + grad_sum = param.grad.detach().clone() + dist.all_reduce(grad_sum, group=self.pg.cp) + _assert_cp_tensor_match( + grad_sum, ref_grad, f"layer={layer_number}:{partition_mode}:param_grad:{name}" + ) + + del cp_attn, ref_attn, full_hidden, local_hidden, ref_hidden, local_out, ref_out, grad + _clear_cuda_test_state() + + @pytest.mark.parametrize( + "layer_number", + [1, 2, 3], + ids=["ratio_0_window_only", "ratio_4_indexer", "ratio_128_compressor"], + ) + @pytest.mark.parametrize("fused", [True, False]) + @pytest.mark.parametrize( + "partition_mode", + [DSV4_CP_PARTITION_CONTIGUOUS, DSV4_CP_PARTITION_TWO_CHUNK], + ids=["contiguous", "two_chunk"], + ) + def test_thd_cp_cuda_graph_matches_eager_forward_backward( + self, layer_number, fused, partition_mode + ): + """CUDA graph replay matches eager THD CP forward/backward. + + Captures the DSv4 attention layer's CP-local forward and backward + graph, replays it with fresh static-buffer contents, and compares the + local output, hidden grad, and parameter grads against an eager module + with identical weights. The unfused graph path is deterministic and + must match bitwise for both forward and backward. Fused forward outputs + are expected to match bitwise. Fused gradients use strict similarity + plus elementwise ``assert_close`` gates because they may be + nondeterministic against eager execution. + """ + context = nullcontext() if fused else _deterministic_torch_algorithms() + mode = "fused" if fused else "unfused" + with context: + packed, padded_tokens, local_idx = _make_ragged_cp_case( + partition_mode, self.cp_size, self.cp_rank + ) + + torch.manual_seed(_SEED + 700 + layer_number) + model_parallel_cuda_manual_seed(_SEED + 700 + layer_number) + config = _make_dsv4_cp_config( + context_parallel_size=self.cp_size, + dsa_indexer_loss_coeff=1.0, + dsa_indexer_use_sparse_loss=True, + apply_dsa_kernel_fusion=fused, + apply_rope_fusion=True, + csa_cp_partition_mode=partition_mode, + ) + graph_attn = _build_attention( + config, layer_number=layer_number, pg_collection=self.pg + ).cuda() + eager_attn = _build_attention( + config, layer_number=layer_number, pg_collection=self.pg + ).cuda() + graph_attn.train() + eager_attn.train() + _copy_module_parameters(graph_attn, eager_attn) + + full_hidden = torch.randn( + padded_tokens, 1, config.hidden_size, dtype=torch.bfloat16, device='cuda' + ) + test_hidden = full_hidden.index_select(0, local_idx).detach().clone() + test_grad = torch.randn_like(test_hidden) + static_hidden = test_hidden.detach().clone().requires_grad_(True) + eager_hidden = test_hidden.detach().clone().requires_grad_(True) + static_grad = test_grad.detach().clone() + + graph, graph_output = _capture_dsv4_attention_forward_backward( + graph_attn, static_hidden, static_grad, packed + ) + with torch.no_grad(): + static_hidden.copy_(test_hidden) + static_grad.copy_(test_grad) + if static_hidden.grad is not None: + static_hidden.grad.zero_() + for param in graph_attn.parameters(): + if param.grad is not None: + param.grad.zero_() + graph.replay() + torch.cuda.synchronize() + graph_out = graph_output.detach().clone() + graph_hidden_grad = static_hidden.grad.detach().clone() + graph_param_grads = { + name: param.grad.detach().clone() + for name, param in graph_attn.named_parameters() + if param.grad is not None + } + + eager_out, eager_hidden_grad, eager_param_grads = _run_dsv4_attention_forward_backward( + eager_attn, eager_hidden, test_grad, packed + ) + torch.cuda.synchronize() + assert graph_param_grads.keys() == eager_param_grads.keys() + output_label = f"layer={layer_number}:{mode}:{partition_mode}:output" + _assert_cp_graph_bitwise_match(graph_out, eager_out, output_label) + bwd_match_fn = ( + _assert_cp_graph_fused_grad_match if fused else _assert_cp_graph_bitwise_match + ) + bwd_match_fn( + graph_hidden_grad, + eager_hidden_grad, + f"layer={layer_number}:{mode}:{partition_mode}:hidden_grad", + ) + for name, graph_grad in graph_param_grads.items(): + bwd_match_fn( + graph_grad, + eager_param_grads[name], + f"layer={layer_number}:{mode}:{partition_mode}:param_grad:{name}", + ) + + del graph, graph_output + del graph_attn, eager_attn, full_hidden, test_hidden, test_grad, static_hidden + del eager_hidden, static_grad, graph_out, graph_hidden_grad + _clear_cuda_test_state() + + @pytest.mark.parametrize( + "layer_number", [2, 3], ids=["ratio_4_indexer", "ratio_128_compressor"] + ) + @pytest.mark.parametrize( + "partition_mode", + [DSV4_CP_PARTITION_CONTIGUOUS, DSV4_CP_PARTITION_TWO_CHUNK], + ids=["contiguous", "two_chunk"], + ) + def test_thd_cp_cuda_graph_replay_accepts_changed_padded_boundaries( + self, layer_number, partition_mode + ): + """CUDA graph replay uses updated device cu_seqlens_padded values. + + The graph is captured with one padded THD layout, then replayed after + overwriting the same PackedSeqParams metadata tensors with another + layout whose tensor shapes, padded total, and max seqlen are unchanged. + Expected: replay matches eager execution with the new metadata. A + failure means the CP path baked capture-time padded boundaries or a + host-derived dynamic shape into the graph. + """ + capture_packed = _make_thd_packed_seq_params( + _DSV4_CP_RAGGED_SEG_LENS, _DSV4_CP_RAGGED_PADDED_SEG_LENS + ) + replay_packed = _make_thd_packed_seq_params( + _DSV4_CP_RAGGED_SEG_LENS, _DSV4_CP_REPLAY_PADDED_SEG_LENS + ) + padded_tokens = sum(_DSV4_CP_RAGGED_PADDED_SEG_LENS) + assert padded_tokens == sum(_DSV4_CP_REPLAY_PADDED_SEG_LENS) + assert capture_packed.cu_seqlens_q.shape == replay_packed.cu_seqlens_q.shape + assert capture_packed.cu_seqlens_q_padded.shape == replay_packed.cu_seqlens_q_padded.shape + assert capture_packed.max_seqlen_q == replay_packed.max_seqlen_q + assert capture_packed.max_seqlen_kv == replay_packed.max_seqlen_kv + assert not torch.equal( + capture_packed.cu_seqlens_q_padded, replay_packed.cu_seqlens_q_padded + ) + local_idx = _make_cp_partition_indices(partition_mode, padded_tokens, self.cp_size)[ + self.cp_rank + ] + + torch.manual_seed(_SEED + 1100 + layer_number) + model_parallel_cuda_manual_seed(_SEED + 1100 + layer_number) + config = _make_dsv4_cp_config( + context_parallel_size=self.cp_size, + dsa_indexer_loss_coeff=1.0, + dsa_indexer_use_sparse_loss=True, + apply_dsa_kernel_fusion=True, + apply_rope_fusion=True, + csa_cp_partition_mode=partition_mode, + ) + graph_attn = _build_attention( + config, layer_number=layer_number, pg_collection=self.pg + ).cuda() + eager_attn = _build_attention( + config, layer_number=layer_number, pg_collection=self.pg + ).cuda() + graph_attn.train() + eager_attn.train() + _copy_module_parameters(graph_attn, eager_attn) + + capture_full_hidden = torch.randn( + padded_tokens, 1, config.hidden_size, dtype=torch.bfloat16, device='cuda' + ) + replay_full_hidden = torch.randn_like(capture_full_hidden) + capture_hidden = capture_full_hidden.index_select(0, local_idx).detach().clone() + replay_hidden = replay_full_hidden.index_select(0, local_idx).detach().clone() + capture_grad = torch.randn_like(capture_hidden) + replay_grad = torch.randn_like(replay_hidden) + static_hidden = capture_hidden.detach().clone().requires_grad_(True) + static_grad = capture_grad.detach().clone() + + graph, graph_output = _capture_dsv4_attention_forward_backward( + graph_attn, static_hidden, static_grad, capture_packed + ) + with torch.no_grad(): + capture_packed.cu_seqlens_q.copy_(replay_packed.cu_seqlens_q) + capture_packed.cu_seqlens_kv.copy_(replay_packed.cu_seqlens_kv) + capture_packed.cu_seqlens_q_padded.copy_(replay_packed.cu_seqlens_q_padded) + capture_packed.cu_seqlens_kv_padded.copy_(replay_packed.cu_seqlens_kv_padded) + static_hidden.copy_(replay_hidden) + static_grad.copy_(replay_grad) + if static_hidden.grad is not None: + static_hidden.grad.zero_() + for param in graph_attn.parameters(): + if param.grad is not None: + param.grad.zero_() + graph.replay() + torch.cuda.synchronize() + graph_out = graph_output.detach().clone() + graph_hidden_grad = static_hidden.grad.detach().clone() + graph_param_grads = { + name: param.grad.detach().clone() + for name, param in graph_attn.named_parameters() + if param.grad is not None + } + + eager_hidden = replay_hidden.detach().clone().requires_grad_(True) + eager_out, eager_hidden_grad, eager_param_grads = _run_dsv4_attention_forward_backward( + eager_attn, eager_hidden, replay_grad, replay_packed + ) + torch.cuda.synchronize() + assert graph_param_grads.keys() == eager_param_grads.keys() + _assert_cp_graph_bitwise_match( + graph_out, eager_out, f"layer={layer_number}:metadata_replay:{partition_mode}:output" + ) + _assert_cp_graph_fused_grad_match( + graph_hidden_grad, + eager_hidden_grad, + f"layer={layer_number}:metadata_replay:{partition_mode}:hidden_grad", + ) + for name, graph_grad in graph_param_grads.items(): + _assert_cp_graph_fused_grad_match( + graph_grad, + eager_param_grads[name], + f"layer={layer_number}:metadata_replay:{partition_mode}:param_grad:{name}", + ) + + del graph, graph_output, graph_attn, eager_attn, capture_full_hidden, replay_full_hidden + del capture_hidden, replay_hidden, capture_grad, replay_grad, static_hidden, static_grad + del graph_out, graph_hidden_grad, eager_hidden + _clear_cuda_test_state() + + @pytest.mark.parametrize( + "layer_number", + [1, 2, 3], + ids=["ratio_0_window_only", "ratio_4_indexer", "ratio_128_compressor"], + ) + @pytest.mark.parametrize( + "partition_mode", + [DSV4_CP_PARTITION_CONTIGUOUS, DSV4_CP_PARTITION_TWO_CHUNK], + ids=["contiguous", "two_chunk"], + ) + def test_thd_cp_peak_allocated_delta_scales_vs_cp1(self, layer_number, partition_mode): + """CP-local THD forward/backward peak allocated delta scales down vs CP1. + + This is a black-box guard against implementations that pass parity but + accidentally gather full hidden states or otherwise allocate global + activation-sized buffers before the layer work. + """ + packed, padded_tokens, local_idx = _make_ragged_cp_case( + partition_mode, self.cp_size, self.cp_rank + ) + + ref_delta = self._measure_cp1_peak_allocated_delta(layer_number, packed, padded_tokens) + + torch.manual_seed(_SEED + 900 + layer_number) + model_parallel_cuda_manual_seed(_SEED + 900 + layer_number) + config_cp = _make_dsv4_cp_config( + context_parallel_size=self.cp_size, + dsa_indexer_loss_coeff=1.0, + dsa_indexer_use_sparse_loss=True, + apply_rope_fusion=True, + csa_cp_partition_mode=partition_mode, + ) + cp_attn = _build_attention( + config_cp, layer_number=layer_number, pg_collection=self.pg + ).cuda() + full_hidden_values, full_grad_values = _make_hidden_and_grad( + padded_tokens, config_cp.hidden_size + ) + local_hidden = ( + full_hidden_values.index_select(0, local_idx).detach().clone().requires_grad_(True) + ) + local_grad = full_grad_values.index_select(0, local_idx).detach().clone() + del full_hidden_values, full_grad_values + cp_delta = _measure_peak_allocated_delta(cp_attn, local_hidden, local_grad, packed) + + limit = _DSV4_CP_MEMORY_RATIO_LIMITS[self.cp_size] + scale = float(cp_delta) / float(ref_delta) + report = _format_scale_report( + "memory_delta_mib", + layer_number, + self.cp_size, + partition_mode, + ref_delta / (1024.0**2), + cp_delta / (1024.0**2), + scale, + limit, + ) + print(report) + assert scale <= limit, report + + del cp_attn, local_hidden, local_grad + _clear_cuda_test_state() + + @pytest.mark.parametrize( + "layer_number", + [1, 2, 3], + ids=["ratio_0_window_only", "ratio_4_indexer", "ratio_128_compressor"], + ) + @pytest.mark.parametrize( + "partition_mode", + [DSV4_CP_PARTITION_CONTIGUOUS, DSV4_CP_PARTITION_TWO_CHUNK], + ids=["contiguous", "two_chunk"], + ) + def test_thd_cp_cuda_graph_time_scales_vs_cp1(self, layer_number, partition_mode): + """CP-local CUDA graph replay time scales down vs CP1 on the same machine. + + This is a black-box performance guard: correctness and memory scaling + are not enough if the implementation falls back to slow dynamic PyTorch + work in the production graph path. + """ + packed, padded_tokens, local_idx = _make_ragged_cp_case( + partition_mode, self.cp_size, self.cp_rank + ) + + ref_time_ms = self._measure_cp1_cuda_graph_time(layer_number, packed, padded_tokens) + + torch.manual_seed(_SEED + 1000 + layer_number) + model_parallel_cuda_manual_seed(_SEED + 1000 + layer_number) + config_cp = _make_dsv4_cp_config( + context_parallel_size=self.cp_size, + dsa_indexer_loss_coeff=1.0, + dsa_indexer_use_sparse_loss=True, + apply_rope_fusion=True, + csa_cp_partition_mode=partition_mode, + ) + cp_attn = _build_attention( + config_cp, layer_number=layer_number, pg_collection=self.pg + ).cuda() + full_hidden_values, full_grad_values = _make_hidden_and_grad( + padded_tokens, config_cp.hidden_size + ) + local_hidden = ( + full_hidden_values.index_select(0, local_idx).detach().clone().requires_grad_(True) + ) + local_grad = full_grad_values.index_select(0, local_idx).detach().clone() + del full_hidden_values, full_grad_values + cp_time_ms = _measure_cuda_graph_time(cp_attn, local_hidden, local_grad, packed) + + limit = _DSV4_CP_GRAPH_TIME_RATIO_LIMITS[self.cp_size] + scale = cp_time_ms / ref_time_ms + report = _format_scale_report( + "graph_time_ms", + layer_number, + self.cp_size, + partition_mode, + ref_time_ms, + cp_time_ms, + scale, + limit, + ) + print(report) + assert scale <= limit, report + + del cp_attn, local_hidden, local_grad + _clear_cuda_test_state() From 8a271beb67cadd9de9d9eebd329ec94db52afb10 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Fri, 26 Jun 2026 05:43:08 -0700 Subject: [PATCH 02/23] Fix DSv4 CP lint formatting Signed-off-by: kunlunl --- megatron/core/datasets/data_schedule_utils.py | 4 +- .../csa_cp_layout_kernels.py | 87 +++++-------------- .../csa_cp_utils.py | 11 +-- .../dsa_kernels.py | 6 +- .../test_csa_cp_layout_kernels.py | 24 +++-- .../test_csa_cp_utils.py | 42 +++++---- 6 files changed, 65 insertions(+), 109 deletions(-) diff --git a/megatron/core/datasets/data_schedule_utils.py b/megatron/core/datasets/data_schedule_utils.py index 0f30a28d73c..0f4125baa08 100644 --- a/megatron/core/datasets/data_schedule_utils.py +++ b/megatron/core/datasets/data_schedule_utils.py @@ -30,8 +30,8 @@ def get_cp_slice_for_thd( keys: Sequence data keys to slice. Defaults to the original THD data tensors. csa_cp_partition_mode: Optional CSA CP partition mode for DSv4 hybrid attention. partition_total_tokens: Optional padded total used only for choosing CP row indices. - When set, tensors selected by ``keys`` are tail-padded to this length before slicing. Existing - cu_seqlens metadata is left unchanged. + When set, tensors selected by ``keys`` are tail-padded to this length before + slicing. Existing cu_seqlens metadata is left unchanged. """ cp_size = cp_group.size() if cp_size <= 1: diff --git a/megatron/core/transformer/experimental_attention_variant/csa_cp_layout_kernels.py b/megatron/core/transformer/experimental_attention_variant/csa_cp_layout_kernels.py index ebf0df3ac6f..61579c34661 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa_cp_layout_kernels.py +++ b/megatron/core/transformer/experimental_attention_variant/csa_cp_layout_kernels.py @@ -172,16 +172,8 @@ def _thd_local_rope_kernel( copy_vecs_per_head = nope_dim // 8 copy_use_u256 = True elif cutlass.const_expr( - ( - x.element_type.width == 16 - and nope_dim % 4 == 0 - and head_dim % 4 == 0 - ) - or ( - x.element_type.width == 32 - and nope_dim % 2 == 0 - and head_dim % 2 == 0 - ) + (x.element_type.width == 16 and nope_dim % 4 == 0 and head_dim % 4 == 0) + or (x.element_type.width == 32 and nope_dim % 2 == 0 and head_dim % 2 == 0) ): copy_use_u256 = False if cutlass.const_expr(x.element_type.width == 16): @@ -201,16 +193,8 @@ def _thd_local_rope_kernel( head = copy_work // copy_vecs_per_head offset = row * row_width + head * head_dim + unit * copy_vec_elems if cutlass.const_expr( - ( - x.element_type.width == 16 - and nope_dim % 4 == 0 - and head_dim % 4 == 0 - ) - or ( - x.element_type.width == 32 - and nope_dim % 2 == 0 - and head_dim % 2 == 0 - ) + (x.element_type.width == 16 and nope_dim % 4 == 0 and head_dim % 4 == 0) + or (x.element_type.width == 32 and nope_dim % 2 == 0 and head_dim % 2 == 0) ): if cutlass.const_expr(copy_use_u256): src_ptr = _ptr_as_i64(x, offset) @@ -226,9 +210,7 @@ def _thd_local_rope_kernel( level1_eviction_priority="evict_no_allocate", ) cute.arch.store( - dst_ptr.llvm_ptr, - value, - level1_eviction_priority="evict_first", + dst_ptr.llvm_ptr, value, level1_eviction_priority="evict_first" ) else: col = head * head_dim + unit @@ -286,12 +268,8 @@ def _thd_local_rope_kernel( value_left = x1_cos_left - x2_sin_left value_right = x2_cos_right + x1_sin_right if adjoint != 0: - x2_sin_right = (x2 * sin_right).to(out.element_type).to( - cutlass.Float32 - ) - x1_sin_left = (x1 * sin_left).to(out.element_type).to( - cutlass.Float32 - ) + x2_sin_right = (x2 * sin_right).to(out.element_type).to(cutlass.Float32) + x1_sin_left = (x1 * sin_left).to(out.element_type).to(cutlass.Float32) value_left = x1_cos_left + x2_sin_right value_right = cutlass.Float32(0.0) - x1_sin_left + x2_cos_right out[row, col] = value_left.to(out.element_type) @@ -577,11 +555,7 @@ def _compressor_input_compact_fwd_kernel( cutlass.Int64, level1_eviction_priority="evict_no_allocate", ) - cute.arch.store( - dst_ptr.llvm_ptr, - value, - level1_eviction_priority="evict_first", - ) + cute.arch.store(dst_ptr.llvm_ptr, value, level1_eviction_priority="evict_first") else: value = cutlass.Float32(0.0).to(hidden_compact.element_type) if src_global >= 0: @@ -743,9 +717,7 @@ def _compressor_input_compact_bwd_kernel( bidx, _, _ = cute.arch.block_idx() out_row = bidx - if cutlass.const_expr( - grad_hidden_compact.element_type.width == 16 and row_width % 4 == 0 - ): + if cutlass.const_expr(grad_hidden_compact.element_type.width == 16 and row_width % 4 == 0): vec_elems = cutlass.Int32(4) vec_cols = row_width // 4 else: @@ -824,22 +796,17 @@ def _compressor_input_compact_bwd_kernel( ) if out_row < d_window: dst_ptr = cute.recast_ptr( - grad_boundary_hidden.iterator + dst_offset, - dtype=cutlass.Int64, + grad_boundary_hidden.iterator + dst_offset, dtype=cutlass.Int64 ) cute.arch.store( - dst_ptr.llvm_ptr, - value, - level1_eviction_priority="evict_first", + dst_ptr.llvm_ptr, value, level1_eviction_priority="evict_first" ) else: dst_ptr = cute.recast_ptr( grad_hidden_local.iterator + dst_offset, dtype=cutlass.Int64 ) cute.arch.store( - dst_ptr.llvm_ptr, - value, - level1_eviction_priority="evict_first", + dst_ptr.llvm_ptr, value, level1_eviction_priority="evict_first" ) else: value = cutlass.Float32(0.0).to(grad_hidden_compact.element_type) @@ -1086,19 +1053,18 @@ def _thd_full_kv_pack_kernel( elif source_kind == 2: src_offset = source_index * row_width + vec_col * vec_elems src_ptr = cute.recast_ptr( - compressed_rank_major.iterator + src_offset, - dtype=cutlass.Int64, + compressed_rank_major.iterator + src_offset, dtype=cutlass.Int64 ) value = cute.arch.load( src_ptr.llvm_ptr, cutlass.Int64, level1_eviction_priority="evict_no_allocate", ) - dst_ptr = cute.recast_ptr(kv_full.iterator + dst_offset, dtype=cutlass.Int64) + dst_ptr = cute.recast_ptr( + kv_full.iterator + dst_offset, dtype=cutlass.Int64 + ) cute.arch.store( - dst_ptr.llvm_ptr, - value, - level1_eviction_priority="evict_first", + dst_ptr.llvm_ptr, value, level1_eviction_priority="evict_first" ) else: value = cutlass.Float32(0.0).to(kv_full.element_type) @@ -1112,7 +1078,9 @@ def _thd_full_kv_pack_kernel( else: if cutlass.const_expr(kv_full.element_type.width == 16 and row_width % 4 == 0): src_offset = out_row * row_width + vec_col * vec_elems - src_ptr = cute.recast_ptr(kv_full.iterator + src_offset, dtype=cutlass.Int64) + src_ptr = cute.recast_ptr( + kv_full.iterator + src_offset, dtype=cutlass.Int64 + ) value = cute.arch.load( src_ptr.llvm_ptr, cutlass.Int64, @@ -1124,9 +1092,7 @@ def _thd_full_kv_pack_kernel( kv_local.iterator + dst_offset, dtype=cutlass.Int64 ) cute.arch.store( - dst_ptr.llvm_ptr, - value, - level1_eviction_priority="evict_first", + dst_ptr.llvm_ptr, value, level1_eviction_priority="evict_first" ) elif source_kind == 1: dst_offset = source_index * row_width + vec_col * vec_elems @@ -1134,20 +1100,15 @@ def _thd_full_kv_pack_kernel( boundary_kv.iterator + dst_offset, dtype=cutlass.Int64 ) cute.arch.store( - dst_ptr.llvm_ptr, - value, - level1_eviction_priority="evict_first", + dst_ptr.llvm_ptr, value, level1_eviction_priority="evict_first" ) elif source_kind == 2: dst_offset = source_index * row_width + vec_col * vec_elems dst_ptr = cute.recast_ptr( - compressed_rank_major.iterator + dst_offset, - dtype=cutlass.Int64, + compressed_rank_major.iterator + dst_offset, dtype=cutlass.Int64 ) cute.arch.store( - dst_ptr.llvm_ptr, - value, - level1_eviction_priority="evict_first", + dst_ptr.llvm_ptr, value, level1_eviction_priority="evict_first" ) else: if source_kind == 0: diff --git a/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py b/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py index c6b24b12a9f..3bb795b68d7 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py +++ b/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py @@ -168,10 +168,7 @@ def build_thd_cp_local_padding_row_mask( chunk_ranges, _, _ = _normalize_row_ranges(chunk_ranges, "DSv4 CP padding-row mask") device = cu_seqlens_padded.device global_rows = torch.cat( - [ - torch.arange(start, end, device=device, dtype=torch.int64) - for start, end in chunk_ranges - ], + [torch.arange(start, end, device=device, dtype=torch.int64) for start, end in chunk_ranges], dim=0, ) @@ -596,11 +593,7 @@ def build_cp_compressor_prep_compact_fused( torch.cumsum(cu_deltas, dim=0), ) ) - return ( - torch.cat(hidden_parts, dim=0), - cu_compact, - torch.cat(comp_parts, dim=0), - ) + return (torch.cat(hidden_parts, dim=0), cu_compact, torch.cat(comp_parts, dim=0)) # ============================================================================= diff --git a/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py b/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py index 60c5eea56bc..214a405e557 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py +++ b/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py @@ -1581,11 +1581,7 @@ def forward( ) raw_local_loss = _kl_loss_from_target_predict( - target, - predict, - indexer_topk_idxs_for_loss, - loss_coeff, - calculate_per_token_loss=True, + target, predict, indexer_topk_idxs_for_loss, loss_coeff, calculate_per_token_loss=True ) if calculate_per_token_loss: indexer_loss = raw_local_loss diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_layout_kernels.py b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_layout_kernels.py index 6518bf8c18f..0039f661185 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_layout_kernels.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_layout_kernels.py @@ -32,7 +32,9 @@ def _make_e2e_like_cu_seqlens(device: str = "cuda", *, padded: bool = True) -> t full e2e hidden sizes in these focused kernel tests. """ lengths = _E2E_RAGGED_PADDED_SEG_LENS if padded else _E2E_RAGGED_SEG_LENS - return torch.tensor([0] + list(torch.tensor(lengths).cumsum(0).tolist()), dtype=torch.int32, device=device) + return torch.tensor( + [0] + list(torch.tensor(lengths).cumsum(0).tolist()), dtype=torch.int32, device=device + ) def _compressed_cu_seqlens(cu_seqlens: torch.Tensor, ratio: int) -> torch.Tensor: @@ -804,9 +806,9 @@ def test_thd_full_kv_pack_matches_native_forward_backward_and_reports_bandwidth( seq_ids, comp_ids, valid = _native_build_compressed_row_metadata( cu, _E2E_CP_SIZE, l_local, ratio, d_comp, c_cap, c_cap, False ) - compressed = torch.randn(seq_ids.shape[0], 4, dtype=torch.bfloat16, device="cuda").requires_grad_( - True - ) + compressed = torch.randn( + seq_ids.shape[0], 4, dtype=torch.bfloat16, device="cuda" + ).requires_grad_(True) dummy = torch.empty((1,), dtype=torch.int32, device="cuda") capacity = l_local + d_window * (cu.shape[0] - 1) + compressed.shape[0] @@ -1153,7 +1155,9 @@ def test_build_indexer_loss_indices_matches_native_and_reports_bandwidth(): window = 16 compressed_width = 8 logical_ids = ( - torch.arange(compressed_width, dtype=torch.int32, device="cuda").unsqueeze(0).repeat(l_local, 1) + torch.arange(compressed_width, dtype=torch.int32, device="cuda") + .unsqueeze(0) + .repeat(l_local, 1) ) logical_ids[1::5, -1] = -1 rank_by_seq = torch.arange(int(cu_comp[-1]), dtype=torch.int32, device="cuda") @@ -1161,15 +1165,7 @@ def test_build_indexer_loss_indices_matches_native_and_reports_bandwidth(): cu, cu_comp, global_start, l_local, d_window, window, ratio, logical_ids, rank_by_seq ) expected = _native_indexer_loss_indices( - cu, - cu_comp, - global_start, - l_local, - d_window, - window, - ratio, - logical_ids, - rank_by_seq, + cu, cu_comp, global_start, l_local, d_window, window, ratio, logical_ids, rank_by_seq ) assert torch.equal(fused[0], expected[0]) assert torch.equal(fused[1], expected[1]) diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py index b9a69a2ce77..69fbf88bda4 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py @@ -291,9 +291,7 @@ def fake_apply( cu = torch.tensor([0, 8], dtype=torch.int32) equal = torch.arange(8, dtype=torch.float32).reshape(4, 2) - out = apply_thd_cp_local_rope_fused( - equal, cos, sin, 1, 1, cu, ((2, 4), (10, 12)), inverse=True - ) + out = apply_thd_cp_local_rope_fused(equal, cos, sin, 1, 1, cu, ((2, 4), (10, 12)), inverse=True) assert torch.equal(out, equal + 10) assert calls[-1] == ((4, 2), 2, 2, 1, 1, 10, True, False) @@ -374,16 +372,15 @@ def fake_build(cu, cp_size, l_local, ratio, d_comp, c_cap, **kwargs): torch.zeros((rows,), dtype=torch.bool), ) - monkeypatch.setattr(csa_cp_utils.csa_cp_layout_kernels, "build_compressed_row_metadata", fake_build) + monkeypatch.setattr( + csa_cp_utils.csa_cp_layout_kernels, "build_compressed_row_metadata", fake_build + ) cu = torch.tensor([0, 32, 64], dtype=torch.int32) build_cp_rank_major_compressed_metadata_fused(cu, ((16, 32),), 4, 4, 8) build_cp_rank_major_compressed_metadata_fused(cu, ((0, 4), (28, 32)), 4, 4, 8) - assert calls == [ - (4, 16, 4, 8, 8, None, False), - (4, 4, 4, 8, 3, 8, True), - ] + assert calls == [(4, 16, 4, 8, 8, None, False), (4, 4, 4, 8, 3, 8, True)] def test_compressor_prep_compact_wrapper_merges_two_chunk_outputs(monkeypatch): @@ -453,7 +450,9 @@ def fake_apply(*args): capacity = int(args[13]) return kv_local.new_full((capacity,) + tuple(kv_local.shape[1:]), len(calls)) - monkeypatch.setattr(csa_cp_utils.csa_cp_layout_kernels.ThdFullKvPack, "apply", staticmethod(fake_apply)) + monkeypatch.setattr( + csa_cp_utils.csa_cp_layout_kernels.ThdFullKvPack, "apply", staticmethod(fake_apply) + ) kv_local = torch.ones(4, 2) boundary = torch.ones(2, 2) compressed = torch.ones(6, 2) @@ -510,7 +509,9 @@ def fake_repack(rank_major, seq_ids, comp_ids, valid, cu_comp, rows): captured.update(rows=int(rows), rank_major=rank_major, seq_ids=seq_ids) return rank_major[:rows].clone(), torch.arange(rows, dtype=torch.int32) - monkeypatch.setattr(csa_cp_utils.csa_cp_layout_kernels, "repack_compressed_kv_to_seq_major", fake_repack) + monkeypatch.setattr( + csa_cp_utils.csa_cp_layout_kernels, "repack_compressed_kv_to_seq_major", fake_repack + ) rank_major = torch.arange(12, dtype=torch.float32).reshape(6, 2) seq_ids = torch.zeros(6, dtype=torch.int32) comp_ids = torch.arange(6, dtype=torch.int32) @@ -577,7 +578,9 @@ def fake_indexer_topk( ) return torch.full((q.shape[0], int(topk)), len(topk_calls), dtype=torch.int32), None - monkeypatch.setattr(csa_cp_utils.csa_cp_layout_kernels, "build_indexer_topk_metadata", fake_metadata) + monkeypatch.setattr( + csa_cp_utils.csa_cp_layout_kernels, "build_indexer_topk_metadata", fake_metadata + ) monkeypatch.setattr(csa_cp_utils, "indexer_topk", fake_indexer_topk) q = torch.randn(5, 2) @@ -604,9 +607,12 @@ def fake_indexer_topk( assert torch.equal(out, torch.tensor([[1, 1], [1, 1], [2, 2], [2, 2], [2, 2]])) assert torch.equal(topk_calls[0][-1], torch.ones(2, dtype=torch.int32)) assert torch.equal(topk_calls[1][-1], torch.full((3,), 2, dtype=torch.int32)) - assert compute_cp_indexer_topk_logical_fused( - q, weights, k_seq[:0], cu_q, cu_comp, ((2, 4),), 4, 2, 1.0, 10, 3 - ) is None + assert ( + compute_cp_indexer_topk_logical_fused( + q, weights, k_seq[:0], cu_q, cu_comp, ((2, 4),), 4, 2, 1.0, 10, 3 + ) + is None + ) def test_final_index_wrappers_lower_two_chunk_arguments(monkeypatch): @@ -622,7 +628,9 @@ def fake_attention(*args, **kwargs): captured["attention"] = (args, kwargs) l_local = int(args[2]) width = int(args[6]) - return torch.zeros(l_local, width, dtype=torch.int32), torch.zeros(l_local, dtype=torch.int32) + return torch.zeros(l_local, width, dtype=torch.int32), torch.zeros( + l_local, dtype=torch.int32 + ) def fake_loss(*args, **kwargs): captured["loss"] = (args, kwargs) @@ -630,7 +638,9 @@ def fake_loss(*args, **kwargs): width = args[7].shape[1] + int(args[5]) return torch.zeros(l_local, width, dtype=torch.int32), torch.zeros_like(args[7]) - monkeypatch.setattr(csa_cp_utils.csa_cp_layout_kernels, "build_attention_indices", fake_attention) + monkeypatch.setattr( + csa_cp_utils.csa_cp_layout_kernels, "build_attention_indices", fake_attention + ) monkeypatch.setattr(csa_cp_utils.csa_cp_layout_kernels, "build_indexer_loss_indices", fake_loss) cu = torch.tensor([0, 16], dtype=torch.int32) From 3aeb5f5bcbf382ef43d10ac9797e0698b1ab8852 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Fri, 26 Jun 2026 07:23:34 -0700 Subject: [PATCH 03/23] Fix SBHD indexer topk length handling Signed-off-by: kunlunl --- .../experimental_attention_variant/dsa_kernels.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py b/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py index 214a405e557..fe23e281a04 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py +++ b/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py @@ -602,10 +602,9 @@ def _indexer_topk_core( pad = torch.full((total_q, topk - topk_k), -1, dtype=torch.int32, device=device) topk_indices = torch.cat([topk_indices, pad], dim=-1) - row_valid = (topk_indices >= 0) & (topk_indices < seq_lens.unsqueeze(1)) - topk_indices = topk_indices.masked_fill(~row_valid, -1) - if is_thd: + row_valid = (topk_indices >= 0) & (topk_indices < seq_lens.unsqueeze(1)) + topk_indices = topk_indices.masked_fill(~row_valid, -1) safe_topk = topk_indices.clamp(min=0, max=max(sk - 1, 0)).to(torch.long) selected_scores = torch.gather(scores_flat, dim=-1, index=safe_topk) selected_valid = (topk_indices >= 0) & (topk_indices < sk) & torch.isfinite(selected_scores) From 91710ac3b460250c6628e5c53819712a89a27f1e Mon Sep 17 00:00:00 2001 From: kunlunl Date: Sat, 27 Jun 2026 01:49:58 -0700 Subject: [PATCH 04/23] Gate DSv4 CP fused tests by availability Signed-off-by: kunlunl --- .../test_dsv4_hybrid_attention_cp.py | 38 +++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py index 040495b1bff..26cfc05616f 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py @@ -74,6 +74,29 @@ _DSV4_CP_REPLAY_PADDED_SEG_LENS = (8, 128, 1000, 32, 132, 904, 64, 260, 804, 96, 512, 156) +def _dsv4_cp_fused_kernels_available(): + """Return whether this host can run the fused DSv4 THD CP test path.""" + if not torch.cuda.is_available(): + return False + try: + sm_major = torch.cuda.get_device_capability()[0] + except RuntimeError: + return False + if sm_major < 10: + return False + try: + from cudnn import DSA # noqa: F401 + from flash_mla import flash_mla_sparse_fwd # noqa: F401 + except ImportError: + return False + return True + + +_DSV4_CP_FUSED_KERNELS_UNAVAILABLE_REASON = ( + "DSv4 fused DSA cases require SM100+, flash_mla, and cudnn.DSA" +) + + class _ReferenceCPGroup: def rank(self): """Return the CP rank used by the CP1 reference path.""" @@ -462,6 +485,7 @@ def setup_method(self, request): model_parallel_cuda_manual_seed(_SEED) cls = request.cls + cls.fused_kernels_available = _dsv4_cp_fused_kernels_available() cls.cp_size = cp_size cls.cp_rank = parallel_state.get_context_parallel_rank() cls.pg = ProcessGroupCollection.use_mpu_process_groups() @@ -564,6 +588,7 @@ def test_thd_cp_matches_full_reference_forward_backward(self, layer_number, part context_parallel_size=self.cp_size, dsa_indexer_loss_coeff=1.0, dsa_indexer_use_sparse_loss=True, + apply_dsa_kernel_fusion=self.fused_kernels_available, apply_rope_fusion=True, csa_cp_partition_mode=partition_mode, ) @@ -573,6 +598,7 @@ def test_thd_cp_matches_full_reference_forward_backward(self, layer_number, part dsa_indexer_use_sparse_loss=True, # Numerical parity uses the unfused CP1 reference so failures point # at the CP path instead of non-CP fused RoPE behavior. + apply_dsa_kernel_fusion=self.fused_kernels_available, apply_rope_fusion=False, ) cp_attn = _build_attention( @@ -649,6 +675,9 @@ def test_thd_cp_cuda_graph_matches_eager_forward_backward( plus elementwise ``assert_close`` gates because they may be nondeterministic against eager execution. """ + if fused and not self.fused_kernels_available: + pytest.skip(_DSV4_CP_FUSED_KERNELS_UNAVAILABLE_REASON) + context = nullcontext() if fused else _deterministic_torch_algorithms() mode = "fused" if fused else "unfused" with context: @@ -753,6 +782,9 @@ def test_thd_cp_cuda_graph_replay_accepts_changed_padded_boundaries( failure means the CP path baked capture-time padded boundaries or a host-derived dynamic shape into the graph. """ + if not self.fused_kernels_available: + pytest.skip(_DSV4_CP_FUSED_KERNELS_UNAVAILABLE_REASON) + capture_packed = _make_thd_packed_seq_params( _DSV4_CP_RAGGED_SEG_LENS, _DSV4_CP_RAGGED_PADDED_SEG_LENS ) @@ -871,6 +903,9 @@ def test_thd_cp_peak_allocated_delta_scales_vs_cp1(self, layer_number, partition accidentally gather full hidden states or otherwise allocate global activation-sized buffers before the layer work. """ + if not self.fused_kernels_available: + pytest.skip(_DSV4_CP_FUSED_KERNELS_UNAVAILABLE_REASON) + packed, padded_tokens, local_idx = _make_ragged_cp_case( partition_mode, self.cp_size, self.cp_rank ) @@ -934,6 +969,9 @@ def test_thd_cp_cuda_graph_time_scales_vs_cp1(self, layer_number, partition_mode are not enough if the implementation falls back to slow dynamic PyTorch work in the production graph path. """ + if not self.fused_kernels_available: + pytest.skip(_DSV4_CP_FUSED_KERNELS_UNAVAILABLE_REASON) + packed, padded_tokens, local_idx = _make_ragged_cp_case( partition_mode, self.cp_size, self.cp_rank ) From 3418b5819643d3f9689a0ce967eeb2f41f8cfea8 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Sat, 27 Jun 2026 04:22:06 -0700 Subject: [PATCH 05/23] Simplify DSv4 CP layout flow Signed-off-by: kunlunl --- .../experimental_attention_variant/csa.py | 321 ++- .../csa_cp_layout_kernels.py | 2315 +++-------------- .../csa_cp_utils.py | 547 +--- .../test_csa_cp_layout_kernels.py | 713 +---- .../test_csa_cp_utils.py | 335 +-- .../test_dsv4_hybrid_attention_cp.py | 44 + 6 files changed, 795 insertions(+), 3480 deletions(-) diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index 535880ff7c9..a00bdc0e88a 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -12,7 +12,9 @@ from megatron.core.models.common.embeddings import RotaryEmbedding, apply_rotary_pos_emb from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.tensor_parallel.mappings import gather_from_sequence_parallel_region from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.experimental_attention_variant import csa_cp_layout_kernels from megatron.core.transformer.experimental_attention_variant import csa_cp_utils as cp_utils from megatron.core.transformer.experimental_attention_variant.dsa import ( DSAIndexerLossAutoScaler, @@ -86,32 +88,6 @@ def get_compress_topk_idxs( return matrix.unsqueeze(0).expand(batch_size, -1, -1) -def _get_csa_compressed_capacity( - packed_seq_params: Optional[PackedSeqParams], ratio: int, total_tokens: int -) -> Optional[int]: - """Return a host-known compressed capacity for THD CUDA graph capture. - - The exact ``sum(seq_len // ratio)`` lives in ``cu_seqlens`` on device. - Reading it on the host would break CUDA graph capture, and - ``PackedSeqParams`` must stay a generic MCore contract rather than - carrying CSA-specific metadata. Use a static upper bound instead; - device-side ``cu_seqlens_compressed`` keeps the true valid rows and - downstream kernels leave the extra rows as tail padding. - """ - if packed_seq_params is None or ratio <= 1: - return None - max_seqlen = packed_seq_params.max_seqlen_q - cu_seqlens = ( - packed_seq_params.cu_seqlens_q_padded - if packed_seq_params.cu_seqlens_q_padded is not None - else packed_seq_params.cu_seqlens_q - ) - if max_seqlen is None or cu_seqlens is None: - return None - num_sequences = max(int(cu_seqlens.shape[0]) - 1, 0) - return min(int(total_tokens) // ratio, num_sequences * (int(max_seqlen) // ratio)) - - # --------------------------------------------------------------------------- # THD (packed) variants of the index helpers above. # @@ -909,7 +885,7 @@ def _forward_thd( max_seqlen_q: Optional[int] = None, compressed_group_ids: Optional[torch.Tensor] = None, fixed_total_comp: Optional[int] = None, - ) -> Tuple[Optional[torch.Tensor], torch.Tensor]: + ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: """THD per-segment compression — fully vectorized. Linear projections are token-wise on the flat input. The gated @@ -939,26 +915,31 @@ def _forward_thd( ``fixed_total_comp`` is not set) and ``cu_seqlens_compressed`` is ``(B+1,)`` int32 with ``cu_seqlens_compressed[b+1] - cu_seqlens_compressed[b] = seqlen_b // ratio``. + Pre-grouped CP inputs return ``None`` for this unused second value. """ ratio = self.compress_ratio device = x.device dtype = x.dtype pre_grouped = compressed_group_ids is not None - # Per-segment compressed lengths (vectorized). - seq_lens = cu_seqlens[1:] - cu_seqlens[:-1] - seg_compressed_lens = seq_lens // ratio - cu_seqlens_compressed = torch.cat( - [ - torch.zeros(1, dtype=cu_seqlens.dtype, device=device), - seg_compressed_lens.cumsum(0).to(cu_seqlens.dtype), - ] - ) - total_comp = ( - int(fixed_total_comp) - if fixed_total_comp is not None - else int(cu_seqlens_compressed[-1].item()) - ) + if pre_grouped: + if fixed_total_comp is None: + raise RuntimeError("Pre-grouped THD compression requires a fixed row count.") + cu_seqlens_compressed = None + total_comp = int(fixed_total_comp) + else: + seg_compressed_lens = (cu_seqlens[1:] - cu_seqlens[:-1]) // ratio + cu_seqlens_compressed = torch.cat( + [ + torch.zeros(1, dtype=cu_seqlens.dtype, device=device), + seg_compressed_lens.cumsum(0).to(cu_seqlens.dtype), + ] + ) + total_comp = ( + int(fixed_total_comp) + if fixed_total_comp is not None + else int(cu_seqlens_compressed[-1].item()) + ) if total_comp == 0: return None, cu_seqlens_compressed @@ -1031,14 +1012,27 @@ def _forward_thd( rotary_pos_cos, rotary_pos_sin = self.rotary_pos_emb.get_cached_cos_sin( int(max_seqlen_q), dtype=compressed_thd.dtype, packed_seq=True, mscale=1.0 ) - compressed_thd = cp_utils.apply_thd_cp_compressed_rope_fused( - compressed_thd, - rotary_pos_cos, - rotary_pos_sin, - compressed_group_ids, - ratio, - self.head_dim - self.qk_pos_emb_head_dim, - self.qk_pos_emb_head_dim, + positions = compressed_group_ids[:total_comp].clamp_min(0).to(torch.long) * ratio + cos = rotary_pos_cos.flatten(1).index_select(0, positions) + sin = rotary_pos_sin.flatten(1).index_select(0, positions) + pos_dim = self.qk_pos_emb_head_dim + table_shape = (total_comp,) + (1,) * (compressed_thd.ndim - 2) + (pos_dim,) + cos = cos[:, :pos_dim].view(table_shape) + sin = sin[:, :pos_dim].view(table_shape) + x_nope, x_pos = torch.split(compressed_thd, [self.head_dim - pos_dim, pos_dim], dim=-1) + x1, x2 = x_pos[..., 0::2], x_pos[..., 1::2] + half = pos_dim // 2 + dtype = compressed_thd.dtype + left = ( + (x1.float() * cos[..., :half].float()).to(dtype).float() + - (x2.float() * sin[..., :half].float()).to(dtype).float() + ).to(dtype) + right = ( + (x2.float() * cos[..., half:].float()).to(dtype).float() + + (x1.float() * sin[..., half:].float()).to(dtype).float() + ).to(dtype) + compressed_thd = torch.cat( + (x_nope, torch.stack((left, right), dim=-1).flatten(-2)), dim=-1 ) if self.rotate: @@ -1075,13 +1069,14 @@ def forward( "to avoid a GPU→CPU sync that breaks CUDA graph capture." ) max_seqlen_q = int(packed_seq_params.max_seqlen_q) + fixed_total_comp = None + if self.compress_ratio > 1: + fixed_total_comp = min( + x.shape[0] // self.compress_ratio, + (cu_seqlens.shape[0] - 1) * (max_seqlen_q // self.compress_ratio), + ) result = self._forward_thd( - x, - cu_seqlens, - max_seqlen_q=max_seqlen_q, - fixed_total_comp=_get_csa_compressed_capacity( - packed_seq_params, self.compress_ratio, x.shape[0] - ), + x, cu_seqlens, max_seqlen_q=max_seqlen_q, fixed_total_comp=fixed_total_comp ) else: result = self._forward_sbhd(x) @@ -1177,39 +1172,6 @@ def __init__( pg_collection=pg_collection, ) - def _forward_thd_query_weights_cp( - self, - x: torch.Tensor, - qr: torch.Tensor, - cu_seqlens_padded: torch.Tensor, - max_position: int, - chunk_ranges: Tuple[Tuple[int, int], ...], - ) -> Tuple[torch.Tensor, torch.Tensor]: - """Compute local THD indexer Q/weights for the current CP row order.""" - local_rows, bsz, _ = x.size() - if bsz != 1: - raise RuntimeError(f"DSv4 THD CP indexer expects bsz=1, got {bsz}.") - - q, _ = self.linear_wq_b(qr) - q = q.reshape(local_rows, bsz, self.index_n_heads, self.index_head_dim).squeeze(1) - rotary_pos_cos, rotary_pos_sin = self.rotary_pos_emb.get_cached_cos_sin( - int(max_position), dtype=q.dtype, packed_seq=True, mscale=1.0 - ) - q = cp_utils.apply_thd_cp_local_rope_fused( - q, - rotary_pos_cos, - rotary_pos_sin, - self.index_head_dim - self.qk_pos_emb_head_dim, - self.qk_pos_emb_head_dim, - cu_seqlens_padded, - chunk_ranges, - ) - q = rotate_activation(q) - - weights, _ = self.linear_weights_proj(x) - weights = weights.squeeze(1) * (self.index_n_heads**-0.5) - return q, weights - def forward_before_topk( self, x: torch.Tensor, qr: torch.Tensor, packed_seq_params: Optional[PackedSeqParams] = None ) -> Union[ @@ -2228,6 +2190,14 @@ def _forward_thd_cp( chunk_ranges = cp_utils.local_q_cp_chunk_ranges( self.config.csa_cp_partition_mode, l_local, cp_size, cp_rank ) + if len(chunk_ranges) == 2: + chunk_starts = tuple(start for start, _ in chunk_ranges) + chunk_len = l_local // 2 + index_global_start = 0 + else: + chunk_starts = None + chunk_len = l_local + index_global_start = chunk_ranges[0][0] kv_local = key.squeeze(-2).squeeze(1) if boundary_hidden is None or boundary_kv is None: raise RuntimeError( @@ -2236,16 +2206,13 @@ def _forward_thd_cp( ) boundary_kv = boundary_kv.squeeze(-2).squeeze(1) d_window = boundary_hidden.shape[0] // len(chunk_ranges) - # Window-only defaults; compression fills rank-major buffers/metadata. + # Window-only defaults; compression fills the rank-major KV buffer. compressed_kv_rank_major = kv_local.new_empty((0, kv_local.shape[-1])) cu_seqlens_compressed = None - seq_ids_rank_major = torch.empty((0,), dtype=torch.int32, device=query.device) - comp_ids_rank_major = torch.empty((0,), dtype=torch.int32, device=query.device) - valid_rank_major = torch.empty((0,), dtype=torch.bool, device=query.device) # Logical top-k uses per-sequence comp ids; the map lowers them to # rank-major gathered rows when needed. - indexer_topk_logical = rank_major_by_seq_major = None + indexer_topk_logical = rank_row_for_seq_row = None ratio = self.compress_ratio indexer = self.indexer indexer_loss_coeff = self.config.dsa_indexer_loss_coeff or 0.0 @@ -2258,23 +2225,29 @@ def _forward_thd_cp( d_comp = 8 if ratio == 4 else ratio # One compressed row per full ratio-sized group; tails are dropped. - cu_seqlens_compressed = cp_utils.build_global_compressed_cu_seqlens( - cu_seqlens_kv, ratio + compressed_lens = torch.div( + cu_seqlens_kv[1:] - cu_seqlens_kv[:-1], ratio, rounding_mode="floor" ) - seq_major_compressed_capacity = (l_local * cp_size) // ratio - - # Compact local+boundary tokens for Compressor. ``cu_compact`` is - # local token prefixes; ``comp_ids_local`` drives compressed RoPE. - hidden_compact, cu_compact, comp_ids_local = ( - cp_utils.build_cp_compressor_prep_compact_fused( - x, boundary_hidden, cu_seqlens_kv, chunk_ranges, ratio, d_comp, d_window + cu_seqlens_compressed = torch.cat( + ( + torch.zeros_like(cu_seqlens_kv[:1]), + torch.cumsum(compressed_lens, dim=0, dtype=torch.int32), ) ) - - # Rank-major metadata: sequence id, comp id, and canonical validity. - seq_ids_rank_major, comp_ids_rank_major, valid_rank_major = ( - cp_utils.build_cp_rank_major_compressed_metadata_fused( - cu_seqlens_kv, chunk_ranges, cp_size, ratio, d_comp + # Compact local+boundary tokens for Compressor. ``comp_ids_local`` + # drives compressed RoPE; the row map directly lowers global + # sequence-major compressed ids into rank-major all-gather rows. + hidden_compact, comp_ids_local, rank_row_for_seq_row = ( + cp_utils.build_cp_compressor_prep_compact_fused( + x, + boundary_hidden, + cu_seqlens_kv, + cu_seqlens_compressed, + chunk_ranges, + cp_size, + ratio, + d_comp, + d_window, ) ) @@ -2290,32 +2263,48 @@ def _forward_thd_cp( "Dense CP-aware indexer loss needs a CP-aware dense score kernel." ) - q_indexer_cp, weights_indexer_cp = indexer._forward_thd_query_weights_cp( - x.detach(), qr.detach(), cu_seqlens_q, max_seqlen_q, chunk_ranges=chunk_ranges + indexer_x, indexer_qr = x.detach(), qr.detach() + if indexer_x.shape[1] != 1: + raise RuntimeError( + f"DSv4 THD CP indexer expects bsz=1, got {indexer_x.shape[1]}." + ) + q_indexer_cp, _ = indexer.linear_wq_b(indexer_qr) + q_indexer_cp = q_indexer_cp.reshape( + l_local, indexer.index_n_heads, indexer.index_head_dim + ) + rotary_pos_cos, rotary_pos_sin = indexer.rotary_pos_emb.get_cached_cos_sin( + max_seqlen_q, dtype=q_indexer_cp.dtype, packed_seq=True, mscale=1.0 + ) + q_indexer_cp = rotate_activation( + cp_utils.apply_thd_cp_local_rope_fused( + q_indexer_cp, + rotary_pos_cos, + rotary_pos_sin, + indexer.index_head_dim - indexer.qk_pos_emb_head_dim, + indexer.qk_pos_emb_head_dim, + cu_seqlens_q, + chunk_ranges, + ) ) + weights_indexer_cp, _ = indexer.linear_weights_proj(indexer_x) + weights_indexer_cp = weights_indexer_cp.squeeze(1) * (indexer.index_n_heads**-0.5) indexer_compressed_local, _ = indexer.compressor._forward_thd( hidden_compact.detach(), - cu_compact, + cu_seqlens_kv, max_seqlen_q=max_seqlen_q, compressed_group_ids=comp_ids_local, fixed_total_comp=hidden_compact.shape[0] // ratio, ) # Gather indexer compressed K in rank-major order. - k_indexer_rank_major = cp_utils.all_gather_fixed_cp_tensor( - indexer_compressed_local.squeeze(1), cp_group + k_indexer_rank_major = gather_from_sequence_parallel_region( + indexer_compressed_local.squeeze(1), group=cp_group ) - # Repack valid compressed K to seq-major order for top-k. - k_indexer_seq_major, rank_major_by_seq_major = ( - cp_utils.repack_rank_major_compressed_to_seq_major_fused( - k_indexer_rank_major, - seq_ids_rank_major, - comp_ids_rank_major, - valid_rank_major, - cu_seqlens_compressed, - seq_major_rows=seq_major_compressed_capacity, - ) + # Indexer top-k consumes sequence-major K rows. Capacity-tail + # map entries are unused by cu_seqlens_compressed. + k_indexer_seq_major = torch.index_select( + k_indexer_rank_major, 0, rank_row_for_seq_row.clamp_min(0) ) # Top-k is still in logical compressed coordinates here. indexer_topk_logical = cp_utils.compute_cp_indexer_topk_logical_fused( @@ -2335,57 +2324,37 @@ def _forward_thd_cp( # ---- Step 5: attention compressed KV path ------------------------- compressed_kv_local, _ = self.compressor._forward_thd( hidden_compact, - cu_compact, + cu_seqlens_kv, max_seqlen_q=max_seqlen_q, compressed_group_ids=comp_ids_local, fixed_total_comp=hidden_compact.shape[0] // ratio, ) # Gather attention compressed KV in the same rank-major layout. - compressed_kv_rank_major = cp_utils.all_gather_fixed_cp_tensor( - compressed_kv_local.squeeze(1), cp_group + compressed_kv_rank_major = gather_from_sequence_parallel_region( + compressed_kv_local.squeeze(1), group=cp_group ) - # Two-chunk lowering needs the seq-major -> rank-major map. - if len(chunk_ranges) == 2 and rank_major_by_seq_major is None: - _, rank_major_by_seq_major = ( - cp_utils.repack_rank_major_compressed_to_seq_major_fused( - compressed_kv_rank_major, - seq_ids_rank_major, - comp_ids_rank_major, - valid_rank_major, - cu_seqlens_compressed, - seq_major_rows=seq_major_compressed_capacity, - ) - ) - # ---- Step 6: pack the local full-KV layout --------------------------- - # Pack the local KV buffer addressed by the final CP attention indices. - kv_full_thd, shared_compressed_base = cp_utils.pack_cp_kv_full_fused( - kv_local, - boundary_kv, - compressed_kv_rank_major, - seq_ids_rank_major, - comp_ids_rank_major, - valid_rank_major, - cu_seqlens_kv, - chunk_ranges, - d_window, - ratio, - rank_major_by_seq_major=rank_major_by_seq_major, - cu_seqlens_compressed=cu_seqlens_compressed, - ) + # ---- Step 6: concatenate the raw local KV sources ------------------- + # Final indices address these source rows directly. This avoids a + # second per-sequence packed KV layout and lets torch.cat own backward. + kv_full_thd = torch.cat((boundary_kv, kv_local, compressed_kv_rank_major), dim=0) + compressed_base = len(chunk_ranges) * d_window + l_local if training_with_grad and indexer_topk_logical is not None: # ---- Step 7a: indexer-loss path ---------------------------------- # Build attention indices plus rank-major indices for indexer loss. - topk_idxs, indexer_topk_rank_major = cp_utils.build_cp_indexer_loss_indices_fused( + topk_idxs, indexer_topk_rank_major = csa_cp_layout_kernels.build_indexer_loss_indices( cu_seqlens_q, cu_seqlens_compressed, - chunk_ranges, + index_global_start, + l_local, d_window, self.window_size, ratio, indexer_topk_logical, - rank_major_by_seq_major, - shared_compressed_base=shared_compressed_base, + rank_row_for_seq_row, + chunk_starts=chunk_starts, + chunk_len=chunk_len, + compressed_base=compressed_base, ) cu_seqlens_q_unpadded = None if ( @@ -2395,12 +2364,20 @@ def _forward_thd_cp( != packed_seq_params.cu_seqlens_q_padded.data_ptr() ): cu_seqlens_q_unpadded = packed_seq_params.cu_seqlens_q - q_padding_mask = cp_utils.build_thd_cp_local_padding_row_mask( - cu_seqlens_q, - cu_seqlens_q_unpadded, - chunk_ranges, - total_padded_rows=l_local * cp_size, - ) + q_padding_mask = None + if cu_seqlens_q_unpadded is not None: + global_rows = torch.cat( + [ + torch.arange(start, end, device=query.device, dtype=torch.int64) + for start, end in chunk_ranges + ] + ) + batch_ids = batch_of_row(cu_seqlens_q, total_q=l_local * cp_size).index_select( + 0, global_rows + ) + real_seqlens = cu_seqlens_q_unpadded[1:] - cu_seqlens_q_unpadded[:-1] + positions = global_rows.to(cu_seqlens_q.dtype) - cu_seqlens_q[batch_ids] + q_padding_mask = positions >= real_seqlens[batch_ids] output, indexer_loss = ( FusedIndexerSparseAttnFromTopkFunc.apply if self.apply_dsa_kernel_fusion @@ -2434,17 +2411,25 @@ def _forward_thd_cp( # ---- Step 7b: sparse attention path ---------------------------------- # No indexer-loss backward: lower final attention indices and run attention only. - topk_idxs, topk_length = cp_utils.build_cp_attention_indices_fused( + compressed_width = ( + indexer_topk_logical.shape[-1] + if indexer_topk_logical is not None + else (max_seqlen_q // ratio if ratio > 1 else 0) + ) + topk_idxs, topk_length = csa_cp_layout_kernels.build_attention_indices( cu_seqlens_q, - chunk_ranges, + index_global_start, + l_local, d_window, self.window_size, ratio, + compressed_width, indexer_topk_logical, - max_n_compressed=max_seqlen_q // ratio if ratio > 1 else 0, - rank_major_by_seq_major=rank_major_by_seq_major, cu_seqlens_compressed=cu_seqlens_compressed, - shared_compressed_base=shared_compressed_base, + rank_row_for_seq_row=rank_row_for_seq_row, + chunk_starts=chunk_starts, + chunk_len=chunk_len, + compressed_base=compressed_base, ) if self.apply_dsa_kernel_fusion: diff --git a/megatron/core/transformer/experimental_attention_variant/csa_cp_layout_kernels.py b/megatron/core/transformer/experimental_attention_variant/csa_cp_layout_kernels.py index 61579c34661..0055666cff1 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa_cp_layout_kernels.py +++ b/megatron/core/transformer/experimental_attention_variant/csa_cp_layout_kernels.py @@ -1,15 +1,13 @@ # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. """CuTeDSL kernels for DSv4 THD context parallel layout work. -This module is intentionally below the MCore-facing CP utility layer. Callers -outside ``csa_cp_utils.py`` should not import this file directly. The functions -here allocate output tensors, compile/cache CuTeDSL launches, and execute small -layout kernels for the DSv4 CP path. +This module contains the few CuTeDSL kernels retained by the DSv4 CP path. +``csa_cp_utils.py`` owns partition semantics and calls the RoPE/compaction +autograd wrappers; ``csa.py`` calls the final-index builders directly. """ import math import os -from functools import lru_cache from typing import Optional, Tuple import torch @@ -328,133 +326,21 @@ def _thd_local_rope_launch( stream=stream, ) - # Kernel contract: - # x/out: row-major compressed tensor views, shape (rows, row_width). - # cos/sin: rotary tables, shape (max_seq_len, pos_dim). - # comp_ids: int32 compressed group id per row, shape (rows,); -1 is treated as 0. - # ratio maps comp_id to original token position comp_id * ratio. - # Copies no-RoPE dims and rotates each non-interleaved pair in the position dims. - @cute.kernel - def _thd_compressed_rope_kernel( - x: cute.Tensor, - out: cute.Tensor, - cos: cute.Tensor, - sin: cute.Tensor, - comp_ids: cute.Tensor, - ratio: cutlass.Int32, - row_width: cutlass.Constexpr, - head_dim: cutlass.Int32, - nope_dim: cutlass.Int32, - pos_dim: cutlass.Int32, - inverse: cutlass.Int32, - adjoint: cutlass.Int32, - total_work: cutlass.Int32, - ): - tidx, _, _ = cute.arch.thread_idx() - bidx, _, _ = cute.arch.block_idx() - linear = bidx * 128 + tidx - if linear < total_work: - row = linear // row_width - col = linear - row * row_width - head_col = col - (col // head_dim) * head_dim - - if head_col < nope_dim: - out[row, col] = x[row, col] - else: - comp_id = comp_ids[row] - if comp_id < 0: - comp_id = 0 - rope_pos = comp_id * ratio - - pos_col = head_col - nope_dim - pair = pos_col // 2 - pair_base = col - head_col + nope_dim + pair * 2 - x1 = x[row, pair_base].to(cutlass.Float32) - x2 = x[row, pair_base + 1].to(cutlass.Float32) - - cos_left = cos[rope_pos, pair].to(cutlass.Float32) - sin_left = sin[rope_pos, pair].to(cutlass.Float32) - cos_right = cos[rope_pos, pos_dim // 2 + pair].to(cutlass.Float32) - sin_right = sin[rope_pos, pos_dim // 2 + pair].to(cutlass.Float32) - if inverse != 0: - sin_left = cutlass.Float32(0.0) - sin_left - sin_right = cutlass.Float32(0.0) - sin_right - - x1_cos_left = (x1 * cos_left).to(out.element_type).to(cutlass.Float32) - x2_sin_left = (x2 * sin_left).to(out.element_type).to(cutlass.Float32) - x2_cos_right = (x2 * cos_right).to(out.element_type).to(cutlass.Float32) - x1_sin_right = (x1 * sin_right).to(out.element_type).to(cutlass.Float32) - value = x1_cos_left - x2_sin_left - if pos_col - pair * 2 != 0: - value = x2_cos_right + x1_sin_right - if adjoint != 0: - x2_sin_right = (x2 * sin_right).to(out.element_type).to(cutlass.Float32) - x1_sin_left = (x1 * sin_left).to(out.element_type).to(cutlass.Float32) - value = x1_cos_left + x2_sin_right - if pos_col - pair * 2 != 0: - value = cutlass.Float32(0.0) - x1_sin_left + x2_cos_right - out[row, col] = value.to(out.element_type) - - # Launch contract for _thd_compressed_rope_kernel. - # row_width is static; total_work is rows * row_width. - @cute.jit - def _thd_compressed_rope_launch( - x: cute.Tensor, - out: cute.Tensor, - cos: cute.Tensor, - sin: cute.Tensor, - comp_ids: cute.Tensor, - ratio: cutlass.Int32, - row_width: cutlass.Constexpr, - head_dim: cutlass.Int32, - nope_dim: cutlass.Int32, - pos_dim: cutlass.Int32, - inverse: cutlass.Int32, - adjoint: cutlass.Int32, - total_work: cutlass.Int32, - stream: cuda.CUstream, - ): - _launch_1d( - _thd_compressed_rope_kernel, - "dsv4_cp_thd_compressed_rope", - ( - x, - out, - cos, - sin, - comp_ids, - ratio, - row_width, - head_dim, - nope_dim, - pos_dim, - inverse, - adjoint, - total_work, - ), - total_work, - stream, - ) - # Kernel contract: # hidden_local: local hidden rows, shape (l_local, row_width), bf16/fp16/fp32. # boundary_hidden: left boundary rows, shape (d_window, row_width). # hidden_compact: output compact rows, shape (compact_len, row_width). - # cu_compact: int32 output prefix, shape (n_seq + 1,). - # seq_ids/comp_ids/valid: compressed metadata outputs, shape (c_cap,). + # comp_ids: original per-sequence compressed group ids, shape (c_cap,). # Enumerates visible full compression groups in [global_start-d_comp, # global_start+l_local), copies their ratio tokens from boundary/local - # hidden into hidden_compact, and emits compact sequence metadata. + # hidden into hidden_compact, and emits compressed group ids. @cute.kernel def _compressor_input_compact_fwd_kernel( hidden_local: cute.Tensor, boundary_hidden: cute.Tensor, hidden_compact: cute.Tensor, cu_seqlens: cute.Tensor, - cu_compact: cute.Tensor, - seq_ids: cute.Tensor, comp_ids: cute.Tensor, - valid: cute.Tensor, n_seq: cutlass.Int32, global_start: cutlass.Int32, l_local: cutlass.Int32, @@ -570,9 +456,7 @@ def _compressor_input_compact_fwd_kernel( if row < c_cap and tidx == 0: slot = row - seq_ids[slot] = -1 comp_ids[slot] = -1 - valid[slot] = False running = 0 for seq in range(n_seq): @@ -602,43 +486,9 @@ def _compressor_input_compact_fwd_kernel( if slot >= running and slot < running + visible_group_count: comp_id = first_visible_group + slot - running - group_end = seq_start + (comp_id + 1) * ratio - seq_ids[slot] = seq comp_ids[slot] = comp_id - valid[slot] = group_end - 1 >= range_start and group_end - 1 < range_end running = running + visible_group_count - if row < n_seq + 1 and tidx == 1: - upto = row - running = 0 - for seq in range(n_seq): - if seq < upto: - seq_start = cu_seqlens[seq] - seq_end = cu_seqlens[seq + 1] - local_seq_start = seq_start - if local_seq_start < range_start: - local_seq_start = range_start - local_seq_end = seq_end - if local_seq_end > range_end: - local_seq_end = range_end - - if local_seq_start < local_seq_end: - seq_full_group_count = (seq_end - seq_start) // ratio - first_visible_numer = first_range_group_start - seq_start - if first_visible_numer < 0: - first_visible_numer = 0 - first_visible_group = 0 - if first_visible_numer > 0: - first_visible_group = (first_visible_numer + ratio - 1) // ratio - stop_visible_group = (local_seq_end - seq_start) // ratio - if stop_visible_group > seq_full_group_count: - stop_visible_group = seq_full_group_count - visible_group_count = stop_visible_group - first_visible_group - if visible_group_count < 0: - visible_group_count = 0 - running = running + visible_group_count * ratio - cu_compact[row] = running - # Launch contract for _compressor_input_compact_fwd_kernel. # total_work covers compact rows plus metadata rows. @cute.jit @@ -647,10 +497,7 @@ def _compressor_input_compact_fwd_launch( boundary_hidden: cute.Tensor, hidden_compact: cute.Tensor, cu_seqlens: cute.Tensor, - cu_compact: cute.Tensor, - seq_ids: cute.Tensor, comp_ids: cute.Tensor, - valid: cute.Tensor, n_seq: cutlass.Int32, global_start: cutlass.Int32, l_local: cutlass.Int32, @@ -671,10 +518,7 @@ def _compressor_input_compact_fwd_launch( boundary_hidden, hidden_compact, cu_seqlens, - cu_compact, - seq_ids, comp_ids, - valid, n_seq, global_start, l_local, @@ -860,314 +704,185 @@ def _compressor_input_compact_bwd_launch( stream=stream, ) - # Kernel contract: - # kv_local/boundary_kv/compressed_rank_major: source tensors for forward, - # or output gradient tensors for backward, each flattened to - # ``(rows, row_width)``. - # kv_full: output tensor for forward, or grad_kv_full input for backward, - # shape ``(kv_full_capacity, row_width)``. - # rank-major metadata inputs describe compressed rows after all-gather. - # direction: 0 packs source tensors into kv_full; 1 scatters kv_full - # gradients back to the source tensors. - # Computes the physical KV-full row mapping inline. No source-map tensor is - # materialized or saved for backward; backward recomputes the same mapping. @cute.kernel - def _thd_full_kv_pack_kernel( - kv_local: cute.Tensor, - boundary_kv: cute.Tensor, - compressed_rank_major: cute.Tensor, - kv_full: cute.Tensor, - seq_ids_rank_major: cute.Tensor, - comp_ids_rank_major: cute.Tensor, - valid_rank_major: cute.Tensor, + def _build_attention_indices_kernel( cu_seqlens: cute.Tensor, cu_seqlens_compressed: cute.Tensor, - rank_major_by_seq_major: cute.Tensor, + indexer_topk: cute.Tensor, + rank_row_for_seq_row: cute.Tensor, + topk_idxs: cute.Tensor, + topk_length: cute.Tensor, n_seq: cutlass.Int32, global_start: cutlass.Int32, l_local: cutlass.Int32, chunk0_start: cutlass.Int32, chunk1_start: cutlass.Int32, - chunk_count: cutlass.Constexpr, + chunk_count: cutlass.Int32, chunk_len: cutlass.Int32, - window_capacity_per_chunk: cutlass.Int32, - d_window: cutlass.Constexpr, - ratio: cutlass.Constexpr, - compressed_rows: cutlass.Int32, - use_seq_major_map: cutlass.Constexpr, + d_window: cutlass.Int32, + window_size: cutlass.Int32, + ratio: cutlass.Int32, + compressed_width: cutlass.Int32, seq_major_rows: cutlass.Int32, - kv_full_capacity: cutlass.Int32, - row_width: cutlass.Constexpr, - direction: cutlass.Constexpr, - total_work: cutlass.Int32, + compressed_base: cutlass.Int32, + has_indexer_topk: cutlass.Int32, + has_compressed: cutlass.Int32, + total_width: cutlass.Int32, + parallel_dense: cutlass.Int32, ): tidx, _, _ = cute.arch.thread_idx() bidx, _, _ = cute.arch.block_idx() - out_row = bidx + linear = bidx * 128 + tidx - if cutlass.const_expr(kv_full.element_type.width == 16 and row_width % 4 == 0): - vec_elems = cutlass.Int32(4) - vec_cols = row_width // 4 - else: - vec_elems = cutlass.Int32(1) - vec_cols = row_width + row = linear + col = cutlass.Int32(0) + if parallel_dense != 0: + if linear < l_local * total_width: + row = linear // total_width + col = linear - row * total_width + else: + row = l_local - if out_row < total_work: - source_kind = cutlass.Int32(-1) - source_index = cutlass.Int32(0) - resolve_row = tidx == 0 or tidx == 32 - - if resolve_row and chunk_count != 0: - shared_compressed_base = chunk_count * window_capacity_per_chunk - if out_row >= shared_compressed_base: - compressed_idx = out_row - shared_compressed_base - if compressed_idx >= 0 and compressed_idx < compressed_rows: - source_kind = 2 - source_index = compressed_idx - else: - chunk_id = out_row // window_capacity_per_chunk - chunk_row = out_row - chunk_id * window_capacity_per_chunk - chunk_start = chunk0_start - if chunk_id == 1: - chunk_start = chunk1_start - chunk_end = chunk_start + chunk_len - - offset = 0 - for seq in range(n_seq): - seq_start = cu_seqlens[seq] - seq_end = cu_seqlens[seq + 1] - local_seq_start = seq_start - if local_seq_start < chunk_start: - local_seq_start = chunk_start - local_seq_end = seq_end - if local_seq_end > chunk_end: - local_seq_end = chunk_end - if local_seq_start < local_seq_end: - window_start = local_seq_start - d_window - if window_start < seq_start: - window_start = seq_start - window_len = local_seq_end - window_start - if chunk_row >= offset and chunk_row < offset + window_len: - src_global = window_start + chunk_row - offset - if src_global < chunk_start: - source_kind = 1 - source_index = ( - chunk_id * d_window + src_global - (chunk_start - d_window) - ) - else: - source_kind = 0 - source_index = chunk_id * chunk_len + src_global - chunk_start - offset = offset + window_len - elif resolve_row: - offset = 0 - for seq in range(n_seq): - seq_start = cu_seqlens[seq] - seq_end = cu_seqlens[seq + 1] - global_end = global_start + l_local - local_seq_start = seq_start - if local_seq_start < global_start: - local_seq_start = global_start - local_seq_end = seq_end - if local_seq_end > global_end: - local_seq_end = global_end + if row < l_local: + local_start = global_start + row_in_chunk = row + chunk_id = cutlass.Int32(0) + boundary_rows = d_window + if chunk_count != 0: + chunk_id = row // chunk_len + row_in_chunk = row - chunk_id * chunk_len + local_start = chunk0_start + if chunk_id == 1: + local_start = chunk1_start + boundary_rows = chunk_count * d_window - if local_seq_start < local_seq_end: - window_start = local_seq_start - d_window - if window_start < seq_start: - window_start = seq_start - window_len = local_seq_end - window_start - comp_len = 0 - comp_start = 0 - comp_end = 0 - if ratio > 1 and use_seq_major_map != 0: - comp_start = cu_seqlens_compressed[seq] - comp_end = cu_seqlens_compressed[seq + 1] - comp_len = comp_end - comp_start - elif ratio > 1: - comp_len = (seq_end - seq_start) // ratio - seq_total = window_len + comp_len - - if out_row >= offset and out_row < offset + seq_total: - inner = out_row - offset - if inner < window_len: - src_global = window_start + inner - if src_global < global_start: - source_kind = 1 - source_index = src_global - (global_start - d_window) - else: - source_kind = 0 - source_index = src_global - global_start + global_q = local_start + row_in_chunk + seq_id = cutlass.Int32(-1) + seq_start_found = cutlass.Int32(0) + seq_comp_start = cutlass.Int32(0) + seq_comp_len = cutlass.Int32(0) + for seq in range(n_seq): + seq_start = cu_seqlens[seq] + seq_end = cu_seqlens[seq + 1] + if global_q >= seq_start and global_q < seq_end: + seq_id = seq + seq_start_found = seq_start + if has_compressed != 0: + seq_comp_start = cu_seqlens_compressed[seq] + seq_comp_len = cu_seqlens_compressed[seq + 1] - seq_comp_start + + if parallel_dense != 0: + topk_value = cutlass.Int32(-1) + length = cutlass.Int32(0) + if seq_id >= 0: + window_start = global_q - window_size + 1 + if window_start < seq_start_found: + window_start = seq_start_found + window_count = global_q - window_start + 1 + comp_count = cutlass.Int32(0) + if has_compressed != 0 and ratio > 1: + comp_count = (global_q - seq_start_found + 1) // ratio + if comp_count > compressed_width: + comp_count = compressed_width + if comp_count > seq_comp_len: + comp_count = seq_comp_len + length = window_count + comp_count + if col < window_count: + pos = window_start + col + if pos < local_start: + topk_value = chunk_id * d_window + pos - (local_start - d_window) + else: + topk_value = boundary_rows + chunk_id * chunk_len + pos - local_start + elif col < length: + seq_major_id = seq_comp_start + col - window_count + if seq_major_id >= 0 and seq_major_id < seq_major_rows: + rank_major_id = rank_row_for_seq_row[seq_major_id] + if rank_major_id >= 0: + topk_value = compressed_base + rank_major_id + elif total_width > 0: + length = 1 + if col == 0: + topk_value = 0 + topk_idxs[row, col] = topk_value + if col == 0: + topk_length[row] = length + else: + for out_col in range(total_width): + topk_idxs[row, out_col] = -1 + topk_length[row] = 0 + + if seq_id >= 0: + write_col = cutlass.Int32(0) + window_start = global_q - window_size + 1 + if window_start < seq_start_found: + window_start = seq_start_found + window_count = global_q - window_start + 1 + for window_col in range(window_size): + if window_col < window_count: + pos = window_start + window_col + if pos < local_start: + topk_idxs[row, write_col] = ( + chunk_id * d_window + pos - (local_start - d_window) + ) else: - comp_id = inner - window_len - if use_seq_major_map != 0: - seq_major = comp_start + comp_id - if seq_major >= comp_start and seq_major < comp_end: - if seq_major >= 0 and seq_major < seq_major_rows: - rank_major = rank_major_by_seq_major[seq_major] - if rank_major >= 0: - source_kind = 2 - source_index = rank_major - else: - # Slow fallback for no-indexer paths: linear scan is - # intentional and not expected to be the hot training path. - for comp_row in range(compressed_rows): - if ( - valid_rank_major[comp_row] - and seq_ids_rank_major[comp_row] == seq - and comp_ids_rank_major[comp_row] == comp_id - ): - source_kind = 2 - source_index = comp_row - offset = offset + seq_total - - source_kind = cute.arch.shuffle_sync(source_kind, 0, mask=-1, mask_and_clamp=31) - source_index = cute.arch.shuffle_sync(source_index, 0, mask=-1, mask_and_clamp=31) + topk_idxs[row, write_col] = ( + boundary_rows + chunk_id * chunk_len + pos - local_start + ) + write_col = write_col + 1 - vec_col = tidx - while vec_col < vec_cols: - col = vec_col - if direction == 0: - if cutlass.const_expr(kv_full.element_type.width == 16 and row_width % 4 == 0): - dst_offset = out_row * row_width + vec_col * vec_elems - value = cutlass.Int64(0) - if source_kind == 0: - src_offset = source_index * row_width + vec_col * vec_elems - src_ptr = cute.recast_ptr( - kv_local.iterator + src_offset, dtype=cutlass.Int64 - ) - value = cute.arch.load( - src_ptr.llvm_ptr, - cutlass.Int64, - level1_eviction_priority="evict_no_allocate", - ) - elif source_kind == 1: - src_offset = source_index * row_width + vec_col * vec_elems - src_ptr = cute.recast_ptr( - boundary_kv.iterator + src_offset, dtype=cutlass.Int64 - ) - value = cute.arch.load( - src_ptr.llvm_ptr, - cutlass.Int64, - level1_eviction_priority="evict_no_allocate", - ) - elif source_kind == 2: - src_offset = source_index * row_width + vec_col * vec_elems - src_ptr = cute.recast_ptr( - compressed_rank_major.iterator + src_offset, dtype=cutlass.Int64 - ) - value = cute.arch.load( - src_ptr.llvm_ptr, - cutlass.Int64, - level1_eviction_priority="evict_no_allocate", - ) - dst_ptr = cute.recast_ptr( - kv_full.iterator + dst_offset, dtype=cutlass.Int64 - ) - cute.arch.store( - dst_ptr.llvm_ptr, value, level1_eviction_priority="evict_first" - ) - else: - value = cutlass.Float32(0.0).to(kv_full.element_type) - if source_kind == 0: - value = kv_local[source_index, col] - elif source_kind == 1: - value = boundary_kv[source_index, col] - elif source_kind == 2: - value = compressed_rank_major[source_index, col] - kv_full[out_row, col] = value - else: - if cutlass.const_expr(kv_full.element_type.width == 16 and row_width % 4 == 0): - src_offset = out_row * row_width + vec_col * vec_elems - src_ptr = cute.recast_ptr( - kv_full.iterator + src_offset, dtype=cutlass.Int64 - ) - value = cute.arch.load( - src_ptr.llvm_ptr, - cutlass.Int64, - level1_eviction_priority="evict_no_allocate", - ) - if source_kind == 0: - dst_offset = source_index * row_width + vec_col * vec_elems - dst_ptr = cute.recast_ptr( - kv_local.iterator + dst_offset, dtype=cutlass.Int64 - ) - cute.arch.store( - dst_ptr.llvm_ptr, value, level1_eviction_priority="evict_first" - ) - elif source_kind == 1: - dst_offset = source_index * row_width + vec_col * vec_elems - dst_ptr = cute.recast_ptr( - boundary_kv.iterator + dst_offset, dtype=cutlass.Int64 - ) - cute.arch.store( - dst_ptr.llvm_ptr, value, level1_eviction_priority="evict_first" - ) - elif source_kind == 2: - dst_offset = source_index * row_width + vec_col * vec_elems - dst_ptr = cute.recast_ptr( - compressed_rank_major.iterator + dst_offset, dtype=cutlass.Int64 - ) - cute.arch.store( - dst_ptr.llvm_ptr, value, level1_eviction_priority="evict_first" - ) - else: - if source_kind == 0: - kv_local[source_index, col] = kv_full[out_row, col] - elif source_kind == 1: - boundary_kv[source_index, col] = kv_full[out_row, col] - elif source_kind == 2: - compressed_rank_major[source_index, col] = kv_full[out_row, col] - vec_col = vec_col + 64 + if has_compressed != 0 and ratio > 1: + for compressed_col in range(compressed_width): + comp_id = indexer_topk[row, compressed_col] + if comp_id >= 0 and comp_id < seq_comp_len: + seq_major_id = seq_comp_start + comp_id + if seq_major_id >= 0 and seq_major_id < seq_major_rows: + rank_major_id = rank_row_for_seq_row[seq_major_id] + if rank_major_id >= 0: + topk_idxs[row, write_col] = compressed_base + rank_major_id + write_col = write_col + 1 + topk_length[row] = write_col + elif total_width > 0: + topk_idxs[row, 0] = 0 + topk_length[row] = 1 - # Launch contract for _thd_full_kv_pack_kernel. - # Supports contiguous CP and two-chunk CP. ``use_seq_major_map`` selects - # whether compressed ids must be translated through rank_major_by_seq_major. - # ``row_width`` is constexpr to preserve the existing compiled layout key. @cute.jit - def _thd_full_kv_pack_launch( - kv_local: cute.Tensor, - boundary_kv: cute.Tensor, - compressed_rank_major: cute.Tensor, - kv_full: cute.Tensor, - seq_ids_rank_major: cute.Tensor, - comp_ids_rank_major: cute.Tensor, - valid_rank_major: cute.Tensor, + def _build_attention_indices_launch( cu_seqlens: cute.Tensor, cu_seqlens_compressed: cute.Tensor, - rank_major_by_seq_major: cute.Tensor, + indexer_topk: cute.Tensor, + rank_row_for_seq_row: cute.Tensor, + topk_idxs: cute.Tensor, + topk_length: cute.Tensor, n_seq: cutlass.Int32, global_start: cutlass.Int32, l_local: cutlass.Int32, chunk0_start: cutlass.Int32, chunk1_start: cutlass.Int32, - chunk_count: cutlass.Constexpr, + chunk_count: cutlass.Int32, chunk_len: cutlass.Int32, - window_capacity_per_chunk: cutlass.Int32, - d_window: cutlass.Constexpr, - ratio: cutlass.Constexpr, - compressed_rows: cutlass.Int32, - use_seq_major_map: cutlass.Constexpr, + d_window: cutlass.Int32, + window_size: cutlass.Int32, + ratio: cutlass.Int32, + compressed_width: cutlass.Int32, seq_major_rows: cutlass.Int32, - kv_full_capacity: cutlass.Int32, - row_width: cutlass.Constexpr, - direction: cutlass.Constexpr, - total_work: cutlass.Int32, + compressed_base: cutlass.Int32, + has_indexer_topk: cutlass.Int32, + has_compressed: cutlass.Int32, + total_width: cutlass.Int32, + parallel_dense: cutlass.Int32, + launch_work: cutlass.Int32, stream: cuda.CUstream, ): - _launch_named( - _thd_full_kv_pack_kernel, - "dsv4_cp_thd_full_kv_pack", + _launch_1d( + _build_attention_indices_kernel, + "dsv4_cp_build_attention_indices", ( - kv_local, - boundary_kv, - compressed_rank_major, - kv_full, - seq_ids_rank_major, - comp_ids_rank_major, - valid_rank_major, cu_seqlens, cu_seqlens_compressed, - rank_major_by_seq_major, + indexer_topk, + rank_row_for_seq_row, + topk_idxs, + topk_length, n_seq, global_start, l_local, @@ -1175,771 +890,24 @@ def _thd_full_kv_pack_launch( chunk1_start, chunk_count, chunk_len, - window_capacity_per_chunk, d_window, + window_size, ratio, - compressed_rows, - use_seq_major_map, + compressed_width, seq_major_rows, - kv_full_capacity, - row_width, - direction, - total_work, + compressed_base, + has_indexer_topk, + has_compressed, + total_width, + parallel_dense, ), - grid=(total_work, 1, 1), - block=(64, 1, 1), - stream=stream, + launch_work, + stream, ) - # Kernel contract: - # rank_by_seq_major: int32 output reverse map, shape (total_seq_major_rows,). - # Initializes the reverse map before the deterministic rank-map build. - @cute.kernel - def _repack_compressed_kv_init_rank_map_kernel( - rank_major_by_seq_major: cute.Tensor, seq_major_rows: cutlass.Int32 - ): - tidx, _, _ = cute.arch.thread_idx() - bidx, _, _ = cute.arch.block_idx() - linear = bidx * 128 + tidx - if linear < seq_major_rows: - rank_major_by_seq_major[linear] = -1 - - # Kernel contract: - # seq_ids/comp_ids/valid: rank-major metadata, shape (rows,). - # cu_compressed: int32 compressed sequence prefixes, shape (n_seq + 1,). - # rank_by_seq_major: int32 output reverse map, shape (total_seq_major_rows,). - # Builds a deterministic seq-major -> rank-major map. Multiple rank-major - # candidates can describe the same compressed group near CP chunk - # boundaries, so the kernel selects the greatest rank-major row with - # atomic max instead of racing concurrent value writes. @cute.kernel - def _repack_compressed_kv_rank_map_kernel( - seq_ids_rank_major: cute.Tensor, - comp_ids_rank_major: cute.Tensor, - valid_rank_major: cute.Tensor, - cu_seqlens_compressed: cute.Tensor, - rank_major_by_seq_major: cute.Tensor, - rank_major_rows: cutlass.Int32, - seq_major_rows: cutlass.Int32, - ): - tidx, _, _ = cute.arch.thread_idx() - bidx, _, _ = cute.arch.block_idx() - rank_row = bidx * 128 + tidx - if rank_row < rank_major_rows and valid_rank_major[rank_row] == 1: - seq = seq_ids_rank_major[rank_row] - comp = comp_ids_rank_major[rank_row] - seq_major = cu_seqlens_compressed[seq] + comp - if seq_major >= 0 and seq_major < seq_major_rows: - rank_map_ptr = rank_major_by_seq_major.iterator + seq_major - cute.arch.atomic_max(rank_map_ptr.llvm_ptr, rank_row, sem="relaxed", scope="gpu") - - # Kernel contract: - # rank_major: compressed rows in rank-major order, shape (rows, row_width). - # rank_by_seq_major: deterministic reverse map built by - # _repack_compressed_kv_rank_map_kernel. - # out: seq-major compressed rows, shape (total_seq_major_rows, row_width). - # Gathers selected rank-major rows into seq-major order and zero-fills - # rows with no valid compressed group. - @cute.kernel - def _repack_compressed_kv_to_seq_major_kernel( - compressed_rank_major: cute.Tensor, - compressed_seq_major: cute.Tensor, - rank_major_by_seq_major: cute.Tensor, - rank_major_rows: cutlass.Int32, - seq_major_rows: cutlass.Int32, - row_width: cutlass.Constexpr, - total_work: cutlass.Int32, - ): - tidx, _, _ = cute.arch.thread_idx() - bidx, _, _ = cute.arch.block_idx() - linear = bidx * 128 + tidx - - if linear < seq_major_rows * row_width: - row = linear // row_width - col = linear - row * row_width - rank_row = rank_major_by_seq_major[row] - value = cutlass.Float32(0.0).to(compressed_seq_major.element_type) - if rank_row >= 0 and rank_row < rank_major_rows: - value = compressed_rank_major[rank_row, col] - compressed_seq_major[row, col] = value - - # Launch contract for _repack_compressed_kv_init_rank_map_kernel. - # launch_work is seq-major rows. - @cute.jit - def _repack_compressed_kv_init_rank_map_launch( - rank_major_by_seq_major: cute.Tensor, - seq_major_rows: cutlass.Int32, - launch_work: cutlass.Int32, - stream: cuda.CUstream, - ): - _launch_1d( - _repack_compressed_kv_init_rank_map_kernel, - "dsv4_cp_repack_compressed_kv_init_rank_map", - (rank_major_by_seq_major, seq_major_rows), - launch_work, - stream, - ) - - # Launch contract for _repack_compressed_kv_rank_map_kernel. - # launch_work is rank-major rows. - @cute.jit - def _repack_compressed_kv_rank_map_launch( - seq_ids_rank_major: cute.Tensor, - comp_ids_rank_major: cute.Tensor, - valid_rank_major: cute.Tensor, - cu_seqlens_compressed: cute.Tensor, - rank_major_by_seq_major: cute.Tensor, - rank_major_rows: cutlass.Int32, - seq_major_rows: cutlass.Int32, - launch_work: cutlass.Int32, - stream: cuda.CUstream, - ): - _launch_1d( - _repack_compressed_kv_rank_map_kernel, - "dsv4_cp_repack_compressed_kv_rank_map", - ( - seq_ids_rank_major, - comp_ids_rank_major, - valid_rank_major, - cu_seqlens_compressed, - rank_major_by_seq_major, - rank_major_rows, - seq_major_rows, - ), - launch_work, - stream, - ) - - # Launch contract for _repack_compressed_kv_to_seq_major_kernel. - # total_work is seq-major rows * row_width. - @cute.jit - def _repack_compressed_kv_to_seq_major_launch( - compressed_rank_major: cute.Tensor, - compressed_seq_major: cute.Tensor, - rank_major_by_seq_major: cute.Tensor, - rank_major_rows: cutlass.Int32, - seq_major_rows: cutlass.Int32, - row_width: cutlass.Constexpr, - total_work: cutlass.Int32, - stream: cuda.CUstream, - ): - _launch_1d( - _repack_compressed_kv_to_seq_major_kernel, - "dsv4_cp_repack_compressed_kv_to_seq_major", - ( - compressed_rank_major, - compressed_seq_major, - rank_major_by_seq_major, - rank_major_rows, - seq_major_rows, - row_width, - total_work, - ), - total_work, - stream, - ) - - # Kernel contract: - # cu_seqlens: int32, shape (n_seq + 1,), global packed sequence prefixes. - # seq_ids/comp_ids: int32 outputs, shape (total_rows,). - # valid: bool output, shape (total_rows,). - # total_rows is cp_size * c_cap_per_rank. For each rank-major compressed - # row, computes (seq_id, compressed group id) and marks whether that group - # ends inside the owning rank/chunk. - @cute.kernel - def _build_compressed_row_metadata_kernel( - cu_seqlens: cute.Tensor, - seq_ids: cute.Tensor, - comp_ids: cute.Tensor, - valid: cute.Tensor, - n_seq: cutlass.Int32, - cp_size: cutlass.Int32, - chunk_len: cutlass.Int32, - ratio: cutlass.Int32, - d_comp: cutlass.Int32, - c_cap_per_chunk: cutlass.Int32, - c_cap_per_rank: cutlass.Int32, - two_chunk: cutlass.Int32, - total_rows: cutlass.Int32, - ): - tidx, _, _ = cute.arch.thread_idx() - bidx, _, _ = cute.arch.block_idx() - row = bidx * 128 + tidx - if row < total_rows: - rank = row // c_cap_per_rank - rank_slot = row - rank * c_cap_per_rank - slot = rank_slot - rank_start = rank * chunk_len - row_is_padding = False - if two_chunk != 0: - local_chunk = rank_slot // c_cap_per_chunk - if local_chunk >= 2: - row_is_padding = True - else: - slot = rank_slot - local_chunk * c_cap_per_chunk - chunk_id = rank - if local_chunk != 0: - chunk_id = cp_size * 2 - 1 - rank - rank_start = chunk_id * chunk_len - rank_end = rank_start + chunk_len - - seq_ids[row] = -1 - comp_ids[row] = -1 - valid[row] = False - - if not row_is_padding: - running = 0 - for seq in range(n_seq): - seq_start = cu_seqlens[seq] - seq_end = cu_seqlens[seq + 1] - local_seq_start = seq_start - if local_seq_start < rank_start: - local_seq_start = rank_start - local_seq_end = seq_end - if local_seq_end > rank_end: - local_seq_end = rank_end - - if local_seq_start < local_seq_end: - n_full_groups = (seq_end - seq_start) // ratio - first_numer = rank_start - d_comp - seq_start - if first_numer < 0: - first_numer = 0 - first_group = 0 - if first_numer > 0: - first_group = (first_numer + ratio - 1) // ratio - stop_group = (local_seq_end - seq_start) // ratio - if stop_group > n_full_groups: - stop_group = n_full_groups - group_count = stop_group - first_group - if group_count < 0: - group_count = 0 - - if slot >= running and slot < running + group_count: - comp_id = first_group + slot - running - group_end = seq_start + (comp_id + 1) * ratio - seq_ids[row] = seq - comp_ids[row] = comp_id - valid[row] = group_end - 1 >= rank_start and group_end - 1 < rank_end - running = running + group_count - - # Launch contract for _build_compressed_row_metadata_kernel. - # two_chunk selects one contiguous chunk per rank or two chunks - # per rank; c_cap_per_rank is the fixed local compressed capacity. - @cute.jit - def _build_compressed_row_metadata_launch( - cu_seqlens: cute.Tensor, - seq_ids: cute.Tensor, - comp_ids: cute.Tensor, - valid: cute.Tensor, - n_seq: cutlass.Int32, - cp_size: cutlass.Int32, - chunk_len: cutlass.Int32, - ratio: cutlass.Int32, - d_comp: cutlass.Int32, - c_cap_per_chunk: cutlass.Int32, - c_cap_per_rank: cutlass.Int32, - two_chunk: cutlass.Int32, - total_rows: cutlass.Int32, - stream: cuda.CUstream, - ): - _launch_1d( - _build_compressed_row_metadata_kernel, - "dsv4_cp_build_compressed_row_metadata", - ( - cu_seqlens, - seq_ids, - comp_ids, - valid, - n_seq, - cp_size, - chunk_len, - ratio, - d_comp, - c_cap_per_chunk, - c_cap_per_rank, - two_chunk, - total_rows, - ), - total_rows, - stream, - ) - - # Kernel contract: - # k_indexer_seq_major: compressed indexer K, shape (seq_major_rows, row_width). - # cu_q: int32 query prefixes, shape (n_seq + 1,). - # cu_compressed: int32 compressed K prefixes, shape (n_seq + 1,). - # k_topk: output packed K for local trapezoid top-k, shape (k_rows, row_width). - # cu_q_topk/cu_k_topk/seq_lens: int32 outputs, shape (n_seq + 1)/(n_seq + 1)/(n_seq). - # Copies the compressed K rows visible to local query rows and emits local - # q/k prefix sums for the indexer kernel. - @cute.kernel - def _build_indexer_topk_metadata_kernel( - k_seq_major: cute.Tensor, - k_topk: cute.Tensor, - cu_seqlens_q: cute.Tensor, - cu_seqlens_compressed: cute.Tensor, - cu_q_topk: cute.Tensor, - cu_k_topk: cute.Tensor, - seq_lens: cute.Tensor, - n_seq: cutlass.Int32, - global_start: cutlass.Int32, - l_local: cutlass.Int32, - ratio: cutlass.Int32, - row_width: cutlass.Constexpr, - seq_major_rows: cutlass.Int32, - total_work: cutlass.Int32, - ): - tidx, _, _ = cute.arch.thread_idx() - bidx, _, _ = cute.arch.block_idx() - if bidx == 0 and tidx == 0: - global_end = global_start + l_local - q_prefix = 0 - k_prefix = 0 - cu_q_topk[0] = cutlass.Int32(0) - cu_k_topk[0] = cutlass.Int32(0) - for seq in range(n_seq): - seq_start = cu_seqlens_q[seq] - seq_end = cu_seqlens_q[seq + 1] - local_start = seq_start - if local_start < global_start: - local_start = global_start - local_end = seq_end - if local_end > global_end: - local_end = global_end - - q_len = 0 - k_len = 0 - if local_start < local_end: - q_len = local_end - local_start - seq_comp_start = cu_seqlens_compressed[seq] - seq_comp_end = cu_seqlens_compressed[seq + 1] - seq_comp_len = seq_comp_end - seq_comp_start - k_len = (local_end - seq_start) // ratio - if k_len > seq_comp_len: - k_len = seq_comp_len - - q_prefix = q_prefix + q_len - k_prefix = k_prefix + k_len - cu_q_topk[seq + 1] = q_prefix - cu_k_topk[seq + 1] = k_prefix - for row in range(q_len): - seq_lens[q_prefix - q_len + row] = k_len - - actual_total = cu_seqlens_q[n_seq] - padding_q = 0 - if global_end > actual_total: - padding_start = actual_total - if padding_start < global_start: - padding_start = global_start - if padding_start < global_end: - padding_q = global_end - padding_start - cu_q_topk[n_seq + 1] = q_prefix + padding_q - cu_k_topk[n_seq + 1] = k_prefix - for row in range(padding_q): - seq_lens[q_prefix + row] = cutlass.Int32(0) - - linear = bidx * 128 + tidx - if linear < total_work: - dst_row = linear // row_width - col = linear - dst_row * row_width - k_topk[dst_row, col] = cutlass.Float32(0.0).to(k_topk.element_type) - - global_end = global_start + l_local - k_prefix = 0 - src_row = -1 - for seq in range(n_seq): - seq_start = cu_seqlens_q[seq] - seq_end = cu_seqlens_q[seq + 1] - local_start = seq_start - if local_start < global_start: - local_start = global_start - local_end = seq_end - if local_end > global_end: - local_end = global_end - - k_len = 0 - if local_start < local_end: - seq_comp_start = cu_seqlens_compressed[seq] - seq_comp_end = cu_seqlens_compressed[seq + 1] - seq_comp_len = seq_comp_end - seq_comp_start - k_len = (local_end - seq_start) // ratio - if k_len > seq_comp_len: - k_len = seq_comp_len - if dst_row >= k_prefix and dst_row < k_prefix + k_len: - src_row = seq_comp_start + dst_row - k_prefix - k_prefix = k_prefix + k_len - - if src_row >= 0 and src_row < seq_major_rows: - k_topk[dst_row, col] = k_seq_major[src_row, col] - - # Launch contract for _build_indexer_topk_metadata_kernel. - # total_work covers max(k_rows * row_width, n_seq + 1). - @cute.jit - def _build_indexer_topk_metadata_launch( - k_seq_major: cute.Tensor, - k_topk: cute.Tensor, - cu_seqlens_q: cute.Tensor, - cu_seqlens_compressed: cute.Tensor, - cu_q_topk: cute.Tensor, - cu_k_topk: cute.Tensor, - seq_lens: cute.Tensor, - n_seq: cutlass.Int32, - global_start: cutlass.Int32, - l_local: cutlass.Int32, - ratio: cutlass.Int32, - row_width: cutlass.Constexpr, - seq_major_rows: cutlass.Int32, - total_work: cutlass.Int32, - launch_work: cutlass.Int32, - stream: cuda.CUstream, - ): - _launch_1d( - _build_indexer_topk_metadata_kernel, - "dsv4_cp_build_indexer_topk_metadata", - ( - k_seq_major, - k_topk, - cu_seqlens_q, - cu_seqlens_compressed, - cu_q_topk, - cu_k_topk, - seq_lens, - n_seq, - global_start, - l_local, - ratio, - row_width, - seq_major_rows, - total_work, - ), - launch_work, - stream, - ) - - # Kernel contract: - # cu_seqlens: int32 global query prefixes, shape (n_seq + 1,). - # optional indexer_topk: int32 compressed logical ids, shape (l_local, compressed_width). - # rank_major_by_seq_major: int32 reverse compressed map when compression is active. - # topk_idxs: int32 output, shape (l_local, window_size + compressed_width). - # topk_length: int32 output valid count, shape (l_local,). - # Lowers each local query row to physical KV-full ids: window ids first, - # then compressed ids from either dense causal range or indexer top-k. - @cute.kernel - def _build_attention_indices_kernel( - cu_seqlens: cute.Tensor, - cu_seqlens_compressed: cute.Tensor, - indexer_topk: cute.Tensor, - rank_major_by_seq_major: cute.Tensor, - topk_idxs: cute.Tensor, - topk_length: cute.Tensor, - n_seq: cutlass.Int32, - global_start: cutlass.Int32, - l_local: cutlass.Int32, - chunk0_start: cutlass.Int32, - chunk1_start: cutlass.Int32, - chunk_count: cutlass.Int32, - chunk_len: cutlass.Int32, - window_capacity_per_chunk: cutlass.Int32, - shared_compressed_base: cutlass.Int32, - d_window: cutlass.Int32, - window_size: cutlass.Int32, - ratio: cutlass.Int32, - compressed_width: cutlass.Int32, - seq_major_rows: cutlass.Int32, - has_indexer_topk: cutlass.Int32, - has_compressed: cutlass.Int32, - total_width: cutlass.Int32, - parallel_dense: cutlass.Int32, - ): - tidx, _, _ = cute.arch.thread_idx() - bidx, _, _ = cute.arch.block_idx() - linear = bidx * 128 + tidx - row_count = l_local - if chunk_count != 0: - row_count = chunk_count * chunk_len - - row = 0 - if parallel_dense != 0: - if linear < row_count * total_width: - row = linear // total_width - col = linear - row * total_width - topk_value = -1 - length = 0 - - global_q = global_start + row - seq_id = -1 - seq_start_found = 0 - seq_end_found = 0 - for seq in range(n_seq): - seq_start = cu_seqlens[seq] - seq_end = cu_seqlens[seq + 1] - if global_q >= seq_start and global_q < seq_end: - seq_id = seq - seq_start_found = seq_start - seq_end_found = seq_end - - if seq_id >= 0: - seq_offset = 0 - seq_window_start = 0 - seq_window_len = 0 - seq_comp_len = 0 - global_end = global_start + row_count - for seq in range(n_seq): - seq_start = cu_seqlens[seq] - seq_end = cu_seqlens[seq + 1] - local_seq_start = seq_start - if local_seq_start < global_start: - local_seq_start = global_start - local_seq_end = seq_end - if local_seq_end > global_end: - local_seq_end = global_end - if local_seq_start < local_seq_end: - window_start = local_seq_start - d_window - if window_start < seq_start: - window_start = seq_start - window_len = local_seq_end - window_start - comp_len = 0 - if has_compressed != 0 and ratio > 1: - comp_len = (seq_end - seq_start) // ratio - if seq == seq_id: - seq_window_start = window_start - seq_window_len = window_len - seq_comp_len = comp_len - if seq < seq_id: - seq_offset = seq_offset + window_len + comp_len - - window_start_for_q = global_q - window_size + 1 - if window_start_for_q < seq_start_found: - window_start_for_q = seq_start_found - if window_start_for_q < seq_window_start: - window_start_for_q = seq_window_start - window_count = global_q - window_start_for_q + 1 - if window_count < 0: - window_count = 0 - if window_count > window_size: - window_count = window_size - - comp_count = 0 - if has_compressed != 0 and compressed_width > 0 and ratio > 1: - comp_count = (global_q - seq_start_found + 1) // ratio - if comp_count > compressed_width: - comp_count = compressed_width - if comp_count > seq_comp_len: - comp_count = seq_comp_len - length = window_count + comp_count - if col < window_count: - topk_value = seq_offset + window_start_for_q - seq_window_start + col - elif col < length: - topk_value = seq_offset + seq_window_len + col - window_count - elif total_width > 0: - # Fused DSA sparse attention expects at least one KV id per - # query row. Padded THD rows carry zero loss gradient, so a - # dummy in-range id is enough to keep the kernel contract. - length = 1 - if col == 0: - topk_value = 0 - - topk_idxs[row, col] = topk_value - if col == 0: - topk_length[row] = length - else: - row = linear - if row >= row_count: - row = row_count - if parallel_dense == 0 and row < row_count: - for col in range(total_width): - topk_idxs[row, col] = -1 - topk_length[row] = 0 - - local_start = global_start - row_in_local = row - chunk_id = 0 - if chunk_count != 0: - chunk_id = row // chunk_len - row_in_local = row - chunk_id * chunk_len - local_start = chunk0_start - if chunk_id == 1: - local_start = chunk1_start - global_q = local_start + row_in_local - seq_id = -1 - seq_start_found = 0 - seq_end_found = 0 - for seq in range(n_seq): - seq_start = cu_seqlens[seq] - seq_end = cu_seqlens[seq + 1] - if global_q >= seq_start and global_q < seq_end: - seq_id = seq - seq_start_found = seq_start - seq_end_found = seq_end - - if seq_id >= 0: - seq_offset = 0 - seq_window_start = 0 - seq_window_len = 0 - seq_comp_len = 0 - seq_comp_start = 0 - seq_comp_end = 0 - global_end = local_start + row_count - if chunk_count != 0: - global_end = local_start + chunk_len - for seq in range(n_seq): - seq_start = cu_seqlens[seq] - seq_end = cu_seqlens[seq + 1] - local_seq_start = seq_start - if local_seq_start < local_start: - local_seq_start = local_start - local_seq_end = seq_end - if local_seq_end > global_end: - local_seq_end = global_end - if local_seq_start < local_seq_end: - window_start = local_seq_start - d_window - if window_start < seq_start: - window_start = seq_start - window_len = local_seq_end - window_start - comp_len = 0 - if has_compressed != 0 and ratio > 1: - if chunk_count != 0: - comp_len = ( - cu_seqlens_compressed[seq + 1] - cu_seqlens_compressed[seq] - ) - else: - comp_len = (seq_end - seq_start) // ratio - if seq == seq_id: - seq_window_start = window_start - seq_window_len = window_len - seq_comp_len = comp_len - if chunk_count != 0 and has_compressed != 0 and ratio > 1: - seq_comp_start = cu_seqlens_compressed[seq] - seq_comp_end = cu_seqlens_compressed[seq + 1] - if seq < seq_id: - seq_offset = seq_offset + window_len - if chunk_count == 0: - seq_offset = seq_offset + comp_len - - write_col = 0 - window_start_for_q = global_q - window_size + 1 - if window_start_for_q < seq_start_found: - window_start_for_q = seq_start_found - window_count = global_q - window_start_for_q + 1 - for w in range(window_size): - pos = window_start_for_q + w - if ( - w < window_count - and pos >= seq_window_start - and pos < seq_window_start + seq_window_len - ): - window_value = seq_offset + pos - seq_window_start - if chunk_count != 0: - window_value = chunk_id * window_capacity_per_chunk + window_value - topk_idxs[row, write_col] = window_value - write_col = write_col + 1 - - if has_compressed != 0 and compressed_width > 0 and ratio > 1: - pos_in_seq = global_q - seq_start_found - for j in range(compressed_width): - comp_id = -1 - if has_indexer_topk != 0: - comp_id = indexer_topk[row, j] - else: - n_visible = (pos_in_seq + 1) // ratio - if n_visible > compressed_width: - n_visible = compressed_width - if j < n_visible: - comp_id = j - if comp_id >= 0 and comp_id < seq_comp_len: - if chunk_count != 0: - seq_major_id = seq_comp_start + comp_id - if seq_major_id >= 0 and seq_major_id < seq_major_rows: - rank_major_id = rank_major_by_seq_major[seq_major_id] - if rank_major_id >= 0: - topk_idxs[row, write_col] = ( - shared_compressed_base + rank_major_id - ) - write_col = write_col + 1 - else: - topk_idxs[row, write_col] = seq_offset + seq_window_len + comp_id - write_col = write_col + 1 - topk_length[row] = write_col - elif total_width > 0: - # Same padding-row contract as the parallel dense path above. - topk_idxs[row, 0] = 0 - topk_length[row] = 1 - - # Launch contract for _build_attention_indices_kernel. - # Supports contiguous and two-chunk CP. total_work is - # l_local * (window_size + compressed_width). - @cute.jit - def _build_attention_indices_launch( - cu_seqlens: cute.Tensor, - cu_seqlens_compressed: cute.Tensor, - indexer_topk: cute.Tensor, - rank_major_by_seq_major: cute.Tensor, - topk_idxs: cute.Tensor, - topk_length: cute.Tensor, - n_seq: cutlass.Int32, - global_start: cutlass.Int32, - l_local: cutlass.Int32, - chunk0_start: cutlass.Int32, - chunk1_start: cutlass.Int32, - chunk_count: cutlass.Int32, - chunk_len: cutlass.Int32, - window_capacity_per_chunk: cutlass.Int32, - shared_compressed_base: cutlass.Int32, - d_window: cutlass.Int32, - window_size: cutlass.Int32, - ratio: cutlass.Int32, - compressed_width: cutlass.Int32, - seq_major_rows: cutlass.Int32, - has_indexer_topk: cutlass.Int32, - has_compressed: cutlass.Int32, - total_width: cutlass.Int32, - parallel_dense: cutlass.Int32, - launch_work: cutlass.Int32, - stream: cuda.CUstream, - ): - _launch_1d( - _build_attention_indices_kernel, - "dsv4_cp_build_attention_indices", - ( - cu_seqlens, - cu_seqlens_compressed, - indexer_topk, - rank_major_by_seq_major, - topk_idxs, - topk_length, - n_seq, - global_start, - l_local, - chunk0_start, - chunk1_start, - chunk_count, - chunk_len, - window_capacity_per_chunk, - shared_compressed_base, - d_window, - window_size, - ratio, - compressed_width, - seq_major_rows, - has_indexer_topk, - has_compressed, - total_width, - parallel_dense, - ), - launch_work, - stream, - ) - - # Kernel contract: - # Same inputs as _build_attention_indices_kernel, but emits compressed-first - # indexer-loss order. - # topk_idxs: int32 output, shape (l_local, compressed_width + window_size), - # compressed ids first then window ids. - # indexer_rank_major: int32 output, shape (l_local, compressed_width), mapping - # each compressed logical top-k id to rank-major indexer-K row. - @cute.kernel - def _build_indexer_loss_indices_kernel( - cu_seqlens: cute.Tensor, + def _build_indexer_loss_indices_kernel( + cu_seqlens: cute.Tensor, cu_seqlens_compressed: cute.Tensor, indexer_topk_logical: cute.Tensor, indexer_rank_by_seq_major: cute.Tensor, @@ -1952,13 +920,11 @@ def _build_indexer_loss_indices_kernel( chunk1_start: cutlass.Int32, chunk_count: cutlass.Int32, chunk_len: cutlass.Int32, - window_capacity_per_chunk: cutlass.Int32, - shared_compressed_base: cutlass.Int32, d_window: cutlass.Int32, window_size: cutlass.Int32, - ratio: cutlass.Int32, compressed_width: cutlass.Int32, seq_major_rows: cutlass.Int32, + compressed_base: cutlass.Int32, total_width: cutlass.Int32, total_work: cutlass.Int32, ): @@ -1968,97 +934,62 @@ def _build_indexer_loss_indices_kernel( if linear < total_work: row = linear // total_width col = linear - row * total_width - topk_value = -1 - rank_major_value = -1 + topk_value = cutlass.Int32(-1) + rank_major_value = cutlass.Int32(-1) local_start = global_start - row_in_local = row - chunk_id = 0 + row_in_chunk = row + chunk_id = cutlass.Int32(0) + boundary_rows = d_window if chunk_count != 0: chunk_id = row // chunk_len - row_in_local = row - chunk_id * chunk_len + row_in_chunk = row - chunk_id * chunk_len local_start = chunk0_start if chunk_id == 1: local_start = chunk1_start - global_q = local_start + row_in_local - seq_id = -1 - seq_start_found = 0 + boundary_rows = chunk_count * d_window + + global_q = local_start + row_in_chunk + seq_id = cutlass.Int32(-1) + seq_start_found = cutlass.Int32(0) + seq_comp_start = cutlass.Int32(0) + seq_comp_len = cutlass.Int32(0) for seq in range(n_seq): seq_start = cu_seqlens[seq] seq_end = cu_seqlens[seq + 1] if global_q >= seq_start and global_q < seq_end: seq_id = seq seq_start_found = seq_start + seq_comp_start = cu_seqlens_compressed[seq] + seq_comp_len = cu_seqlens_compressed[seq + 1] - seq_comp_start if seq_id >= 0: - seq_offset = 0 - seq_window_start = 0 - seq_window_len = 0 - seq_comp_len = 0 - global_end = local_start + l_local - if chunk_count != 0: - global_end = local_start + chunk_len - for seq in range(n_seq): - seq_start = cu_seqlens[seq] - seq_end = cu_seqlens[seq + 1] - local_seq_start = seq_start - if local_seq_start < local_start: - local_seq_start = local_start - local_seq_end = seq_end - if local_seq_end > global_end: - local_seq_end = global_end - if local_seq_start < local_seq_end: - window_start = local_seq_start - d_window - if window_start < seq_start: - window_start = seq_start - window_len = local_seq_end - window_start - comp_len = 0 - if ratio > 1: - comp_len = cu_seqlens_compressed[seq + 1] - cu_seqlens_compressed[seq] - if seq == seq_id: - seq_window_start = window_start - seq_window_len = window_len - seq_comp_len = comp_len - if seq < seq_id: - seq_offset = seq_offset + window_len - if chunk_count == 0: - seq_offset = seq_offset + comp_len - if col < compressed_width: comp_id = indexer_topk_logical[row, col] if comp_id >= 0 and comp_id < seq_comp_len: - seq_major_id = cu_seqlens_compressed[seq_id] + comp_id + seq_major_id = seq_comp_start + comp_id if seq_major_id >= 0 and seq_major_id < seq_major_rows: rank_major_id = indexer_rank_by_seq_major[seq_major_id] if rank_major_id >= 0: - if chunk_count != 0: - topk_value = shared_compressed_base + rank_major_id - else: - topk_value = seq_offset + seq_window_len + comp_id + topk_value = compressed_base + rank_major_id rank_major_value = rank_major_id else: window_col = col - compressed_width - window_start_for_q = global_q - window_size + 1 - if window_start_for_q < seq_start_found: - window_start_for_q = seq_start_found - window_count = global_q - window_start_for_q + 1 - pos = window_start_for_q + window_col - if ( - window_col < window_count - and pos >= seq_window_start - and pos < seq_window_start + seq_window_len - ): - topk_value = seq_offset + pos - seq_window_start - if chunk_count != 0: - topk_value = chunk_id * window_capacity_per_chunk + topk_value + window_start = global_q - window_size + 1 + if window_start < seq_start_found: + window_start = seq_start_found + window_count = global_q - window_start + 1 + if window_col < window_count: + pos = window_start + window_col + if pos < local_start: + topk_value = chunk_id * d_window + pos - (local_start - d_window) + else: + topk_value = boundary_rows + chunk_id * chunk_len + pos - local_start topk_idxs[row, col] = topk_value if col < compressed_width: indexer_rank_major[row, col] = rank_major_value - # Launch contract for _build_indexer_loss_indices_kernel. - # Supports contiguous and two-chunk CP. total_work is - # l_local * (compressed_width + window_size). @cute.jit def _build_indexer_loss_indices_launch( cu_seqlens: cute.Tensor, @@ -2074,13 +1005,11 @@ def _build_indexer_loss_indices_launch( chunk1_start: cutlass.Int32, chunk_count: cutlass.Int32, chunk_len: cutlass.Int32, - window_capacity_per_chunk: cutlass.Int32, - shared_compressed_base: cutlass.Int32, d_window: cutlass.Int32, window_size: cutlass.Int32, - ratio: cutlass.Int32, compressed_width: cutlass.Int32, seq_major_rows: cutlass.Int32, + compressed_base: cutlass.Int32, total_width: cutlass.Int32, total_work: cutlass.Int32, stream: cuda.CUstream, @@ -2102,13 +1031,11 @@ def _build_indexer_loss_indices_launch( chunk1_start, chunk_count, chunk_len, - window_capacity_per_chunk, - shared_compressed_base, d_window, window_size, - ratio, compressed_width, seq_major_rows, + compressed_base, total_width, total_work, ), @@ -2120,8 +1047,8 @@ def _build_indexer_loss_indices_launch( # ============================================================================= # Torch Wrapper Functions # ============================================================================= -# Torch-facing entry points called by csa_cp_utils.py. They validate inputs, -# allocate outputs, and dispatch the CuTeDSL kernels defined above. +# Torch-facing entry points called by csa.py and csa_cp_utils.py. They validate +# inputs, allocate outputs, and dispatch the CuTeDSL kernels defined above. # ============================================================================= @@ -2135,127 +1062,69 @@ def _require_cute(message: str, *tensors: Optional[torch.Tensor]) -> None: raise RuntimeError(message) -def _make_compiled_launch_runner(): - """Build the CuTe launch helper with private helpers and persistent caches.""" - compiled_launch_cache = {} +_COMPILED_LAUNCH_CACHE = {} - def _cute_scalar(value): - if isinstance(value, bool): - return cutlass.Int32(1 if value else 0) - if isinstance(value, int): - return cutlass.Int32(value) - if isinstance(value, float): - return cutlass.Float32(value) - return value - @lru_cache(maxsize=None) - def _gpu_arch_flag() -> str: +def _cute_scalar(value): + if isinstance(value, bool): + return cutlass.Int32(1 if value else 0) + if isinstance(value, int): + return cutlass.Int32(value) + if isinstance(value, float): + return cutlass.Float32(value) + return value + + +def _run_compiled_launch( + launch_fn, + tensor_args: Tuple[torch.Tensor, ...], + scalar_args: Tuple, + key_arg_indices: Optional[Tuple[int, ...]] = None, + static_arg_indices: Optional[Tuple[int, ...]] = None, +) -> None: + """Compile/cache a CuTe launch function and invoke it on the current stream.""" + key_arg_indices = ( + tuple(range(len(scalar_args))) if key_arg_indices is None else tuple(key_arg_indices) + ) + static_arg_indices = () if static_arg_indices is None else tuple(static_arg_indices) + static_arg_set = set(static_arg_indices) + key = ( + launch_fn.__name__, + tuple( + (tensor.dtype, tuple(tensor.shape), tuple(tensor.stride())) for tensor in tensor_args + ), + key_arg_indices, + tuple(scalar_args[i] for i in key_arg_indices), + static_arg_indices, + ) + compiled = _COMPILED_LAUNCH_CACHE.get(key) + if compiled is None: cap = torch.cuda.get_device_capability() - arch_map = {(9, 0): "sm_90a", (10, 0): "sm_100a", (10, 3): "sm_103a"} - arch = arch_map.get(cap) + arch = {(9, 0): "sm_90a", (10, 0): "sm_100a", (10, 3): "sm_103a"}.get(cap) if arch is None: raise RuntimeError(f"Unsupported GPU compute capability {cap} for CSA CP CuTe kernels.") - return arch - - def _run_compiled_launch( - launch_fn, - tensor_args: Tuple[torch.Tensor, ...], - scalar_args: Tuple, - key_arg_indices: Optional[Tuple[int, ...]] = None, - static_arg_indices: Optional[Tuple[int, ...]] = None, - ) -> None: - """Compile/cache a CuTe launch function and invoke it on the current stream. - - Inputs: - launch_fn: ``@cute.jit`` launcher. - tensor_args: Torch CUDA tensors passed through DLPack with dynamic leading - layouts. - scalar_args: Python scalar launch parameters. - key_arg_indices: Scalar indices that participate in the compile cache key. - static_arg_indices: Scalar indices passed as compile-time constexpr values. - - Output: - None. The launch writes into output tensors passed in ``tensor_args``. - """ - key_arg_indices = ( - tuple(range(len(scalar_args))) if key_arg_indices is None else tuple(key_arg_indices) - ) - static_arg_indices = () if static_arg_indices is None else tuple(static_arg_indices) - static_arg_set = set(static_arg_indices) - key = ( - launch_fn.__name__, - tuple( - (tensor.dtype, tuple(tensor.shape), tuple(tensor.stride())) - for tensor in tensor_args + cute_tensor_args = [] + for tensor in tensor_args: + cute_tensor = from_dlpack(tensor.detach(), assumed_align=16, enable_tvm_ffi=True) + if tensor.ndim != 0: + cute_tensor = cute_tensor.mark_layout_dynamic(leading_dim=tensor.ndim - 1) + cute_tensor_args.append(cute_tensor) + compiled = cute.compile( + launch_fn, + *cute_tensor_args, + *( + arg if i in static_arg_set else _cute_scalar(arg) + for i, arg in enumerate(scalar_args) ), - key_arg_indices, - tuple(scalar_args[i] for i in key_arg_indices), - static_arg_indices, - ) - compiled = compiled_launch_cache.get(key) - if compiled is None: - fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) - cute_tensor_args = [] - for tensor in tensor_args: - cute_tensor = from_dlpack(tensor.detach(), assumed_align=16, enable_tvm_ffi=True) - if tensor.ndim != 0: - cute_tensor = cute_tensor.mark_layout_dynamic(leading_dim=tensor.ndim - 1) - cute_tensor_args.append(cute_tensor) - compiled = cute.compile( - launch_fn, - *cute_tensor_args, - *( - arg if i in static_arg_set else _cute_scalar(arg) - for i, arg in enumerate(scalar_args) - ), - fake_stream, - options=f"--enable-tvm-ffi --gpu-arch {_gpu_arch_flag()}", - ) - compiled_launch_cache[key] = compiled - compiled( - *tensor_args, - *(_cute_scalar(arg) for i, arg in enumerate(scalar_args) if i not in static_arg_set), - cuda.CUstream(torch.cuda.current_stream(tensor_args[0].device).cuda_stream), - ) - - return _run_compiled_launch - - -_run_compiled_launch = _make_compiled_launch_runner() - - -def _validate_rope_inputs( - x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, nope_dim: int, pos_dim: int -) -> Tuple[int, int, int]: - """Validate common RoPE tensor dimensions and return flattened row metadata. - - Inputs: - x: Tensor with first dimension as rows and remaining dims flattened per row. - cos/sin: Rotary tables with first dimension as sequence position and at - least ``pos_dim`` values per row. - nope_dim/pos_dim: Per-head no-RoPE and RoPE dimensions. - - Output: - ``(rows, row_width, head_dim)`` where ``head_dim = nope_dim + pos_dim``. - """ - rows = x.shape[0] - row_width = math.prod(x.shape[1:]) - head_dim = nope_dim + pos_dim - if pos_dim <= 0 or pos_dim % 2 != 0: - raise RuntimeError(f"DSv4 CP RoPE expects an even positive pos_dim, got {pos_dim}.") - if head_dim <= 0 or row_width % head_dim != 0: - raise RuntimeError( - "DSv4 CP RoPE expects the flattened row width to be a multiple of " - f"head_dim={head_dim}, got row_width={row_width}." - ) - cos_width = math.prod(cos.shape[1:]) - sin_width = math.prod(sin.shape[1:]) - if cos_width < pos_dim or sin_width < pos_dim: - raise RuntimeError( - "DSv4 CP RoPE cos/sin tables are narrower than pos_dim: " - f"cos={cos_width}, sin={sin_width}, pos_dim={pos_dim}." + make_fake_stream(use_tvm_ffi_env_stream=False), + options=f"--enable-tvm-ffi --gpu-arch {arch}", ) - return rows, row_width, head_dim + _COMPILED_LAUNCH_CACHE[key] = compiled + compiled( + *tensor_args, + *(_cute_scalar(arg) for i, arg in enumerate(scalar_args) if i not in static_arg_set), + cuda.CUstream(torch.cuda.current_stream(tensor_args[0].device).cuda_stream), + ) class ThdLocalRope(torch.autograd.Function): @@ -2298,7 +1167,18 @@ def _run( sin, cu_seqlens_padded, ) - rows, row_width, head_dim = _validate_rope_inputs(x, cos, sin, nope_dim, pos_dim) + rows = x.shape[0] + row_width = math.prod(x.shape[1:]) + head_dim = nope_dim + pos_dim + if pos_dim <= 0 or pos_dim % 2 != 0: + raise RuntimeError(f"DSv4 CP RoPE expects an even positive pos_dim, got {pos_dim}.") + if head_dim <= 0 or row_width % head_dim != 0: + raise RuntimeError( + "DSv4 CP RoPE expects row width to be a multiple of " + f"head_dim={head_dim}, got {row_width}." + ) + if math.prod(cos.shape[1:]) < pos_dim or math.prod(sin.shape[1:]) < pos_dim: + raise RuntimeError("DSv4 CP RoPE cos/sin tables are narrower than pos_dim.") chunk_len = int(chunk_len) if chunk_len <= 0: raise RuntimeError(f"DSv4 THD CP RoPE expects positive chunk_len, got {chunk_len}.") @@ -2384,102 +1264,6 @@ def backward(ctx, grad_output: torch.Tensor): return (grad_x, *([None] * 10)) -class ThdCompressedRope(torch.autograd.Function): - """Apply RoPE to compressed rows using per-row compressed group ids. - - Inputs: - x: CUDA tensor, shape ``(rows, ...)``; flattened row width must be a - multiple of ``nope_dim + pos_dim``. - cos/sin: CUDA tensors, shape ``(max_seq_len, >= pos_dim)``. - comp_ids_local: int32 CUDA tensor, shape at least ``(rows,)``; -1 maps to 0. - ratio: compression ratio, used as ``rope_position = comp_id * ratio``. - - Output: - Tensor with same shape and dtype as ``x``. No-RoPE dims are copied; RoPE - dims are rotated pairwise, with inverse/adjoint variants controlled by - flags. - """ - - @staticmethod - def _run( - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - comp_ids_local: torch.Tensor, - ratio: int, - nope_dim: int, - pos_dim: int, - inverse: bool = False, - adjoint: bool = False, - ) -> torch.Tensor: - _require_cute( - "DSv4 CP compressed RoPE requires CUDA tensors and available CuTe kernels.", - x, - cos, - sin, - comp_ids_local, - ) - rows, row_width, head_dim = _validate_rope_inputs(x, cos, sin, nope_dim, pos_dim) - if comp_ids_local.shape[0] < rows: - raise RuntimeError( - "DSv4 CP compressed RoPE received too few comp ids: " - f"comp_ids={comp_ids_local.shape[0]}, rows={rows}." - ) - if rows == 0: - return x - out = torch.empty_like(x) - total_work = rows * row_width - _run_compiled_launch( - _thd_compressed_rope_launch, - ( - x.reshape(rows, row_width), - out.reshape(rows, row_width), - cos.reshape(cos.shape[0], math.prod(cos.shape[1:])), - sin.reshape(sin.shape[0], math.prod(sin.shape[1:])), - comp_ids_local[:rows], - ), - ( - int(ratio), - row_width, - head_dim, - int(nope_dim), - int(pos_dim), - 1 if inverse else 0, - 1 if adjoint else 0, - total_work, - ), - key_arg_indices=(1,), - static_arg_indices=(1,), - ) - return out - - @staticmethod - def forward( - ctx, - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - comp_ids_local: torch.Tensor, - ratio: int, - nope_dim: int, - pos_dim: int, - inverse: bool, - ) -> torch.Tensor: - """Apply compressed-row RoPE and save metadata for backward.""" - ctx.rope_args = (int(ratio), int(nope_dim), int(pos_dim), bool(inverse)) - ctx.save_for_backward(cos, sin, comp_ids_local) - return ThdCompressedRope._run(x, cos, sin, comp_ids_local, *ctx.rope_args) - - @staticmethod - def backward(ctx, grad_output: torch.Tensor): - """Apply the adjoint compressed-row RoPE transform to gradients.""" - cos, sin, comp_ids_local = ctx.saved_tensors - grad_x = ThdCompressedRope._run( - grad_output.contiguous(), cos, sin, comp_ids_local, *ctx.rope_args, adjoint=True - ) - return (grad_x, *([None] * 7)) - - class CompressorInputCompact(torch.autograd.Function): """Compact local and boundary hidden rows for compressor input. @@ -2491,15 +1275,15 @@ class CompressorInputCompact(torch.autograd.Function): ratio/d_comp/d_window/c_cap: fixed compressor window and capacity. Outputs: - ``hidden_compact`` shape ``(c_cap * ratio, ...)``, same dtype as hidden; - ``cu_compact`` int32 shape ``(n_seq + 1,)``; ``seq_ids`` int32, - ``comp_ids`` int32, and ``valid`` bool each shape ``(c_cap,)``. + ``hidden_compact`` shape ``(c_cap * ratio, ...)``, same dtype as hidden, + and ``comp_ids`` int32 shape ``(c_cap,)``. The kernel copies the ratio source tokens for each visible compressed - group and emits metadata for the compressed rows. + group and emits each row's compressed group id. """ @staticmethod - def _forward_kernel( + def forward( + ctx, hidden_local: torch.Tensor, boundary_hidden: torch.Tensor, cu_seqlens: torch.Tensor, @@ -2509,522 +1293,76 @@ def _forward_kernel( d_comp: int, d_window: int, c_cap: int, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - compact_len = c_cap * ratio - output_shape = (compact_len,) + tuple(hidden_local.shape[1:]) - hidden_compact = hidden_local.new_empty(output_shape) - cu_compact = torch.empty_like(cu_seqlens) - seq_ids = torch.empty((c_cap,), dtype=torch.int32, device=hidden_local.device) - comp_ids = torch.empty((c_cap,), dtype=torch.int32, device=hidden_local.device) - valid = torch.empty((c_cap,), dtype=torch.bool, device=hidden_local.device) + ): + """Compact local and boundary hidden rows into compressor input rows.""" + ctx.hidden_shape = tuple(hidden_local.shape) + ctx.boundary_shape = tuple(boundary_hidden.shape) + ctx.compact_args = (int(global_start), int(l_local), int(ratio), int(d_comp), int(d_window)) + ctx.save_for_backward(cu_seqlens) + compact_len = int(c_cap) * int(ratio) + hidden_compact = hidden_local.new_empty((compact_len,) + tuple(hidden_local.shape[1:])) + comp_ids = torch.empty((int(c_cap),), dtype=torch.int32, device=hidden_local.device) row_width = math.prod(hidden_local.shape[1:]) - hidden_local_flat = hidden_local.reshape(l_local, row_width) - boundary_flat = boundary_hidden.reshape(d_window, row_width) - compact_flat = hidden_compact.reshape(compact_len, row_width) - total_work = max(compact_len, c_cap, cu_seqlens.shape[0]) _run_compiled_launch( _compressor_input_compact_fwd_launch, ( - hidden_local_flat, - boundary_flat, - compact_flat, + hidden_local.reshape(int(l_local), row_width), + boundary_hidden.reshape(int(d_window), row_width), + hidden_compact.reshape(compact_len, row_width), cu_seqlens, - cu_compact, - seq_ids, comp_ids, - valid, ), ( cu_seqlens.shape[0] - 1, - global_start, - l_local, - ratio, - d_comp, - d_window, - c_cap, + int(global_start), + int(l_local), + int(ratio), + int(d_comp), + int(d_window), + int(c_cap), compact_len, row_width, - total_work, + max(compact_len, int(c_cap)), ), key_arg_indices=(3, 4, 5, 8), static_arg_indices=(3, 4, 5, 8), ) - return hidden_compact, cu_compact, seq_ids, comp_ids, valid - - @staticmethod - def _backward_kernel( - grad_hidden_compact: torch.Tensor, - hidden_shape: Tuple[int, ...], - boundary_shape: Tuple[int, ...], - cu_seqlens: torch.Tensor, - global_start: int, - l_local: int, - ratio: int, - d_comp: int, - d_window: int, - ) -> Tuple[torch.Tensor, torch.Tensor]: - grad_hidden = grad_hidden_compact.new_empty(hidden_shape) - grad_boundary = grad_hidden_compact.new_empty(boundary_shape) - compact_len = grad_hidden_compact.shape[0] - row_width = math.prod(hidden_shape[1:]) - - total_work = l_local + d_window - if total_work > 0: - _run_compiled_launch( - _compressor_input_compact_bwd_launch, - ( - grad_hidden_compact.reshape(compact_len, row_width), - grad_hidden.reshape(l_local, row_width), - grad_boundary.reshape(d_window, row_width), - cu_seqlens, - ), - ( - cu_seqlens.shape[0] - 1, - global_start, - l_local, - ratio, - d_comp, - d_window, - compact_len, - row_width, - total_work, - ), - key_arg_indices=(3, 4, 5, 7), - static_arg_indices=(3, 4, 5, 7), - ) - return grad_hidden, grad_boundary - - @staticmethod - def forward( - ctx, - hidden_local: torch.Tensor, - boundary_hidden: torch.Tensor, - cu_seqlens: torch.Tensor, - global_start: int, - l_local: int, - ratio: int, - d_comp: int, - d_window: int, - c_cap: int, - ): - """Compact local and boundary hidden rows into compressor input rows.""" - ctx.hidden_shape = tuple(hidden_local.shape) - ctx.boundary_shape = tuple(boundary_hidden.shape) - ctx.compact_args = (int(global_start), int(l_local), int(ratio), int(d_comp), int(d_window)) - ctx.save_for_backward(cu_seqlens) - return CompressorInputCompact._forward_kernel( - hidden_local, - boundary_hidden, - cu_seqlens, - int(global_start), - int(l_local), - int(ratio), - int(d_comp), - int(d_window), - int(c_cap), - ) + return hidden_compact, comp_ids @staticmethod - def backward( - ctx, - grad_hidden_compact: torch.Tensor, - _grad_cu_compact: torch.Tensor, - _grad_seq_ids: torch.Tensor, - _grad_comp_ids: torch.Tensor, - _grad_valid: torch.Tensor, - ): + def backward(ctx, grad_hidden_compact: torch.Tensor, _grad_comp_ids: torch.Tensor): """Scatter compacted compressor gradients to local and boundary rows.""" (cu_seqlens,) = ctx.saved_tensors - grad_hidden, grad_boundary = CompressorInputCompact._backward_kernel( - grad_hidden_compact.contiguous(), - ctx.hidden_shape, - ctx.boundary_shape, - cu_seqlens, - *ctx.compact_args, - ) - return (grad_hidden, grad_boundary, *([None] * 7)) - - -class ThdFullKvPack(torch.autograd.Function): - """Autograd wrapper for static-shape KV-full packing. - - Forward builds a physical ``kv_full`` row space from local window rows, - boundary rows, and all-gathered compressed rows. Backward recomputes the - same row mapping and scatters gradients back to those three sources. - """ - - @staticmethod - def forward( - ctx, - kv_local: torch.Tensor, - boundary_kv: torch.Tensor, - compressed_rank_major: torch.Tensor, - seq_ids_rank_major: torch.Tensor, - comp_ids_rank_major: torch.Tensor, - valid_rank_major: torch.Tensor, - cu_seqlens: torch.Tensor, - rank_major_by_seq_major: torch.Tensor, - cu_seqlens_compressed: torch.Tensor, - global_start: int, - l_local: int, - d_window: int, - ratio: int, - kv_full_capacity: int, - chunk0_start: int, - chunk1_start: int, - chunk_count: int, - chunk_len: int, - window_capacity_per_chunk: int, - ) -> torch.Tensor: - """Pack KV rows into the fixed physical sparse-attention row order. - - Inputs: - kv_local: CUDA tensor, shape ``(local_rows, ...)``. - boundary_kv: CUDA tensor, shape ``(d_window or 2*d_window, ...)``. - compressed_rank_major: CUDA tensor, shape ``(compressed_rows, ...)``. - seq_ids/comp_ids/valid: rank-major compressed metadata. - cu_seqlens: int32 CUDA tensor, shape ``(n_seq + 1,)``. - optional compressed reverse-map metadata for contiguous CP. - chunk fields select contiguous or two-chunk layout. - - Output: - ``kv_full`` CUDA tensor, shape ``(kv_full_capacity, ...)``, same dtype - as ``kv_local``. Invalid tail rows are zero-filled. - """ - ctx.pack_shapes = ( - tuple(kv_local.shape), - tuple(boundary_kv.shape), - tuple(compressed_rank_major.shape), - ) - ctx.kv_full_capacity = int(kv_full_capacity) - - output_shape = (int(kv_full_capacity),) + tuple(kv_local.shape[1:]) - kv_full = kv_local.new_empty(output_shape) - row_width = math.prod(kv_local.shape[1:]) - local_rows = kv_local.shape[0] - boundary_rows = boundary_kv.shape[0] - kv_local_flat = kv_local.reshape(local_rows, row_width) - boundary_flat = boundary_kv.reshape(boundary_rows, row_width) - compressed_rows = compressed_rank_major.shape[0] - compressed_flat = compressed_rank_major.reshape(compressed_rows, row_width) - kv_full_flat = kv_full.reshape(int(kv_full_capacity), row_width) - total_work = int(kv_full_capacity) - use_seq_major_map = ( - int(chunk_count) == 0 - and ratio > 1 - and compressed_rows > 0 - and rank_major_by_seq_major.shape[0] > 1 - ) + global_start, l_local, ratio, d_comp, d_window = ctx.compact_args + grad_hidden_compact = grad_hidden_compact.contiguous() + grad_hidden = grad_hidden_compact.new_empty(ctx.hidden_shape) + grad_boundary = grad_hidden_compact.new_empty(ctx.boundary_shape) + compact_len = grad_hidden_compact.shape[0] + row_width = math.prod(ctx.hidden_shape[1:]) _run_compiled_launch( - _thd_full_kv_pack_launch, + _compressor_input_compact_bwd_launch, ( - kv_local_flat, - boundary_flat, - compressed_flat, - kv_full_flat, - seq_ids_rank_major, - comp_ids_rank_major, - valid_rank_major, + grad_hidden_compact.reshape(compact_len, row_width), + grad_hidden.reshape(l_local, row_width), + grad_boundary.reshape(d_window, row_width), cu_seqlens, - cu_seqlens_compressed, - rank_major_by_seq_major, ), ( cu_seqlens.shape[0] - 1, - int(global_start), - int(l_local), - int(chunk0_start), - int(chunk1_start), - int(chunk_count), - int(chunk_len), - int(window_capacity_per_chunk), - int(d_window), - int(ratio), - compressed_rows, - 1 if use_seq_major_map else 0, - rank_major_by_seq_major.shape[0] if use_seq_major_map else 0, - int(kv_full_capacity), + global_start, + l_local, + ratio, + d_comp, + d_window, + compact_len, row_width, - 0, - total_work, + l_local + d_window, ), - key_arg_indices=(5, 8, 9, 11, 14, 15), - static_arg_indices=(5, 8, 9, 11, 14, 15), - ) - ctx.save_for_backward( - seq_ids_rank_major, - comp_ids_rank_major, - valid_rank_major, - cu_seqlens, - cu_seqlens_compressed, - rank_major_by_seq_major, + key_arg_indices=(3, 4, 5, 7), + static_arg_indices=(3, 4, 5, 7), ) - ctx.pack_args = ( - int(global_start), - int(l_local), - int(chunk0_start), - int(chunk1_start), - int(chunk_count), - int(chunk_len), - int(window_capacity_per_chunk), - int(d_window), - int(ratio), - compressed_rows, - 1 if use_seq_major_map else 0, - rank_major_by_seq_major.shape[0] if use_seq_major_map else 0, - ) - return kv_full - - @staticmethod - def backward(ctx, grad_kv_full: torch.Tensor): - """Scatter ``kv_full`` gradients back to local, boundary, and compressed KV.""" - local_shape, boundary_shape, compressed_shape = ctx.pack_shapes - grad_kv_local = grad_kv_full.new_empty(local_shape) - grad_boundary_kv = grad_kv_full.new_empty(boundary_shape) - grad_compressed = grad_kv_full.new_empty(compressed_shape) - - compressed_rows = compressed_shape[0] - kv_full_capacity = ctx.kv_full_capacity - row_width = math.prod(local_shape[1:]) - ( - global_start, - l_local, - chunk0_start, - chunk1_start, - chunk_count, - chunk_len, - window_capacity_per_chunk, - d_window, - ratio, - compressed_rows, - use_seq_major_map, - seq_major_rows, - ) = ctx.pack_args - grad_kv_local.zero_() - grad_boundary_kv.zero_() - grad_compressed.zero_() - - total_work = kv_full_capacity - if total_work > 0: - ( - seq_ids_rank_major, - comp_ids_rank_major, - valid_rank_major, - cu_seqlens, - cu_seqlens_compressed, - rank_major_by_seq_major, - ) = ctx.saved_tensors - _run_compiled_launch( - _thd_full_kv_pack_launch, - ( - grad_kv_local.reshape(local_shape[0], row_width), - grad_boundary_kv.reshape(boundary_shape[0], row_width), - grad_compressed.reshape(compressed_rows, row_width), - grad_kv_full.reshape(kv_full_capacity, row_width), - seq_ids_rank_major, - comp_ids_rank_major, - valid_rank_major, - cu_seqlens, - cu_seqlens_compressed, - rank_major_by_seq_major, - ), - ( - cu_seqlens.shape[0] - 1, - global_start, - l_local, - chunk0_start, - chunk1_start, - chunk_count, - chunk_len, - window_capacity_per_chunk, - d_window, - ratio, - compressed_rows, - use_seq_major_map, - seq_major_rows, - kv_full_capacity, - row_width, - 1, - total_work, - ), - key_arg_indices=(5, 8, 9, 11, 14, 15), - static_arg_indices=(5, 8, 9, 11, 14, 15), - ) - return (grad_kv_local, grad_boundary_kv, grad_compressed, *([None] * 16)) - - -def repack_compressed_kv_to_seq_major( - compressed_rank_major: torch.Tensor, - seq_ids_rank_major: torch.Tensor, - comp_ids_rank_major: torch.Tensor, - valid_rank_major: torch.Tensor, - cu_seqlens_compressed: torch.Tensor, - seq_major_rows: int, -) -> Tuple[torch.Tensor, torch.Tensor]: - """Repack all-gathered compressed rows from rank-major to seq-major order. - - Inputs: - compressed_rank_major: CUDA tensor, shape ``(rank_major_rows, ...)``. - seq_ids_rank_major/comp_ids_rank_major: int32 metadata, shape - ``(rank_major_rows,)``. - valid_rank_major: bool metadata, shape ``(rank_major_rows,)``. - cu_seqlens_compressed: int32 prefixes, shape ``(n_seq + 1,)``. - seq_major_rows: total compressed seq-major rows. - - Outputs: - ``compressed_seq_major`` shape ``(seq_major_rows, ...)`` and - ``rank_major_by_seq_major`` int32 shape ``(seq_major_rows,)``. - """ - output_shape = (seq_major_rows,) + tuple(compressed_rank_major.shape[1:]) - compressed_seq_major = compressed_rank_major.new_empty(output_shape) - rank_major_by_seq_major = torch.empty( - (seq_major_rows,), dtype=torch.int32, device=compressed_rank_major.device - ) - rank_major_rows = compressed_rank_major.shape[0] - row_width = math.prod(compressed_rank_major.shape[1:]) - map_work = int(rank_major_rows) - gather_work = int(seq_major_rows) * int(row_width) - _run_compiled_launch( - _repack_compressed_kv_init_rank_map_launch, - (rank_major_by_seq_major,), - (seq_major_rows, max(1, seq_major_rows)), - ) - _run_compiled_launch( - _repack_compressed_kv_rank_map_launch, - ( - seq_ids_rank_major, - comp_ids_rank_major, - valid_rank_major, - cu_seqlens_compressed, - rank_major_by_seq_major, - ), - (rank_major_rows, seq_major_rows, max(1, map_work)), - ) - _run_compiled_launch( - _repack_compressed_kv_to_seq_major_launch, - ( - compressed_rank_major.reshape(rank_major_rows, row_width), - compressed_seq_major.reshape(seq_major_rows, row_width), - rank_major_by_seq_major, - ), - (rank_major_rows, seq_major_rows, row_width, max(1, gather_work)), - key_arg_indices=(2,), - static_arg_indices=(2,), - ) - return compressed_seq_major, rank_major_by_seq_major - - -def build_compressed_row_metadata( - cu_seqlens: torch.Tensor, - cp_size: int, - l_local: int, - ratio: int, - d_comp: int, - c_cap: int, - *, - c_cap_per_rank: Optional[int] = None, - use_two_chunk: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Build rank-major compressed metadata for all CP ranks. - - Inputs: - cu_seqlens: int32 CUDA tensor, shape ``(n_seq + 1,)``. - cp_size: CP world size. - l_local: rows per contiguous rank, or chunk length when ``use_two_chunk``. - ratio/d_comp: compression parameters. - c_cap: fixed compressed capacity per chunk. - c_cap_per_rank: optional fixed compressed capacity per rank. - use_two_chunk: whether each rank owns two chunks. - - Outputs: - ``seq_ids`` int32, ``comp_ids`` int32, and ``valid`` bool, each shape - ``(cp_size * c_cap_per_rank,)``. Rows are ordered rank-major and describe - which compressed group each all-gathered row represents. - """ - c_cap_per_rank = int(c_cap if c_cap_per_rank is None else c_cap_per_rank) - total_rows = int(cp_size) * c_cap_per_rank - seq_ids = torch.empty((total_rows,), dtype=torch.int32, device=cu_seqlens.device) - comp_ids = torch.empty((total_rows,), dtype=torch.int32, device=cu_seqlens.device) - valid = torch.empty((total_rows,), dtype=torch.bool, device=cu_seqlens.device) - _run_compiled_launch( - _build_compressed_row_metadata_launch, - (cu_seqlens, seq_ids, comp_ids, valid), - ( - cu_seqlens.shape[0] - 1, - int(cp_size), - int(l_local), - int(ratio), - int(d_comp), - int(c_cap), - c_cap_per_rank, - 1 if use_two_chunk else 0, - total_rows, - ), - key_arg_indices=(), - ) - return seq_ids, comp_ids, valid - - -def build_indexer_topk_metadata( - k_indexer_seq_major: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_compressed: torch.Tensor, - global_start: int, - l_local: int, - ratio: int, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Build local trapezoid inputs for the indexer top-k kernel. - - Inputs: - k_indexer_seq_major: CUDA tensor, shape ``(seq_major_rows, ...)``. - cu_seqlens_q/cu_seqlens_compressed: int32 prefixes, shape ``(n_seq + 1,)``. - global_start/l_local: local query range. - ratio: compression ratio. - - Outputs: - ``k_topk`` tensor with same shape as ``k_indexer_seq_major``; - ``cu_q_topk`` and ``cu_k_topk`` int32 tensors, shape ``(n_seq + 2,)``; - ``seq_lens`` int32 tensor, shape ``(l_local,)``. The metadata represents - the local Q/K trapezoid that the indexer kernel consumes. - """ - seq_major_rows = k_indexer_seq_major.shape[0] - row_width = math.prod(k_indexer_seq_major.shape[1:]) - k_topk = k_indexer_seq_major.new_empty(k_indexer_seq_major.shape) - cu_q_topk = torch.empty( - (cu_seqlens_q.shape[0] + 1,), dtype=cu_seqlens_q.dtype, device=cu_seqlens_q.device - ) - cu_k_topk = torch.empty( - (cu_seqlens_q.shape[0] + 1,), dtype=cu_seqlens_compressed.dtype, device=cu_seqlens_q.device - ) - seq_lens = torch.empty((l_local,), dtype=torch.int32, device=cu_seqlens_q.device) - n_seq = cu_seqlens_q.shape[0] - 1 - total_work = seq_major_rows * row_width - _run_compiled_launch( - _build_indexer_topk_metadata_launch, - ( - k_indexer_seq_major.reshape(seq_major_rows, row_width), - k_topk.reshape(seq_major_rows, row_width), - cu_seqlens_q, - cu_seqlens_compressed, - cu_q_topk, - cu_k_topk, - seq_lens, - ), - ( - n_seq, - global_start, - l_local, - ratio, - row_width, - seq_major_rows, - total_work, - max(1, total_work), - ), - key_arg_indices=(4,), - static_arg_indices=(4,), - ) - return k_topk, cu_q_topk, cu_k_topk, seq_lens + return (grad_hidden, grad_boundary, *([None] * 7)) def build_attention_indices( @@ -3037,11 +1375,10 @@ def build_attention_indices( compressed_width: int, indexer_topk_compressed_logical_ids: Optional[torch.Tensor] = None, cu_seqlens_compressed: Optional[torch.Tensor] = None, - rank_major_by_seq_major: Optional[torch.Tensor] = None, + rank_row_for_seq_row: Optional[torch.Tensor] = None, chunk_starts: Optional[Tuple[int, ...]] = None, chunk_len: int = 0, - window_capacity_per_chunk: int = 0, - shared_compressed_base: int = 0, + compressed_base: int = 0, ) -> Tuple[torch.Tensor, torch.Tensor]: """Build final sparse-attention indices in physical ``kv_full`` row space. @@ -3052,8 +1389,9 @@ def build_attention_indices( ratio/compressed_width: compression mode and number of compressed columns. indexer_topk_compressed_logical_ids: optional int32 tensor, shape ``(l_local, compressed_width)``. - cu_seqlens_compressed/rank_major_by_seq_major/chunk metadata: required for - two-chunk compressed lowering. + cu_seqlens_compressed/rank_row_for_seq_row: compressed row mapping. + chunk metadata: optional two-chunk row layout. + compressed_base: first compressed row in the raw KV tensor. Outputs: ``topk_idxs`` int32, shape ``(l_local, window_size + compressed_width)``; @@ -3064,22 +1402,23 @@ def build_attention_indices( chunk1_start = 0 if chunk_starts is not None: _require_cute( - "DSv4 two-chunk shared final idx requires CUDA tensors and CuTeDSL.", + "DSv4 two-chunk final indices require CUDA tensors and CuTeDSL.", cu_seqlens_compressed, - rank_major_by_seq_major, + rank_row_for_seq_row, ) if not chunk_starts or len(chunk_starts) > 2: raise RuntimeError( - "DSv4 two-chunk shared final idx supports one or two chunks, " + "DSv4 two-chunk final indices support one or two chunks, " f"got {len(chunk_starts)}." ) chunk_count = len(chunk_starts) l_local = int(chunk_count) * int(chunk_len) global_start = 0 chunk1_start = int(chunk_starts[1]) if chunk_count > 1 else 0 - else: + if cu_seqlens_compressed is None: cu_seqlens_compressed = cu_seqlens - rank_major_by_seq_major = torch.empty((1,), dtype=torch.int32, device=cu_seqlens.device) + if rank_row_for_seq_row is None: + rank_row_for_seq_row = torch.empty((1,), dtype=torch.int32, device=cu_seqlens.device) total_width = window_size + compressed_width topk_idxs = torch.empty((l_local, total_width), dtype=torch.int32, device=cu_seqlens.device) @@ -3090,11 +1429,11 @@ def build_attention_indices( else: dummy = indexer_topk_compressed_logical_ids has_indexer = 1 - parallel_dense = 1 if has_indexer == 0 and chunk_count == 0 else 0 + parallel_dense = 1 if has_indexer == 0 else 0 launch_work = l_local * total_width if parallel_dense else l_local _run_compiled_launch( _build_attention_indices_launch, - (cu_seqlens, cu_seqlens_compressed, dummy, rank_major_by_seq_major, topk_idxs, topk_length), + (cu_seqlens, cu_seqlens_compressed, dummy, rank_row_for_seq_row, topk_idxs, topk_length), ( cu_seqlens.shape[0] - 1, global_start, @@ -3103,13 +1442,12 @@ def build_attention_indices( chunk1_start, chunk_count, int(chunk_len), - int(window_capacity_per_chunk), - int(shared_compressed_base), d_window, window_size, ratio, compressed_width, - rank_major_by_seq_major.shape[0], + rank_row_for_seq_row.shape[0], + int(compressed_base), has_indexer, 1 if ratio > 1 and compressed_width > 0 else 0, total_width, @@ -3133,8 +1471,7 @@ def build_indexer_loss_indices( indexer_rank_by_seq_major: torch.Tensor, chunk_starts: Optional[Tuple[int, ...]] = None, chunk_len: int = 0, - window_capacity_per_chunk: int = 0, - shared_compressed_base: int = 0, + compressed_base: int = 0, ) -> Tuple[torch.Tensor, torch.Tensor]: """Build compressed-first sparse-attention ids and indexer-loss rank-major ids. @@ -3164,7 +1501,7 @@ def build_indexer_loss_indices( if chunk_starts is not None: if not chunk_starts or len(chunk_starts) > 2: raise RuntimeError( - "DSv4 two-chunk shared indexer-loss final idx supports one or two chunks, " + "DSv4 two-chunk indexer-loss indices support one or two chunks, " f"got {len(chunk_starts)}." ) chunk_count = len(chunk_starts) @@ -3199,13 +1536,11 @@ def build_indexer_loss_indices( chunk1_start, chunk_count, int(chunk_len), - int(window_capacity_per_chunk), - int(shared_compressed_base), d_window, window_size, - ratio, compressed_width, indexer_rank_by_seq_major.shape[0], + int(compressed_base), total_width, total_work, ), diff --git a/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py b/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py index 3bb795b68d7..7b9681c5e84 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py +++ b/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py @@ -1,10 +1,9 @@ # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. """MCore-facing utilities for the DSv4 THD context-parallel path. -This module is the boundary between attention code and the small CP layout -kernels in ``csa_cp_layout_kernels.py``. It owns CP partition metadata, boundary -communication, fixed-shape collectives, and the typed wrappers used by CSA/DSv4 -attention modules. +This module owns CP partition semantics, boundary exchange, compressor-input +layout, and indexer top-k metadata. It calls the retained RoPE/compaction +kernels; ``csa.py`` calls the two final-index kernels directly. """ import math @@ -14,10 +13,7 @@ import torch.distributed as dist from megatron.core.transformer.experimental_attention_variant import csa_cp_layout_kernels -from megatron.core.transformer.experimental_attention_variant.dsa_kernels import ( - batch_of_row, - indexer_topk, -) +from megatron.core.transformer.experimental_attention_variant.dsa_kernels import indexer_topk # ============================================================================= # CP Partition Modes And Layout @@ -141,45 +137,6 @@ def _normalize_row_ranges( return tuple(normalized), tuple(lengths), sum(lengths) -def _two_chunk_layout( - chunk_ranges: Sequence[Tuple[int, int]], op_name: str -) -> Tuple[Tuple[int, int], int, int]: - ranges, lengths, l_local = _normalize_row_ranges(chunk_ranges, op_name) - if len(ranges) != 2 or lengths[0] != lengths[1]: - raise RuntimeError(f"{op_name} expects exactly two equal-length chunks.") - return (int(ranges[0][0]), int(ranges[1][0])), int(lengths[0]), int(l_local) - - -def build_thd_cp_local_padding_row_mask( - cu_seqlens_padded: torch.Tensor, - cu_seqlens_unpadded: Optional[torch.Tensor], - chunk_ranges: Sequence[Tuple[int, int]], - total_padded_rows: int, -) -> Optional[torch.Tensor]: - """Return this CP rank's local query rows that are CUDA-graph padding rows.""" - if cu_seqlens_unpadded is None: - return None - if cu_seqlens_padded.shape != cu_seqlens_unpadded.shape: - raise RuntimeError( - "padded and unpadded THD cu_seqlens must have the same shape, got " - f"{tuple(cu_seqlens_padded.shape)} and {tuple(cu_seqlens_unpadded.shape)}." - ) - - chunk_ranges, _, _ = _normalize_row_ranges(chunk_ranges, "DSv4 CP padding-row mask") - device = cu_seqlens_padded.device - global_rows = torch.cat( - [torch.arange(start, end, device=device, dtype=torch.int64) for start, end in chunk_ranges], - dim=0, - ) - - row_batch_ids_full = batch_of_row(cu_seqlens_padded, total_q=int(total_padded_rows)) - row_batch_ids = row_batch_ids_full.index_select(0, global_rows) - - real_seg_lens = cu_seqlens_unpadded[1:] - cu_seqlens_unpadded[:-1] - pos_in_seg = global_rows.to(cu_seqlens_padded.dtype) - cu_seqlens_padded[row_batch_ids] - return pos_in_seg >= real_seg_lens[row_batch_ids] - - # ============================================================================= # RoPE Wrappers # ============================================================================= @@ -239,22 +196,6 @@ def apply_thd_cp_local_rope_fused( return torch.cat(parts, dim=0) -def apply_thd_cp_compressed_rope_fused( - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - comp_ids_local: torch.Tensor, - ratio: int, - nope_dim: int, - pos_dim: int, - inverse: bool = False, -) -> torch.Tensor: - """Apply fused RoPE to compressed rows using their original compression ids.""" - return csa_cp_layout_kernels.ThdCompressedRope.apply( - x, cos, sin, comp_ids_local, ratio, nope_dim, pos_dim, bool(inverse) - ) - - # ============================================================================= # Boundary Hidden Exchange # ============================================================================= @@ -440,63 +381,13 @@ def exchange_cp_boundary_hidden( # ============================================================================= -def build_global_compressed_cu_seqlens(cu_seqlens_padded: torch.Tensor, ratio: int) -> torch.Tensor: - """Build fixed-shape compressed sequence prefix sums. - - Inputs: - cu_seqlens_padded: int32 CUDA tensor, shape ``(n_seq + 1,)``. - ratio: positive compression ratio. - - Output: - int32 CUDA tensor, shape ``(n_seq + 1,)`` where each prefix increments by - ``floor(seq_len / ratio)``. This is intentionally a PyTorch utility, not - a CuTe kernel, because it is fixed-shape and does not need a host sync. - """ - if ratio <= 1: - raise RuntimeError(f"DSv4 CP compressed cu_seqlens expects ratio > 1, got {ratio}.") - out = torch.empty_like(cu_seqlens_padded) - out[:1].zero_() - compressed_lens = torch.div( - cu_seqlens_padded[1:] - cu_seqlens_padded[:-1], int(ratio), rounding_mode="floor" - ) - out[1:].copy_(torch.cumsum(compressed_lens, dim=0, dtype=torch.int32)) - return out - - -def build_cp_rank_major_compressed_metadata_fused( - cu_seqlens: torch.Tensor, - chunk_ranges: Sequence[Tuple[int, int]], - cp_size: int, - ratio: int, - d_comp: int, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Build rank-major compressed metadata for the selected CP partition mode.""" - if len(chunk_ranges) == 2: - _, chunk_len, _ = _two_chunk_layout(chunk_ranges, "DSv4 two-chunk compressed metadata") - c_cap_per_chunk = (chunk_len + int(d_comp)) // int(ratio) - c_cap_per_rank = _compressed_group_capacity(2 * chunk_len, ratio, 2 * d_comp) - return csa_cp_layout_kernels.build_compressed_row_metadata( - cu_seqlens, - int(cp_size), - chunk_len, - ratio, - d_comp, - c_cap_per_chunk, - c_cap_per_rank=c_cap_per_rank, - use_two_chunk=True, - ) - _, _, l_local = _normalize_row_ranges(chunk_ranges, "DSv4 CP compressed metadata") - c_cap = _compressed_group_capacity(l_local, ratio, d_comp) - return csa_cp_layout_kernels.build_compressed_row_metadata( - cu_seqlens, int(cp_size), l_local, ratio, d_comp, c_cap - ) - - def build_cp_compressor_prep_compact_fused( hidden_local: torch.Tensor, boundary_hidden: torch.Tensor, cu_seqlens: torch.Tensor, + cu_seqlens_compressed: torch.Tensor, chunk_ranges: Sequence[Tuple[int, int]], + cp_size: int, ratio: int, d_comp: int, d_window: int, @@ -506,20 +397,25 @@ def build_cp_compressor_prep_compact_fused( Returns: ``hidden_compact``: rank-local compressor input, shape ``(compact_group_capacity * ratio, ...)``. - ``cu_compact``: rank-local token-row prefix sums over global sequence - ids, shape ``(n_seq + 1,)``. The compressor divides these deltas - by ``ratio`` to recover this rank's visible compact groups. ``comp_ids_local``: original per-sequence compressed group id for each compact group, shape ``(compact_group_capacity,)``. For example, with ``ratio=4``, ``comp_id=3`` maps to RoPE position ``12``. + ``rank_row_for_seq_row``: map from global sequence-major compressed rows + to their canonical rank-major all-gather rows. """ d_window = int(d_window) two_chunk = len(chunk_ranges) == 2 if two_chunk: - chunk_starts, chunk_len, l_local = _two_chunk_layout( + chunk_ranges, chunk_lengths, l_local = _normalize_row_ranges( chunk_ranges, "DSv4 two-chunk compressor-prep" ) + if chunk_lengths[0] != chunk_lengths[1]: + raise RuntimeError( + "DSv4 two-chunk compressor-prep expects exactly two equal-length chunks." + ) + chunk_starts = (chunk_ranges[0][0], chunk_ranges[1][0]) + chunk_len = chunk_lengths[0] else: start, end = chunk_ranges[0] chunk_ranges = ((int(start), int(end)),) @@ -564,208 +460,74 @@ def build_cp_compressor_prep_compact_fused( c_cap, ) ) - if not two_chunk: - hidden_compact, cu_compact, _seq_ids, comp_ids, _valid = parts[0] - return hidden_compact, cu_compact, comp_ids - hidden_parts, cu_parts, _seq_parts, comp_parts, _valid_parts = zip(*parts) - target_c_cap = _compressed_group_capacity(l_local, ratio, len(chunk_ranges) * d_comp) - current_c_cap = sum(part.shape[0] for part in comp_parts) - if current_c_cap > target_c_cap: - raise RuntimeError( - "DSv4 two-chunk compressor-prep produced more groups than its aligned capacity: " - f"current={current_c_cap}, capacity={target_c_cap}." - ) - if current_c_cap < target_c_cap: - pad_groups = target_c_cap - current_c_cap - hidden_pad_rows = pad_groups * int(ratio) - hidden_parts = hidden_parts + ( - hidden_local.new_zeros((hidden_pad_rows,) + tuple(hidden_local.shape[1:])), - ) - comp_parts = comp_parts + ( - torch.full((pad_groups,), -1, dtype=torch.int32, device=hidden_local.device), - ) - cu_deltas = cu_parts[0][1:] - cu_parts[0][:-1] - for cu in cu_parts[1:]: - cu_deltas = cu_deltas + (cu[1:] - cu[:-1]) - cu_compact = torch.cat( - ( - torch.zeros((1,), dtype=cu_seqlens.dtype, device=cu_seqlens.device), - torch.cumsum(cu_deltas, dim=0), - ) - ) - return (torch.cat(hidden_parts, dim=0), cu_compact, torch.cat(comp_parts, dim=0)) - - -# ============================================================================= -# Fixed-Shape CP Collectives -# ============================================================================= - - -class _AllGatherFixedCPTensor(torch.autograd.Function): - """All-gather a fixed local tensor and reduce-scatter gradients in backward.""" - - @staticmethod - def forward(ctx, tensor: torch.Tensor, cp_group) -> torch.Tensor: - """Gather fixed-size CP-local rows into rank-major order.""" - cp_size = 1 if cp_group is None else cp_group.size() - ctx.cp_group = cp_group - ctx.input_shape = tuple(tensor.shape) - ctx.cp_size = cp_size - - if cp_size == 1: - return tensor - - local = tensor.contiguous() - output_shape = (cp_size * local.shape[0],) + tuple(local.shape[1:]) - output = local.new_empty(output_shape) - if hasattr(dist, "all_gather_into_tensor"): - dist.all_gather_into_tensor(output, local, group=cp_group) - elif hasattr(dist, "_all_gather_base"): - dist._all_gather_base(output, local, group=cp_group) - else: - raise RuntimeError("DSv4 CP fixed all-gather requires all_gather_into_tensor support") - return output - - @staticmethod - def backward(ctx, grad_output: torch.Tensor): - """Reduce-scatter rank-major gradients back to the local CP shard.""" - cp_size = ctx.cp_size - if cp_size == 1: - return grad_output, None - - expected_shape = (cp_size * ctx.input_shape[0],) + tuple(ctx.input_shape[1:]) - if tuple(grad_output.shape) != expected_shape: + if two_chunk: + hidden_parts, comp_parts = zip(*parts) + target_c_cap = _compressed_group_capacity(l_local, ratio, len(chunk_ranges) * d_comp) + current_c_cap = sum(part.shape[0] for part in comp_parts) + if current_c_cap > target_c_cap: raise RuntimeError( - "DSv4 CP fixed all-gather backward received an unexpected gradient shape: " - f"grad={tuple(grad_output.shape)}, expected={expected_shape}" + "DSv4 two-chunk compressor-prep produced more groups than its aligned capacity: " + f"current={current_c_cap}, capacity={target_c_cap}." ) - - grad_output = grad_output.contiguous() - grad_input = grad_output.new_empty(ctx.input_shape) - if hasattr(dist, "reduce_scatter_tensor"): - dist.reduce_scatter_tensor( - grad_input, grad_output, op=dist.ReduceOp.SUM, group=ctx.cp_group + if current_c_cap < target_c_cap: + pad_groups = target_c_cap - current_c_cap + hidden_parts = hidden_parts + ( + hidden_local.new_zeros((pad_groups * int(ratio),) + tuple(hidden_local.shape[1:])), ) - elif hasattr(dist, "_reduce_scatter_base"): - dist._reduce_scatter_base( - grad_input, grad_output, op=dist.ReduceOp.SUM, group=ctx.cp_group + comp_parts = comp_parts + ( + torch.full((pad_groups,), -1, dtype=torch.int32, device=hidden_local.device), ) - else: - raise RuntimeError( - "DSv4 CP fixed all-gather backward requires reduce_scatter_tensor support" - ) - return grad_input, None - - -def all_gather_fixed_cp_tensor( - tensor: torch.Tensor, cp_group: torch.distributed.ProcessGroup -) -> torch.Tensor: - """All-gather fixed-capacity CP tensors along dim 0 with autograd support.""" - return _AllGatherFixedCPTensor.apply(tensor, cp_group) - - -# ============================================================================= -# KV Packing And Compressed KV Repacking -# ============================================================================= - - -def pack_cp_kv_full_fused( - kv_local: torch.Tensor, - boundary_kv: torch.Tensor, - compressed_kv_rank_major: torch.Tensor, - seq_ids_rank_major: torch.Tensor, - comp_ids_rank_major: torch.Tensor, - valid_rank_major: torch.Tensor, - cu_seqlens: torch.Tensor, - chunk_ranges: Sequence[Tuple[int, int]], - d_window: int, - ratio: int, - rank_major_by_seq_major: Optional[torch.Tensor] = None, - cu_seqlens_compressed: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, Optional[int]]: - """Pack full KV for the selected CP partition mode.""" - d_window = int(d_window) - shared_compressed_base = None - if len(chunk_ranges) == 2: - chunk_starts, chunk_len, l_local = _two_chunk_layout(chunk_ranges, "DSv4 two-chunk KV pack") - window_capacity = max(1, int(chunk_len) + d_window * (cu_seqlens.shape[0] - 1)) - shared_compressed_base = 2 * window_capacity - total_capacity = max(1, shared_compressed_base + int(compressed_kv_rank_major.shape[0])) - dummy_ids = torch.empty((1,), dtype=torch.int32, device=kv_local.device) - dummy_valid = torch.empty((1,), dtype=torch.bool, device=kv_local.device) - seq_ids_rank_major = comp_ids_rank_major = rank_major_by_seq_major = dummy_ids - valid_rank_major = dummy_valid - cu_seqlens_compressed = cu_seqlens - global_start = local_rows = kernel_ratio = 0 - chunk0_start, chunk1_start, chunk_count = int(chunk_starts[0]), int(chunk_starts[1]), 2 + hidden_compact = torch.cat(hidden_parts, dim=0) + comp_ids_local = torch.cat(comp_parts, dim=0) else: - chunk_ranges, _, l_local = _normalize_row_ranges(chunk_ranges, "DSv4 CP KV pack") - total_capacity = max( - 1, - int(l_local) - + d_window * (cu_seqlens.shape[0] - 1) - + int(compressed_kv_rank_major.shape[0]), - ) - if rank_major_by_seq_major is None: - rank_major_by_seq_major = torch.empty((1,), dtype=torch.int32, device=kv_local.device) - if cu_seqlens_compressed is None: - cu_seqlens_compressed = cu_seqlens - global_start = int(chunk_ranges[0][0]) - local_rows = int(l_local) - kernel_ratio = int(ratio) - chunk0_start = chunk1_start = chunk_count = chunk_len = window_capacity = 0 - if kv_local.shape[0] != l_local: - raise RuntimeError( - f"DSv4 CP KV pack expects local KV rows to be {l_local}, got {kv_local.shape[0]}." - ) - expected_boundary_rows = len(chunk_ranges) * d_window - if boundary_kv.shape[0] != expected_boundary_rows: - raise RuntimeError( - "DSv4 CP KV pack expects one boundary window per chunk: " - f"boundary={boundary_kv.shape[0]}, expected={expected_boundary_rows}." - ) - - kv_full = csa_cp_layout_kernels.ThdFullKvPack.apply( - kv_local, - boundary_kv, - compressed_kv_rank_major, - seq_ids_rank_major, - comp_ids_rank_major, - valid_rank_major, - cu_seqlens, - rank_major_by_seq_major, - cu_seqlens_compressed, - global_start, - local_rows, - d_window, - kernel_ratio, - int(total_capacity), - chunk0_start, - chunk1_start, - chunk_count, - int(chunk_len), - window_capacity, + hidden_compact, comp_ids_local = parts[0] + + # A compressed group belongs canonically to the chunk containing its last + # token. From that chunk's first visible compressed row, its fixed-capacity + # rank-major slot follows directly; no (seq, comp, valid) tensors or repack + # kernel are needed. + cp_size = int(cp_size) + ratio = int(ratio) + total_chunks = cp_size * (2 if two_chunk else 1) + seq_major_rows = (l_local * cp_size) // ratio + logical_rows = torch.arange(seq_major_rows, dtype=cu_seqlens.dtype, device=cu_seqlens.device) + n_seq = cu_seqlens.shape[0] - 1 + seq_ids = torch.bucketize( + logical_rows, cu_seqlens_compressed[1:], out_int32=True, right=True + ).clamp_max(n_seq - 1) + comp_ids = logical_rows - cu_seqlens_compressed[seq_ids] + group_last_rows = cu_seqlens[seq_ids] + (comp_ids + 1) * ratio - 1 + canonical_chunks = torch.div( + group_last_rows, chunk_len if two_chunk else l_local, rounding_mode="floor" + ).clamp_(0, total_chunks - 1) + + chunk_length = chunk_len if two_chunk else l_local + chunk_starts = torch.arange( + total_chunks, dtype=cu_seqlens.dtype, device=cu_seqlens.device + ) * int(chunk_length) + first_seq_ids = torch.bucketize( + chunk_starts, cu_seqlens[1:], out_int32=True, right=True + ).clamp_max(n_seq - 1) + first_comp_ids = torch.div( + (chunk_starts - int(d_comp) - cu_seqlens[first_seq_ids]).clamp_min_(0) + ratio - 1, + ratio, + rounding_mode="floor", ) - return kv_full, shared_compressed_base - - -def repack_rank_major_compressed_to_seq_major_fused( - rank_major: torch.Tensor, - seq_ids_rank_major: torch.Tensor, - comp_ids_rank_major: torch.Tensor, - valid_rank_major: torch.Tensor, - cu_seqlens_compressed: torch.Tensor, - seq_major_rows: int, -) -> Tuple[torch.Tensor, torch.Tensor]: - """Repack gathered compressed KV from rank-major rows to sequence-major rows.""" - return csa_cp_layout_kernels.repack_compressed_kv_to_seq_major( - rank_major, - seq_ids_rank_major, - comp_ids_rank_major, - valid_rank_major, - cu_seqlens_compressed, - seq_major_rows, + first_logical_rows = cu_seqlens_compressed[first_seq_ids] + first_comp_ids + rank_slots = logical_rows - first_logical_rows[canonical_chunks] + if two_chunk: + late_chunks = canonical_chunks >= cp_size + owner_ranks = torch.where( + late_chunks, total_chunks - 1 - canonical_chunks, canonical_chunks + ) + rank_slots = rank_slots + late_chunks * int(c_cap) + else: + owner_ranks = canonical_chunks + rank_rows = owner_ranks * comp_ids_local.shape[0] + rank_slots + rank_row_for_seq_row = torch.where(logical_rows < cu_seqlens_compressed[-1], rank_rows, -1).to( + torch.int32 ) + return hidden_compact, comp_ids_local, rank_row_for_seq_row def compute_cp_indexer_topk_logical_fused( @@ -817,17 +579,43 @@ def compute_cp_indexer_topk_logical_fused( outputs = [] row_start = 0 for (global_start, _), length in zip(chunk_ranges, lengths): + global_start, length = int(global_start), int(length) + global_end = global_start + length q_for_topk = q_indexer_local.narrow(0, row_start, int(length)) weights_for_topk = weights_indexer_local.narrow(0, row_start, int(length)) - k_for_topk, cu_q_topk, cu_k_topk, seq_lens_topk = ( - csa_cp_layout_kernels.build_indexer_topk_metadata( - k_indexer_seq_major, - cu_seqlens_q, - cu_seqlens_compressed, - int(global_start), - int(length), - ratio, - ) + local_starts = cu_seqlens_q[:-1].clamp_min(global_start) + local_ends = cu_seqlens_q[1:].clamp_max(global_end) + q_lens = (local_ends - local_starts).clamp_min(0) + k_lens = torch.minimum( + torch.div( + (local_ends - cu_seqlens_q[:-1]).clamp_min(0), int(ratio), rounding_mode="floor" + ), + cu_seqlens_compressed[1:] - cu_seqlens_compressed[:-1], + ) + k_lens = torch.where(q_lens > 0, k_lens, 0) + q_prefix = torch.cumsum(q_lens, dim=0, dtype=torch.int32) + k_prefix = torch.cumsum(k_lens, dim=0, dtype=torch.int32) + zero = torch.zeros((1,), dtype=cu_seqlens_q.dtype, device=cu_seqlens_q.device) + padding_q = (global_end - cu_seqlens_q[-1].clamp_min(global_start)).clamp_min(0) + cu_q_topk = torch.cat((zero, q_prefix, (q_prefix[-1] + padding_q).view(1))) + cu_k_topk = torch.cat((zero, k_prefix, k_prefix[-1:])) + + k_rows = torch.arange( + k_indexer_seq_major.shape[0], + dtype=cu_seqlens_compressed.dtype, + device=cu_seqlens_compressed.device, + ) + k_seq_ids = torch.bucketize(k_rows, cu_k_topk[1:-1], out_int32=True, right=True).clamp_max( + cu_seqlens_q.shape[0] - 2 + ) + valid_k = k_rows < k_prefix[-1] + source_rows = cu_seqlens_compressed[k_seq_ids] + k_rows - cu_k_topk[k_seq_ids] + k_for_topk = torch.index_select( + k_indexer_seq_major, 0, torch.where(valid_k, source_rows, 0) + ) + k_for_topk = k_for_topk * valid_k.view((-1,) + (1,) * (k_for_topk.ndim - 1)) + seq_lens_topk = torch.repeat_interleave( + torch.cat((k_lens, zero)), torch.cat((q_lens, padding_q.view(1))), output_size=length ) topk_chunk, _ = indexer_topk( q_for_topk, @@ -843,116 +631,5 @@ def compute_cp_indexer_topk_logical_fused( visible_k_lengths=seq_lens_topk, ) outputs.append(topk_chunk) - row_start += int(length) + row_start += length return torch.cat(outputs, dim=0) - - -# ============================================================================= -# Final Attention And Indexer-Loss Indices -# ============================================================================= - - -def build_cp_attention_indices_fused( - cu_seqlens: torch.Tensor, - chunk_ranges: Sequence[Tuple[int, int]], - d_window: int, - window_size: int, - ratio: int, - indexer_topk_compressed_logical_ids: Optional[torch.Tensor] = None, - max_n_compressed: int = 0, - rank_major_by_seq_major: Optional[torch.Tensor] = None, - cu_seqlens_compressed: Optional[torch.Tensor] = None, - shared_compressed_base: Optional[int] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """Build final sparse-attention indices for the selected CP partition mode.""" - compressed_width = max_n_compressed if ratio > 1 else 0 - if indexer_topk_compressed_logical_ids is not None: - compressed_width = indexer_topk_compressed_logical_ids.shape[-1] - - if len(chunk_ranges) == 2: - chunk_starts, chunk_len, l_local = _two_chunk_layout( - chunk_ranges, "DSv4 two-chunk final idx" - ) - if shared_compressed_base is None: - raise RuntimeError("DSv4 two-chunk final idx requires a shared compressed base.") - if compressed_width > 0 and ( - rank_major_by_seq_major is None or cu_seqlens_compressed is None - ): - raise RuntimeError("DSv4 two-chunk final idx with compression requires metadata.") - if compressed_width == 0: - cu_seqlens_compressed = cu_seqlens - rank_major_by_seq_major = torch.empty((1,), dtype=torch.int32, device=cu_seqlens.device) - global_start, window_capacity_per_chunk = 0, int(shared_compressed_base) // 2 - else: - chunk_ranges, _, l_local = _normalize_row_ranges(chunk_ranges, "DSv4 CP final idx") - global_start = chunk_ranges[0][0] - chunk_starts, chunk_len, window_capacity_per_chunk, shared_compressed_base = None, 0, 0, 0 - if ( - indexer_topk_compressed_logical_ids is not None - and indexer_topk_compressed_logical_ids.shape[0] != l_local - ): - raise RuntimeError( - "DSv4 CP final idx expects indexer rows to be " - f"{l_local}, got {indexer_topk_compressed_logical_ids.shape[0]}." - ) - return csa_cp_layout_kernels.build_attention_indices( - cu_seqlens, - global_start, - l_local, - d_window, - window_size, - ratio, - compressed_width, - indexer_topk_compressed_logical_ids, - cu_seqlens_compressed=cu_seqlens_compressed, - rank_major_by_seq_major=rank_major_by_seq_major, - chunk_starts=chunk_starts, - chunk_len=chunk_len, - window_capacity_per_chunk=window_capacity_per_chunk, - shared_compressed_base=shared_compressed_base, - ) - - -def build_cp_indexer_loss_indices_fused( - cu_seqlens: torch.Tensor, - cu_seqlens_compressed: torch.Tensor, - chunk_ranges: Sequence[Tuple[int, int]], - d_window: int, - window_size: int, - ratio: int, - indexer_topk_compressed_logical_ids: torch.Tensor, - rank_major_by_seq_major: torch.Tensor, - shared_compressed_base: Optional[int] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """Build indexer-loss indices for the selected CP partition mode.""" - if len(chunk_ranges) == 2: - chunk_starts, chunk_len, l_local = _two_chunk_layout( - chunk_ranges, "DSv4 two-chunk indexer-loss idx" - ) - if shared_compressed_base is None: - raise RuntimeError("DSv4 two-chunk indexer-loss idx requires a shared compressed base.") - global_start, window_capacity_per_chunk = 0, int(shared_compressed_base) // 2 - else: - chunk_ranges, _, l_local = _normalize_row_ranges(chunk_ranges, "DSv4 CP indexer-loss idx") - global_start = chunk_ranges[0][0] - chunk_starts, chunk_len, window_capacity_per_chunk, shared_compressed_base = None, 0, 0, 0 - if indexer_topk_compressed_logical_ids.shape[0] != l_local: - raise RuntimeError( - "DSv4 CP indexer-loss idx expects indexer rows to be " - f"{l_local}, got {indexer_topk_compressed_logical_ids.shape[0]}." - ) - return csa_cp_layout_kernels.build_indexer_loss_indices( - cu_seqlens, - cu_seqlens_compressed, - global_start, - l_local, - d_window, - window_size, - ratio, - indexer_topk_compressed_logical_ids, - rank_major_by_seq_major, - chunk_starts=chunk_starts, - chunk_len=chunk_len, - window_capacity_per_chunk=window_capacity_per_chunk, - shared_compressed_base=shared_compressed_base, - ) diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_layout_kernels.py b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_layout_kernels.py index 0039f661185..31ecaa5f502 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_layout_kernels.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_layout_kernels.py @@ -228,205 +228,32 @@ def _native_compressor_input_compact( dim=0, ) - cu_compact = torch.zeros_like(cu_seqlens) - running = 0 - for seq in range(len(cu) - 1): - running += sum(1 for group_seq, _ in groups if group_seq == seq) * int(ratio) - cu_compact[seq + 1] = running - - seq_ids = torch.full((c_cap,), -1, dtype=torch.int32, device=hidden_local.device) comp_ids = torch.full((c_cap,), -1, dtype=torch.int32, device=hidden_local.device) - valid = torch.zeros((c_cap,), dtype=torch.bool, device=hidden_local.device) - for slot, (seq, comp_id) in enumerate(groups[:c_cap]): - seq_start = cu[seq] - group_end = seq_start + (comp_id + 1) * int(ratio) - seq_ids[slot] = seq + for slot, (_, comp_id) in enumerate(groups[:c_cap]): comp_ids[slot] = comp_id - valid[slot] = group_end - 1 >= global_start and group_end - 1 < global_start + l_local - return hidden_compact, cu_compact, seq_ids, comp_ids, valid + return hidden_compact, comp_ids -def _native_build_compressed_row_metadata( - cu_seqlens: torch.Tensor, - cp_size: int, - chunk_len: int, - ratio: int, - d_comp: int, - c_cap_per_chunk: int, - c_cap_per_rank: int, - use_two_chunk: bool, -): - total_rows = int(cp_size) * int(c_cap_per_rank) - seq_ids = torch.full((total_rows,), -1, dtype=torch.int32, device=cu_seqlens.device) - comp_ids = torch.full((total_rows,), -1, dtype=torch.int32, device=cu_seqlens.device) - valid = torch.zeros((total_rows,), dtype=torch.bool, device=cu_seqlens.device) - for row in range(total_rows): - rank = row // c_cap_per_rank - rank_slot = row - rank * c_cap_per_rank - local_chunk = 0 - slot = rank_slot - chunk_id = rank - if use_two_chunk: - local_chunk = rank_slot // c_cap_per_chunk - if local_chunk >= 2: - continue - slot = rank_slot - local_chunk * c_cap_per_chunk - chunk_id = rank if local_chunk == 0 else cp_size * 2 - 1 - rank - rank_start = chunk_id * chunk_len - groups = _compressed_groups(cu_seqlens, rank_start, chunk_len, ratio, d_comp) - if slot < len(groups): - seq, comp = groups[slot] - seq_start = int(cu_seqlens[seq].item()) - group_end = seq_start + (comp + 1) * ratio - seq_ids[row] = seq - comp_ids[row] = comp - valid[row] = group_end - 1 >= rank_start and group_end - 1 < rank_start + chunk_len - return seq_ids, comp_ids, valid - - -def _native_repack_compressed_kv_to_seq_major( - compressed_rank_major: torch.Tensor, - seq_ids: torch.Tensor, - comp_ids: torch.Tensor, - valid: torch.Tensor, - cu_seqlens_compressed: torch.Tensor, - seq_major_rows: int, -): - out = compressed_rank_major.new_zeros( - (seq_major_rows,) + tuple(compressed_rank_major.shape[1:]) - ) - rank_by_seq = torch.full( - (seq_major_rows,), -1, dtype=torch.int32, device=compressed_rank_major.device - ) - for rank_row in range(compressed_rank_major.shape[0]): - if bool(valid[rank_row]): - seq = int(seq_ids[rank_row]) - comp = int(comp_ids[rank_row]) - seq_major = int(cu_seqlens_compressed[seq]) + comp - if 0 <= seq_major < seq_major_rows: - out[seq_major] = compressed_rank_major[rank_row] - rank_by_seq[seq_major] = rank_row - return out, rank_by_seq - - -def _native_indexer_topk_metadata( - k_seq_major: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_compressed: torch.Tensor, - global_start: int, - l_local: int, - ratio: int, -): - n_seq = cu_seqlens_q.shape[0] - 1 - k_topk = torch.zeros_like(k_seq_major) - cu_q = torch.zeros((n_seq + 2,), dtype=torch.int32, device=k_seq_major.device) - cu_k = torch.zeros((n_seq + 2,), dtype=torch.int32, device=k_seq_major.device) - seq_lens = torch.empty((l_local,), dtype=torch.int32, device=k_seq_major.device) - q_prefix = 0 - k_prefix = 0 - for seq in range(n_seq): - seq_start = int(cu_seqlens_q[seq]) - seq_end = int(cu_seqlens_q[seq + 1]) - local_start = max(seq_start, int(global_start)) - local_end = min(seq_end, int(global_start) + int(l_local)) - q_len = 0 - k_len = 0 - if local_start < local_end: - q_len = local_end - local_start - seq_comp_start = int(cu_seqlens_compressed[seq]) - seq_comp_end = int(cu_seqlens_compressed[seq + 1]) - k_len = min((local_end - seq_start) // int(ratio), seq_comp_end - seq_comp_start) - if k_len: - k_topk[k_prefix : k_prefix + k_len] = k_seq_major[ - seq_comp_start : seq_comp_start + k_len - ] - if q_len: - seq_lens[q_prefix : q_prefix + q_len] = k_len - q_prefix += q_len - k_prefix += k_len - cu_q[seq + 1] = q_prefix - cu_k[seq + 1] = k_prefix - actual_total = int(cu_seqlens_q[-1]) - padding_start = max(actual_total, int(global_start)) - padding_q = max(0, int(global_start) + int(l_local) - padding_start) - if padding_q: - seq_lens[q_prefix : q_prefix + padding_q] = 0 - cu_q[n_seq + 1] = q_prefix + padding_q - cu_k[n_seq + 1] = k_prefix - return k_topk, cu_q, cu_k, seq_lens - - -def _native_thd_full_kv_pack( - kv_local: torch.Tensor, - boundary_kv: torch.Tensor, - compressed_rank_major: torch.Tensor, - seq_ids: torch.Tensor, - comp_ids: torch.Tensor, - valid: torch.Tensor, - cu_seqlens: torch.Tensor, - global_start: int, +def _native_attention_indices_one_seq( l_local: int, d_window: int, + window_size: int, ratio: int, - capacity: int, -): - rows = [] - n_seq = cu_seqlens.shape[0] - 1 - for seq in range(n_seq): - seq_start = int(cu_seqlens[seq]) - seq_end = int(cu_seqlens[seq + 1]) - local_start = max(seq_start, int(global_start)) - local_end = min(seq_end, int(global_start) + int(l_local)) - if local_start >= local_end: - continue - window_start = max(seq_start, local_start - int(d_window)) - for src_global in range(window_start, local_end): - if src_global < global_start: - rows.append( - boundary_kv[ - src_global - - (global_start - d_window) : src_global - - (global_start - d_window) - + 1 - ] - ) - else: - rows.append(kv_local[src_global - global_start : src_global - global_start + 1]) - if ratio > 1: - for comp_id in range((seq_end - seq_start) // int(ratio)): - for rank_row in range(compressed_rank_major.shape[0]): - if ( - bool(valid[rank_row]) - and int(seq_ids[rank_row]) == seq - and int(comp_ids[rank_row]) == comp_id - ): - rows.append(compressed_rank_major[rank_row : rank_row + 1]) - if rows: - out = torch.cat(rows, dim=0) - else: - out = kv_local.new_empty((0,) + tuple(kv_local.shape[1:])) - if out.shape[0] < capacity: - out = torch.cat( - (out, kv_local.new_zeros((capacity - out.shape[0],) + tuple(kv_local.shape[1:]))), dim=0 - ) - return out - - -def _native_attention_indices_one_seq( - l_local: int, window_size: int, ratio: int, compressed_width: int, device: torch.device + compressed_width: int, + device: torch.device, ): rows = torch.arange(l_local, dtype=torch.int32, device=device).unsqueeze(1) window_cols = torch.arange(window_size, dtype=torch.int32, device=device).unsqueeze(0) window_start = (rows - int(window_size) + 1).clamp_min(0) window_count = rows - window_start + 1 - window_values = window_start + window_cols + window_values = int(d_window) + window_start + window_cols window_values = torch.where( window_cols < window_count, window_values, torch.full_like(window_values, -1) ) if ratio > 1 and compressed_width > 0: comp_cols = torch.arange(compressed_width, dtype=torch.int32, device=device).unsqueeze(0) visible = ((rows + 1) // int(ratio)).clamp(max=int(compressed_width)) - comp_values = int(l_local) + comp_cols + comp_values = int(d_window) + int(l_local) + comp_cols comp_values = torch.where( comp_cols < visible, comp_values, torch.full_like(comp_values, -1) ) @@ -440,12 +267,15 @@ def _native_attention_indices_one_seq( def _native_attention_indices( cu_seqlens: torch.Tensor, + cu_seqlens_compressed: torch.Tensor, global_start: int, l_local: int, d_window: int, window_size: int, ratio: int, compressed_width: int, + rank_by_seq_major: torch.Tensor, + compressed_base: int, indexer_topk: torch.Tensor = None, ): """Reference contiguous final-index lowering for ragged THD CP rows.""" @@ -454,36 +284,32 @@ def _native_attention_indices( topk = torch.full((int(l_local), total_width), -1, dtype=torch.int32, device=device) lengths = torch.zeros((int(l_local),), dtype=torch.int32, device=device) cu = [int(x) for x in cu_seqlens.cpu().tolist()] - seq_infos = [] - running = 0 - global_end = int(global_start) + int(l_local) - for seq, (seq_start, seq_end) in enumerate(zip(cu[:-1], cu[1:])): - local_start = max(seq_start, int(global_start)) - local_end = min(seq_end, global_end) - if local_start >= local_end: - continue - window_start = max(seq_start, local_start - int(d_window)) - window_len = local_end - window_start - comp_len = (seq_end - seq_start) // int(ratio) if int(ratio) > 1 else 0 - seq_infos.append((seq, seq_start, seq_end, window_start, window_len, comp_len, running)) - running += window_len + comp_len + cu_comp = [int(x) for x in cu_seqlens_compressed.cpu().tolist()] for row in range(int(l_local)): global_q = int(global_start) + row - info = next((item for item in seq_infos if item[1] <= global_q < item[2]), None) - if info is None: + seq_id = next( + ( + seq + for seq, (start, end) in enumerate(zip(cu[:-1], cu[1:])) + if start <= global_q < end + ), + None, + ) + if seq_id is None: if total_width > 0: topk[row, 0] = 0 lengths[row] = 1 continue - _seq, seq_start, _seq_end, seq_window_start, seq_window_len, seq_comp_len, seq_offset = info + seq_start = cu[seq_id] + seq_comp_len = cu_comp[seq_id + 1] - cu_comp[seq_id] write_col = 0 - window_start_for_q = max(global_q - int(window_size) + 1, seq_start, seq_window_start) + window_start_for_q = max(global_q - int(window_size) + 1, seq_start) window_count = max(0, global_q - window_start_for_q + 1) for w in range(int(window_size)): pos = window_start_for_q + w - if w < window_count and seq_window_start <= pos < seq_window_start + seq_window_len: - topk[row, write_col] = seq_offset + pos - seq_window_start + if w < window_count: + topk[row, write_col] = pos - (int(global_start) - int(d_window)) write_col += 1 if int(ratio) > 1 and int(compressed_width) > 0: pos_in_seq = global_q - seq_start @@ -494,14 +320,17 @@ def _native_attention_indices( n_visible = min((pos_in_seq + 1) // int(ratio), int(compressed_width)) comp_id = j if j < n_visible else -1 if 0 <= comp_id < seq_comp_len: - topk[row, write_col] = seq_offset + seq_window_len + comp_id - write_col += 1 + rank_id = int(rank_by_seq_major[cu_comp[seq_id] + comp_id]) + if rank_id >= 0: + topk[row, write_col] = int(compressed_base) + rank_id + write_col += 1 lengths[row] = write_col return topk, lengths def _native_indexer_loss_indices_one_seq( l_local: int, + d_window: int, window_size: int, ratio: int, logical_ids: torch.Tensor, @@ -509,12 +338,14 @@ def _native_indexer_loss_indices_one_seq( ): device = logical_ids.device compressed_width = logical_ids.shape[1] - attention_topk, _ = _native_attention_indices_one_seq(l_local, window_size, ratio, 0, device) + attention_topk, _ = _native_attention_indices_one_seq( + l_local, d_window, window_size, ratio, 0, device + ) compressed = torch.full((l_local, compressed_width), -1, dtype=torch.int32, device=device) rank_major = torch.full_like(compressed, -1) valid = (logical_ids >= 0) & (logical_ids < rank_by_seq_major.shape[0]) if valid.any(): - lowered = int(l_local) + logical_ids.clamp(min=0) + lowered = int(d_window) + int(l_local) + logical_ids.clamp(min=0) rank_ids = rank_by_seq_major.index_select(0, logical_ids.clamp(min=0).reshape(-1)).view_as( logical_ids ) @@ -533,6 +364,7 @@ def _native_indexer_loss_indices( ratio: int, logical_ids: torch.Tensor, rank_by_seq_major: torch.Tensor, + compressed_base: int, ): """Reference contiguous compressed-first indexer-loss lowering for ragged THD rows.""" device = logical_ids.device @@ -542,26 +374,20 @@ def _native_indexer_loss_indices( rank_major = torch.full((int(l_local), compressed_width), -1, dtype=torch.int32, device=device) cu = [int(x) for x in cu_seqlens.cpu().tolist()] cu_comp = [int(x) for x in cu_seqlens_compressed.cpu().tolist()] - seq_infos = [] - running = 0 - global_end = int(global_start) + int(l_local) - for seq, (seq_start, seq_end) in enumerate(zip(cu[:-1], cu[1:])): - local_start = max(seq_start, int(global_start)) - local_end = min(seq_end, global_end) - if local_start >= local_end: - continue - window_start = max(seq_start, local_start - int(d_window)) - window_len = local_end - window_start - comp_len = cu_comp[seq + 1] - cu_comp[seq] if int(ratio) > 1 else 0 - seq_infos.append((seq, seq_start, seq_end, window_start, window_len, comp_len, running)) - running += window_len + comp_len - for row in range(int(l_local)): global_q = int(global_start) + row - info = next((item for item in seq_infos if item[1] <= global_q < item[2]), None) - if info is None: + seq_id = next( + ( + seq + for seq, (start, end) in enumerate(zip(cu[:-1], cu[1:])) + if start <= global_q < end + ), + None, + ) + if seq_id is None: continue - seq_id, seq_start, _seq_end, window_start, window_len, comp_len, offset = info + seq_start = cu[seq_id] + comp_len = cu_comp[seq_id + 1] - cu_comp[seq_id] for col in range(compressed_width): comp_id = int(logical_ids[row, col]) @@ -569,19 +395,15 @@ def _native_indexer_loss_indices( seq_major = cu_comp[seq_id] + comp_id rank_id = int(rank_by_seq_major[seq_major]) if rank_id >= 0: - topk[row, col] = offset + window_len + comp_id + topk[row, col] = int(compressed_base) + rank_id rank_major[row, col] = rank_id window_start_for_q = max(global_q - int(window_size) + 1, seq_start) window_count = global_q - window_start_for_q + 1 for window_col in range(int(window_size)): pos = window_start_for_q + window_col - if ( - window_col < window_count - and pos >= window_start - and pos < window_start + window_len - ): - topk[row, compressed_width + window_col] = offset + pos - window_start + if window_col < window_count: + topk[row, compressed_width + window_col] = pos - (int(global_start) - int(d_window)) return topk, rank_major @@ -649,53 +471,6 @@ def run_torch(): ) -def test_thd_compressed_rope_matches_native_forward_backward_and_reports_bandwidth(): - _require_cute_cuda() - torch.manual_seed(22) - comp_ids = torch.tensor([0, 2, -1, 17, 31, 127], dtype=torch.int32, device="cuda") - ratio = 4 - nope_dim = 3 - pos_dim = 4 - x = torch.randn(comp_ids.numel(), 1, 2, nope_dim + pos_dim, device="cuda", requires_grad=True) - cos = torch.randn(max(_E2E_RAGGED_PADDED_SEG_LENS), 1, 1, pos_dim, device="cuda") - sin = torch.randn(max(_E2E_RAGGED_PADDED_SEG_LENS), 1, 1, pos_dim, device="cuda") - positions = comp_ids.clamp(min=0).to(torch.long) * ratio - ref_x = x.detach().clone().requires_grad_(True) - ref = _rope_reference(ref_x, cos, sin, positions, nope_dim, pos_dim) - grad = torch.randn_like(ref) - ref.backward(grad) - - fused_x = x.detach().clone().requires_grad_(True) - fused = csa_cp_layout_kernels.ThdCompressedRope.apply( - fused_x, cos, sin, comp_ids, ratio, nope_dim, pos_dim, False - ) - fused.backward(grad) - _assert_rope_close(fused, ref, "thd_compressed_rope.forward") - _assert_rope_close(fused_x.grad, ref_x.grad, "thd_compressed_rope.backward") - - bench_x = torch.randn(1024, 4, nope_dim + pos_dim, device="cuda") - bench_grad = torch.randn_like(bench_x) - bench_ids = torch.arange(1024, dtype=torch.int32, device="cuda") % 8 - bench_pos = bench_ids.to(torch.long) * ratio - - def run_cute(): - bx = bench_x.detach().clone().requires_grad_(True) - y = csa_cp_layout_kernels.ThdCompressedRope.apply( - bx, cos, sin, bench_ids, ratio, nope_dim, pos_dim, False - ) - y.backward(bench_grad) - - def run_torch(): - bx = bench_x.detach().clone().requires_grad_(True) - y = _rope_reference(bx, cos, sin, bench_pos, nope_dim, pos_dim) - y.backward(bench_grad) - - effective_bytes = 4 * _num_bytes(bench_x) - _print_perf( - "thd_compressed_rope_fwd_bwd", _time_cuda(run_cute), _time_cuda(run_torch), effective_bytes - ) - - def test_compressor_input_compact_matches_native_forward_backward_and_reports_bandwidth(): _require_cute_cuda() cu = _make_e2e_like_cu_seqlens() @@ -736,8 +511,7 @@ def test_compressor_input_compact_matches_native_forward_backward_and_reports_ba ) fused[0].backward(grad) assert torch.equal(fused[0], ref[0]) - for actual, expected in zip(fused[1:], ref[1:]): - assert torch.equal(actual, expected) + assert torch.equal(fused[1], ref[1]) assert torch.equal(hidden.grad, ref_hidden_grad) assert torch.equal(boundary.grad, ref_boundary_grad) @@ -793,317 +567,73 @@ def run_torch(): ) -def test_thd_full_kv_pack_matches_native_forward_backward_and_reports_bandwidth(): +def test_build_attention_indices_matches_native_and_reports_bandwidth(): _require_cute_cuda() cu = _make_e2e_like_cu_seqlens() global_start, l_local = _e2e_like_contiguous_range() d_window = 128 - ratio = 128 - d_comp = 128 - c_cap = (l_local + d_comp) // ratio - kv_local = torch.randn(l_local, 4, dtype=torch.bfloat16, device="cuda").requires_grad_(True) - boundary = torch.randn(d_window, 4, dtype=torch.bfloat16, device="cuda").requires_grad_(True) - seq_ids, comp_ids, valid = _native_build_compressed_row_metadata( - cu, _E2E_CP_SIZE, l_local, ratio, d_comp, c_cap, c_cap, False - ) - compressed = torch.randn( - seq_ids.shape[0], 4, dtype=torch.bfloat16, device="cuda" - ).requires_grad_(True) - dummy = torch.empty((1,), dtype=torch.int32, device="cuda") - capacity = l_local + d_window * (cu.shape[0] - 1) + compressed.shape[0] - - ref = _native_thd_full_kv_pack( - kv_local, - boundary, - compressed, - seq_ids, - comp_ids, - valid, + window = 16 + ratio = 4 + compressed_width = 16 + cu_comp = _compressed_cu_seqlens(cu, ratio) + rank_by_seq = torch.arange(int(cu_comp[-1]), dtype=torch.int32, device="cuda").flip(0) + compressed_base = d_window + l_local + fused = csa_cp_layout_kernels.build_attention_indices( cu, global_start, l_local, d_window, + window, ratio, - capacity, - ) - grad = torch.randn_like(ref) - ref.backward(grad) - ref_grads = ( - kv_local.grad.detach().clone(), - boundary.grad.detach().clone(), - compressed.grad.detach().clone(), + compressed_width, + cu_seqlens_compressed=cu_comp, + rank_row_for_seq_row=rank_by_seq, + compressed_base=compressed_base, ) - kv_local.grad.zero_() - boundary.grad.zero_() - compressed.grad.zero_() - - fused = csa_cp_layout_kernels.ThdFullKvPack.apply( - kv_local, - boundary, - compressed, - seq_ids, - comp_ids, - valid, - cu, - dummy, + expected = _native_attention_indices( cu, + cu_comp, global_start, l_local, d_window, + window, ratio, - capacity, - 0, - 0, - 0, - 0, - 0, + compressed_width, + rank_by_seq, + compressed_base, ) - fused.backward(grad) - assert torch.equal(fused, ref) - assert torch.equal(kv_local.grad, ref_grads[0]) - assert torch.equal(boundary.grad, ref_grads[1]) - assert torch.equal(compressed.grad, ref_grads[2]) - - bench_cu = torch.tensor([0, 8192], dtype=torch.int32, device="cuda") - bench_global_start = 2048 - bench_l_local = 4096 - bench_d_window = 128 - bench_ratio = 128 - bench_compressed_rows = int(bench_cu[-1]) // bench_ratio - bench_capacity = bench_d_window + bench_l_local + bench_compressed_rows - bench_kv = torch.randn(bench_l_local, 64, dtype=torch.bfloat16, device="cuda") - bench_boundary = torch.randn(bench_d_window, 64, dtype=torch.bfloat16, device="cuda") - bench_compressed = torch.randn(bench_compressed_rows, 64, dtype=torch.bfloat16, device="cuda") - bench_seq_ids = torch.zeros(bench_compressed_rows, dtype=torch.int32, device="cuda") - bench_comp_ids = torch.arange(bench_compressed_rows, dtype=torch.int32, device="cuda") - bench_valid = torch.ones(bench_compressed_rows, dtype=torch.bool, device="cuda") - bench_dummy = torch.empty((1,), dtype=torch.int32, device="cuda") - bench_grad = torch.randn(bench_capacity, 64, dtype=torch.bfloat16, device="cuda") - - def run_cute(): - k = bench_kv.detach().clone().requires_grad_(True) - b = bench_boundary.detach().clone().requires_grad_(True) - c = bench_compressed.detach().clone().requires_grad_(True) - y = csa_cp_layout_kernels.ThdFullKvPack.apply( - k, - b, - c, - bench_seq_ids, - bench_comp_ids, - bench_valid, - bench_cu, - bench_dummy, - bench_cu, - bench_global_start, - bench_l_local, - bench_d_window, - bench_ratio, - bench_capacity, - 0, - 0, - 0, - 0, - 0, - ) - y.backward(bench_grad) - - def run_torch(): - k = bench_kv.detach().clone().requires_grad_(True) - b = bench_boundary.detach().clone().requires_grad_(True) - c = bench_compressed.detach().clone().requires_grad_(True) - y = _native_thd_full_kv_pack( - k, - b, - c, - bench_seq_ids, - bench_comp_ids, - bench_valid, - bench_cu, - bench_global_start, - bench_l_local, - bench_d_window, - bench_ratio, - bench_capacity, - ) - y.backward(bench_grad) - - effective_bytes = 2 * _num_bytes(bench_kv, bench_boundary, bench_compressed, bench_grad) - _print_perf( - "thd_full_kv_pack_fwd_bwd", _time_cuda(run_cute), _time_cuda(run_torch), effective_bytes - ) - - -def test_repack_compressed_kv_to_seq_major_matches_native_and_reports_bandwidth(): - _require_cute_cuda() - cu = _make_e2e_like_cu_seqlens() - ratio = 128 - d_comp = 128 - global_start, l_local = _e2e_like_contiguous_range() - c_cap = (l_local + d_comp) // ratio - seq_ids, comp_ids, valid = _native_build_compressed_row_metadata( - cu, _E2E_CP_SIZE, l_local, ratio, d_comp, c_cap, c_cap, False - ) - rank_major = torch.arange(seq_ids.numel() * 4, dtype=torch.bfloat16, device="cuda").reshape( - seq_ids.numel(), 4 - ) - cu_comp = _compressed_cu_seqlens(cu, ratio) - ref = _native_repack_compressed_kv_to_seq_major( - rank_major, seq_ids, comp_ids, valid, cu_comp, int(cu_comp[-1]) - ) - fused = csa_cp_layout_kernels.repack_compressed_kv_to_seq_major( - rank_major, seq_ids, comp_ids, valid, cu_comp, int(cu_comp[-1]) - ) - assert torch.equal(fused[0], ref[0]) - assert torch.equal(fused[1], ref[1]) - - bench_rows = 4096 - bench_width = 64 - bench_rank_major = torch.randn(bench_rows, bench_width, dtype=torch.bfloat16, device="cuda") - bench_seq_ids = torch.zeros(bench_rows, dtype=torch.int32, device="cuda") - bench_comp_ids = torch.arange(bench_rows, dtype=torch.int32, device="cuda") - bench_valid = torch.ones(bench_rows, dtype=torch.bool, device="cuda") - bench_cu_comp = torch.tensor([0, bench_rows], dtype=torch.int32, device="cuda") - - def run_cute(): - csa_cp_layout_kernels.repack_compressed_kv_to_seq_major( - bench_rank_major, bench_seq_ids, bench_comp_ids, bench_valid, bench_cu_comp, bench_rows - ) - - def run_torch(): - _native_repack_compressed_kv_to_seq_major( - bench_rank_major, bench_seq_ids, bench_comp_ids, bench_valid, bench_cu_comp, bench_rows - ) - - effective_bytes = _num_bytes(bench_rank_major) * 2 + bench_rows * 4 - _print_perf( - "repack_compressed_kv_to_seq_major", - _time_cuda(run_cute), - _time_cuda(run_torch), - effective_bytes, - ) - + assert torch.equal(fused[0], expected[0]) + assert torch.equal(fused[1], expected[1]) -def test_build_compressed_row_metadata_matches_native_and_reports_bandwidth(): - _require_cute_cuda() - cu = _make_e2e_like_cu_seqlens() - cp_size = _E2E_CP_SIZE - chunk_len = (sum(_E2E_RAGGED_PADDED_SEG_LENS) // cp_size) // 2 - ratio = 4 - d_comp = 8 - c_cap_per_chunk = (chunk_len + d_comp) // ratio - c_cap_per_rank = 2 * c_cap_per_chunk - ref = _native_build_compressed_row_metadata( - cu, cp_size, chunk_len, ratio, d_comp, c_cap_per_chunk, c_cap_per_rank, True - ) - fused = csa_cp_layout_kernels.build_compressed_row_metadata( + logical_ids = torch.tensor([2, 0, -1], dtype=torch.int32, device="cuda").expand(l_local, -1) + fused_indexer = csa_cp_layout_kernels.build_attention_indices( cu, - cp_size, - chunk_len, + global_start, + l_local, + d_window, + window, ratio, - d_comp, - c_cap_per_chunk, - c_cap_per_rank=c_cap_per_rank, - use_two_chunk=True, - ) - for actual, expected in zip(fused, ref): - assert torch.equal(actual, expected) - - bench_cu = torch.tensor([0, 65536], dtype=torch.int32, device="cuda") - bench_cp_size = 4 - bench_chunk_len = 8192 - bench_ratio = 128 - bench_d_comp = 128 - bench_c_cap = (bench_chunk_len + bench_d_comp) // bench_ratio - bench_c_cap_per_rank = bench_c_cap * 2 - - def run_cute(): - csa_cp_layout_kernels.build_compressed_row_metadata( - bench_cu, - bench_cp_size, - bench_chunk_len, - bench_ratio, - bench_d_comp, - bench_c_cap, - c_cap_per_rank=bench_c_cap_per_rank, - use_two_chunk=True, - ) - - def run_torch(): - _native_build_compressed_row_metadata( - bench_cu, - bench_cp_size, - bench_chunk_len, - bench_ratio, - bench_d_comp, - bench_c_cap, - bench_c_cap_per_rank, - True, - ) - - effective_bytes = bench_cp_size * bench_c_cap_per_rank * (4 + 4 + 1) - _print_perf( - "build_compressed_row_metadata", - _time_cuda(run_cute), - _time_cuda(run_torch), - effective_bytes, - ) - - -def test_build_indexer_topk_metadata_matches_native_and_reports_bandwidth(): - _require_cute_cuda() - ratio = 4 - cu_q = _make_e2e_like_cu_seqlens() - cu_comp = _compressed_cu_seqlens(cu_q, ratio) - k_seq = torch.arange(int(cu_comp[-1]) * 2, dtype=torch.bfloat16, device="cuda").reshape( - int(cu_comp[-1]), 2 - ) - global_start, l_local = _e2e_like_contiguous_range() - ref = _native_indexer_topk_metadata(k_seq, cu_q, cu_comp, global_start, l_local, ratio) - fused = csa_cp_layout_kernels.build_indexer_topk_metadata( - k_seq, cu_q, cu_comp, global_start, l_local, ratio + logical_ids.shape[1], + logical_ids, + cu_seqlens_compressed=cu_comp, + rank_row_for_seq_row=rank_by_seq, + compressed_base=compressed_base, ) - for actual, expected in zip(fused, ref): - assert torch.equal(actual, expected) - - bench_ratio = 128 - bench_l_local = 8192 - bench_rows = 512 - bench_k_seq = torch.randn(bench_rows, 64, dtype=torch.bfloat16, device="cuda") - bench_cu_q = torch.tensor([0, bench_l_local], dtype=torch.int32, device="cuda") - bench_cu_comp = torch.tensor([0, bench_rows], dtype=torch.int32, device="cuda") - - def run_cute(): - csa_cp_layout_kernels.build_indexer_topk_metadata( - bench_k_seq, bench_cu_q, bench_cu_comp, 0, bench_l_local, bench_ratio - ) - - def run_torch(): - _native_indexer_topk_metadata( - bench_k_seq, bench_cu_q, bench_cu_comp, 0, bench_l_local, bench_ratio - ) - - effective_bytes = _num_bytes(bench_k_seq) * 2 + bench_l_local * 4 - _print_perf( - "build_indexer_topk_metadata", _time_cuda(run_cute), _time_cuda(run_torch), effective_bytes - ) - - -def test_build_attention_indices_matches_native_and_reports_bandwidth(): - _require_cute_cuda() - cu = _make_e2e_like_cu_seqlens() - global_start, l_local = _e2e_like_contiguous_range() - d_window = 128 - window = 16 - ratio = 4 - compressed_width = 16 - fused = csa_cp_layout_kernels.build_attention_indices( - cu, global_start, l_local, d_window, window, ratio, compressed_width - ) - expected = _native_attention_indices( - cu, global_start, l_local, d_window, window, ratio, compressed_width + expected_indexer = _native_attention_indices( + cu, + cu_comp, + global_start, + l_local, + d_window, + window, + ratio, + logical_ids.shape[1], + rank_by_seq, + compressed_base, + logical_ids, ) - assert torch.equal(fused[0], expected[0]) - assert torch.equal(fused[1], expected[1]) + assert torch.equal(fused_indexer[0], expected_indexer[0]) + assert torch.equal(fused_indexer[1], expected_indexer[1]) padded = csa_cp_layout_kernels.build_attention_indices( torch.tensor([0, 8], dtype=torch.int32, device="cuda"), 0, 10, 2, 2, 0, 0 @@ -1119,8 +649,15 @@ def test_build_attention_indices_matches_native_and_reports_bandwidth(): bench_ratio = 4 bench_compressed_width = bench_l_local // bench_ratio bench_cu = torch.tensor([0, bench_l_local], dtype=torch.int32, device="cuda") + bench_cu_comp = torch.tensor([0, bench_compressed_width], dtype=torch.int32, device="cuda") + bench_rank_by_seq = torch.arange(bench_compressed_width, dtype=torch.int32, device="cuda") bench_ref = _native_attention_indices_one_seq( - bench_l_local, bench_window, bench_ratio, bench_compressed_width, torch.device("cuda") + bench_l_local, + bench_window, + bench_window, + bench_ratio, + bench_compressed_width, + torch.device("cuda"), ) def run_cute(): @@ -1132,11 +669,19 @@ def run_cute(): bench_window, bench_ratio, bench_compressed_width, + cu_seqlens_compressed=bench_cu_comp, + rank_row_for_seq_row=bench_rank_by_seq, + compressed_base=bench_window + bench_l_local, ) def run_torch(): _native_attention_indices_one_seq( - bench_l_local, bench_window, bench_ratio, bench_compressed_width, torch.device("cuda") + bench_l_local, + bench_window, + bench_window, + bench_ratio, + bench_compressed_width, + torch.device("cuda"), ) effective_bytes = _num_bytes(*bench_ref) @@ -1160,12 +705,31 @@ def test_build_indexer_loss_indices_matches_native_and_reports_bandwidth(): .repeat(l_local, 1) ) logical_ids[1::5, -1] = -1 - rank_by_seq = torch.arange(int(cu_comp[-1]), dtype=torch.int32, device="cuda") + rank_by_seq = torch.arange(int(cu_comp[-1]), dtype=torch.int32, device="cuda").flip(0) + compressed_base = d_window + l_local fused = csa_cp_layout_kernels.build_indexer_loss_indices( - cu, cu_comp, global_start, l_local, d_window, window, ratio, logical_ids, rank_by_seq + cu, + cu_comp, + global_start, + l_local, + d_window, + window, + ratio, + logical_ids, + rank_by_seq, + compressed_base=compressed_base, ) expected = _native_indexer_loss_indices( - cu, cu_comp, global_start, l_local, d_window, window, ratio, logical_ids, rank_by_seq + cu, + cu_comp, + global_start, + l_local, + d_window, + window, + ratio, + logical_ids, + rank_by_seq, + compressed_base, ) assert torch.equal(fused[0], expected[0]) assert torch.equal(fused[1], expected[1]) @@ -1183,7 +747,7 @@ def test_build_indexer_loss_indices_matches_native_and_reports_bandwidth(): ) bench_rank_by_seq = torch.arange(bench_compressed_width, dtype=torch.int32, device="cuda") bench_ref = _native_indexer_loss_indices_one_seq( - bench_l_local, bench_window, bench_ratio, bench_logical, bench_rank_by_seq + bench_l_local, bench_window, bench_window, bench_ratio, bench_logical, bench_rank_by_seq ) def run_cute(): @@ -1197,11 +761,12 @@ def run_cute(): bench_ratio, bench_logical, bench_rank_by_seq, + compressed_base=bench_window + bench_l_local, ) def run_torch(): _native_indexer_loss_indices_one_seq( - bench_l_local, bench_window, bench_ratio, bench_logical, bench_rank_by_seq + bench_l_local, bench_window, bench_window, bench_ratio, bench_logical, bench_rank_by_seq ) effective_bytes = _num_bytes(bench_logical, bench_rank_by_seq, *bench_ref) diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py index 69fbf88bda4..698422265bf 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py @@ -9,20 +9,12 @@ from megatron.core.transformer.experimental_attention_variant.csa_cp_utils import ( DSV4_CP_PARTITION_CONTIGUOUS, DSV4_CP_PARTITION_TWO_CHUNK, - all_gather_fixed_cp_tensor, - apply_thd_cp_compressed_rope_fused, apply_thd_cp_local_rope_fused, - build_cp_attention_indices_fused, build_cp_compressor_prep_compact_fused, - build_cp_indexer_loss_indices_fused, - build_cp_rank_major_compressed_metadata_fused, - build_global_compressed_cu_seqlens, compute_cp_indexer_topk_logical_fused, exchange_cp_boundary_hidden, local_kv_cp_chunk_ranges, local_q_cp_chunk_ranges, - pack_cp_kv_full_fused, - repack_rank_major_compressed_to_seq_major_fused, thd_cp_local_row_indices, ) from tests.unit_tests.test_utilities import Utils @@ -232,22 +224,6 @@ def test_thd_cp_left_boundary_exchange_forward_backward(): Utils.destroy_model_parallel() -def test_global_compressed_cu_seqlens_matches_reference(): - """Validate compressed-prefix metadata for ragged padded sequences. - - Expected: each sequence contributes ``padded_seq_len // ratio`` compressed - rows to the global seq-major prefix. A failure here means rank-major - compressed rows could be repacked with the wrong sequence offsets. - """ - _require_cuda() - cu_seqlens_cpu = torch.tensor([0, 7, 128, 381, 512], dtype=torch.int32) - ratio = 4 - - actual = build_global_compressed_cu_seqlens(cu_seqlens_cpu.cuda(), ratio) - - assert torch.equal(actual.cpu(), torch.tensor([0, 1, 31, 94, 126], dtype=torch.int32)) - - def test_apply_thd_cp_local_rope_fused_dispatches_equal_and_ragged_chunks(monkeypatch): """Validate the public local-RoPE wrapper dispatch contract. @@ -309,80 +285,6 @@ def fake_apply( apply_thd_cp_local_rope_fused(equal[:3], cos, sin, 1, 1, cu, ((2, 4), (10, 12))) -def test_apply_thd_cp_compressed_rope_fused_forwards_kernel_arguments(monkeypatch): - """Validate the public compressed-RoPE wrapper is a stable typed facade. - - Expected: the wrapper forwards tensor arguments, dimensions, ratio, and - inverse flag unchanged to the fused kernel entry point. - """ - captured = {} - - def fake_apply(x, cos, sin, comp_ids, ratio, nope_dim, pos_dim, inverse): - captured.update( - ratio=int(ratio), - nope_dim=int(nope_dim), - pos_dim=int(pos_dim), - inverse=bool(inverse), - comp_ids=comp_ids, - ) - return x + 1 - - monkeypatch.setattr( - csa_cp_utils.csa_cp_layout_kernels.ThdCompressedRope, "apply", staticmethod(fake_apply) - ) - x = torch.arange(6, dtype=torch.float32).reshape(3, 2) - comp_ids = torch.tensor([0, 2, -1], dtype=torch.int32) - out = apply_thd_cp_compressed_rope_fused( - x, torch.empty(16, 1), torch.empty(16, 1), comp_ids, 4, 1, 1, inverse=True - ) - - assert torch.equal(out, x + 1) - assert captured["ratio"] == 4 - assert captured["nope_dim"] == 1 - assert captured["pos_dim"] == 1 - assert captured["inverse"] is True - assert captured["comp_ids"] is comp_ids - - -def test_rank_major_compressed_metadata_wrapper_computes_capacities(monkeypatch): - """Validate public compressed-metadata capacity and partition dispatch. - - Expected: contiguous and two-chunk modes compute aligned capacities before - delegating to the kernel wrapper, so all ranks use fixed-size all-gather - rows independent of the number of valid compressed groups. - """ - calls = [] - - def fake_build(cu, cp_size, l_local, ratio, d_comp, c_cap, **kwargs): - calls.append( - ( - int(cp_size), - int(l_local), - int(ratio), - int(d_comp), - int(c_cap), - kwargs.get("c_cap_per_rank"), - kwargs.get("use_two_chunk", False), - ) - ) - rows = int(cp_size) * int(kwargs.get("c_cap_per_rank", c_cap)) - return ( - torch.full((rows,), -1, dtype=torch.int32), - torch.full((rows,), -1, dtype=torch.int32), - torch.zeros((rows,), dtype=torch.bool), - ) - - monkeypatch.setattr( - csa_cp_utils.csa_cp_layout_kernels, "build_compressed_row_metadata", fake_build - ) - cu = torch.tensor([0, 32, 64], dtype=torch.int32) - - build_cp_rank_major_compressed_metadata_fused(cu, ((16, 32),), 4, 4, 8) - build_cp_rank_major_compressed_metadata_fused(cu, ((0, 4), (28, 32)), 4, 4, 8) - - assert calls == [(4, 16, 4, 8, 8, None, False), (4, 4, 4, 8, 3, 8, True)] - - def test_compressor_prep_compact_wrapper_merges_two_chunk_outputs(monkeypatch): """Validate public compressor-prep wrapper splitting, merge, and padding. @@ -396,137 +298,25 @@ def fake_apply(hidden, boundary, cu, global_start, rows, ratio, d_comp, d_window calls.append((int(global_start), int(rows), int(c_cap), tuple(boundary.shape))) marker = len(calls) hidden_compact = torch.full((int(c_cap) * int(ratio), 1), marker, dtype=hidden.dtype) - cu_compact = torch.tensor([0, int(c_cap) * int(ratio)], dtype=cu.dtype) - seq_ids = torch.full((int(c_cap),), marker, dtype=torch.int32) comp_ids = torch.arange(int(c_cap), dtype=torch.int32) + marker * 10 - valid = torch.ones((int(c_cap),), dtype=torch.bool) - return hidden_compact, cu_compact, seq_ids, comp_ids, valid + return hidden_compact, comp_ids monkeypatch.setattr( csa_cp_utils.csa_cp_layout_kernels.CompressorInputCompact, "apply", staticmethod(fake_apply) ) - hidden = torch.arange(6, dtype=torch.float32).reshape(6, 1) + hidden = torch.arange(16, dtype=torch.float32).reshape(16, 1) boundary = torch.arange(4, dtype=torch.float32).reshape(4, 1) - cu = torch.tensor([0, 6], dtype=torch.int32) + cu = torch.tensor([0, 32], dtype=torch.int32) + cu_comp = torch.tensor([0, 8], dtype=torch.int32) - hidden_compact, cu_compact, comp_ids = build_cp_compressor_prep_compact_fused( - hidden, boundary, cu, ((0, 3), (10, 13)), ratio=2, d_comp=2, d_window=2 + hidden_compact, comp_ids, rank_rows = build_cp_compressor_prep_compact_fused( + hidden, boundary, cu, cu_comp, ((0, 8), (24, 32)), cp_size=2, ratio=4, d_comp=8, d_window=2 ) - assert calls == [(0, 3, 2, (2, 1)), (10, 3, 2, (2, 1))] + assert calls == [(0, 8, 4, (2, 1)), (24, 8, 4, (2, 1))] assert hidden_compact.shape == (32, 1) - assert torch.equal(cu_compact, torch.tensor([0, 8], dtype=torch.int32)) - assert torch.equal(comp_ids[:4], torch.tensor([10, 11, 20, 21], dtype=torch.int32)) - assert torch.equal(comp_ids[4:], torch.full((12,), -1, dtype=torch.int32)) - - -def test_all_gather_fixed_cp_tensor_single_rank_preserves_autograd(): - """Validate the public fixed CP all-gather no-CP fallback. - - Expected: with cp_group=None, the helper returns the local tensor and - backward is the identity, matching the disabled-CP execution path. - """ - x = torch.arange(6, dtype=torch.float32).reshape(3, 2).requires_grad_(True) - - out = all_gather_fixed_cp_tensor(x, None) - - assert torch.equal(out, x) - out.square().sum().backward() - assert torch.equal(x.grad, 2 * x.detach()) - - -def test_pack_cp_kv_full_fused_wrapper_selects_contiguous_and_two_chunk_layouts(monkeypatch): - """Validate public KV-pack wrapper argument lowering for both layouts. - - Expected: contiguous CP passes the physical local range through directly, - while two-chunk CP switches to shared-window layout and returns the shared - compressed base used by final index lowering. - """ - calls = [] - - def fake_apply(*args): - calls.append(args) - kv_local = args[0] - capacity = int(args[13]) - return kv_local.new_full((capacity,) + tuple(kv_local.shape[1:]), len(calls)) - - monkeypatch.setattr( - csa_cp_utils.csa_cp_layout_kernels.ThdFullKvPack, "apply", staticmethod(fake_apply) - ) - kv_local = torch.ones(4, 2) - boundary = torch.ones(2, 2) - compressed = torch.ones(6, 2) - ids = torch.zeros(6, dtype=torch.int32) - valid = torch.ones(6, dtype=torch.bool) - cu = torch.tensor([0, 8], dtype=torch.int32) - rank_by_seq = torch.arange(2, dtype=torch.int32) - cu_comp = torch.tensor([0, 2], dtype=torch.int32) - - out, shared = pack_cp_kv_full_fused( - kv_local, - boundary, - compressed, - ids, - ids, - valid, - cu, - ((4, 8),), - d_window=2, - ratio=4, - rank_major_by_seq_major=rank_by_seq, - cu_seqlens_compressed=cu_comp, - ) - assert shared is None - assert out.shape == (12, 2) - assert calls[-1][9:19] == (4, 4, 2, 4, 12, 0, 0, 0, 0, 0) - - out, shared = pack_cp_kv_full_fused( - kv_local, - torch.ones(4, 2), - compressed, - ids, - ids, - valid, - cu, - ((0, 2), (14, 16)), - d_window=2, - ratio=4, - ) - assert shared == 8 - assert out.shape == (14, 2) - assert calls[-1][9:19] == (0, 0, 2, 0, 14, 0, 14, 2, 2, 4) - - -def test_repack_rank_major_compressed_to_seq_major_fused_forwards_arguments(monkeypatch): - """Validate the public compressed-KV repack wrapper. - - Expected: the wrapper forwards rank-major tensors and requested seq-major - capacity unchanged to the fused repack kernel. - """ - captured = {} - - def fake_repack(rank_major, seq_ids, comp_ids, valid, cu_comp, rows): - captured.update(rows=int(rows), rank_major=rank_major, seq_ids=seq_ids) - return rank_major[:rows].clone(), torch.arange(rows, dtype=torch.int32) - - monkeypatch.setattr( - csa_cp_utils.csa_cp_layout_kernels, "repack_compressed_kv_to_seq_major", fake_repack - ) - rank_major = torch.arange(12, dtype=torch.float32).reshape(6, 2) - seq_ids = torch.zeros(6, dtype=torch.int32) - comp_ids = torch.arange(6, dtype=torch.int32) - valid = torch.ones(6, dtype=torch.bool) - cu_comp = torch.tensor([0, 3], dtype=torch.int32) - - out, rank_by_seq = repack_rank_major_compressed_to_seq_major_fused( - rank_major, seq_ids, comp_ids, valid, cu_comp, 3 - ) - - assert captured["rows"] == 3 - assert captured["rank_major"] is rank_major - assert captured["seq_ids"] is seq_ids - assert torch.equal(out, rank_major[:3]) - assert torch.equal(rank_by_seq, torch.arange(3, dtype=torch.int32)) + assert torch.equal(comp_ids, torch.tensor([10, 11, 12, 13, 20, 21, 22, 23])) + assert torch.equal(rank_rows, torch.tensor([0, 1, 10, 11, 14, 15, 6, 7], dtype=torch.int32)) def test_compute_cp_indexer_topk_logical_fused_splits_chunks_and_passes_visible_lengths( @@ -538,19 +328,8 @@ def test_compute_cp_indexer_topk_logical_fused_splits_chunks_and_passes_visible_ call, visible K lengths are passed through, and chunk outputs are concatenated in local row order. """ - metadata_calls = [] topk_calls = [] - def fake_metadata(k_seq, cu_q, cu_comp, global_start, length, ratio): - metadata_calls.append((int(global_start), int(length), int(ratio))) - marker = len(metadata_calls) - return ( - k_seq + marker, - torch.tensor([0, int(length)], dtype=torch.int32), - torch.tensor([0, 1], dtype=torch.int32), - torch.full((int(length),), marker, dtype=torch.int32), - ) - def fake_indexer_topk( q, k, @@ -573,19 +352,19 @@ def fake_indexer_topk( float(indexer_softmax_scale), int(max_seqlen_q), int(max_seqlen_kv), + k.clone(), + cu_seqlens_q.clone(), + cu_seqlens_kv.clone(), visible_k_lengths.clone(), ) ) return torch.full((q.shape[0], int(topk)), len(topk_calls), dtype=torch.int32), None - monkeypatch.setattr( - csa_cp_utils.csa_cp_layout_kernels, "build_indexer_topk_metadata", fake_metadata - ) monkeypatch.setattr(csa_cp_utils, "indexer_topk", fake_indexer_topk) q = torch.randn(5, 2) weights = torch.randn(5, 1) - k_seq = torch.randn(4, 2) + k_seq = torch.arange(8, dtype=torch.float32).reshape(4, 2) cu_q = torch.tensor([0, 10], dtype=torch.int32) cu_comp = torch.tensor([0, 3], dtype=torch.int32) @@ -603,90 +382,20 @@ def fake_indexer_topk( max_seqlen_kv=3, ) - assert metadata_calls == [(2, 2, 4), (8, 3, 4)] assert torch.equal(out, torch.tensor([[1, 1], [1, 1], [2, 2], [2, 2], [2, 2]])) + assert torch.equal(topk_calls[0][-4][0], k_seq[0]) + assert torch.count_nonzero(topk_calls[0][-4][1:]) == 0 + assert torch.equal(topk_calls[0][-3], torch.tensor([0, 2, 2], dtype=torch.int32)) + assert torch.equal(topk_calls[0][-2], torch.tensor([0, 1, 1], dtype=torch.int32)) assert torch.equal(topk_calls[0][-1], torch.ones(2, dtype=torch.int32)) - assert torch.equal(topk_calls[1][-1], torch.full((3,), 2, dtype=torch.int32)) + assert torch.equal(topk_calls[1][-4][:2], k_seq[:2]) + assert torch.count_nonzero(topk_calls[1][-4][2:]) == 0 + assert torch.equal(topk_calls[1][-3], torch.tensor([0, 2, 3], dtype=torch.int32)) + assert torch.equal(topk_calls[1][-2], torch.tensor([0, 2, 2], dtype=torch.int32)) + assert torch.equal(topk_calls[1][-1], torch.tensor([2, 2, 0], dtype=torch.int32)) assert ( compute_cp_indexer_topk_logical_fused( q, weights, k_seq[:0], cu_q, cu_comp, ((2, 4),), 4, 2, 1.0, 10, 3 ) is None ) - - -def test_final_index_wrappers_lower_two_chunk_arguments(monkeypatch): - """Validate public final attention and indexer-loss index wrappers. - - Expected: two-chunk wrappers require shared compressed metadata and pass - chunk starts, chunk length, shared base, and local row count to the fused - kernel entry points. - """ - captured = {} - - def fake_attention(*args, **kwargs): - captured["attention"] = (args, kwargs) - l_local = int(args[2]) - width = int(args[6]) - return torch.zeros(l_local, width, dtype=torch.int32), torch.zeros( - l_local, dtype=torch.int32 - ) - - def fake_loss(*args, **kwargs): - captured["loss"] = (args, kwargs) - l_local = int(args[3]) - width = args[7].shape[1] + int(args[5]) - return torch.zeros(l_local, width, dtype=torch.int32), torch.zeros_like(args[7]) - - monkeypatch.setattr( - csa_cp_utils.csa_cp_layout_kernels, "build_attention_indices", fake_attention - ) - monkeypatch.setattr(csa_cp_utils.csa_cp_layout_kernels, "build_indexer_loss_indices", fake_loss) - - cu = torch.tensor([0, 16], dtype=torch.int32) - cu_comp = torch.tensor([0, 4], dtype=torch.int32) - rank_by_seq = torch.arange(4, dtype=torch.int32) - logical = torch.zeros(4, 2, dtype=torch.int32) - chunks = ((0, 2), (14, 16)) - - with pytest.raises(RuntimeError, match="requires a shared compressed base"): - build_cp_attention_indices_fused(cu, chunks, 2, 2, 4, logical, 2) - - build_cp_attention_indices_fused( - cu, - chunks, - d_window=2, - window_size=2, - ratio=4, - indexer_topk_compressed_logical_ids=logical, - max_n_compressed=2, - rank_major_by_seq_major=rank_by_seq, - cu_seqlens_compressed=cu_comp, - shared_compressed_base=8, - ) - args, kwargs = captured["attention"] - assert args[0] is cu - assert args[1:7] == (0, 4, 2, 2, 4, 2) - assert kwargs["chunk_starts"] == (0, 14) - assert kwargs["chunk_len"] == 2 - assert kwargs["shared_compressed_base"] == 8 - - build_cp_indexer_loss_indices_fused( - cu, - cu_comp, - chunks, - d_window=2, - window_size=2, - ratio=4, - indexer_topk_compressed_logical_ids=logical, - rank_major_by_seq_major=rank_by_seq, - shared_compressed_base=8, - ) - args, kwargs = captured["loss"] - assert args[0] is cu - assert args[1] is cu_comp - assert args[2:7] == (0, 4, 2, 2, 4) - assert args[7] is logical - assert kwargs["chunk_starts"] == (0, 14) - assert kwargs["chunk_len"] == 2 - assert kwargs["shared_compressed_base"] == 8 diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py index 26cfc05616f..2a7468ace7e 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py @@ -650,6 +650,50 @@ def test_thd_cp_matches_full_reference_forward_backward(self, layer_number, part del cp_attn, ref_attn, full_hidden, local_hidden, ref_hidden, local_out, ref_out, grad _clear_cuda_test_state() + def test_thd_cp_ratio4_eval_matches_full_reference(self): + """Ratio-4 two-chunk inference lowers logical indexer top-k rows correctly.""" + partition_mode = DSV4_CP_PARTITION_TWO_CHUNK + packed, padded_tokens, local_idx = _make_ragged_cp_case( + partition_mode, self.cp_size, self.cp_rank + ) + + torch.manual_seed(_SEED + 1202) + model_parallel_cuda_manual_seed(_SEED + 1202) + config_cp = _make_dsv4_cp_config( + context_parallel_size=self.cp_size, + apply_dsa_kernel_fusion=self.fused_kernels_available, + apply_rope_fusion=True, + csa_cp_partition_mode=partition_mode, + ) + config_ref = _make_dsv4_cp_config( + context_parallel_size=1, + apply_dsa_kernel_fusion=self.fused_kernels_available, + apply_rope_fusion=False, + ) + cp_attn = _build_attention(config_cp, layer_number=2, pg_collection=self.pg).cuda().eval() + ref_attn = ( + _build_attention(config_ref, layer_number=2, pg_collection=self.ref_pg).cuda().eval() + ) + _copy_module_parameters(cp_attn, ref_attn) + + full_hidden = torch.randn( + padded_tokens, 1, config_cp.hidden_size, dtype=torch.bfloat16, device='cuda' + ) + local_hidden = full_hidden.index_select(0, local_idx) + with torch.no_grad(): + local_out, _ = cp_attn( + hidden_states=local_hidden, attention_mask=None, packed_seq_params=packed + ) + ref_out, _ = ref_attn( + hidden_states=full_hidden, attention_mask=None, packed_seq_params=packed + ) + _assert_cp_tensor_match( + local_out, ref_out.index_select(0, local_idx), f"layer=2:eval:{partition_mode}:output" + ) + + del cp_attn, ref_attn, full_hidden, local_hidden, local_out, ref_out + _clear_cuda_test_state() + @pytest.mark.parametrize( "layer_number", [1, 2, 3], From 4afb6fe79eb2191afb4930ebc274582a89a22bba Mon Sep 17 00:00:00 2001 From: kunlunl Date: Sat, 27 Jun 2026 09:50:15 -0700 Subject: [PATCH 06/23] Simplify DSv4 THD context parallel flow Signed-off-by: kunlunl --- megatron/core/datasets/data_schedule.py | 60 +- megatron/core/datasets/data_schedule_utils.py | 41 +- .../core/fusions/fused_mla_yarn_rope_apply.py | 32 +- .../experimental_attention_variant/csa.py | 239 +--- .../csa_cp_layout_kernels.py | 1160 +++-------------- .../csa_cp_utils.py | 657 +++------- .../deepseek_v4_hybrid_attention.py | 50 +- .../dsa_kernels.py | 145 +-- .../core/transformer/transformer_config.py | 5 - tests/unit_tests/test_sequence_packing.py | 131 +- .../test_csa_cp_layout_kernels.py | 430 +----- .../test_csa_cp_utils.py | 344 ++--- .../test_dsa_kernels.py | 49 +- .../test_dsv4_hybrid_attention.py | 6 +- .../test_dsv4_hybrid_attention_cp.py | 140 +- 15 files changed, 853 insertions(+), 2636 deletions(-) diff --git a/megatron/core/datasets/data_schedule.py b/megatron/core/datasets/data_schedule.py index 6471d7465ff..00ca7e59cc2 100644 --- a/megatron/core/datasets/data_schedule.py +++ b/megatron/core/datasets/data_schedule.py @@ -613,15 +613,6 @@ def get_batch_on_this_rank_for_sequence_packing( is_first_or_last_stage = is_first_stage or is_last_stage dev = torch.cuda.current_device() - static_cp_target_len = None - if ( - config is not None - and getattr(config, 'experimental_attention_variant', None) == "dsv4_hybrid" - and cp_group.size() > 1 - and getattr(config, 'pad_packed_seq_alignment', None) == "max" - and getattr(config, 'max_seqlen_per_dp_cp_rank', None) is not None - ): - static_cp_target_len = int(config.max_seqlen_per_dp_cp_rank) # data_iterator should return a batch including the following keys. batch_keys = ['cu_seqlens', 'cu_seqlens_padded', 'max_seqlen'] @@ -653,6 +644,32 @@ def get_batch_on_this_rank_for_sequence_packing( group_size=local_cp_size_val ) + use_dsv4_cp_slice = ( + config is not None + and getattr(config, 'experimental_attention_variant', None) == "dsv4_hybrid" + ) + dsv4_cp_local_target_len = None + pad_alignment = ( + getattr(config, 'pad_packed_seq_alignment', None) if config is not None else None + ) + alignment = target_len = max_num_seqs = None + if pad_alignment is not None: + alignment, target_len, max_num_seqs = get_thd_padding_kwargs( + pad_alignment, + getattr(config, 'max_seqlen_per_dp_cp_rank', None), + getattr(config, 'thd_max_packed_sequences', None), + getattr(config, 'cuda_graph_impl', 'none') != 'none', + ) + if is_tp_rank_0 and use_dsv4_cp_slice and pad_alignment is not None: + if target_len is not None: + dsv4_cp_local_target_len = target_len + else: + # Fix the local width before slicing so later padding cannot shift rank origins. + assert alignment is not None + total_rows = int(batch['cu_seqlens_padded'][-1].item()) + local_rows = (total_rows + cp_group.size() - 1) // cp_group.size() + dsv4_cp_local_target_len = ((local_rows + alignment - 1) // alignment) * alignment + # Build padding_mask before CP slicing while tensors still have the full # packed length represented by cu_seqlens_padded[-1]. if is_tp_rank_0: @@ -667,20 +684,16 @@ def get_batch_on_this_rank_for_sequence_packing( cp_slice_keys = ['padding_mask'] if is_first_or_last_stage or mtp_on_this_rank: cp_slice_keys.extend(['tokens', 'position_ids', 'labels', 'loss_mask']) - csa_cp_partition_mode = None - if ( - config is not None - and getattr(config, 'experimental_attention_variant', None) == "dsv4_hybrid" - ): - csa_cp_partition_mode = config.csa_cp_partition_mode partition_total_tokens = ( - static_cp_target_len * cp_group.size() if static_cp_target_len is not None else None + dsv4_cp_local_target_len * cp_group.size() + if dsv4_cp_local_target_len is not None + else None ) get_cp_slice_for_thd( batch, cp_group, keys=cp_slice_keys, - csa_cp_partition_mode=csa_cp_partition_mode, + use_dsv4_cp_slice=use_dsv4_cp_slice, partition_total_tokens=partition_total_tokens, ) @@ -696,8 +709,8 @@ def get_batch_on_this_rank_for_sequence_packing( # Broadcast total_tokens because padding_mask is prepared on every PP stage. # Tokens/labels/loss_mask/position_ids use the same length on stages that own them. if is_tp_rank_0: - if static_cp_target_len is not None: - total_tokens = torch.tensor([static_cp_target_len], dtype=torch.int32, device=dev) + if dsv4_cp_local_target_len is not None: + total_tokens = torch.tensor([dsv4_cp_local_target_len], dtype=torch.int32, device=dev) else: # Under VPP, the last PP stage has labels but no tokens, so derive # total_tokens from cu_seqlens_padded, which is present on every @@ -824,16 +837,7 @@ def get_batch_on_this_rank_for_sequence_packing( # Pad the already-packed THD tensors at the end when requested. CUDA Graph # additionally pads cu_seqlens tensors to thd_max_packed_sequences + 1 entries. - pad_alignment = ( - getattr(config, 'pad_packed_seq_alignment', None) if config is not None else None - ) if pad_alignment is not None and packed_seq_params is not None: - alignment, target_len, max_num_seqs = get_thd_padding_kwargs( - pad_alignment, - getattr(config, 'max_seqlen_per_dp_cp_rank', None), - getattr(config, 'thd_max_packed_sequences', None), - getattr(config, 'cuda_graph_impl', 'none') != 'none', - ) tokens, labels, loss_mask, position_ids, packed_seq_params, padding_mask = ( pad_sequence_for_thd( tokens, diff --git a/megatron/core/datasets/data_schedule_utils.py b/megatron/core/datasets/data_schedule_utils.py index 0f4125baa08..88e4b78394d 100644 --- a/megatron/core/datasets/data_schedule_utils.py +++ b/megatron/core/datasets/data_schedule_utils.py @@ -15,26 +15,24 @@ def get_cp_slice_for_thd( batch, cp_group, keys: Optional[Sequence[str]] = None, - csa_cp_partition_mode: Optional[str] = None, + use_dsv4_cp_slice: bool = False, partition_total_tokens: Optional[int] = None, ): """Partition sequence data for context parallelism in THD format. - Uses TE's THD partitioned indices by default. When ``csa_cp_partition_mode`` - is provided, uses the CSA partition helper so data row order matches the - DSv4 attention layer's CP row order. Only keys present in the batch are sliced. + Uses TE's THD partitioned indices by default. DSv4 instead assigns each rank one + ordered range of the global rows. Only keys present in the batch are sliced. Args: batch: Dict with packed sequence data. cp_group: Context parallel process group. keys: Sequence data keys to slice. Defaults to the original THD data tensors. - csa_cp_partition_mode: Optional CSA CP partition mode for DSv4 hybrid attention. - partition_total_tokens: Optional padded total used only for choosing CP row indices. - When set, tensors selected by ``keys`` are tail-padded to this length before - slicing. Existing cu_seqlens metadata is left unchanged. + use_dsv4_cp_slice: Whether to use the DSv4 row assignment. + partition_total_tokens: Optional total used to tail-pad tensors selected by ``keys`` + before slicing. Existing cu_seqlens metadata is left unchanged. """ cp_size = cp_group.size() - if cp_size <= 1: + if cp_size <= 1 and partition_total_tokens is None: return cp_rank = cp_group.rank() # Partition with padded cumulative lengths so CP slices match the THD @@ -63,16 +61,23 @@ def get_cp_slice_for_thd( if pad_len > 0: pad_value = True if key == 'padding_mask' else 0 batch[key] = torch.cat([batch[key], batch[key].new_full((pad_len,), pad_value)]) - if csa_cp_partition_mode is None: - index = get_thd_partitioned_indices(cu_seqlens, total_tokens, cp_size, cp_rank) - else: - from megatron.core.transformer.experimental_attention_variant.csa_cp_utils import ( - thd_cp_local_row_indices, - ) + if cp_size <= 1: + return + if use_dsv4_cp_slice: + if total_tokens % cp_size != 0: + raise RuntimeError( + f"DSv4 CP slicing requires total_tokens={total_tokens} to be divisible by " + f"cp_size={cp_size}." + ) + local_rows = total_tokens // cp_size + # Rank r takes [r * local_rows, (r + 1) * local_rows). + row_slice = slice(cp_rank * local_rows, (cp_rank + 1) * local_rows) + for key in keys: + if key in batch and batch[key] is not None: + batch[key] = batch[key][row_slice] + return - index = thd_cp_local_row_indices( - csa_cp_partition_mode, total_tokens, cp_size, cp_rank, cu_seqlens.device - ) + index = get_thd_partitioned_indices(cu_seqlens, total_tokens, cp_size, cp_rank) for key in keys: if key in batch and batch[key] is not None: batch[key] = batch[key].index_select(0, index) diff --git a/megatron/core/fusions/fused_mla_yarn_rope_apply.py b/megatron/core/fusions/fused_mla_yarn_rope_apply.py index a679162020c..b9e13fcf7cb 100644 --- a/megatron/core/fusions/fused_mla_yarn_rope_apply.py +++ b/megatron/core/fusions/fused_mla_yarn_rope_apply.py @@ -86,6 +86,7 @@ def _mla_rope_fwd_inplace_kernel( batch_size, seq_num, cu_seqlens_q, + position_ids, stride_x_seq, stride_x_nheads, stride_cos_seq, @@ -112,7 +113,9 @@ def _mla_rope_fwd_inplace_kernel( pid_m = tl.program_id(axis=0) pid_head = tl.program_id(axis=1) - if cu_seqlens_q is None: + 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) @@ -181,6 +184,7 @@ def _mla_rope_bwd_inplace_kernel( batch_size, seq_num, cu_seqlens_q, + position_ids, stride_x_seq, stride_x_nheads, stride_cos_seq, @@ -205,7 +209,9 @@ def _mla_rope_bwd_inplace_kernel( pid_m = tl.program_id(axis=0) pid_head = tl.program_id(axis=1) - if cu_seqlens_q is None: + 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) @@ -270,6 +276,7 @@ def forward( rotary_interleaved=False, inverse=False, remove_interleaving=False, + position_ids=None, ): """ Forward function for _FusedMLARoPEInplace. @@ -288,6 +295,7 @@ def forward( seq_num = None if cu_seqlens_q is None: # sbhd + assert position_ids is None max_seqlen, batch_size, nheads, headdim = q.shape q = q.view(-1, nheads, headdim) total_seqlen = q.shape[0] @@ -295,6 +303,8 @@ def forward( # thd total_seqlen, nheads, headdim = q.shape seq_num = len(cu_seqlens_q) - 1 + if position_ids is not None: + assert position_ids.shape == (total_seqlen,) assert q.stride(-1) == 1 assert cos.stride(-1) == 1 assert sin.stride(-1) == 1 @@ -312,6 +322,7 @@ def forward( batch_size, seq_num, cu_seqlens_q, + position_ids, q.stride(0), q.stride(1), cos.stride(0), @@ -321,7 +332,8 @@ def forward( INVERSE=inverse, REMOVE_INTERLEAVING=remove_interleaving, ) - ctx.save_for_backward(cos, sin) + ctx.save_for_backward(cos, sin, *(() if position_ids is None else (position_ids,))) + ctx.has_position_ids = position_ids is not None ctx.nope_dim = nope_dim ctx.emb_dim = emb_dim ctx.cu_seqlens_q = cu_seqlens_q @@ -343,7 +355,11 @@ def backward(ctx, grad): grad: [seq_len, batch_size, head_num, nope_dim + emb_dim] or [total_seq_len, head_num, nope_dim + emb_dim] """ - cos, sin = ctx.saved_tensors + if ctx.has_position_ids: + cos, sin, position_ids = ctx.saved_tensors + else: + cos, sin = ctx.saved_tensors + position_ids = None max_seqlen = None batch_size = None seq_num = None @@ -353,6 +369,7 @@ def backward(ctx, grad): total_seqlen = grad.shape[0] else: seq_num = len(ctx.cu_seqlens_q) - 1 + grad = grad.contiguous() total_seqlen, nheads, headdim = grad.shape assert grad.stride(-1) == 1 @@ -367,6 +384,7 @@ def backward(ctx, grad): batch_size, seq_num, ctx.cu_seqlens_q, + position_ids, grad.stride(0), grad.stride(1), cos.stride(0), @@ -378,7 +396,7 @@ def backward(ctx, grad): ) if ctx.cu_seqlens_q is None: grad = grad.view(max_seqlen, batch_size, nheads, headdim) - return grad, None, None, None, None, None, None, None, None, None, None + return grad, None, None, None, None, None, None, None, None, None, None, None def fused_mla_rope_inplace( @@ -393,6 +411,7 @@ def fused_mla_rope_inplace( rotary_interleaved: bool = False, inverse: bool = False, remove_interleaving: bool = False, + position_ids: Optional[torch.Tensor] = None, ): """ Fused RoPE applied inplace to the trailing emb_dim elements of a tensor, @@ -414,6 +433,8 @@ def fused_mla_rope_inplace( rotary_interleaved: whether to apply RoPE interleaved, only supports False for now inverse: if True, apply the inverse rotation remove_interleaving: if True, output RoPE dims in non-interleaved layout + position_ids: optional THD row positions. When supplied, these positions + replace the built-in CP row-to-position mapping. Returns: t: inplace modified input tensor @@ -430,6 +451,7 @@ def fused_mla_rope_inplace( rotary_interleaved, inverse, remove_interleaving, + position_ids, ) diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index a00bdc0e88a..3482bcb3896 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -608,8 +608,7 @@ def _unfused_indexer_sparse_attn_from_topk( softmax_scale: float, indexer_softmax_scale: float, loss_coeff: float, - calculate_per_token_loss: bool, - global_query_rows: int, + loss_divisor: float, q_padding_mask: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """PyTorch sparse attention plus caller-supplied top-k indexer loss for THD CP. @@ -622,21 +621,8 @@ def _unfused_indexer_sparse_attn_from_topk( total_q, np_, hn = query.shape indexer_topk = indexer_topk_indices.shape[-1] - if indexer_topk == 0: - return output, query.new_zeros((), dtype=torch.float32) if q_padding_mask is not None: - if q_padding_mask.shape != (total_q,): - raise RuntimeError( - "q_padding_mask must have shape " - f"({total_q},), got {tuple(q_padding_mask.shape)}." - ) - if q_padding_mask.dtype != torch.bool: - raise RuntimeError(f"q_padding_mask must be bool, got {q_padding_mask.dtype}.") - if q_padding_mask.device != query.device: - raise RuntimeError( - f"q_padding_mask must be on {query.device}, got {q_padding_mask.device}." - ) indexer_topk_indices = indexer_topk_indices.masked_fill(q_padding_mask.unsqueeze(-1), -1) valid = indexer_topk_indices >= 0 @@ -679,12 +665,7 @@ def _unfused_indexer_sparse_attn_from_topk( kl_per_row = torch.where(row_valid.squeeze(-1), kl_per_row, torch.zeros_like(kl_per_row)) raw_local_loss = kl_per_row.sum() - if calculate_per_token_loss: - indexer_loss = raw_local_loss * loss_coeff - else: - if global_query_rows <= 0: - raise RuntimeError(f"global_query_rows must be positive, got {global_query_rows}.") - indexer_loss = raw_local_loss * (float(loss_coeff) / float(global_query_rows)) + indexer_loss = raw_local_loss * float(loss_coeff) / float(loss_divisor) return output, indexer_loss @@ -892,10 +873,9 @@ def _forward_thd( softmax + reduce is batched across ALL compressed entries from all segments via a ``(total_compressed, ratio)`` gather index. - When ``fixed_total_comp`` is supplied the output is padded to that - static capacity so tensor shapes are host-known and CUDA-graph - capturable. Rows beyond the true ``cu_seqlens_compressed[-1]`` - gather from position 0 and are left as tail padding. + Non-CP callers may supply ``fixed_total_comp`` to pad the output to a + host-known CUDA-graph capacity. Rows beyond the true + ``cu_seqlens_compressed[-1]`` gather from position 0 and remain tail padding. Args: x: ``(total, 1, hidden_size)`` packed bf16. @@ -905,14 +885,14 @@ def _forward_thd( sync when building the rotary table for ``ratio > 1``). compressed_group_ids: CP compressor-prep group ids. When supplied, ``x`` is already packed into ``ratio``-sized groups. - fixed_total_comp: when set, overrides ``cu_seqlens_compressed[-1]`` - as the output row count. Must be >= the true compressed count. + fixed_total_comp: optional non-CP output row capacity. Must be at + least the true compressed count. Returns: ``(compressed_thd, cu_seqlens_compressed)`` where ``compressed_thd`` is ``(total_compressed, 1, head_dim)`` bf16 - (or ``None`` when no sequence has ``seg_len >= ratio`` and - ``fixed_total_comp`` is not set) and + (or ``None`` when no sequence has ``seg_len >= ratio`` and no + fixed capacity is set) and ``cu_seqlens_compressed`` is ``(B+1,)`` int32 with ``cu_seqlens_compressed[b+1] - cu_seqlens_compressed[b] = seqlen_b // ratio``. Pre-grouped CP inputs return ``None`` for this unused second value. @@ -923,10 +903,8 @@ def _forward_thd( pre_grouped = compressed_group_ids is not None if pre_grouped: - if fixed_total_comp is None: - raise RuntimeError("Pre-grouped THD compression requires a fixed row count.") cu_seqlens_compressed = None - total_comp = int(fixed_total_comp) + total_comp = compressed_group_ids.shape[0] else: seg_compressed_lens = (cu_seqlens[1:] - cu_seqlens[:-1]) // ratio cu_seqlens_compressed = torch.cat( @@ -1428,31 +1406,6 @@ def __init__( # Private helpers – each owns one logical slice of the forward pass. # ------------------------------------------------------------------ - def _build_kv_full( - self, kv: torch.Tensor, x: torch.Tensor - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], int]: - """Concatenate original KV with compressed KV (if applicable). - - Returns: - kv_full: [n_kv, b, v_head_dim] original + compressed KV. - compressed_kv: [n_compressed, b, v_head_dim] or None. - n_compressed: number of compressed positions (0 when unused). - """ - if self.compressor is not None and self.compress_ratio > 1: - compressed_kv = self.compressor(x) - if compressed_kv is not None: - kv_full = torch.cat([kv, compressed_kv], dim=0) - n_compressed = compressed_kv.size(0) - else: - kv_full = kv - compressed_kv = None - n_compressed = 0 - else: - kv_full = kv - compressed_kv = None - n_compressed = 0 - return kv_full, compressed_kv, n_compressed - def _forward_unfused_csa( self, query: torch.Tensor, @@ -1582,7 +1535,6 @@ def _forward_fused_indexer_inference( x: torch.Tensor, qr: torch.Tensor, kv_full: torch.Tensor, - n_compressed: int, offset: int, window_idxs: torch.Tensor, packed_seq_params: Optional[PackedSeqParams], @@ -1628,7 +1580,6 @@ def _forward_fused_indexer_training( x: torch.Tensor, qr: torch.Tensor, kv_full: torch.Tensor, - n_compressed: int, offset: int, window_idxs: torch.Tensor, packed_seq_params: Optional[PackedSeqParams], @@ -1721,7 +1672,11 @@ def forward( sq, b, np, hn = query.size() kv = key.squeeze(-2) # [sq, b, 1, v_head_dim] -> [sq, b, v_head_dim] - kv_full, compressed_kv, n_compressed = self._build_kv_full(kv, x) + compressed_kv = ( + self.compressor(x) if self.compressor is not None and self.compress_ratio > 1 else None + ) + n_compressed = 0 if compressed_kv is None else compressed_kv.shape[0] + kv_full = kv if compressed_kv is None else torch.cat((kv, compressed_kv), dim=0) offset = sq # compressed indices start after original positions window_idxs = get_window_topk_idxs(self.window_size, b, sq, query.device) @@ -1745,11 +1700,11 @@ def forward( ) elif has_indexer_compressed and self.training and torch.is_grad_enabled(): output, indexer_loss = self._forward_fused_indexer_training( - query, x, qr, kv_full, n_compressed, offset, window_idxs, packed_seq_params + query, x, qr, kv_full, offset, window_idxs, packed_seq_params ) elif has_indexer_compressed: output = self._forward_fused_indexer_inference( - query, x, qr, kv_full, n_compressed, offset, window_idxs, packed_seq_params + query, x, qr, kv_full, offset, window_idxs, packed_seq_params ) else: output = self._forward_fused_no_indexer( @@ -1781,7 +1736,6 @@ def _forward_unfused_csa_thd( cu_seqlens_kv_full: torch.Tensor, cu_seqlens_compressed: torch.Tensor, window_idxs: torch.Tensor, - max_seqlen_q: int, max_seqlen_compressed_idx: int, packed_seq_params: PackedSeqParams, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: @@ -1999,7 +1953,6 @@ def _forward_fused_indexer_inference_thd( window_idxs: torch.Tensor, max_seqlen_q: int, max_seqlen_compressed_idx: int, - max_seqlen_kv: int, ) -> torch.Tensor: """Path C (THD): separate indexer forward (no loss) + fused sparse attn (compact). @@ -2068,8 +2021,6 @@ def _forward_fused_indexer_training_thd( x: torch.Tensor, qr: torch.Tensor, packed_seq_params: PackedSeqParams, - total_q: int, - np_: int, cu_seqlens_q: torch.Tensor, cu_seqlens_kv: torch.Tensor, cu_seqlens_kv_full: torch.Tensor, @@ -2174,30 +2125,15 @@ def _forward_thd_cp( l_local = query.shape[0] if l_local != key.shape[0]: raise RuntimeError("DSv4 THD CP path currently supports self-attention only.") - cu_seqlens_q = ( + cu_seqlens = ( packed_seq_params.cu_seqlens_q_padded if packed_seq_params.cu_seqlens_q_padded is not None else packed_seq_params.cu_seqlens_q ) - cu_seqlens_kv = ( - packed_seq_params.cu_seqlens_kv_padded - if packed_seq_params.cu_seqlens_kv_padded is not None - else packed_seq_params.cu_seqlens_kv - ) max_seqlen_q = int(packed_seq_params.max_seqlen_q) # ---- Step 2: local CP rows, local KV, and boundary tensors ------------ - chunk_ranges = cp_utils.local_q_cp_chunk_ranges( - self.config.csa_cp_partition_mode, l_local, cp_size, cp_rank - ) - if len(chunk_ranges) == 2: - chunk_starts = tuple(start for start, _ in chunk_ranges) - chunk_len = l_local // 2 - index_global_start = 0 - else: - chunk_starts = None - chunk_len = l_local - index_global_start = chunk_ranges[0][0] + global_start = cp_rank * l_local kv_local = key.squeeze(-2).squeeze(1) if boundary_hidden is None or boundary_kv is None: raise RuntimeError( @@ -2205,14 +2141,14 @@ def _forward_thd_cp( "the hidden-only boundary exchange and boundary KV projection path." ) boundary_kv = boundary_kv.squeeze(-2).squeeze(1) - d_window = boundary_hidden.shape[0] // len(chunk_ranges) + d_window = boundary_hidden.shape[0] # Window-only defaults; compression fills the rank-major KV buffer. compressed_kv_rank_major = kv_local.new_empty((0, kv_local.shape[-1])) cu_seqlens_compressed = None # Logical top-k uses per-sequence comp ids; the map lowers them to # rank-major gathered rows when needed. - indexer_topk_logical = rank_row_for_seq_row = None + compressed_topk = seq_to_rank_row = None ratio = self.compress_ratio indexer = self.indexer indexer_loss_coeff = self.config.dsa_indexer_loss_coeff or 0.0 @@ -2221,33 +2157,28 @@ def _forward_thd_cp( if self.compressor is not None and ratio > 1: # ---- Step 3: build fixed-capacity compressor input ---------------- - # ratio=4 overlaps one previous group; ratio=128 is non-overlap. - d_comp = 8 if ratio == 4 else ratio - # One compressed row per full ratio-sized group; tails are dropped. compressed_lens = torch.div( - cu_seqlens_kv[1:] - cu_seqlens_kv[:-1], ratio, rounding_mode="floor" + cu_seqlens[1:] - cu_seqlens[:-1], ratio, rounding_mode="floor" ) cu_seqlens_compressed = torch.cat( ( - torch.zeros_like(cu_seqlens_kv[:1]), + torch.zeros_like(cu_seqlens[:1]), torch.cumsum(compressed_lens, dim=0, dtype=torch.int32), ) ) - # Compact local+boundary tokens for Compressor. ``comp_ids_local`` + # Compact local+boundary tokens for Compressor. ``compressed_group_ids`` # drives compressed RoPE; the row map directly lowers global # sequence-major compressed ids into rank-major all-gather rows. - hidden_compact, comp_ids_local, rank_row_for_seq_row = ( - cp_utils.build_cp_compressor_prep_compact_fused( + hidden_compact, compressed_group_ids, seq_to_rank_row = ( + cp_utils.prepare_cp_compressor_input( x, boundary_hidden, - cu_seqlens_kv, + cu_seqlens, cu_seqlens_compressed, - chunk_ranges, + global_start, cp_size, ratio, - d_comp, - d_window, ) ) @@ -2282,8 +2213,8 @@ def _forward_thd_cp( rotary_pos_sin, indexer.index_head_dim - indexer.qk_pos_emb_head_dim, indexer.qk_pos_emb_head_dim, - cu_seqlens_q, - chunk_ranges, + cu_seqlens, + global_start, ) ) weights_indexer_cp, _ = indexer.linear_weights_proj(indexer_x) @@ -2291,10 +2222,9 @@ def _forward_thd_cp( indexer_compressed_local, _ = indexer.compressor._forward_thd( hidden_compact.detach(), - cu_seqlens_kv, + cu_seqlens, max_seqlen_q=max_seqlen_q, - compressed_group_ids=comp_ids_local, - fixed_total_comp=hidden_compact.shape[0] // ratio, + compressed_group_ids=compressed_group_ids, ) # Gather indexer compressed K in rank-major order. k_indexer_rank_major = gather_from_sequence_parallel_region( @@ -2304,30 +2234,28 @@ def _forward_thd_cp( # Indexer top-k consumes sequence-major K rows. Capacity-tail # map entries are unused by cu_seqlens_compressed. k_indexer_seq_major = torch.index_select( - k_indexer_rank_major, 0, rank_row_for_seq_row.clamp_min(0) + k_indexer_rank_major, 0, seq_to_rank_row.clamp_min(0) ) # Top-k is still in logical compressed coordinates here. - indexer_topk_logical = cp_utils.compute_cp_indexer_topk_logical_fused( + compressed_topk = cp_utils.compute_cp_indexer_topk( q_indexer_cp, weights_indexer_cp, k_indexer_seq_major, - cu_seqlens_q, + cu_seqlens, cu_seqlens_compressed, - chunk_ranges, + global_start, ratio, indexer.index_topk, indexer.softmax_scale, max_seqlen_q=max_seqlen_q, - max_seqlen_kv=max_seqlen_q // ratio, ) # ---- Step 5: attention compressed KV path ------------------------- compressed_kv_local, _ = self.compressor._forward_thd( hidden_compact, - cu_seqlens_kv, + cu_seqlens, max_seqlen_q=max_seqlen_q, - compressed_group_ids=comp_ids_local, - fixed_total_comp=hidden_compact.shape[0] // ratio, + compressed_group_ids=compressed_group_ids, ) # Gather attention compressed KV in the same rank-major layout. compressed_kv_rank_major = gather_from_sequence_parallel_region( @@ -2338,24 +2266,29 @@ def _forward_thd_cp( # Final indices address these source rows directly. This avoids a # second per-sequence packed KV layout and lets torch.cat own backward. kv_full_thd = torch.cat((boundary_kv, kv_local, compressed_kv_rank_major), dim=0) - compressed_base = len(chunk_ranges) * d_window + l_local - if training_with_grad and indexer_topk_logical is not None: - # ---- Step 7a: indexer-loss path ---------------------------------- - # Build attention indices plus rank-major indices for indexer loss. - topk_idxs, indexer_topk_rank_major = csa_cp_layout_kernels.build_indexer_loss_indices( - cu_seqlens_q, - cu_seqlens_compressed, - index_global_start, + use_indexer_loss = training_with_grad and compressed_topk is not None + compressed_width = ( + compressed_topk.shape[-1] + if compressed_topk is not None + else (max_seqlen_q // ratio if ratio > 1 else 0) + ) + topk_idxs, topk_length, indexer_topk_rank_major = ( + csa_cp_layout_kernels.build_attention_indices( + cu_seqlens, + global_start, l_local, d_window, self.window_size, ratio, - indexer_topk_logical, - rank_row_for_seq_row, - chunk_starts=chunk_starts, - chunk_len=chunk_len, - compressed_base=compressed_base, + compressed_width, + compressed_topk, + cu_seqlens_compressed=cu_seqlens_compressed, + seq_to_rank_row=seq_to_rank_row, + for_indexer_loss=use_indexer_loss, ) + ) + if use_indexer_loss: + # ---- Step 7a: indexer-loss path ---------------------------------- cu_seqlens_q_unpadded = None if ( packed_seq_params.cu_seqlens_q is not None @@ -2366,17 +2299,17 @@ def _forward_thd_cp( cu_seqlens_q_unpadded = packed_seq_params.cu_seqlens_q q_padding_mask = None if cu_seqlens_q_unpadded is not None: - global_rows = torch.cat( - [ - torch.arange(start, end, device=query.device, dtype=torch.int64) - for start, end in chunk_ranges - ] - ) - batch_ids = batch_of_row(cu_seqlens_q, total_q=l_local * cp_size).index_select( - 0, global_rows + global_rows = torch.arange( + global_start, + global_start + l_local, + device=query.device, + dtype=cu_seqlens.dtype, ) + batch_ids = torch.bucketize( + global_rows, cu_seqlens[1:], out_int32=True, right=True + ).clamp_max(cu_seqlens.shape[0] - 2) real_seqlens = cu_seqlens_q_unpadded[1:] - cu_seqlens_q_unpadded[:-1] - positions = global_rows.to(cu_seqlens_q.dtype) - cu_seqlens_q[batch_ids] + positions = global_rows - cu_seqlens[batch_ids] q_padding_mask = positions >= real_seqlens[batch_ids] output, indexer_loss = ( FusedIndexerSparseAttnFromTopkFunc.apply @@ -2386,17 +2319,16 @@ def _forward_thd_cp( query, kv_full_thd, self.attn_sink.float(), - topk_idxs.int(), + topk_idxs, q_indexer_cp, k_indexer_rank_major, weights_indexer_cp, - indexer_topk_rank_major.int(), + indexer_topk_rank_major, compressed_kv_rank_major, self.softmax_scale, indexer.softmax_scale, indexer_loss_coeff, - self.config.calculate_per_token_loss, - l_local * cp_size, + 1 if self.config.calculate_per_token_loss else l_local * cp_size, q_padding_mask, ) if indexer_loss_coeff > 0: @@ -2410,41 +2342,19 @@ def _forward_thd_cp( return output.unsqueeze(1) # ---- Step 7b: sparse attention path ---------------------------------- - # No indexer-loss backward: lower final attention indices and run attention only. - compressed_width = ( - indexer_topk_logical.shape[-1] - if indexer_topk_logical is not None - else (max_seqlen_q // ratio if ratio > 1 else 0) - ) - topk_idxs, topk_length = csa_cp_layout_kernels.build_attention_indices( - cu_seqlens_q, - index_global_start, - l_local, - d_window, - self.window_size, - ratio, - compressed_width, - indexer_topk_logical, - cu_seqlens_compressed=cu_seqlens_compressed, - rank_row_for_seq_row=rank_row_for_seq_row, - chunk_starts=chunk_starts, - chunk_len=chunk_len, - compressed_base=compressed_base, - ) - if self.apply_dsa_kernel_fusion: output = dsa_sparse_attn( query, kv_full_thd, self.attn_sink.float(), - topk_idxs.int(), + topk_idxs, self.softmax_scale, topk_length=topk_length, is_thd=True, ) else: output = unfused_compressed_sparse_attn( - query, kv_full_thd, self.attn_sink.float(), topk_idxs.int(), self.softmax_scale + query, kv_full_thd, self.attn_sink.float(), topk_idxs, self.softmax_scale ) return output.unsqueeze(1) @@ -2476,7 +2386,7 @@ def _forward_thd( # DSv4 hybrid adds a dummy batch dim to keep the MQA-head # unsqueeze symmetric with SBHD) # x, qr : (total_q, 1, *) - total_q, _np, _ = query.shape + total_q, np_, _ = query.shape cu_seqlens_q = ( packed_seq_params.cu_seqlens_q_padded @@ -2489,7 +2399,6 @@ def _forward_thd( else packed_seq_params.cu_seqlens_kv ) max_seqlen_q = int(packed_seq_params.max_seqlen_q) - max_seqlen_kv = int(packed_seq_params.max_seqlen_kv) # Squeeze the dummy b=1 and MQA head-dim to get the KV-flat layout. # (key arrives as (total_kv, 1, 1, hn) for MQA.) @@ -2545,14 +2454,13 @@ def _forward_thd( kv_full_thd, compressed_kv, n_compressed_total, - _np, + np_, total_q, cu_seqlens_q, cu_seqlens_kv, cu_seqlens_kv_full, cu_seqlens_compressed, window_idxs, - max_seqlen_q, max_seqlen_compressed_idx, packed_seq_params, ) @@ -2562,8 +2470,6 @@ def _forward_thd( x, qr, packed_seq_params, - total_q, - _np, cu_seqlens_q, cu_seqlens_kv, cu_seqlens_kv_full, @@ -2587,7 +2493,6 @@ def _forward_thd( window_idxs, max_seqlen_q, max_seqlen_compressed_idx, - max_seqlen_kv, ) else: output = self._forward_fused_no_indexer_thd( diff --git a/megatron/core/transformer/experimental_attention_variant/csa_cp_layout_kernels.py b/megatron/core/transformer/experimental_attention_variant/csa_cp_layout_kernels.py index 0055666cff1..a8ef4df4cb9 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa_cp_layout_kernels.py +++ b/megatron/core/transformer/experimental_attention_variant/csa_cp_layout_kernels.py @@ -1,9 +1,10 @@ # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. """CuTeDSL kernels for DSv4 THD context parallel layout work. -This module contains the few CuTeDSL kernels retained by the DSv4 CP path. -``csa_cp_utils.py`` owns partition semantics and calls the RoPE/compaction -autograd wrappers; ``csa.py`` calls the final-index builders directly. +This module contains the CuTeDSL compaction and final-index kernels retained by +the DSv4 CP path. ``csa_cp_utils.py`` owns row mapping and compressor +input layout; ``csa.py`` calls final-index lowering directly. Local THD RoPE +reuses MCore's fused MLA implementation. """ import math @@ -16,18 +17,13 @@ import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute - from cutlass._mlir.dialects import llvm from cutlass.cute.runtime import from_dlpack, make_fake_stream - from cutlass.cutlass_dsl import T, dsl_user_op _CUTE_AVAILABLE = True except ImportError: cuda = None cutlass = None cute = None - llvm = None - T = None - dsl_user_op = None from_dlpack = None make_fake_stream = None _CUTE_AVAILABLE = False @@ -41,63 +37,6 @@ if _CUTE_AVAILABLE: - @dsl_user_op - def _ptr_as_i64( - tensor: cute.Tensor, offset: cutlass.Int32, *, loc=None, ip=None - ) -> cutlass.Int64: - elem_ptr = tensor.iterator + cutlass.Int32(offset) - return cutlass.Int64(llvm.ptrtoint(T.i64(), elem_ptr.llvm_ptr, loc=loc, ip=ip)) - - @dsl_user_op - def _ld_global_v4_u64( - base_ptr: cutlass.Int64, *, loc=None, ip=None - ) -> Tuple[cutlass.Uint64, cutlass.Uint64, cutlass.Uint64, cutlass.Uint64]: - result = llvm.inline_asm( - llvm.StructType.get_literal([T.i64(), T.i64(), T.i64(), T.i64()]), - [cutlass.Int64(base_ptr).ir_value(loc=loc, ip=ip)], - "ld.global.L1::no_allocate.v4.u64 {$0, $1, $2, $3}, [$4];", - "=l,=l,=l,=l,l", - has_side_effects=False, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - v0 = llvm.extractvalue(T.i64(), result, [0], loc=loc, ip=ip) - v1 = llvm.extractvalue(T.i64(), result, [1], loc=loc, ip=ip) - v2 = llvm.extractvalue(T.i64(), result, [2], loc=loc, ip=ip) - v3 = llvm.extractvalue(T.i64(), result, [3], loc=loc, ip=ip) - return cutlass.Uint64(v0), cutlass.Uint64(v1), cutlass.Uint64(v2), cutlass.Uint64(v3) - - @dsl_user_op - def _st_global_v4_u64( - base_ptr: cutlass.Int64, - v0: cutlass.Uint64, - v1: cutlass.Uint64, - v2: cutlass.Uint64, - v3: cutlass.Uint64, - *, - loc=None, - ip=None, - ): - llvm.inline_asm( - None, - [ - cutlass.Int64(base_ptr).ir_value(loc=loc, ip=ip), - cutlass.Uint64(v0).ir_value(loc=loc, ip=ip), - cutlass.Uint64(v1).ir_value(loc=loc, ip=ip), - cutlass.Uint64(v2).ir_value(loc=loc, ip=ip), - cutlass.Uint64(v3).ir_value(loc=loc, ip=ip), - ], - "st.global.L1::evict_first.v4.u64 [$0], {$1, $2, $3, $4};", - "l,l,l,l,l", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - def _launch_named( kernel, name: str, args: Tuple, grid: Tuple, block: Tuple, stream: cuda.CUstream ): @@ -113,219 +52,6 @@ def _launch_named( if dsl is not None and hasattr(dsl, "_name_prefix"): dsl._name_prefix = None - def _launch_1d(kernel, name: str, args: Tuple, work: cutlass.Int32, stream: cuda.CUstream): - """Launch a 1-D CuTe kernel with 128 threads per block.""" - _launch_named( - kernel, - name, - args, - grid=(cute.ceil_div(work, 128), 1, 1), - block=(128, 1, 1), - stream=stream, - ) - - # Kernel contract: - # x/out: row-major tensor views, shape (local_rows, row_width), bf16/fp16/fp32. - # cos/sin: rotary tables, shape (max_seq_len, pos_dim), floating dtype. - # cu_seqlens: int32 global sequence prefix sums, shape (n_seq + 1,). - # chunk0/1_global_start and chunk_len map local rows back to two global - # packed-token chunks. Applies non-interleaved RoPE to each head's position - # dims; if adjoint is set, applies the backward/adjoint transform. - @cute.kernel - def _thd_local_rope_kernel( - x: cute.Tensor, - out: cute.Tensor, - cos: cute.Tensor, - sin: cute.Tensor, - cu_seqlens: cute.Tensor, - n_seq: cutlass.Constexpr, - chunk0_global_start: cutlass.Int32, - chunk1_global_start: cutlass.Int32, - chunk_len: cutlass.Int32, - clamp_to_valid_token: cutlass.Constexpr, - row_width: cutlass.Constexpr, - head_dim: cutlass.Constexpr, - nope_dim: cutlass.Constexpr, - pos_dim: cutlass.Constexpr, - inverse: cutlass.Constexpr, - adjoint: cutlass.Constexpr, - head_group: cutlass.Constexpr, - total_work: cutlass.Int32, - ): - tidx, _, _ = cute.arch.thread_idx() - bidx, _, _ = cute.arch.block_idx() - row = bidx - if row < total_work: - n_heads = row_width // head_dim - if cutlass.const_expr( - x.element_type.width == 16 and nope_dim % 16 == 0 and head_dim % 16 == 0 - ): - copy_vec_elems = cutlass.Int32(16) - copy_vecs_per_head = nope_dim // 16 - copy_use_u256 = True - elif cutlass.const_expr( - x.element_type.width == 32 and nope_dim % 8 == 0 and head_dim % 8 == 0 - ): - copy_vec_elems = cutlass.Int32(8) - copy_vecs_per_head = nope_dim // 8 - copy_use_u256 = True - elif cutlass.const_expr( - (x.element_type.width == 16 and nope_dim % 4 == 0 and head_dim % 4 == 0) - or (x.element_type.width == 32 and nope_dim % 2 == 0 and head_dim % 2 == 0) - ): - copy_use_u256 = False - if cutlass.const_expr(x.element_type.width == 16): - copy_vec_elems = cutlass.Int32(4) - copy_vecs_per_head = nope_dim // 4 - else: - copy_vec_elems = cutlass.Int32(2) - copy_vecs_per_head = nope_dim // 2 - else: - copy_vec_elems = cutlass.Int32(1) - copy_vecs_per_head = nope_dim - copy_use_u256 = False - - copy_work = tidx - while copy_work < n_heads * copy_vecs_per_head: - unit = copy_work % copy_vecs_per_head - head = copy_work // copy_vecs_per_head - offset = row * row_width + head * head_dim + unit * copy_vec_elems - if cutlass.const_expr( - (x.element_type.width == 16 and nope_dim % 4 == 0 and head_dim % 4 == 0) - or (x.element_type.width == 32 and nope_dim % 2 == 0 and head_dim % 2 == 0) - ): - if cutlass.const_expr(copy_use_u256): - src_ptr = _ptr_as_i64(x, offset) - dst_ptr = _ptr_as_i64(out, offset) - v0, v1, v2, v3 = _ld_global_v4_u64(src_ptr) - _st_global_v4_u64(dst_ptr, v0, v1, v2, v3) - else: - src_ptr = cute.recast_ptr(x.iterator + offset, dtype=cutlass.Int64) - dst_ptr = cute.recast_ptr(out.iterator + offset, dtype=cutlass.Int64) - value = cute.arch.load( - src_ptr.llvm_ptr, - cutlass.Int64, - level1_eviction_priority="evict_no_allocate", - ) - cute.arch.store( - dst_ptr.llvm_ptr, value, level1_eviction_priority="evict_first" - ) - else: - col = head * head_dim + unit - out[row, col] = x[row, col] - copy_work = copy_work + 128 - - rope_pos = 0 - if tidx == 0 or tidx == 32 or tidx == 64 or tidx == 96: - local_row = row - global_token = chunk0_global_start + local_row - if row >= chunk_len: - local_row = row - chunk_len - global_token = chunk1_global_start + local_row - - if clamp_to_valid_token != 0: - if global_token < 0: - global_token = 0 - last_token = cu_seqlens[n_seq] - 1 - if global_token > last_token: - global_token = last_token - - for seq in range(n_seq): - seq_start = cu_seqlens[seq] - seq_end = cu_seqlens[seq + 1] - if global_token >= seq_start and global_token < seq_end: - rope_pos = global_token - seq_start - - rope_pos = cute.arch.shuffle_sync(rope_pos, 0, mask=-1, mask_and_clamp=31) - - pos_pairs = pos_dim // 2 - grouped_heads = n_heads // head_group - work = tidx - while work < grouped_heads * pos_pairs: - pair = work % pos_pairs - head_block = work // pos_pairs - cos_left = cos[rope_pos, pair].to(cutlass.Float32) - sin_left = sin[rope_pos, pair].to(cutlass.Float32) - cos_right = cos[rope_pos, pos_pairs + pair].to(cutlass.Float32) - sin_right = sin[rope_pos, pos_pairs + pair].to(cutlass.Float32) - if inverse != 0: - sin_left = cutlass.Float32(0.0) - sin_left - sin_right = cutlass.Float32(0.0) - sin_right - - for slot in cutlass.range_constexpr(16): - if slot < head_group: - head = head_block * head_group + slot - col = head * head_dim + nope_dim + pair * 2 - x1 = x[row, col].to(cutlass.Float32) - x2 = x[row, col + 1].to(cutlass.Float32) - - x1_cos_left = (x1 * cos_left).to(out.element_type).to(cutlass.Float32) - x2_sin_left = (x2 * sin_left).to(out.element_type).to(cutlass.Float32) - x2_cos_right = (x2 * cos_right).to(out.element_type).to(cutlass.Float32) - x1_sin_right = (x1 * sin_right).to(out.element_type).to(cutlass.Float32) - value_left = x1_cos_left - x2_sin_left - value_right = x2_cos_right + x1_sin_right - if adjoint != 0: - x2_sin_right = (x2 * sin_right).to(out.element_type).to(cutlass.Float32) - x1_sin_left = (x1 * sin_left).to(out.element_type).to(cutlass.Float32) - value_left = x1_cos_left + x2_sin_right - value_right = cutlass.Float32(0.0) - x1_sin_left + x2_cos_right - out[row, col] = value_left.to(out.element_type) - out[row, col + 1] = value_right.to(out.element_type) - work = work + 128 - - # Launch contract for _thd_local_rope_kernel. - # row_width is static; total_work is the local row count. - @cute.jit - def _thd_local_rope_launch( - x: cute.Tensor, - out: cute.Tensor, - cos: cute.Tensor, - sin: cute.Tensor, - cu_seqlens: cute.Tensor, - n_seq: cutlass.Constexpr, - chunk0_global_start: cutlass.Int32, - chunk1_global_start: cutlass.Int32, - chunk_len: cutlass.Int32, - clamp_to_valid_token: cutlass.Constexpr, - row_width: cutlass.Constexpr, - head_dim: cutlass.Constexpr, - nope_dim: cutlass.Constexpr, - pos_dim: cutlass.Constexpr, - inverse: cutlass.Constexpr, - adjoint: cutlass.Constexpr, - head_group: cutlass.Constexpr, - total_work: cutlass.Int32, - stream: cuda.CUstream, - ): - _launch_named( - _thd_local_rope_kernel, - "dsv4_cp_thd_local_rope", - ( - x, - out, - cos, - sin, - cu_seqlens, - n_seq, - chunk0_global_start, - chunk1_global_start, - chunk_len, - clamp_to_valid_token, - row_width, - head_dim, - nope_dim, - pos_dim, - inverse, - adjoint, - head_group, - total_work, - ), - grid=(total_work, 1, 1), - block=(128, 1, 1), - stream=stream, - ) - # Kernel contract: # hidden_local: local hidden rows, shape (l_local, row_width), bf16/fp16/fp32. # boundary_hidden: left boundary rows, shape (d_window, row_width). @@ -347,10 +73,8 @@ def _compressor_input_compact_fwd_kernel( ratio: cutlass.Constexpr, d_comp: cutlass.Constexpr, d_window: cutlass.Constexpr, - c_cap: cutlass.Int32, compact_len: cutlass.Int32, row_width: cutlass.Constexpr, - total_work: cutlass.Int32, ): tidx, _, _ = cute.arch.thread_idx() bidx, _, _ = cute.arch.block_idx() @@ -368,30 +92,23 @@ def _compressor_input_compact_fwd_kernel( first_range_group_start = range_start - d_comp src_global = cutlass.Int32(-1) + visible_comp_id = cutlass.Int32(-1) if row < compact_len: if tidx == 0 or tidx == 32: running_tokens = 0 for seq in range(n_seq): seq_start = cu_seqlens[seq] seq_end = cu_seqlens[seq + 1] - local_seq_start = seq_start - if local_seq_start < range_start: - local_seq_start = range_start local_seq_end = seq_end if local_seq_end > range_end: local_seq_end = range_end - if local_seq_start < local_seq_end: - seq_full_group_count = (seq_end - seq_start) // ratio + if seq_start < local_seq_end and range_start < local_seq_end: first_visible_numer = first_range_group_start - seq_start if first_visible_numer < 0: first_visible_numer = 0 - first_visible_group = 0 - if first_visible_numer > 0: - first_visible_group = (first_visible_numer + ratio - 1) // ratio + first_visible_group = (first_visible_numer + ratio - 1) // ratio stop_visible_group = (local_seq_end - seq_start) // ratio - if stop_visible_group > seq_full_group_count: - stop_visible_group = seq_full_group_count visible_group_count = stop_visible_group - first_visible_group if visible_group_count < 0: visible_group_count = 0 @@ -404,9 +121,12 @@ def _compressor_input_compact_fwd_kernel( local_visible_token - (local_visible_token // ratio) * ratio ) src_global = seq_start + comp_id * ratio + token_in_group + visible_comp_id = comp_id running_tokens = running_tokens + visible_token_count src_global = cute.arch.shuffle_sync(src_global, 0, mask=-1, mask_and_clamp=31) + if row % ratio == 0 and tidx == 0: + comp_ids[row // ratio] = visible_comp_id vec_col = tidx while vec_col < vec_cols: @@ -454,43 +174,8 @@ def _compressor_input_compact_fwd_kernel( hidden_compact[row, vec_col] = value vec_col = vec_col + 64 - if row < c_cap and tidx == 0: - slot = row - comp_ids[slot] = -1 - - running = 0 - for seq in range(n_seq): - seq_start = cu_seqlens[seq] - seq_end = cu_seqlens[seq + 1] - local_seq_start = seq_start - if local_seq_start < range_start: - local_seq_start = range_start - local_seq_end = seq_end - if local_seq_end > range_end: - local_seq_end = range_end - - if local_seq_start < local_seq_end: - seq_full_group_count = (seq_end - seq_start) // ratio - first_visible_numer = first_range_group_start - seq_start - if first_visible_numer < 0: - first_visible_numer = 0 - first_visible_group = 0 - if first_visible_numer > 0: - first_visible_group = (first_visible_numer + ratio - 1) // ratio - stop_visible_group = (local_seq_end - seq_start) // ratio - if stop_visible_group > seq_full_group_count: - stop_visible_group = seq_full_group_count - visible_group_count = stop_visible_group - first_visible_group - if visible_group_count < 0: - visible_group_count = 0 - - if slot >= running and slot < running + visible_group_count: - comp_id = first_visible_group + slot - running - comp_ids[slot] = comp_id - running = running + visible_group_count - # Launch contract for _compressor_input_compact_fwd_kernel. - # total_work covers compact rows plus metadata rows. + # One block copies each compact row; group-leading blocks also emit comp_ids. @cute.jit def _compressor_input_compact_fwd_launch( hidden_local: cute.Tensor, @@ -504,10 +189,8 @@ def _compressor_input_compact_fwd_launch( ratio: cutlass.Constexpr, d_comp: cutlass.Constexpr, d_window: cutlass.Constexpr, - c_cap: cutlass.Int32, compact_len: cutlass.Int32, row_width: cutlass.Constexpr, - total_work: cutlass.Int32, stream: cuda.CUstream, ): _launch_named( @@ -525,12 +208,10 @@ def _compressor_input_compact_fwd_launch( ratio, d_comp, d_window, - c_cap, compact_len, row_width, - total_work, ), - grid=(total_work, 1, 1), + grid=(compact_len, 1, 1), block=(64, 1, 1), stream=stream, ) @@ -555,7 +236,6 @@ def _compressor_input_compact_bwd_kernel( d_window: cutlass.Constexpr, compact_len: cutlass.Int32, row_width: cutlass.Constexpr, - total_work: cutlass.Int32, ): tidx, _, _ = cute.arch.thread_idx() bidx, _, _ = cute.arch.block_idx() @@ -569,98 +249,84 @@ def _compressor_input_compact_bwd_kernel( vec_cols = row_width compact_row = cutlass.Int32(-1) - if out_row < total_work: - range_start = global_start - range_end = global_start + l_local - first_range_group_start = range_start - d_comp - src_global = range_start + (out_row - d_window) - dst_row = out_row - d_window - if out_row < d_window: - src_global = range_start - d_window + out_row - dst_row = out_row - - if tidx == 0 or tidx == 32: - running_tokens = 0 - for seq in range(n_seq): - seq_start = cu_seqlens[seq] - seq_end = cu_seqlens[seq + 1] - local_seq_start = seq_start - if local_seq_start < range_start: - local_seq_start = range_start - local_seq_end = seq_end - if local_seq_end > range_end: - local_seq_end = range_end + range_start = global_start + range_end = global_start + l_local + first_range_group_start = range_start - d_comp + src_global = range_start + (out_row - d_window) + dst_row = out_row - d_window + if out_row < d_window: + dst_row = out_row - if local_seq_start < local_seq_end: - seq_full_group_count = (seq_end - seq_start) // ratio - first_visible_numer = first_range_group_start - seq_start - if first_visible_numer < 0: - first_visible_numer = 0 - first_visible_group = 0 - if first_visible_numer > 0: - first_visible_group = (first_visible_numer + ratio - 1) // ratio - stop_visible_group = (local_seq_end - seq_start) // ratio - if stop_visible_group > seq_full_group_count: - stop_visible_group = seq_full_group_count - visible_group_count = stop_visible_group - first_visible_group - if visible_group_count < 0: - visible_group_count = 0 - visible_token_count = visible_group_count * ratio + if tidx == 0 or tidx == 32: + running_tokens = 0 + for seq in range(n_seq): + seq_start = cu_seqlens[seq] + seq_end = cu_seqlens[seq + 1] + local_seq_end = seq_end + if local_seq_end > range_end: + local_seq_end = range_end - if src_global >= seq_start and src_global < seq_end: - seq_offset = src_global - seq_start - comp_id = seq_offset // ratio - token_in_group = seq_offset - (seq_offset // ratio) * ratio - if comp_id >= first_visible_group and comp_id < stop_visible_group: - compact_row = ( - running_tokens - + (comp_id - first_visible_group) * ratio - + token_in_group - ) - running_tokens = running_tokens + visible_token_count + if seq_start < local_seq_end and range_start < local_seq_end: + first_visible_numer = first_range_group_start - seq_start + if first_visible_numer < 0: + first_visible_numer = 0 + first_visible_group = (first_visible_numer + ratio - 1) // ratio + stop_visible_group = (local_seq_end - seq_start) // ratio + visible_group_count = stop_visible_group - first_visible_group + if visible_group_count < 0: + visible_group_count = 0 + visible_token_count = visible_group_count * ratio + + if src_global >= seq_start and src_global < seq_end: + seq_offset = src_global - seq_start + comp_id = seq_offset // ratio + token_in_group = seq_offset - comp_id * ratio + if comp_id >= first_visible_group and comp_id < stop_visible_group: + compact_row = ( + running_tokens + + (comp_id - first_visible_group) * ratio + + token_in_group + ) + running_tokens = running_tokens + visible_token_count - compact_row = cute.arch.shuffle_sync(compact_row, 0, mask=-1, mask_and_clamp=31) + compact_row = cute.arch.shuffle_sync(compact_row, 0, mask=-1, mask_and_clamp=31) - vec_col = tidx - while vec_col < vec_cols: - dst_offset = dst_row * row_width + vec_col * vec_elems - if cutlass.const_expr( - grad_hidden_compact.element_type.width == 16 and row_width % 4 == 0 - ): - value = cutlass.Int64(0) - if compact_row >= 0 and compact_row < compact_len: - src_offset = compact_row * row_width + vec_col * vec_elems - src_ptr = cute.recast_ptr( - grad_hidden_compact.iterator + src_offset, dtype=cutlass.Int64 - ) - value = cute.arch.load( - src_ptr.llvm_ptr, - cutlass.Int64, - level1_eviction_priority="evict_no_allocate", - ) - if out_row < d_window: - dst_ptr = cute.recast_ptr( - grad_boundary_hidden.iterator + dst_offset, dtype=cutlass.Int64 - ) - cute.arch.store( - dst_ptr.llvm_ptr, value, level1_eviction_priority="evict_first" - ) - else: - dst_ptr = cute.recast_ptr( - grad_hidden_local.iterator + dst_offset, dtype=cutlass.Int64 - ) - cute.arch.store( - dst_ptr.llvm_ptr, value, level1_eviction_priority="evict_first" - ) + vec_col = tidx + while vec_col < vec_cols: + dst_offset = dst_row * row_width + vec_col * vec_elems + if cutlass.const_expr( + grad_hidden_compact.element_type.width == 16 and row_width % 4 == 0 + ): + value = cutlass.Int64(0) + if compact_row >= 0 and compact_row < compact_len: + src_offset = compact_row * row_width + vec_col * vec_elems + src_ptr = cute.recast_ptr( + grad_hidden_compact.iterator + src_offset, dtype=cutlass.Int64 + ) + value = cute.arch.load( + src_ptr.llvm_ptr, + cutlass.Int64, + level1_eviction_priority="evict_no_allocate", + ) + if out_row < d_window: + dst_ptr = cute.recast_ptr( + grad_boundary_hidden.iterator + dst_offset, dtype=cutlass.Int64 + ) + cute.arch.store(dst_ptr.llvm_ptr, value, level1_eviction_priority="evict_first") else: - value = cutlass.Float32(0.0).to(grad_hidden_compact.element_type) - if compact_row >= 0 and compact_row < compact_len: - value = grad_hidden_compact[compact_row, vec_col] - if out_row < d_window: - grad_boundary_hidden[dst_row, vec_col] = value - else: - grad_hidden_local[dst_row, vec_col] = value - vec_col = vec_col + 64 + dst_ptr = cute.recast_ptr( + grad_hidden_local.iterator + dst_offset, dtype=cutlass.Int64 + ) + cute.arch.store(dst_ptr.llvm_ptr, value, level1_eviction_priority="evict_first") + else: + value = cutlass.Float32(0.0).to(grad_hidden_compact.element_type) + if compact_row >= 0 and compact_row < compact_len: + value = grad_hidden_compact[compact_row, vec_col] + if out_row < d_window: + grad_boundary_hidden[dst_row, vec_col] = value + else: + grad_hidden_local[dst_row, vec_col] = value + vec_col = vec_col + 64 # Launch contract for _compressor_input_compact_bwd_kernel. # total_work is the number of boundary + local gradient rows. @@ -697,7 +363,6 @@ def _compressor_input_compact_bwd_launch( d_window, compact_len, row_width, - total_work, ), grid=(total_work, 1, 1), block=(64, 1, 1), @@ -709,34 +374,30 @@ def _build_attention_indices_kernel( cu_seqlens: cute.Tensor, cu_seqlens_compressed: cute.Tensor, indexer_topk: cute.Tensor, - rank_row_for_seq_row: cute.Tensor, + seq_to_rank_row: cute.Tensor, topk_idxs: cute.Tensor, topk_length: cute.Tensor, + indexer_rank_major: cute.Tensor, n_seq: cutlass.Int32, global_start: cutlass.Int32, l_local: cutlass.Int32, - chunk0_start: cutlass.Int32, - chunk1_start: cutlass.Int32, - chunk_count: cutlass.Int32, - chunk_len: cutlass.Int32, d_window: cutlass.Int32, window_size: cutlass.Int32, ratio: cutlass.Int32, compressed_width: cutlass.Int32, seq_major_rows: cutlass.Int32, compressed_base: cutlass.Int32, - has_indexer_topk: cutlass.Int32, - has_compressed: cutlass.Int32, total_width: cutlass.Int32, - parallel_dense: cutlass.Int32, + index_mode: cutlass.Int32, ): tidx, _, _ = cute.arch.thread_idx() bidx, _, _ = cute.arch.block_idx() linear = bidx * 128 + tidx + # 0: selected top-k, 1: all visible compressed rows, 2: indexer loss. row = linear col = cutlass.Int32(0) - if parallel_dense != 0: + if index_mode != 0: if linear < l_local * total_width: row = linear // total_width col = linear - row * total_width @@ -744,74 +405,81 @@ def _build_attention_indices_kernel( row = l_local if row < l_local: - local_start = global_start - row_in_chunk = row - chunk_id = cutlass.Int32(0) - boundary_rows = d_window - if chunk_count != 0: - chunk_id = row // chunk_len - row_in_chunk = row - chunk_id * chunk_len - local_start = chunk0_start - if chunk_id == 1: - local_start = chunk1_start - boundary_rows = chunk_count * d_window - - global_q = local_start + row_in_chunk - seq_id = cutlass.Int32(-1) - seq_start_found = cutlass.Int32(0) + global_q = global_start + row + seq_start_found = cutlass.Int32(-1) seq_comp_start = cutlass.Int32(0) seq_comp_len = cutlass.Int32(0) for seq in range(n_seq): seq_start = cu_seqlens[seq] seq_end = cu_seqlens[seq + 1] if global_q >= seq_start and global_q < seq_end: - seq_id = seq seq_start_found = seq_start - if has_compressed != 0: + if ratio > 1 and compressed_width > 0: seq_comp_start = cu_seqlens_compressed[seq] seq_comp_len = cu_seqlens_compressed[seq + 1] - seq_comp_start - if parallel_dense != 0: + if index_mode != 0: topk_value = cutlass.Int32(-1) + rank_major_value = cutlass.Int32(-1) length = cutlass.Int32(0) - if seq_id >= 0: + if seq_start_found >= 0: window_start = global_q - window_size + 1 if window_start < seq_start_found: window_start = seq_start_found window_count = global_q - window_start + 1 - comp_count = cutlass.Int32(0) - if has_compressed != 0 and ratio > 1: - comp_count = (global_q - seq_start_found + 1) // ratio - if comp_count > compressed_width: - comp_count = compressed_width - if comp_count > seq_comp_len: - comp_count = seq_comp_len - length = window_count + comp_count - if col < window_count: - pos = window_start + col - if pos < local_start: - topk_value = chunk_id * d_window + pos - (local_start - d_window) + if index_mode == 2: + if col < compressed_width: + comp_id = indexer_topk[row, col] + if comp_id >= 0 and comp_id < seq_comp_len: + seq_major_id = seq_comp_start + comp_id + if seq_major_id < seq_major_rows: + rank_major_value = seq_to_rank_row[seq_major_id] + if rank_major_value >= 0: + topk_value = compressed_base + rank_major_value else: - topk_value = boundary_rows + chunk_id * chunk_len + pos - local_start - elif col < length: - seq_major_id = seq_comp_start + col - window_count - if seq_major_id >= 0 and seq_major_id < seq_major_rows: - rank_major_id = rank_row_for_seq_row[seq_major_id] - if rank_major_id >= 0: - topk_value = compressed_base + rank_major_id - elif total_width > 0: + window_col = col - compressed_width + if window_col < window_count: + pos = window_start + window_col + if pos < global_start: + topk_value = pos - (global_start - d_window) + else: + topk_value = d_window + pos - global_start + else: + comp_count = cutlass.Int32(0) + if ratio > 1 and compressed_width > 0: + comp_count = (global_q - seq_start_found + 1) // ratio + if comp_count > compressed_width: + comp_count = compressed_width + if comp_count > seq_comp_len: + comp_count = seq_comp_len + length = window_count + comp_count + if col < window_count: + pos = window_start + col + if pos < global_start: + topk_value = pos - (global_start - d_window) + else: + topk_value = d_window + pos - global_start + elif col < length: + seq_major_id = seq_comp_start + col - window_count + if seq_major_id < seq_major_rows: + rank_major_id = seq_to_rank_row[seq_major_id] + if rank_major_id >= 0: + topk_value = compressed_base + rank_major_id + elif total_width > 0 and index_mode != 2: length = 1 if col == 0: topk_value = 0 topk_idxs[row, col] = topk_value - if col == 0: + if index_mode == 2 and col < compressed_width: + indexer_rank_major[row, col] = rank_major_value + elif col == 0: topk_length[row] = length else: for out_col in range(total_width): topk_idxs[row, out_col] = -1 topk_length[row] = 0 - if seq_id >= 0: + if seq_start_found >= 0: write_col = cutlass.Int32(0) window_start = global_q - window_size + 1 if window_start < seq_start_found: @@ -820,23 +488,19 @@ def _build_attention_indices_kernel( for window_col in range(window_size): if window_col < window_count: pos = window_start + window_col - if pos < local_start: - topk_idxs[row, write_col] = ( - chunk_id * d_window + pos - (local_start - d_window) - ) + if pos < global_start: + topk_idxs[row, write_col] = pos - (global_start - d_window) else: - topk_idxs[row, write_col] = ( - boundary_rows + chunk_id * chunk_len + pos - local_start - ) + topk_idxs[row, write_col] = d_window + pos - global_start write_col = write_col + 1 - if has_compressed != 0 and ratio > 1: + if ratio > 1 and compressed_width > 0: for compressed_col in range(compressed_width): comp_id = indexer_topk[row, compressed_col] if comp_id >= 0 and comp_id < seq_comp_len: seq_major_id = seq_comp_start + comp_id - if seq_major_id >= 0 and seq_major_id < seq_major_rows: - rank_major_id = rank_row_for_seq_row[seq_major_id] + if seq_major_id < seq_major_rows: + rank_major_id = seq_to_rank_row[seq_major_id] if rank_major_id >= 0: topk_idxs[row, write_col] = compressed_base + rank_major_id write_col = write_col + 1 @@ -850,197 +514,50 @@ def _build_attention_indices_launch( cu_seqlens: cute.Tensor, cu_seqlens_compressed: cute.Tensor, indexer_topk: cute.Tensor, - rank_row_for_seq_row: cute.Tensor, + seq_to_rank_row: cute.Tensor, topk_idxs: cute.Tensor, topk_length: cute.Tensor, + indexer_rank_major: cute.Tensor, n_seq: cutlass.Int32, global_start: cutlass.Int32, l_local: cutlass.Int32, - chunk0_start: cutlass.Int32, - chunk1_start: cutlass.Int32, - chunk_count: cutlass.Int32, - chunk_len: cutlass.Int32, d_window: cutlass.Int32, window_size: cutlass.Int32, ratio: cutlass.Int32, compressed_width: cutlass.Int32, seq_major_rows: cutlass.Int32, compressed_base: cutlass.Int32, - has_indexer_topk: cutlass.Int32, - has_compressed: cutlass.Int32, total_width: cutlass.Int32, - parallel_dense: cutlass.Int32, + index_mode: cutlass.Int32, launch_work: cutlass.Int32, stream: cuda.CUstream, ): - _launch_1d( + _launch_named( _build_attention_indices_kernel, "dsv4_cp_build_attention_indices", ( cu_seqlens, cu_seqlens_compressed, indexer_topk, - rank_row_for_seq_row, + seq_to_rank_row, topk_idxs, topk_length, - n_seq, - global_start, - l_local, - chunk0_start, - chunk1_start, - chunk_count, - chunk_len, - d_window, - window_size, - ratio, - compressed_width, - seq_major_rows, - compressed_base, - has_indexer_topk, - has_compressed, - total_width, - parallel_dense, - ), - launch_work, - stream, - ) - - @cute.kernel - def _build_indexer_loss_indices_kernel( - cu_seqlens: cute.Tensor, - cu_seqlens_compressed: cute.Tensor, - indexer_topk_logical: cute.Tensor, - indexer_rank_by_seq_major: cute.Tensor, - topk_idxs: cute.Tensor, - indexer_rank_major: cute.Tensor, - n_seq: cutlass.Int32, - global_start: cutlass.Int32, - l_local: cutlass.Int32, - chunk0_start: cutlass.Int32, - chunk1_start: cutlass.Int32, - chunk_count: cutlass.Int32, - chunk_len: cutlass.Int32, - d_window: cutlass.Int32, - window_size: cutlass.Int32, - compressed_width: cutlass.Int32, - seq_major_rows: cutlass.Int32, - compressed_base: cutlass.Int32, - total_width: cutlass.Int32, - total_work: cutlass.Int32, - ): - tidx, _, _ = cute.arch.thread_idx() - bidx, _, _ = cute.arch.block_idx() - linear = bidx * 128 + tidx - if linear < total_work: - row = linear // total_width - col = linear - row * total_width - topk_value = cutlass.Int32(-1) - rank_major_value = cutlass.Int32(-1) - - local_start = global_start - row_in_chunk = row - chunk_id = cutlass.Int32(0) - boundary_rows = d_window - if chunk_count != 0: - chunk_id = row // chunk_len - row_in_chunk = row - chunk_id * chunk_len - local_start = chunk0_start - if chunk_id == 1: - local_start = chunk1_start - boundary_rows = chunk_count * d_window - - global_q = local_start + row_in_chunk - seq_id = cutlass.Int32(-1) - seq_start_found = cutlass.Int32(0) - seq_comp_start = cutlass.Int32(0) - seq_comp_len = cutlass.Int32(0) - for seq in range(n_seq): - seq_start = cu_seqlens[seq] - seq_end = cu_seqlens[seq + 1] - if global_q >= seq_start and global_q < seq_end: - seq_id = seq - seq_start_found = seq_start - seq_comp_start = cu_seqlens_compressed[seq] - seq_comp_len = cu_seqlens_compressed[seq + 1] - seq_comp_start - - if seq_id >= 0: - if col < compressed_width: - comp_id = indexer_topk_logical[row, col] - if comp_id >= 0 and comp_id < seq_comp_len: - seq_major_id = seq_comp_start + comp_id - if seq_major_id >= 0 and seq_major_id < seq_major_rows: - rank_major_id = indexer_rank_by_seq_major[seq_major_id] - if rank_major_id >= 0: - topk_value = compressed_base + rank_major_id - rank_major_value = rank_major_id - else: - window_col = col - compressed_width - window_start = global_q - window_size + 1 - if window_start < seq_start_found: - window_start = seq_start_found - window_count = global_q - window_start + 1 - if window_col < window_count: - pos = window_start + window_col - if pos < local_start: - topk_value = chunk_id * d_window + pos - (local_start - d_window) - else: - topk_value = boundary_rows + chunk_id * chunk_len + pos - local_start - - topk_idxs[row, col] = topk_value - if col < compressed_width: - indexer_rank_major[row, col] = rank_major_value - - @cute.jit - def _build_indexer_loss_indices_launch( - cu_seqlens: cute.Tensor, - cu_seqlens_compressed: cute.Tensor, - indexer_topk_logical: cute.Tensor, - indexer_rank_by_seq_major: cute.Tensor, - topk_idxs: cute.Tensor, - indexer_rank_major: cute.Tensor, - n_seq: cutlass.Int32, - global_start: cutlass.Int32, - l_local: cutlass.Int32, - chunk0_start: cutlass.Int32, - chunk1_start: cutlass.Int32, - chunk_count: cutlass.Int32, - chunk_len: cutlass.Int32, - d_window: cutlass.Int32, - window_size: cutlass.Int32, - compressed_width: cutlass.Int32, - seq_major_rows: cutlass.Int32, - compressed_base: cutlass.Int32, - total_width: cutlass.Int32, - total_work: cutlass.Int32, - stream: cuda.CUstream, - ): - _launch_1d( - _build_indexer_loss_indices_kernel, - "dsv4_cp_build_indexer_loss_indices", - ( - cu_seqlens, - cu_seqlens_compressed, - indexer_topk_logical, - indexer_rank_by_seq_major, - topk_idxs, indexer_rank_major, n_seq, global_start, l_local, - chunk0_start, - chunk1_start, - chunk_count, - chunk_len, d_window, window_size, + ratio, compressed_width, seq_major_rows, compressed_base, total_width, - total_work, + index_mode, ), - total_work, - stream, + grid=(cute.ceil_div(launch_work, 128), 1, 1), + block=(128, 1, 1), + stream=stream, ) @@ -1065,37 +582,20 @@ def _require_cute(message: str, *tensors: Optional[torch.Tensor]) -> None: _COMPILED_LAUNCH_CACHE = {} -def _cute_scalar(value): - if isinstance(value, bool): - return cutlass.Int32(1 if value else 0) - if isinstance(value, int): - return cutlass.Int32(value) - if isinstance(value, float): - return cutlass.Float32(value) - return value - - def _run_compiled_launch( launch_fn, tensor_args: Tuple[torch.Tensor, ...], scalar_args: Tuple, - key_arg_indices: Optional[Tuple[int, ...]] = None, - static_arg_indices: Optional[Tuple[int, ...]] = None, + static_arg_indices: Tuple[int, ...] = (), ) -> None: """Compile/cache a CuTe launch function and invoke it on the current stream.""" - key_arg_indices = ( - tuple(range(len(scalar_args))) if key_arg_indices is None else tuple(key_arg_indices) - ) - static_arg_indices = () if static_arg_indices is None else tuple(static_arg_indices) static_arg_set = set(static_arg_indices) key = ( launch_fn.__name__, tuple( (tensor.dtype, tuple(tensor.shape), tuple(tensor.stride())) for tensor in tensor_args ), - key_arg_indices, - tuple(scalar_args[i] for i in key_arg_indices), - static_arg_indices, + tuple((i, scalar_args[i]) for i in static_arg_indices), ) compiled = _COMPILED_LAUNCH_CACHE.get(key) if compiled is None: @@ -1113,7 +613,7 @@ def _run_compiled_launch( launch_fn, *cute_tensor_args, *( - arg if i in static_arg_set else _cute_scalar(arg) + arg if i in static_arg_set else cutlass.Int32(arg) for i, arg in enumerate(scalar_args) ), make_fake_stream(use_tvm_ffi_env_stream=False), @@ -1122,148 +622,11 @@ def _run_compiled_launch( _COMPILED_LAUNCH_CACHE[key] = compiled compiled( *tensor_args, - *(_cute_scalar(arg) for i, arg in enumerate(scalar_args) if i not in static_arg_set), + *(cutlass.Int32(arg) for i, arg in enumerate(scalar_args) if i not in static_arg_set), cuda.CUstream(torch.cuda.current_stream(tensor_args[0].device).cuda_stream), ) -class ThdLocalRope(torch.autograd.Function): - """Apply non-interleaved RoPE for THD rows owned by one or two CP chunks. - - Inputs: - x: CUDA tensor, shape ``(local_rows, ...)``; each row flattens to - ``num_heads * (nope_dim + pos_dim)`` and dtype is preserved. - cos/sin: CUDA tensors, shape ``(max_seq_len, >= pos_dim)``. - cu_seqlens_padded: int32 CUDA tensor, shape ``(n_seq + 1,)``. - chunk0_global_start/chunk1_global_start: global packed-token row starts. - chunk_len: positive row count per chunk. ``local_rows <= 2 * chunk_len``. - inverse/adjoint: select inverse RoPE or backward adjoint math. - - Output: - Tensor with same shape and dtype as ``x``. The wrapper maps each local - row back to a global THD token position, derives position within its - sequence from ``cu_seqlens_padded``, and launches the fused RoPE kernel. - """ - - @staticmethod - def _run( - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - cu_seqlens_padded: torch.Tensor, - chunk0_global_start: int, - chunk_len: int, - nope_dim: int, - pos_dim: int, - chunk1_global_start: int = 0, - inverse: bool = False, - clamp_to_valid_token: bool = False, - adjoint: bool = False, - ) -> torch.Tensor: - _require_cute( - "DSv4 THD CP RoPE requires CUDA tensors and available CuTe kernels.", - x, - cos, - sin, - cu_seqlens_padded, - ) - rows = x.shape[0] - row_width = math.prod(x.shape[1:]) - head_dim = nope_dim + pos_dim - if pos_dim <= 0 or pos_dim % 2 != 0: - raise RuntimeError(f"DSv4 CP RoPE expects an even positive pos_dim, got {pos_dim}.") - if head_dim <= 0 or row_width % head_dim != 0: - raise RuntimeError( - "DSv4 CP RoPE expects row width to be a multiple of " - f"head_dim={head_dim}, got {row_width}." - ) - if math.prod(cos.shape[1:]) < pos_dim or math.prod(sin.shape[1:]) < pos_dim: - raise RuntimeError("DSv4 CP RoPE cos/sin tables are narrower than pos_dim.") - chunk_len = int(chunk_len) - if chunk_len <= 0: - raise RuntimeError(f"DSv4 THD CP RoPE expects positive chunk_len, got {chunk_len}.") - if rows > 2 * chunk_len: - raise RuntimeError( - "DSv4 THD CP RoPE supports at most two chunks: " - f"rows={rows}, chunk_len={chunk_len}." - ) - if rows == 0: - return x - out = torch.empty_like(x) - n_heads = row_width // head_dim - head_group = 1 - for candidate in (16, 8, 4, 2): - if n_heads % candidate == 0: - head_group = candidate - break - total_work = rows - _run_compiled_launch( - _thd_local_rope_launch, - ( - x.reshape(rows, row_width), - out.reshape(rows, row_width), - cos.reshape(cos.shape[0], math.prod(cos.shape[1:])), - sin.reshape(sin.shape[0], math.prod(sin.shape[1:])), - cu_seqlens_padded, - ), - ( - cu_seqlens_padded.shape[0] - 1, - int(chunk0_global_start), - int(chunk1_global_start), - chunk_len, - 1 if clamp_to_valid_token else 0, - row_width, - head_dim, - int(nope_dim), - int(pos_dim), - 1 if inverse else 0, - 1 if adjoint else 0, - head_group, - total_work, - ), - key_arg_indices=(0, 4, 5, 6, 7, 8, 9, 10, 11), - static_arg_indices=(0, 4, 5, 6, 7, 8, 9, 10, 11), - ) - return out - - @staticmethod - def forward( - ctx, - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - cu_seqlens_padded: torch.Tensor, - chunk0_global_start: int, - chunk_len: int, - nope_dim: int, - pos_dim: int, - chunk1_global_start: int, - inverse: bool, - clamp_to_valid_token: bool, - ) -> torch.Tensor: - """Apply local THD CP RoPE and save metadata for backward.""" - ctx.rope_args = ( - int(chunk0_global_start), - int(chunk_len), - int(nope_dim), - int(pos_dim), - int(chunk1_global_start), - bool(inverse), - bool(clamp_to_valid_token), - ) - ctx.save_for_backward(cos, sin, cu_seqlens_padded) - return ThdLocalRope._run(x, cos, sin, cu_seqlens_padded, *ctx.rope_args) - - @staticmethod - def backward(ctx, grad_output: torch.Tensor): - """Apply the adjoint local THD CP RoPE transform to gradients.""" - cos, sin, cu_seqlens_padded = ctx.saved_tensors - grad_x = ThdLocalRope._run( - grad_output.contiguous(), cos, sin, cu_seqlens_padded, *ctx.rope_args, adjoint=True - ) - return (grad_x, *([None] * 10)) - - class CompressorInputCompact(torch.autograd.Function): """Compact local and boundary hidden rows for compressor input. @@ -1271,8 +634,8 @@ class CompressorInputCompact(torch.autograd.Function): hidden_local: CUDA tensor, shape ``(l_local, ...)``. boundary_hidden: CUDA tensor, shape ``(d_window, ...)``. cu_seqlens: int32 CUDA tensor, shape ``(n_seq + 1,)``. - global_start/l_local: local contiguous global row range. - ratio/d_comp/d_window/c_cap: fixed compressor window and capacity. + global_start: first global row in ``hidden_local``. + ratio/d_comp/c_cap: compressor window and fixed group capacity. Outputs: ``hidden_compact`` shape ``(c_cap * ratio, ...)``, same dtype as hidden, @@ -1288,13 +651,19 @@ def forward( boundary_hidden: torch.Tensor, cu_seqlens: torch.Tensor, global_start: int, - l_local: int, ratio: int, d_comp: int, - d_window: int, c_cap: int, ): """Compact local and boundary hidden rows into compressor input rows.""" + _require_cute( + "DSv4 CP compressor compaction requires CUDA tensors and CuTeDSL.", + hidden_local, + boundary_hidden, + cu_seqlens, + ) + l_local = hidden_local.shape[0] + d_window = boundary_hidden.shape[0] ctx.hidden_shape = tuple(hidden_local.shape) ctx.boundary_shape = tuple(boundary_hidden.shape) ctx.compact_args = (int(global_start), int(l_local), int(ratio), int(d_comp), int(d_window)) @@ -1320,13 +689,10 @@ def forward( int(ratio), int(d_comp), int(d_window), - int(c_cap), compact_len, row_width, - max(compact_len, int(c_cap)), ), - key_arg_indices=(3, 4, 5, 8), - static_arg_indices=(3, 4, 5, 8), + static_arg_indices=(3, 4, 5, 7), ) return hidden_compact, comp_ids @@ -1359,10 +725,9 @@ def backward(ctx, grad_hidden_compact: torch.Tensor, _grad_comp_ids: torch.Tenso row_width, l_local + d_window, ), - key_arg_indices=(3, 4, 5, 7), static_arg_indices=(3, 4, 5, 7), ) - return (grad_hidden, grad_boundary, *([None] * 7)) + return (grad_hidden, grad_boundary, *([None] * 5)) def build_attention_indices( @@ -1373,177 +738,84 @@ def build_attention_indices( window_size: int, ratio: int, compressed_width: int, - indexer_topk_compressed_logical_ids: Optional[torch.Tensor] = None, + compressed_topk: Optional[torch.Tensor] = None, cu_seqlens_compressed: Optional[torch.Tensor] = None, - rank_row_for_seq_row: Optional[torch.Tensor] = None, - chunk_starts: Optional[Tuple[int, ...]] = None, - chunk_len: int = 0, - compressed_base: int = 0, -) -> Tuple[torch.Tensor, torch.Tensor]: - """Build final sparse-attention indices in physical ``kv_full`` row space. + seq_to_rank_row: Optional[torch.Tensor] = None, + for_indexer_loss: bool = False, +) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: + """Build final sparse-attention and optional indexer-loss indices. Inputs: cu_seqlens: int32 CUDA tensor, shape ``(n_seq + 1,)``. - global_start/l_local: contiguous local query range, ignored for two-chunk mode. + global_start/l_local: this rank's global query start and row count. d_window/window_size: physical left-window capacity and per-query window width. ratio/compressed_width: compression mode and number of compressed columns. - indexer_topk_compressed_logical_ids: optional int32 tensor, shape + compressed_topk: optional per-sequence compressed ids, shape ``(l_local, compressed_width)``. - cu_seqlens_compressed/rank_row_for_seq_row: compressed row mapping. - chunk metadata: optional two-chunk row layout. - compressed_base: first compressed row in the raw KV tensor. + cu_seqlens_compressed/seq_to_rank_row: compressed row mapping. + for_indexer_loss: put compressed ids first and return their rank-major rows. Outputs: ``topk_idxs`` int32, shape ``(l_local, window_size + compressed_width)``; - ``topk_length`` int32, shape ``(l_local,)``. Window ids appear before - compressed ids and are already lowered into ``kv_full`` rows. + optional ``topk_length`` int32, shape ``(l_local,)``; and optional + ``indexer_rank_major`` int32, shape ``(l_local, compressed_width)``. + Normal attention puts window ids first. Indexer-loss mode puts selected + compressed ids first and returns their rank-major rows. """ - chunk_count = 0 - chunk1_start = 0 - if chunk_starts is not None: - _require_cute( - "DSv4 two-chunk final indices require CUDA tensors and CuTeDSL.", - cu_seqlens_compressed, - rank_row_for_seq_row, - ) - if not chunk_starts or len(chunk_starts) > 2: - raise RuntimeError( - "DSv4 two-chunk final indices support one or two chunks, " - f"got {len(chunk_starts)}." - ) - chunk_count = len(chunk_starts) - l_local = int(chunk_count) * int(chunk_len) - global_start = 0 - chunk1_start = int(chunk_starts[1]) if chunk_count > 1 else 0 + if for_indexer_loss and compressed_topk is None: + raise RuntimeError("DSv4 CP indexer-loss indices require compressed_topk.") + _require_cute( + "DSv4 CP final indices require CUDA tensors and CuTeDSL.", + cu_seqlens, + compressed_topk, + cu_seqlens_compressed, + seq_to_rank_row, + ) + global_start, l_local = int(global_start), int(l_local) if cu_seqlens_compressed is None: cu_seqlens_compressed = cu_seqlens - if rank_row_for_seq_row is None: - rank_row_for_seq_row = torch.empty((1,), dtype=torch.int32, device=cu_seqlens.device) + if seq_to_rank_row is None: + seq_to_rank_row = torch.empty((1,), dtype=torch.int32, device=cu_seqlens.device) total_width = window_size + compressed_width topk_idxs = torch.empty((l_local, total_width), dtype=torch.int32, device=cu_seqlens.device) - topk_length = torch.empty((l_local,), dtype=torch.int32, device=cu_seqlens.device) - if indexer_topk_compressed_logical_ids is None: - dummy = torch.empty((1, 1), dtype=torch.int32, device=cu_seqlens.device) - has_indexer = 0 + index_mode = 2 if for_indexer_loss else int(compressed_topk is None) + if compressed_topk is None: + compressed_topk = torch.empty((1, 1), dtype=torch.int32, device=cu_seqlens.device) + if for_indexer_loss: + topk_length_kernel = torch.empty((1,), dtype=torch.int32, device=cu_seqlens.device) + indexer_rank_major = torch.empty_like(compressed_topk) else: - dummy = indexer_topk_compressed_logical_ids - has_indexer = 1 - parallel_dense = 1 if has_indexer == 0 else 0 - launch_work = l_local * total_width if parallel_dense else l_local + topk_length_kernel = torch.empty((l_local,), dtype=torch.int32, device=cu_seqlens.device) + indexer_rank_major = torch.empty((1, 1), dtype=torch.int32, device=cu_seqlens.device) + launch_work = l_local * total_width if index_mode else l_local + compressed_base = int(d_window) + l_local _run_compiled_launch( _build_attention_indices_launch, - (cu_seqlens, cu_seqlens_compressed, dummy, rank_row_for_seq_row, topk_idxs, topk_length), - ( - cu_seqlens.shape[0] - 1, - global_start, - l_local, - int(chunk_starts[0]) if chunk_starts is not None else 0, - chunk1_start, - chunk_count, - int(chunk_len), - d_window, - window_size, - ratio, - compressed_width, - rank_row_for_seq_row.shape[0], - int(compressed_base), - has_indexer, - 1 if ratio > 1 and compressed_width > 0 else 0, - total_width, - parallel_dense, - launch_work, - ), - key_arg_indices=(), - ) - return topk_idxs, topk_length - - -def build_indexer_loss_indices( - cu_seqlens: torch.Tensor, - cu_seqlens_compressed: torch.Tensor, - global_start: int, - l_local: int, - d_window: int, - window_size: int, - ratio: int, - indexer_topk_compressed_logical_ids: torch.Tensor, - indexer_rank_by_seq_major: torch.Tensor, - chunk_starts: Optional[Tuple[int, ...]] = None, - chunk_len: int = 0, - compressed_base: int = 0, -) -> Tuple[torch.Tensor, torch.Tensor]: - """Build compressed-first sparse-attention ids and indexer-loss rank-major ids. - - Inputs: - cu_seqlens/cu_seqlens_compressed: int32 CUDA prefixes, shape ``(n_seq + 1,)``. - global_start/l_local: contiguous local query range, ignored for two-chunk mode. - d_window/window_size/ratio: final-id lowering parameters. - indexer_topk_compressed_logical_ids: int32 tensor, shape - ``(l_local, compressed_width)``. - indexer_rank_by_seq_major: int32 reverse map, shape ``(seq_major_rows,)``. - optional chunk metadata selects two-chunk lowering. - - Outputs: - ``topk_idxs`` int32, shape ``(l_local, compressed_width + window_size)``, - with compressed ids first; ``indexer_rank_major`` int32, shape - ``(l_local, compressed_width)`` for indexer-loss K rows. - """ - _require_cute( - "DSv4 CP indexer-loss final idx generation requires CUDA tensors and CuTeDSL.", - cu_seqlens, - cu_seqlens_compressed, - indexer_topk_compressed_logical_ids, - indexer_rank_by_seq_major, - ) - chunk_count = 0 - chunk1_start = 0 - if chunk_starts is not None: - if not chunk_starts or len(chunk_starts) > 2: - raise RuntimeError( - "DSv4 two-chunk indexer-loss indices support one or two chunks, " - f"got {len(chunk_starts)}." - ) - chunk_count = len(chunk_starts) - l_local = int(chunk_count) * int(chunk_len) - global_start = 0 - chunk1_start = int(chunk_starts[1]) if chunk_count > 1 else 0 - - compressed_width = indexer_topk_compressed_logical_ids.shape[-1] - total_width = compressed_width + window_size - topk_idxs = torch.empty((l_local, total_width), dtype=torch.int32, device=cu_seqlens.device) - indexer_rank_major = torch.empty( - (l_local, compressed_width), dtype=torch.int32, device=cu_seqlens.device - ) - total_work = l_local * total_width - if total_work == 0: - return topk_idxs, indexer_rank_major - _run_compiled_launch( - _build_indexer_loss_indices_launch, ( cu_seqlens, cu_seqlens_compressed, - indexer_topk_compressed_logical_ids, - indexer_rank_by_seq_major, + compressed_topk, + seq_to_rank_row, topk_idxs, + topk_length_kernel, indexer_rank_major, ), ( cu_seqlens.shape[0] - 1, global_start, l_local, - int(chunk_starts[0]) if chunk_starts is not None else 0, - chunk1_start, - chunk_count, - int(chunk_len), d_window, window_size, + ratio, compressed_width, - indexer_rank_by_seq_major.shape[0], - int(compressed_base), + seq_to_rank_row.shape[0], + compressed_base, total_width, - total_work, + index_mode, + launch_work, ), - key_arg_indices=(), ) - return topk_idxs, indexer_rank_major + if for_indexer_loss: + return topk_idxs, None, indexer_rank_major + return topk_idxs, topk_length_kernel, None diff --git a/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py b/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py index 7b9681c5e84..2e1703a0470 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py +++ b/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py @@ -1,142 +1,21 @@ # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. """MCore-facing utilities for the DSv4 THD context-parallel path. -This module owns CP partition semantics, boundary exchange, compressor-input -layout, and indexer top-k metadata. It calls the retained RoPE/compaction -kernels; ``csa.py`` calls the two final-index kernels directly. +This module owns CP row mapping, boundary exchange, compressor-input layout, +and indexer top-k metadata. It reuses MCore's fused MLA RoPE and calls the +retained compaction kernel; ``csa.py`` calls final-index lowering directly. """ import math -from typing import Optional, Sequence, Tuple +from typing import Optional, Tuple import torch import torch.distributed as dist +from megatron.core.fusions.fused_mla_yarn_rope_apply import fused_mla_rope_inplace from megatron.core.transformer.experimental_attention_variant import csa_cp_layout_kernels from megatron.core.transformer.experimental_attention_variant.dsa_kernels import indexer_topk -# ============================================================================= -# CP Partition Modes And Layout -# ============================================================================= - -DSV4_CP_PARTITION_CONTIGUOUS = "contiguous" -DSV4_CP_PARTITION_TWO_CHUNK = "two_chunk" - - -def local_q_cp_chunk_ranges( - partition_mode: Optional[str], local_rows: int, cp_size: int, cp_rank: int -) -> Tuple[Tuple[int, int], ...]: - """Return this rank's local rows as global packed-token ranges.""" - mode = DSV4_CP_PARTITION_CONTIGUOUS if partition_mode is None else partition_mode - if mode not in (DSV4_CP_PARTITION_CONTIGUOUS, DSV4_CP_PARTITION_TWO_CHUNK): - raise RuntimeError( - "Unsupported CSA CP partition mode: " f"{mode!r}. Expected contiguous or two_chunk." - ) - local_rows, cp_size, cp_rank = int(local_rows), int(cp_size), int(cp_rank) - if local_rows <= 0: - raise RuntimeError(f"local_rows must be positive, got {local_rows}.") - if cp_size < 1 or cp_rank < 0 or cp_rank >= cp_size: - raise RuntimeError(f"Invalid CP rank/size: cp_rank={cp_rank}, cp_size={cp_size}.") - - if mode == DSV4_CP_PARTITION_CONTIGUOUS: - start = cp_rank * local_rows - return ((start, start + local_rows),) - - if cp_size == 1: - return ((0, local_rows),) - if local_rows % 2 != 0: - raise RuntimeError( - "DSv4 two-chunk CP partition expects even local_rows: " - f"local_rows={local_rows}, cp_size={cp_size}." - ) - chunk_len = local_rows // 2 - total_chunks = 2 * cp_size - chunk_ids = (cp_rank, total_chunks - 1 - cp_rank) - return tuple((chunk_id * chunk_len, (chunk_id + 1) * chunk_len) for chunk_id in chunk_ids) - - -def local_kv_cp_chunk_ranges( - partition_mode: Optional[str], local_rows: int, boundary_rows: int, cp_size: int, cp_rank: int -) -> Tuple[Tuple[int, int], ...]: - """Return boundary rows followed by local rows as global packed-token ranges.""" - local_ranges = local_q_cp_chunk_ranges(partition_mode, local_rows, cp_size, cp_rank) - boundary_rows = int(boundary_rows) - if boundary_rows <= 0: - raise RuntimeError(f"boundary_rows must be positive, got {boundary_rows}.") - if boundary_rows % len(local_ranges) != 0: - raise RuntimeError( - "CP boundary rows must be divisible by local chunk count: " - f"boundary_rows={boundary_rows}, chunks={len(local_ranges)}." - ) - d_window = boundary_rows // len(local_ranges) - boundary_ranges = tuple((int(start) - d_window, int(start)) for start, _ in local_ranges) - return boundary_ranges + local_ranges - - -def _compressed_group_capacity(rows: int, ratio: int, d_comp: int) -> int: - """Return compressed group capacity with row alignment for TE FP8 linear inputs.""" - c_cap = max(1, (int(rows) + int(d_comp)) // int(ratio)) - group_alignment = 32 // math.gcd(32, int(ratio)) - return ((c_cap + group_alignment - 1) // group_alignment) * group_alignment - - -def thd_cp_local_row_indices( - partition_mode: Optional[str], - padded_total_rows: int, - cp_size: int, - cp_rank: int, - device: torch.device, -) -> torch.Tensor: - """Return global THD row indices selected by this CP rank. - - ``padded_total_rows`` is the final padded packed-token length. The output is - an int64 tensor of shape ``(padded_total_rows / cp_size,)``. Its row order is - exactly the order expected by ``local_q_cp_chunk_ranges`` for the selected - CSA CP partition mode. - """ - padded_total_rows = int(padded_total_rows) - cp_size, cp_rank = int(cp_size), int(cp_rank) - if padded_total_rows <= 0: - raise RuntimeError(f"padded_total_rows must be positive, got {padded_total_rows}.") - if cp_size < 1 or cp_rank < 0 or cp_rank >= cp_size: - raise RuntimeError(f"Invalid CP rank/size: cp_rank={cp_rank}, cp_size={cp_size}.") - if padded_total_rows % cp_size != 0: - raise RuntimeError( - "DSv4 THD CP partition expects padded_total_rows divisible by cp_size: " - f"padded_total_rows={padded_total_rows}, cp_size={cp_size}." - ) - - local_rows = padded_total_rows // cp_size - return torch.cat( - [ - torch.arange(start, end, dtype=torch.long, device=device) - for start, end in local_q_cp_chunk_ranges( - partition_mode, local_rows=local_rows, cp_size=cp_size, cp_rank=cp_rank - ) - ], - dim=0, - ) - - -def _normalize_row_ranges( - chunk_ranges: Sequence[Tuple[int, int]], op_name: str -) -> Tuple[Tuple[Tuple[int, int], ...], Tuple[int, ...], int]: - """Validate row ranges and return normalized ranges, lengths, and total rows.""" - if not chunk_ranges: - raise RuntimeError(f"{op_name} expects at least one chunk range.") - normalized = [] - lengths = [] - for start, end in chunk_ranges: - start = int(start) - end = int(end) - length = end - start - if length <= 0: - raise RuntimeError(f"{op_name} expects positive chunk lengths, got {length}.") - normalized.append((start, end)) - lengths.append(length) - return tuple(normalized), tuple(lengths), sum(lengths) - - # ============================================================================= # RoPE Wrappers # ============================================================================= @@ -149,51 +28,49 @@ def apply_thd_cp_local_rope_fused( nope_dim: int, pos_dim: int, cu_seqlens_padded: torch.Tensor, - chunk_ranges: Sequence[Tuple[int, int]], + global_start: int, inverse: bool = False, clamp_to_valid_token: bool = False, ) -> torch.Tensor: """Apply fused non-interleaved RoPE to local THD CP rows.""" - chunk_ranges, lengths, l_local = _normalize_row_ranges(chunk_ranges, "DSv4 THD CP local RoPE") - if x.shape[0] != l_local: - raise RuntimeError( - f"DSv4 THD CP local RoPE expects x rows to be {l_local}, got {x.shape[0]}." - ) - if len(chunk_ranges) == 1 or (len(chunk_ranges) == 2 and lengths[0] == lengths[1]): - return csa_cp_layout_kernels.ThdLocalRope.apply( - x, - cos, - sin, - cu_seqlens_padded, - int(chunk_ranges[0][0]), - int(lengths[0]), - nope_dim, - pos_dim, - int(chunk_ranges[1][0]) if len(chunk_ranges) == 2 else 0, - bool(inverse), - bool(clamp_to_valid_token), - ) + global_rows = torch.arange( + int(global_start), + int(global_start) + x.shape[0], + dtype=cu_seqlens_padded.dtype, + device=x.device, + ) + if clamp_to_valid_token: + global_rows = torch.minimum(global_rows.clamp_min(0), cu_seqlens_padded[-1] - 1) + sequence_ids = torch.bucketize( + global_rows, cu_seqlens_padded[1:], out_int32=True, right=True + ).clamp_max(cu_seqlens_padded.shape[0] - 2) + sequence_starts = cu_seqlens_padded[sequence_ids] + sequence_ends = cu_seqlens_padded[sequence_ids + 1] + valid_rows = (global_rows >= sequence_starts) & (global_rows < sequence_ends) + position_ids = torch.where( + valid_rows, global_rows - sequence_starts, torch.zeros_like(global_rows) + ) - parts = [] - row_start = 0 - for chunk_range, length in zip(chunk_ranges, lengths): - parts.append( - csa_cp_layout_kernels.ThdLocalRope.apply( - x.narrow(0, row_start, int(length)), - cos, - sin, - cu_seqlens_padded, - int(chunk_range[0]), - int(length), - nope_dim, - pos_dim, - 0, - bool(inverse), - bool(clamp_to_valid_token), - ) - ) - row_start += int(length) - return torch.cat(parts, dim=0) + squeezed_batch = x.ndim == 4 and x.shape[1] == 1 + squeezed_head = x.ndim == 2 + rope_input = x.squeeze(1) if squeezed_batch else x + rope_input = rope_input.unsqueeze(1) if squeezed_head else rope_input + output = fused_mla_rope_inplace( + rope_input.clone(), + cos, + sin, + nope_dim, + pos_dim, + cu_seqlens_q=cu_seqlens_padded, + inverse=inverse, + remove_interleaving=True, + position_ids=position_ids, + ) + if squeezed_batch: + return output.unsqueeze(1) + if squeezed_head: + return output.squeeze(1) + return output # ============================================================================= @@ -205,90 +82,39 @@ class _LeftBoundaryExchange(torch.autograd.Function): """Exchange fixed left-boundary windows and scatter gradients back to senders.""" @staticmethod - def forward( - ctx, - tensor: torch.Tensor, - d_window: int, - cp_group: torch.distributed.ProcessGroup, - two_chunk: bool, - ): + def forward(ctx, tensor: torch.Tensor, d_window: int, cp_group: torch.distributed.ProcessGroup): """Receive fixed left-boundary hidden rows needed by this CP rank.""" cp_size = 1 if cp_group is None else cp_group.size() cp_rank = 0 if cp_group is None else cp_group.rank() - if cp_size <= 1: - total_chunks, local_chunk_ids = 1, (0,) - owner_ranks = (0,) - elif two_chunk: - total_chunks = 2 * cp_size - local_chunk_ids = (cp_rank, total_chunks - 1 - cp_rank) - owner_ranks = tuple( - min(chunk_id, total_chunks - 1 - chunk_id) for chunk_id in range(total_chunks) - ) - else: - total_chunks, local_chunk_ids = cp_size, (cp_rank,) - owner_ranks = tuple(range(cp_size)) ctx.cp_group = cp_group ctx.d_window = d_window ctx.input_shape = tensor.shape - ctx.local_chunk_ids = local_chunk_ids - ctx.owner_ranks = owner_ranks - - local_chunks = len(local_chunk_ids) - if tensor.shape[0] % local_chunks != 0: - raise RuntimeError( - "DSv4 CP boundary exchange expects equal local chunks: " - f"local={tensor.shape[0]}, chunks={local_chunks}." - ) - chunk_len = tensor.shape[0] // local_chunks - ctx.chunk_len = chunk_len - if chunk_len < d_window: + if tensor.shape[0] < d_window: raise RuntimeError( - "DSv4 CP boundary exchange requires chunk_len >= D_window: " - f"chunk_len={chunk_len}, D_window={d_window}." + "DSv4 CP boundary exchange requires local rows >= D_window: " + f"local_rows={tensor.shape[0]}, D_window={d_window}." ) + boundary = tensor.new_zeros((d_window,) + tuple(tensor.shape[1:])) if cp_size <= 1: - return tensor.new_zeros((local_chunks * d_window,) + tuple(tensor.shape[1:])) + return boundary - recv = tensor.new_zeros((local_chunks, d_window) + tuple(tensor.shape[1:])) ops = [] - send_tensors = [] - to_peer = dist.get_global_rank if hasattr(dist, "get_global_rank") else lambda _, rank: rank - - for local_idx, chunk_id in enumerate(local_chunk_ids): - if chunk_id > 0: - prev_chunk_id = chunk_id - 1 - source_rank = owner_ranks[prev_chunk_id] - if source_rank == cp_rank: - source_idx = local_chunk_ids.index(prev_chunk_id) - source_start = source_idx * chunk_len - recv[local_idx].copy_( - tensor[source_start + chunk_len - d_window : source_start + chunk_len] - ) - else: - ops.append( - dist.P2POp( - dist.irecv, recv[local_idx], to_peer(cp_group, source_rank), cp_group - ) - ) - - next_chunk_id = chunk_id + 1 - if next_chunk_id < total_chunks: - target_rank = owner_ranks[next_chunk_id] - if target_rank != cp_rank: - start = local_idx * chunk_len - send_tensors.append( - tensor[start + chunk_len - d_window : start + chunk_len].contiguous() - ) - ops.append( - dist.P2POp( - dist.isend, send_tensors[-1], to_peer(cp_group, target_rank), cp_group - ) - ) - - if ops: - for req in dist.batch_isend_irecv(ops): - req.wait() - return recv.reshape((local_chunks * d_window,) + tuple(tensor.shape[1:])) + if cp_rank > 0: + ops.append( + dist.P2POp( + dist.irecv, boundary, dist.get_global_rank(cp_group, cp_rank - 1), cp_group + ) + ) + if cp_rank + 1 < cp_size: + send_tail = tensor[-d_window:].contiguous() + ops.append( + dist.P2POp( + dist.isend, send_tail, dist.get_global_rank(cp_group, cp_rank + 1), cp_group + ) + ) + for req in dist.batch_isend_irecv(ops): + req.wait() + return boundary @staticmethod def backward(ctx, grad_boundary: torch.Tensor): @@ -297,83 +123,44 @@ def backward(ctx, grad_boundary: torch.Tensor): cp_size = 1 if cp_group is None else cp_group.size() cp_rank = 0 if cp_group is None else cp_group.rank() d_window = ctx.d_window - chunk_len = ctx.chunk_len grad_input = grad_boundary.new_zeros(ctx.input_shape) if cp_size <= 1: - return grad_input, None, None, None - - local_chunk_ids = ctx.local_chunk_ids - owner_ranks = ctx.owner_ranks - total_chunks = len(owner_ranks) - local_chunks = len(local_chunk_ids) - grad_chunks = grad_boundary.reshape( - (local_chunks, d_window) + tuple(grad_boundary.shape[1:]) - ) + return grad_input, None, None + ops = [] - send_tensors = [] - recv_specs = [] - to_peer = dist.get_global_rank if hasattr(dist, "get_global_rank") else lambda _, rank: rank - - for local_idx, chunk_id in enumerate(local_chunk_ids): - if chunk_id > 0: - prev_chunk_id = chunk_id - 1 - source_rank = owner_ranks[prev_chunk_id] - if source_rank == cp_rank: - source_idx = local_chunk_ids.index(prev_chunk_id) - source_start = source_idx * chunk_len - grad_input[ - source_start + chunk_len - d_window : source_start + chunk_len - ] += grad_chunks[local_idx] - else: - send_tensors.append(grad_chunks[local_idx].contiguous()) - ops.append( - dist.P2POp( - dist.isend, send_tensors[-1], to_peer(cp_group, source_rank), cp_group - ) - ) - - next_chunk_id = chunk_id + 1 - if next_chunk_id < total_chunks: - target_rank = owner_ranks[next_chunk_id] - if target_rank != cp_rank: - recv = grad_chunks.new_zeros((d_window,) + tuple(grad_chunks.shape[2:])) - recv_specs.append((local_idx, recv)) - ops.append( - dist.P2POp(dist.irecv, recv, to_peer(cp_group, target_rank), cp_group) - ) - - if ops: - for req in dist.batch_isend_irecv(ops): - req.wait() - for local_idx, recv in recv_specs: - start = local_idx * chunk_len - grad_input[start + chunk_len - d_window : start + chunk_len] += recv - return grad_input, None, None, None + if cp_rank > 0: + send_grad = grad_boundary.contiguous() + ops.append( + dist.P2POp( + dist.isend, send_grad, dist.get_global_rank(cp_group, cp_rank - 1), cp_group + ) + ) + if cp_rank + 1 < cp_size: + recv_grad = grad_boundary.new_empty(grad_boundary.shape) + ops.append( + dist.P2POp( + dist.irecv, recv_grad, dist.get_global_rank(cp_group, cp_rank + 1), cp_group + ) + ) + for req in dist.batch_isend_irecv(ops): + req.wait() + if cp_rank + 1 < cp_size: + grad_input[-d_window:] = recv_grad + return grad_input, None, None def exchange_cp_boundary_hidden( hidden_states: torch.Tensor, - csa_compress_ratios: Sequence[int], + compress_ratio: int, csa_window_size: int, - partition_mode: Optional[str], cp_group: torch.distributed.ProcessGroup, ) -> torch.Tensor: - """Exchange hidden-state boundary rows for the selected CSA CP partition mode.""" - d_comp = max( - (8 if ratio == 4 else ratio for ratio in csa_compress_ratios if ratio and ratio > 1), - default=0, - ) + """Exchange hidden-state rows immediately left of this rank's token block.""" + d_comp = 8 if compress_ratio == 4 else compress_ratio if compress_ratio > 1 else 0 d_window = max(int(csa_window_size), d_comp) hidden_flat = hidden_states.view(hidden_states.shape[0], -1) - cp_size = 1 if cp_group is None else cp_group.size() - cp_rank = 0 if cp_group is None else cp_group.rank() - local_chunks = len( - local_q_cp_chunk_ranges(partition_mode, hidden_flat.shape[0], cp_size, cp_rank) - ) - boundary_hidden = _LeftBoundaryExchange.apply( - hidden_flat, d_window, cp_group, local_chunks == 2 - ) - return boundary_hidden.reshape((local_chunks * d_window,) + tuple(hidden_states.shape[1:])) + boundary_hidden = _LeftBoundaryExchange.apply(hidden_flat, d_window, cp_group) + return boundary_hidden.reshape((d_window,) + tuple(hidden_states.shape[1:])) # ============================================================================= @@ -381,114 +168,42 @@ def exchange_cp_boundary_hidden( # ============================================================================= -def build_cp_compressor_prep_compact_fused( +def prepare_cp_compressor_input( hidden_local: torch.Tensor, boundary_hidden: torch.Tensor, cu_seqlens: torch.Tensor, cu_seqlens_compressed: torch.Tensor, - chunk_ranges: Sequence[Tuple[int, int]], + global_start: int, cp_size: int, ratio: int, - d_comp: int, - d_window: int, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Build fixed-capacity compressor input for the selected CP partition mode. + """Build fixed-capacity compressor input for this rank's token block. Returns: ``hidden_compact``: rank-local compressor input, shape ``(compact_group_capacity * ratio, ...)``. - ``comp_ids_local``: original per-sequence compressed group id for each + ``compressed_group_ids``: original per-sequence compressed group id for each compact group, shape ``(compact_group_capacity,)``. For example, with ``ratio=4``, ``comp_id=3`` maps to RoPE position ``12``. - ``rank_row_for_seq_row``: map from global sequence-major compressed rows + ``seq_to_rank_row``: map from global sequence-major compressed rows to their canonical rank-major all-gather rows. """ - d_window = int(d_window) - - two_chunk = len(chunk_ranges) == 2 - if two_chunk: - chunk_ranges, chunk_lengths, l_local = _normalize_row_ranges( - chunk_ranges, "DSv4 two-chunk compressor-prep" - ) - if chunk_lengths[0] != chunk_lengths[1]: - raise RuntimeError( - "DSv4 two-chunk compressor-prep expects exactly two equal-length chunks." - ) - chunk_starts = (chunk_ranges[0][0], chunk_ranges[1][0]) - chunk_len = chunk_lengths[0] - else: - start, end = chunk_ranges[0] - chunk_ranges = ((int(start), int(end)),) - l_local = chunk_ranges[0][1] - chunk_ranges[0][0] - if hidden_local.shape[0] != l_local: - raise RuntimeError( - "DSv4 CP compressor-prep expects hidden rows to be " - f"{l_local}, got {hidden_local.shape[0]}." - ) - expected_boundary_rows = len(chunk_ranges) * d_window - if boundary_hidden.shape[0] != expected_boundary_rows: - raise RuntimeError( - "DSv4 CP compressor-prep expects one boundary window per chunk: " - f"boundary={boundary_hidden.shape[0]}, expected={expected_boundary_rows}." - ) - - if two_chunk: - chunk_specs = [ - (0, chunk_len, 0, chunk_starts[0]), - (chunk_len, chunk_len, d_window, chunk_starts[1]), - ] - else: - chunk_specs = [(0, l_local, 0, chunk_ranges[0][0])] - - parts = [] - for row_start, rows, boundary_start, global_start in chunk_specs: - c_cap = ( - (int(rows) + int(d_comp)) // int(ratio) - if two_chunk - else _compressed_group_capacity(rows, ratio, d_comp) - ) - parts.append( - csa_cp_layout_kernels.CompressorInputCompact.apply( - hidden_local.narrow(0, row_start, rows), - boundary_hidden.narrow(0, boundary_start, d_window), - cu_seqlens, - global_start, - rows, - ratio, - d_comp, - d_window, - c_cap, - ) - ) - if two_chunk: - hidden_parts, comp_parts = zip(*parts) - target_c_cap = _compressed_group_capacity(l_local, ratio, len(chunk_ranges) * d_comp) - current_c_cap = sum(part.shape[0] for part in comp_parts) - if current_c_cap > target_c_cap: - raise RuntimeError( - "DSv4 two-chunk compressor-prep produced more groups than its aligned capacity: " - f"current={current_c_cap}, capacity={target_c_cap}." - ) - if current_c_cap < target_c_cap: - pad_groups = target_c_cap - current_c_cap - hidden_parts = hidden_parts + ( - hidden_local.new_zeros((pad_groups * int(ratio),) + tuple(hidden_local.shape[1:])), - ) - comp_parts = comp_parts + ( - torch.full((pad_groups,), -1, dtype=torch.int32, device=hidden_local.device), - ) - hidden_compact = torch.cat(hidden_parts, dim=0) - comp_ids_local = torch.cat(comp_parts, dim=0) - else: - hidden_compact, comp_ids_local = parts[0] + cp_size = int(cp_size) + ratio = int(ratio) + d_comp = 8 if ratio == 4 else ratio + global_start = int(global_start) + l_local = hidden_local.shape[0] + group_alignment = 32 // math.gcd(32, ratio) + c_cap = max(1, (l_local + d_comp) // ratio) + c_cap = ((c_cap + group_alignment - 1) // group_alignment) * group_alignment + hidden_compact, compressed_group_ids = csa_cp_layout_kernels.CompressorInputCompact.apply( + hidden_local, boundary_hidden, cu_seqlens, global_start, ratio, d_comp, c_cap + ) - # A compressed group belongs canonically to the chunk containing its last - # token. From that chunk's first visible compressed row, its fixed-capacity + # A compressed group belongs to the rank containing its last token. From + # that rank's first visible compressed row, its fixed-capacity # rank-major slot follows directly; no (seq, comp, valid) tensors or repack # kernel are needed. - cp_size = int(cp_size) - ratio = int(ratio) - total_chunks = cp_size * (2 if two_chunk else 1) seq_major_rows = (l_local * cp_size) // ratio logical_rows = torch.arange(seq_major_rows, dtype=cu_seqlens.dtype, device=cu_seqlens.device) n_seq = cu_seqlens.shape[0] - 1 @@ -497,56 +212,42 @@ def build_cp_compressor_prep_compact_fused( ).clamp_max(n_seq - 1) comp_ids = logical_rows - cu_seqlens_compressed[seq_ids] group_last_rows = cu_seqlens[seq_ids] + (comp_ids + 1) * ratio - 1 - canonical_chunks = torch.div( - group_last_rows, chunk_len if two_chunk else l_local, rounding_mode="floor" - ).clamp_(0, total_chunks - 1) - - chunk_length = chunk_len if two_chunk else l_local - chunk_starts = torch.arange( - total_chunks, dtype=cu_seqlens.dtype, device=cu_seqlens.device - ) * int(chunk_length) + owner_ranks = torch.div(group_last_rows, l_local, rounding_mode="floor").clamp_(0, cp_size - 1) + + rank_starts = torch.arange(cp_size, dtype=cu_seqlens.dtype, device=cu_seqlens.device) * l_local first_seq_ids = torch.bucketize( - chunk_starts, cu_seqlens[1:], out_int32=True, right=True + rank_starts, cu_seqlens[1:], out_int32=True, right=True ).clamp_max(n_seq - 1) first_comp_ids = torch.div( - (chunk_starts - int(d_comp) - cu_seqlens[first_seq_ids]).clamp_min_(0) + ratio - 1, + (rank_starts - d_comp - cu_seqlens[first_seq_ids]).clamp_min_(0) + ratio - 1, ratio, rounding_mode="floor", ) first_logical_rows = cu_seqlens_compressed[first_seq_ids] + first_comp_ids - rank_slots = logical_rows - first_logical_rows[canonical_chunks] - if two_chunk: - late_chunks = canonical_chunks >= cp_size - owner_ranks = torch.where( - late_chunks, total_chunks - 1 - canonical_chunks, canonical_chunks - ) - rank_slots = rank_slots + late_chunks * int(c_cap) - else: - owner_ranks = canonical_chunks - rank_rows = owner_ranks * comp_ids_local.shape[0] + rank_slots - rank_row_for_seq_row = torch.where(logical_rows < cu_seqlens_compressed[-1], rank_rows, -1).to( + rank_slots = logical_rows - first_logical_rows[owner_ranks] + rank_rows = owner_ranks * compressed_group_ids.shape[0] + rank_slots + seq_to_rank_row = torch.where(logical_rows < cu_seqlens_compressed[-1], rank_rows, -1).to( torch.int32 ) - return hidden_compact, comp_ids_local, rank_row_for_seq_row + return hidden_compact, compressed_group_ids, seq_to_rank_row -def compute_cp_indexer_topk_logical_fused( +def compute_cp_indexer_topk( q_indexer_local: torch.Tensor, weights_indexer_local: torch.Tensor, k_indexer_seq_major: torch.Tensor, cu_seqlens_q: torch.Tensor, cu_seqlens_compressed: torch.Tensor, - chunk_ranges: Sequence[Tuple[int, int]], + global_start: int, ratio: int, topk_width: int, indexer_softmax_scale: float, max_seqlen_q: int, - max_seqlen_kv: int, ) -> Optional[torch.Tensor]: - """Run CP-aware indexer top-k for the selected CP partition mode. + """Run indexer top-k for this rank's query block. - This is a workaround for the DSA fused indexer forward mask contract. A CP - chunk may need a mask whose first local query already sees a compressed K + This is a workaround for the DSA fused indexer forward mask contract. A + rank may need a mask whose first local query already sees a compressed K prefix, for example: ``` @@ -556,7 +257,7 @@ def compute_cp_indexer_topk_logical_fused( ``` The fused kernel assumes the last Q row can see all K rows, then derives - earlier rows' causal masks from that endpoint. A CP chunk's last Q row may + earlier rows' causal masks from that endpoint. A rank's last Q row may only see a prefix of the global K rows, so this helper copies that visible prefix before calling the kernel. """ @@ -564,72 +265,56 @@ def compute_cp_indexer_topk_logical_fused( if topk_width == 0 or k_indexer_seq_major.shape[0] == 0: return None - chunk_ranges, lengths, l_local = _normalize_row_ranges(chunk_ranges, "DSv4 CP indexer top-k") - if q_indexer_local.shape[0] != l_local: - raise RuntimeError( - "DSv4 CP indexer top-k expects q_indexer_local rows to be " - f"{l_local}, got {q_indexer_local.shape[0]}." - ) + global_start = int(global_start) + l_local = q_indexer_local.shape[0] if weights_indexer_local.shape[0] != l_local: raise RuntimeError( "DSv4 CP indexer top-k expects weights rows to be " f"{l_local}, got {weights_indexer_local.shape[0]}." ) - outputs = [] - row_start = 0 - for (global_start, _), length in zip(chunk_ranges, lengths): - global_start, length = int(global_start), int(length) - global_end = global_start + length - q_for_topk = q_indexer_local.narrow(0, row_start, int(length)) - weights_for_topk = weights_indexer_local.narrow(0, row_start, int(length)) - local_starts = cu_seqlens_q[:-1].clamp_min(global_start) - local_ends = cu_seqlens_q[1:].clamp_max(global_end) - q_lens = (local_ends - local_starts).clamp_min(0) - k_lens = torch.minimum( - torch.div( - (local_ends - cu_seqlens_q[:-1]).clamp_min(0), int(ratio), rounding_mode="floor" - ), - cu_seqlens_compressed[1:] - cu_seqlens_compressed[:-1], - ) - k_lens = torch.where(q_lens > 0, k_lens, 0) - q_prefix = torch.cumsum(q_lens, dim=0, dtype=torch.int32) - k_prefix = torch.cumsum(k_lens, dim=0, dtype=torch.int32) - zero = torch.zeros((1,), dtype=cu_seqlens_q.dtype, device=cu_seqlens_q.device) - padding_q = (global_end - cu_seqlens_q[-1].clamp_min(global_start)).clamp_min(0) - cu_q_topk = torch.cat((zero, q_prefix, (q_prefix[-1] + padding_q).view(1))) - cu_k_topk = torch.cat((zero, k_prefix, k_prefix[-1:])) - - k_rows = torch.arange( - k_indexer_seq_major.shape[0], - dtype=cu_seqlens_compressed.dtype, - device=cu_seqlens_compressed.device, - ) - k_seq_ids = torch.bucketize(k_rows, cu_k_topk[1:-1], out_int32=True, right=True).clamp_max( - cu_seqlens_q.shape[0] - 2 - ) - valid_k = k_rows < k_prefix[-1] - source_rows = cu_seqlens_compressed[k_seq_ids] + k_rows - cu_k_topk[k_seq_ids] - k_for_topk = torch.index_select( - k_indexer_seq_major, 0, torch.where(valid_k, source_rows, 0) - ) - k_for_topk = k_for_topk * valid_k.view((-1,) + (1,) * (k_for_topk.ndim - 1)) - seq_lens_topk = torch.repeat_interleave( - torch.cat((k_lens, zero)), torch.cat((q_lens, padding_q.view(1))), output_size=length - ) - topk_chunk, _ = indexer_topk( - q_for_topk, - k_for_topk, - weights_for_topk, - topk=topk_width, - ratio=ratio, - indexer_softmax_scale=indexer_softmax_scale, - cu_seqlens_q=cu_q_topk, - cu_seqlens_kv=cu_k_topk, - max_seqlen_q=int(max_seqlen_q), - max_seqlen_kv=int(max_seqlen_kv), - visible_k_lengths=seq_lens_topk, - ) - outputs.append(topk_chunk) - row_start += length - return torch.cat(outputs, dim=0) + global_end = global_start + l_local + zero = torch.zeros((1,), dtype=cu_seqlens_q.dtype, device=cu_seqlens_q.device) + k_rows = torch.arange( + k_indexer_seq_major.shape[0], + dtype=cu_seqlens_compressed.dtype, + device=cu_seqlens_compressed.device, + ) + local_starts = cu_seqlens_q[:-1].clamp_min(global_start) + local_ends = cu_seqlens_q[1:].clamp_max(global_end) + q_lens = (local_ends - local_starts).clamp_min(0) + k_lens = torch.minimum( + torch.div((local_ends - cu_seqlens_q[:-1]).clamp_min(0), int(ratio), rounding_mode="floor"), + cu_seqlens_compressed[1:] - cu_seqlens_compressed[:-1], + ) + k_lens = torch.where(q_lens > 0, k_lens, 0) + q_prefix = torch.cumsum(q_lens, dim=0, dtype=torch.int32) + k_prefix = torch.cumsum(k_lens, dim=0, dtype=torch.int32) + padding_q = (global_end - cu_seqlens_q[-1].clamp_min(global_start)).clamp_min(0) + cu_q_topk = torch.cat((zero, q_prefix, (q_prefix[-1] + padding_q).view(1))) + cu_k_topk = torch.cat((zero, k_prefix, k_prefix[-1:])) + + k_seq_ids = torch.bucketize(k_rows, cu_k_topk[1:-1], out_int32=True, right=True).clamp_max( + cu_seqlens_q.shape[0] - 2 + ) + valid_k = k_rows < k_prefix[-1] + source_rows = cu_seqlens_compressed[k_seq_ids] + k_rows - cu_k_topk[k_seq_ids] + k_for_topk = torch.index_select(k_indexer_seq_major, 0, torch.where(valid_k, source_rows, 0)) + k_for_topk = k_for_topk * valid_k.view((-1,) + (1,) * (k_for_topk.ndim - 1)) + valid_k_lengths = torch.repeat_interleave( + torch.cat((k_lens, zero)), torch.cat((q_lens, padding_q.view(1))), output_size=l_local + ) + topk, _ = indexer_topk( + q_indexer_local, + k_for_topk, + weights_indexer_local, + topk=topk_width, + ratio=ratio, + indexer_softmax_scale=indexer_softmax_scale, + cu_seqlens_q=cu_q_topk, + cu_seqlens_kv=cu_k_topk, + max_seqlen_q=int(max_seqlen_q), + max_seqlen_kv=int(max_seqlen_q) // int(ratio), + valid_k_lengths=valid_k_lengths, + ) + return topk diff --git a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py index 0a281b732df..daccda6198b 100644 --- a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py +++ b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py @@ -128,7 +128,6 @@ def __init__( self.config.csa_compress_rotary_base if use_compressed_yarn else self.config.rotary_base ) self._dsv4_compress_ratio = compress_ratio - self._dsv4_rope_base = rope_base self._dsv4_uses_yarn_rope = use_compressed_yarn if not use_compressed_yarn: self.rotary_pos_emb = RotaryEmbedding( @@ -275,9 +274,8 @@ def forward( if cp_size > 1: boundary_hidden = cp_utils.exchange_cp_boundary_hidden( hidden_states, - self.config.csa_compress_ratios, + self._dsv4_compress_ratio, self.config.csa_window_size, - self.config.csa_cp_partition_mode, self.pg_collection.cp, ) @@ -286,7 +284,7 @@ def forward( # ===================== # Get the query, key and value tensors based on the type of attention - # self or cross attn. - query, key, value, q_compressed, _, boundary_kv = self.get_query_key_value_tensors( + query, key, value, q_compressed, boundary_kv = self.get_query_key_value_tensors( hidden_states, key_value_states, position_ids, @@ -383,12 +381,7 @@ def forward( if cp_size > 1: if rope_seqlen is None: raise RuntimeError("DSv4 THD CP inverse RoPE requires max_seqlen_kv.") - chunk_ranges = cp_utils.local_q_cp_chunk_ranges( - self.config.csa_cp_partition_mode, - core_attn_out.shape[0], - cp_size, - self.pg_collection.cp.rank(), - ) + global_start = self.pg_collection.cp.rank() * core_attn_out.shape[0] core_attn_out = cp_utils.apply_thd_cp_local_rope_fused( core_attn_out, rotary_pos_cos, @@ -396,7 +389,7 @@ def forward( nope_dim, pos_dim, cu_seqlens_kv, - chunk_ranges, + global_start, inverse=True, ) else: @@ -583,7 +576,7 @@ def get_query_key_value_tensors( Derives `query`, `key` and `value` tensors from `hidden_states`. Returns: - Tuple of ``(query, key, value, q_compressed, kv_compressed, boundary_kv)``. + Tuple of ``(query, key, value, q_compressed, boundary_kv)``. ``boundary_kv`` carries CP boundary rows when the DSv4 CP path is active. """ # s = sequence length, b = batch size, h = hidden size, n = num attention heads @@ -651,7 +644,6 @@ def get_query_key_value_tensors( kv_compressed = hidden_states boundary_kv_compressed = boundary_hidden - k_pos_emb = None if packed_seq_params is not None: # If sequence packing, TE expect [t, h, d] shaped qkv input. @@ -675,7 +667,7 @@ def get_query_key_value_tensors( # ========================================= def qkv_up_proj_and_rope_apply( - q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb, boundary_kv_compressed + q_compressed, kv_compressed, rotary_pos_emb, boundary_kv_compressed ): """ Apply the up projection and RoPE to the query and key. @@ -702,19 +694,13 @@ def qkv_up_proj_and_rope_apply( kv = self.kv_layernorm(kv) boundary_kv = None - # [num_tokens, qk_pos_emb_head_dim] -> [num_tokens, 1, qk_pos_emb_head_dim] - if k_pos_emb is not None: - k_pos_emb = torch.unsqueeze(k_pos_emb, -2) - cp_size = self.pg_collection.cp.size() if self.config.apply_rope_fusion: if cp_size > 1: cp_rank = self.pg_collection.cp.rank() - # CP local rows may start mid-sequence; use global THD ranges for RoPE. + # Rank r owns global rows [r * local_rows, (r + 1) * local_rows). assert packed_seq, "DSv4 CP fused RoPE expects THD packed sequence input." - chunk_ranges = cp_utils.local_q_cp_chunk_ranges( - self.config.csa_cp_partition_mode, q.shape[0], cp_size, cp_rank - ) + global_start = cp_rank * q.shape[0] query = cp_utils.apply_thd_cp_local_rope_fused( q, rotary_pos_cos, @@ -722,18 +708,7 @@ def qkv_up_proj_and_rope_apply( self.config.qk_head_dim, self.config.qk_pos_emb_head_dim, cu_seqlens_q, - chunk_ranges, - ) - kv_chunk_ranges = ( - cp_utils.local_kv_cp_chunk_ranges( - self.config.csa_cp_partition_mode, - kv_compressed.shape[0], - boundary_rows, - cp_size, - cp_rank, - ) - if boundary_rows - else chunk_ranges + global_start, ) kv = kv.unsqueeze(-2) kv = cp_utils.apply_thd_cp_local_rope_fused( @@ -743,7 +718,7 @@ def qkv_up_proj_and_rope_apply( self.config.qk_head_dim, self.config.qk_pos_emb_head_dim, cu_seqlens_q, - kv_chunk_ranges, + global_start - boundary_rows, clamp_to_valid_token=bool(boundary_rows), ) if boundary_rows: @@ -852,16 +827,15 @@ def qkv_up_proj_and_rope_apply( qkv_up_proj_and_rope_apply, q_compressed, kv_compressed, - k_pos_emb, rotary_pos_emb, boundary_kv_compressed, ) else: query, key, value, boundary_kv = qkv_up_proj_and_rope_apply( - q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb, boundary_kv_compressed + q_compressed, kv_compressed, rotary_pos_emb, boundary_kv_compressed ) - return query, key, value, q_compressed, kv_compressed, boundary_kv + return query, key, value, q_compressed, boundary_kv def backward_dw(self) -> NoReturn: """Execute weight gradient computation""" diff --git a/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py b/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py index fe23e281a04..69dbf53d542 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py +++ b/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py @@ -479,7 +479,7 @@ def _indexer_topk_core( cu_seqlens_kv: Optional[Tensor] = None, max_seqlen_q: Optional[int] = None, max_seqlen_kv: Optional[int] = None, - visible_k_lengths: Optional[Tensor] = None, + valid_k_lengths: Optional[Tensor] = None, ) -> Tuple[Tensor, Tensor, Tensor]: """Layout-agnostic core for :func:`indexer_topk`. @@ -545,23 +545,23 @@ def _indexer_topk_core( sk = int(max_seqlen_kv) total_q = q.shape[0] - if visible_k_lengths is not None: - if visible_k_lengths.shape[0] != total_q: + if valid_k_lengths is not None: + if valid_k_lengths.shape[0] != total_q: raise ValueError( - "visible_k_lengths must have one entry per THD query row: " - f"got={visible_k_lengths.shape[0]}, expected={total_q}." + "valid_k_lengths must have one entry per THD query row: " + f"got={valid_k_lengths.shape[0]}, expected={total_q}." ) - if visible_k_lengths.dtype != torch.int32: - raise ValueError(f"visible_k_lengths must be int32, got {visible_k_lengths.dtype}.") - if visible_k_lengths.device != device: - raise ValueError("visible_k_lengths must be on the same device as q_indexer.") - seq_lens = visible_k_lengths + if valid_k_lengths.dtype != torch.int32: + raise ValueError(f"valid_k_lengths must be int32, got {valid_k_lengths.dtype}.") + if valid_k_lengths.device != device: + raise ValueError("valid_k_lengths must be on the same device as q_indexer.") + seq_lens = valid_k_lengths else: # Per-row valid KV length: for token ``i`` in batch ``b``, # pos_in_seq = i - cu_seqlens_q[b] # valid = min((pos_in_seq + 1) // ratio, seqlen_kv[b]) - # CP callers with a trapezoid mask pass ``visible_k_lengths`` - # directly because q/k local ranges can differ from this formula. + # Callers with nonstandard Q/K alignment pass ``valid_k_lengths`` + # when this default bottom-right mask formula does not apply. row_idx = torch.arange(total_q, device=device, dtype=torch.int32) row_batch_ids = batch_of_row(cu_seqlens_q, total_q=total_q) row_valid = row_idx < cu_seqlens_q[-1] @@ -631,7 +631,7 @@ def indexer_topk( cu_seqlens_kv: Optional[Tensor] = None, max_seqlen_q: Optional[int] = None, max_seqlen_kv: Optional[int] = None, - visible_k_lengths: Optional[Tensor] = None, + valid_k_lengths: Optional[Tensor] = None, ) -> Tuple[Tensor, Tensor]: """Score + top-K selection for inference (no KL loss, no backward). @@ -653,10 +653,10 @@ def indexer_topk( cu_seqlens_kv: THD only — ``(B+1,)`` int32 CUDA cumulative KV lens. max_seqlen_q: THD only — per-batch max Q length. max_seqlen_kv: THD only — per-batch max KV length. - visible_k_lengths: THD only — optional ``(total_q,)`` int32 CUDA - tensor with the per-row visible K length for the radix top-K wrapper. - Callers pass this from a layout kernel to avoid rebuilding the - same metadata with generic PyTorch indexing during CUDA graph replay. + valid_k_lengths: THD only — optional ``(total_q,)`` int32 CUDA tensor + overriding the per-row valid K count passed to the radix top-K + wrapper when the default bottom-right mask does not describe the + Q/K alignment. Returns: SBHD: ``(topk_indices (b, sq, topk), topk_length (b, sq))`` int32 @@ -670,8 +670,8 @@ def indexer_topk( "indexer_topk THD mode requires cu_seqlens_q, cu_seqlens_kv, " "max_seqlen_q, and max_seqlen_kv to all be supplied." ) - if not is_thd and visible_k_lengths is not None: - raise ValueError("visible_k_lengths is only supported in THD mode.") + if not is_thd and valid_k_lengths is not None: + raise ValueError("valid_k_lengths is only supported in THD mode.") # ``indexer_softmax_scale`` is applied via the # ``relu(c·x) = c·relu(x)`` trick (the cudnn kernel does the relu), @@ -699,7 +699,7 @@ def indexer_topk( cu_seqlens_kv=cu_seqlens_kv, max_seqlen_q=int(max_seqlen_q) if max_seqlen_q is not None else None, max_seqlen_kv=int(max_seqlen_kv) if max_seqlen_kv is not None else None, - visible_k_lengths=visible_k_lengths, + valid_k_lengths=valid_k_lengths, ) return topk_indices, topk_length @@ -717,63 +717,6 @@ def _thd_to_fake_bshd(*tensors: Tensor) -> Tuple[Tensor, ...]: return tuple(t.unsqueeze(0) for t in tensors) -def _compute_indexer_predict( - q_indexer: Tensor, - k_indexer: Tensor, - weights: Tensor, - topk_indices: Tensor, - qhead_per_kv_head: int, - *, - topk_indices_global: bool = False, -) -> Tensor: - """Compute ``predict`` distribution (softmax over top-K of indexer scores). - - Wraps `cudnn.DSA.sparse_indexer_score_recompute_wrapper`. - This function is not used now, but it is kept for potential future use. - - Two layouts: - - * **BSHD** (default; 4-D q): ``q (B, S_q, H, D)``, ``k (B, S_k, D)``, - ``w (B, S_q, H)``, ``topk (B, S_q, topk)``. - * **THD packed** (3-D q): ``q (total_q, H, D)``, ``k (total_k, D)``, - ``w (total_q, H)``, ``topk (total_q, topk)``. Internally - fake-BSHD'd with ``B=1`` so the wrapper's 4-D-Q shape check - passes; ``topk_indices_global=True`` is required (and enforced) so - the kernel decodes the flat ids directly as positions into the - ``(1*total_k, D)`` view. - - Output shape matches the layout: BSHD ``(B, S_q, topk)`` or - THD ``(total_q, topk)``, fp32 softmax over the top-K axis. - """ - _ensure_dsa_namespace() - is_thd = q_indexer.ndim == 3 - if is_thd: - if not topk_indices_global: - raise ValueError( - "THD ``_compute_indexer_predict`` requires " - "``topk_indices_global=True`` so the kernel addresses K " - "by flat ids over the packed ``(total_k, D)`` buffer." - ) - q_bshd, k_bsd, w_bsh, topk_bst = _thd_to_fake_bshd( - q_indexer, k_indexer, weights, topk_indices - ) - else: - q_bshd, k_bsd, w_bsh, topk_bst = q_indexer, k_indexer, weights, topk_indices - - result = _DSA.sparse_indexer_score_recompute_wrapper( - q_bshd, - k_bsd, - w_bsh, - topk_bst, - qhead_per_kv_head=qhead_per_kv_head, - topk_indices_global=topk_indices_global, - ) - predict = result["predict"] - if is_thd: - predict = predict.squeeze(0) - return predict - - def _compute_attn_target( q_attn: Tensor, k_attn: Tensor, @@ -786,10 +729,9 @@ def _compute_attn_target( ) -> Tensor: """Compute ``target`` distribution (L1-normalised head-sum softmax). - Wraps :attr:`cudnn.DSA.sparse_attn_score_recompute_wrapper`. Same - layout convention as :func:`_compute_indexer_predict`: 4-D q is - BSHD; 3-D q is THD and gets fake-BSHD'd with ``B=1`` before the - wrapper call (so the 4-D-Q shape check passes). + Wraps :attr:`cudnn.DSA.sparse_attn_score_recompute_wrapper`. A 4-D q + is BSHD; a 3-D q is THD and gets fake-BSHD'd with ``B=1`` before the + wrapper call so the 4-D-Q shape check passes. """ _ensure_dsa_namespace() is_thd = q_attn.ndim == 3 @@ -1068,7 +1010,6 @@ def forward( # THD: skip the permute; tensors are already flat. if is_thd: total_q = q_indexer.shape[0] - idx_nh = q_indexer.shape[1] np_, d = query.shape[1], query.shape[2] q_indexer_flat = q_indexer @@ -1077,7 +1018,6 @@ def forward( else: sq, b, np_, d = query.shape skv = kv_full.shape[0] - idx_nh = q_indexer.shape[2] q_indexer_flat = q_indexer.permute(1, 0, 2, 3).contiguous() k_indexer_flat = k_indexer.permute(1, 0, 2).contiguous() @@ -1516,8 +1456,7 @@ def forward( softmax_scale: float, indexer_softmax_scale: float, loss_coeff: float, - calculate_per_token_loss: bool, - global_query_rows: int, + loss_divisor: float, q_padding_mask: Optional[Tensor] = None, ) -> Tuple[Tensor, Tensor]: """Run fused sparse attention using caller-supplied top-k indices.""" @@ -1530,17 +1469,6 @@ def forward( indexer_topk_idxs_for_loss = indexer_topk_idxs if q_padding_mask is not None: - if q_padding_mask.shape != (total_q,): - raise RuntimeError( - "q_padding_mask must have shape " - f"({total_q},), got {tuple(q_padding_mask.shape)}." - ) - if q_padding_mask.dtype != torch.bool: - raise RuntimeError(f"q_padding_mask must be bool, got {q_padding_mask.dtype}.") - if q_padding_mask.device != query.device: - raise RuntimeError( - f"q_padding_mask must be on {query.device}, got {q_padding_mask.device}." - ) indexer_topk_idxs_for_loss = indexer_topk_idxs.masked_fill( q_padding_mask.unsqueeze(-1), -1 ) @@ -1554,21 +1482,17 @@ def forward( topk_length=None, indexer_topk=indexer_topk, ) - if lse_indexer is None: - raise RuntimeError("Indexer sparse attention from top-k requires lse_indexer.") if indexer_softmax_scale != 1.0: weights_scaled = (weights.float() * indexer_softmax_scale).to(weights.dtype) else: weights_scaled = weights - predict = _compute_indexer_predict( - q_indexer, - k_indexer, - weights_scaled, - indexer_topk_idxs_for_loss, - qhead_per_kv_head=idx_nh, - topk_indices_global=True, + q_bshd, k_bsd, w_bsh, topk_bst = _thd_to_fake_bshd( + q_indexer, k_indexer, weights_scaled, indexer_topk_idxs_for_loss ) + predict = _DSA.sparse_indexer_score_recompute_wrapper( + q_bshd, k_bsd, w_bsh, topk_bst, qhead_per_kv_head=idx_nh, topk_indices_global=True + )["predict"].squeeze(0) target = _compute_attn_target( query.detach(), compressed_kv.detach(), @@ -1582,14 +1506,8 @@ def forward( raw_local_loss = _kl_loss_from_target_predict( target, predict, indexer_topk_idxs_for_loss, loss_coeff, calculate_per_token_loss=True ) - if calculate_per_token_loss: - indexer_loss = raw_local_loss - bwd_loss_coeff = loss_coeff * total_q - else: - if global_query_rows <= 0: - raise RuntimeError(f"global_query_rows must be positive, got {global_query_rows}.") - indexer_loss = raw_local_loss / float(global_query_rows) - bwd_loss_coeff = loss_coeff * float(total_q) / float(global_query_rows) + indexer_loss = raw_local_loss / loss_divisor + bwd_loss_coeff = loss_coeff * total_q / loss_divisor if loss_coeff > 0: ig = _DSA.indexer_backward_wrapper( @@ -1673,7 +1591,6 @@ def backward(ctx, grad_output, grad_loss): None, None, None, - None, ) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index a8f88aafda8..1a44adcb0d7 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -343,11 +343,6 @@ class TransformerConfig(ModelParallelConfig): """Whether to use dense mode for compressed sparse attention. If True, the CSA indexer will be disabled.""" - csa_cp_partition_mode: str = "contiguous" - """CSA CP token partition mode. ``contiguous`` keeps the usual one-block-per-rank - split. ``two_chunk`` gives each rank two separated token ranges, which helps balance - attention work for a single long packed sequence.""" - apply_dsa_kernel_fusion: bool = False """If True, use fused DSA sparse-attention kernels (FlashMLA forward + cuDNN DSA backward, indexer scoring, and top-K selection). Requires ``flash_mla`` and ``nvidia-cudnn-frontend`` diff --git a/tests/unit_tests/test_sequence_packing.py b/tests/unit_tests/test_sequence_packing.py index 643e1bf02f1..3e123f5a49f 100644 --- a/tests/unit_tests/test_sequence_packing.py +++ b/tests/unit_tests/test_sequence_packing.py @@ -8,6 +8,7 @@ import torch from megatron.core import parallel_state +from megatron.core.datasets import data_schedule from megatron.core.datasets.data_schedule import ( _build_thd_padding_mask, _get_scheduler_max_real_num_seqs, @@ -183,33 +184,123 @@ def rank(self): def test_dsv4_thd_cp_slice_uses_static_partition_total(): from megatron.core.datasets.data_schedule_utils import get_cp_slice_for_thd - from megatron.core.transformer.experimental_attention_variant.csa_cp_utils import ( - DSV4_CP_PARTITION_CONTIGUOUS, - ) batch = { - "tokens": torch.arange(8, dtype=torch.int64), - "position_ids": torch.arange(8, dtype=torch.int64), - "labels": torch.arange(100, 108, dtype=torch.int64), - "loss_mask": torch.ones(8, dtype=torch.float32), - "cu_seqlens": torch.tensor([0, 8], dtype=torch.int32), - "cu_seqlens_padded": torch.tensor([0, 8], dtype=torch.int32), - "max_seqlen": torch.tensor([8], dtype=torch.int32), + "tokens": torch.arange(12, dtype=torch.int64), + "position_ids": torch.arange(12, dtype=torch.int64), + "labels": torch.arange(100, 112, dtype=torch.int64), + "loss_mask": torch.ones(12, dtype=torch.float32), + "cu_seqlens": torch.tensor([0, 12], dtype=torch.int32), + "cu_seqlens_padded": torch.tensor([0, 12], dtype=torch.int32), + "max_seqlen": torch.tensor([12], dtype=torch.int32), } get_cp_slice_for_thd( - batch, - _MockCPGroup(size=4, rank=2), - csa_cp_partition_mode=DSV4_CP_PARTITION_CONTIGUOUS, - partition_total_tokens=16, + batch, _MockCPGroup(size=4, rank=2), use_dsv4_cp_slice=True, partition_total_tokens=16 + ) + + assert torch.equal(batch["tokens"], torch.arange(8, 12, dtype=torch.int64)) + assert torch.equal(batch["position_ids"], torch.arange(8, 12, dtype=torch.int64)) + assert torch.equal(batch["labels"], torch.arange(108, 112, dtype=torch.int64)) + assert torch.equal(batch["loss_mask"], torch.ones(4, dtype=torch.float32)) + assert torch.equal(batch["cu_seqlens"], torch.tensor([0, 12], dtype=torch.int32)) + assert torch.equal(batch["cu_seqlens_padded"], torch.tensor([0, 12], dtype=torch.int32)) + + +@pytest.mark.parametrize( + ("alignment", "cuda_graph_impl", "local_cp_size", "total_tokens", "local_target"), + [("max", "none", 1, 6, 8), (4, "none", 2, 10, 8), (8, "transformer_engine", 2, 10, 8)], +) +def test_dsv4_thd_dynamic_cp_pads_before_slicing( + monkeypatch, alignment, cuda_graph_impl, local_cp_size, total_tokens, local_target +): + """DSv4 padding must not change rank-local row origins after CP slicing.""" + cp_rank = local_cp_size - 1 + dynamic_cp_group = _MockCPGroup(size=local_cp_size, rank=cp_rank) + pg_collection = SimpleNamespace( + tp=_MockCPGroup(size=1, rank=0), + pp=_MockCPGroup(size=1, rank=0), + cp=_MockCPGroup(size=4, rank=cp_rank), + ) + config = SimpleNamespace( + experimental_attention_variant="dsv4_hybrid", + pad_packed_seq_alignment=alignment, + max_seqlen_per_dp_cp_rank=8, + thd_max_packed_sequences=None, + cuda_graph_impl=cuda_graph_impl, + pad_packed_seq_by_appending_dummy_seq=True, ) + tokens = torch.arange(1, total_tokens + 1, dtype=torch.int64) + batch = { + "tokens": tokens.clone(), + "position_ids": torch.arange(total_tokens, dtype=torch.int64), + "labels": tokens + 100, + "loss_mask": torch.ones(total_tokens, dtype=torch.float32), + "cu_seqlens": torch.tensor([0, total_tokens], dtype=torch.int32), + "cu_seqlens_padded": torch.tensor([0, total_tokens], dtype=torch.int32), + "max_seqlen": torch.tensor([total_tokens], dtype=torch.int32), + "local_cp_size": torch.tensor([local_cp_size], dtype=torch.int32), + } + pad_input_lengths = [] + real_pad_sequence_for_thd = data_schedule.pad_sequence_for_thd - assert torch.equal(batch["tokens"], torch.zeros(4, dtype=torch.int64)) - assert torch.equal(batch["position_ids"], torch.zeros(4, dtype=torch.int64)) - assert torch.equal(batch["labels"], torch.zeros(4, dtype=torch.int64)) - assert torch.equal(batch["loss_mask"], torch.zeros(4, dtype=torch.float32)) - assert torch.equal(batch["cu_seqlens"], torch.tensor([0, 8], dtype=torch.int32)) - assert torch.equal(batch["cu_seqlens_padded"], torch.tensor([0, 8], dtype=torch.int32)) + def record_padding( + tokens, labels, loss_mask, position_ids, packed_seq_params, *, padding_mask, **_ + ): + pad_input_lengths.append(tokens.shape[-1]) + assert tokens.shape[-1] == local_target + assert padding_mask.shape[-1] == local_target + result = real_pad_sequence_for_thd( + tokens, + labels, + loss_mask, + position_ids, + packed_seq_params, + padding_mask=padding_mask, + **_, + ) + assert result[0].shape[-1] == local_target + assert result[-1].shape[-1] == local_target + return result + + monkeypatch.setattr(torch.cuda, "current_device", lambda: torch.device("cpu")) + monkeypatch.setattr(torch.distributed, "get_process_group_ranks", lambda _: [0]) + monkeypatch.setattr( + torch.distributed, + "get_world_size", + lambda group=None: group.size() if group is not None else 1, + ) + monkeypatch.setattr( + torch.distributed, "get_rank", lambda group=None: group.rank() if group is not None else 0 + ) + monkeypatch.setattr(data_schedule, "broadcast_tensor", lambda *_: None) + monkeypatch.setattr(data_schedule, "pad_sequence_for_thd", record_padding) + monkeypatch.setattr( + parallel_state, + "get_dynamic_data_context_parallel_groups", + lambda group_size: dynamic_cp_group, + ) + + result = get_batch_on_this_rank_for_sequence_packing( + data_iterator=iter([batch]), dynamic_cp=True, pg_collection=pg_collection, config=config + ) + + local_tokens, _, _, _, _, packed_seq_params, padding_mask = result + global_start = cp_rank * local_target + expected = torch.zeros(local_target, dtype=torch.int64) + copied = max(0, min(total_tokens - global_start, local_target)) + if copied: + expected[:copied] = torch.arange( + global_start + 1, global_start + copied + 1, dtype=torch.int64 + ) + assert torch.equal(local_tokens.squeeze(0), expected) + assert padding_mask.shape == (1, local_target) + assert torch.equal( + padding_mask.squeeze(0), + torch.arange(global_start, global_start + local_target) >= total_tokens, + ) + assert packed_seq_params.local_cp_size == local_cp_size + assert pad_input_lengths == [local_target] def _gather_tensor_from_tp_group(tensor): diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_layout_kernels.py b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_layout_kernels.py index 31ecaa5f502..b24513ca714 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_layout_kernels.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_layout_kernels.py @@ -1,6 +1,5 @@ # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -import os from typing import List, Tuple import pytest @@ -11,10 +10,8 @@ # This file guards only DSv4 CP layout/metadata kernels. Layer-level CUDA graph # tests guard graph capture/replay behavior. -_E2E_RAGGED_SEG_LENS = (1, 127, 1000, 23, 129, 900, 55, 257, 800, 95, 509, 148) _E2E_RAGGED_PADDED_SEG_LENS = (1, 127, 1000, 23, 129, 900, 55, 257, 800, 95, 509, 200) _E2E_CP_SIZE = 4 -_E2E_CP_RANK = 1 def _require_cute_cuda(): @@ -24,16 +21,17 @@ def _require_cute_cuda(): pytest.skip("DSv4 CP CuTe kernels are not available in this environment.") -def _make_e2e_like_cu_seqlens(device: str = "cuda", *, padded: bool = True) -> torch.Tensor: +def _make_e2e_like_cu_seqlens(device: str = "cuda") -> torch.Tensor: """Return the ragged THD prefix pattern used by DSv4 CP e2e tests. - The padded variant sums to 4096, so CP4 owns 1024 local rows. It contains + The lengths sum to 4096, so CP4 owns 1024 local rows. They contain short, long, boundary-crossing, and padded-tail sequences without allocating full e2e hidden sizes in these focused kernel tests. """ - lengths = _E2E_RAGGED_PADDED_SEG_LENS if padded else _E2E_RAGGED_SEG_LENS return torch.tensor( - [0] + list(torch.tensor(lengths).cumsum(0).tolist()), dtype=torch.int32, device=device + [0] + list(torch.tensor(_E2E_RAGGED_PADDED_SEG_LENS).cumsum(0).tolist()), + dtype=torch.int32, + device=device, ) @@ -51,114 +49,11 @@ def _compressed_cu_seqlens(cu_seqlens: torch.Tensor, ratio: int) -> torch.Tensor ) -def _e2e_like_contiguous_range(cp_rank: int = _E2E_CP_RANK) -> Tuple[int, int]: - """Return one CP4 contiguous local range for the e2e-like padded THD layout.""" +def _e2e_like_local_range() -> Tuple[int, int]: + """Return rank 1's CP4 global start and local row count.""" total = sum(_E2E_RAGGED_PADDED_SEG_LENS) local = total // _E2E_CP_SIZE - start = int(cp_rank) * local - return start, local - - -def _e2e_like_two_chunk_ranges(cp_rank: int = _E2E_CP_RANK) -> Tuple[Tuple[int, int], ...]: - """Return one CP4 two-chunk local row order for the e2e-like THD layout.""" - total = sum(_E2E_RAGGED_PADDED_SEG_LENS) - local = total // _E2E_CP_SIZE - chunk_len = local // 2 - total_chunks = 2 * _E2E_CP_SIZE - chunk_ids = (int(cp_rank), total_chunks - 1 - int(cp_rank)) - return tuple((chunk_id * chunk_len, (chunk_id + 1) * chunk_len) for chunk_id in chunk_ids) - - -def _cosine_sim(actual: torch.Tensor, expected: torch.Tensor) -> float: - a = actual.detach().float().reshape(-1) - b = expected.detach().float().reshape(-1) - denom = a.norm() * b.norm() - if denom == 0: - return 1.0 if torch.equal(a, b) else 0.0 - return float(torch.dot(a, b) / denom) - - -def _tensor_sim(actual: torch.Tensor, expected: torch.Tensor) -> float: - a = actual.detach().float() - b = expected.detach().float() - denom = b.norm().clamp_min(1e-12) - return float(1.0 - (a - b).norm() / denom) - - -def _assert_rope_close(actual: torch.Tensor, expected: torch.Tensor, name: str) -> None: - torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) - cos = _cosine_sim(actual, expected) - sim = _tensor_sim(actual, expected) - assert cos > 0.9999, f"{name} cosine similarity too low: {cos}" - assert sim > 0.9999, f"{name} tensor similarity too low: {sim}" - - -def _num_bytes(*tensors: torch.Tensor) -> int: - return sum(t.numel() * t.element_size() for t in tensors) - - -def _time_cuda(fn, *, warmup: int = 3, iters: int = 5) -> float: - iters = int(os.environ.get("DSV4_CP_KERNEL_BENCH_ITERS", iters)) - for _ in range(warmup): - fn() - torch.cuda.synchronize() - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - for _ in range(iters): - fn() - end.record() - torch.cuda.synchronize() - return float(start.elapsed_time(end) / iters) - - -def _print_perf(name: str, cute_ms: float, torch_ms: float, effective_bytes: int) -> None: - cute_gbps = effective_bytes / max(cute_ms, 1e-9) / 1.0e6 - torch_gbps = effective_bytes / max(torch_ms, 1e-9) / 1.0e6 - print( - f"{name}: cute={cute_ms:.4f} ms ({cute_gbps:.2f} GB/s effective), " - f"torch={torch_ms:.4f} ms ({torch_gbps:.2f} GB/s effective), " - f"speedup={torch_ms / max(cute_ms, 1e-9):.2f}x" - ) - - -def _rope_reference( - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - positions: torch.Tensor, - nope_dim: int, - pos_dim: int, - inverse: bool = False, -) -> torch.Tensor: - x_nope, x_pos = torch.split(x, [nope_dim, pos_dim], dim=-1) - cos_flat = cos.reshape(cos.shape[0], -1).index_select(0, positions.to(torch.long)) - sin_flat = sin.reshape(sin.shape[0], -1).index_select(0, positions.to(torch.long)) - view_shape = (x.shape[0],) + (1,) * (x.ndim - 2) + (pos_dim,) - cos_pos = cos_flat[:, :pos_dim].view(view_shape) - sin_pos = sin_flat[:, :pos_dim].view(view_shape) - if inverse: - sin_pos = -sin_pos - - half = pos_dim // 2 - x1 = x_pos[..., 0::2] - x2 = x_pos[..., 1::2] - left = x1 * cos_pos[..., :half] - x2 * sin_pos[..., :half] - right = x2 * cos_pos[..., half:] + x1 * sin_pos[..., half:] - x_rot = torch.stack((left, right), dim=-1).flatten(-2) - return torch.cat((x_nope, x_rot), dim=-1) - - -def _seq_positions_from_rows( - cu_seqlens: torch.Tensor, global_rows: torch.Tensor, clamp_to_valid_token: bool = False -) -> torch.Tensor: - cu = cu_seqlens.cpu().to(torch.long) - rows = global_rows.cpu().to(torch.long) - if clamp_to_valid_token: - rows = rows.clamp(min=0, max=int(cu[-1]) - 1) - seq_ids = torch.searchsorted(cu, rows, right=True) - 1 - seq_ids = seq_ids.clamp(min=0, max=cu.numel() - 2) - return (rows - cu[seq_ids]).to(global_rows.device) + return local, local def _compressed_groups( @@ -234,37 +129,6 @@ def _native_compressor_input_compact( return hidden_compact, comp_ids -def _native_attention_indices_one_seq( - l_local: int, - d_window: int, - window_size: int, - ratio: int, - compressed_width: int, - device: torch.device, -): - rows = torch.arange(l_local, dtype=torch.int32, device=device).unsqueeze(1) - window_cols = torch.arange(window_size, dtype=torch.int32, device=device).unsqueeze(0) - window_start = (rows - int(window_size) + 1).clamp_min(0) - window_count = rows - window_start + 1 - window_values = int(d_window) + window_start + window_cols - window_values = torch.where( - window_cols < window_count, window_values, torch.full_like(window_values, -1) - ) - if ratio > 1 and compressed_width > 0: - comp_cols = torch.arange(compressed_width, dtype=torch.int32, device=device).unsqueeze(0) - visible = ((rows + 1) // int(ratio)).clamp(max=int(compressed_width)) - comp_values = int(d_window) + int(l_local) + comp_cols - comp_values = torch.where( - comp_cols < visible, comp_values, torch.full_like(comp_values, -1) - ) - topk = torch.cat((window_values, comp_values.expand(l_local, -1)), dim=1) - topk_length = window_count.squeeze(1) + visible.squeeze(1) - else: - topk = window_values - topk_length = window_count.squeeze(1) - return topk, topk_length - - def _native_attention_indices( cu_seqlens: torch.Tensor, cu_seqlens_compressed: torch.Tensor, @@ -274,11 +138,11 @@ def _native_attention_indices( window_size: int, ratio: int, compressed_width: int, - rank_by_seq_major: torch.Tensor, + seq_to_rank_row: torch.Tensor, compressed_base: int, indexer_topk: torch.Tensor = None, ): - """Reference contiguous final-index lowering for ragged THD CP rows.""" + """Reference final-index lowering for ragged THD CP rows.""" device = cu_seqlens.device total_width = int(window_size) + int(compressed_width) topk = torch.full((int(l_local), total_width), -1, dtype=torch.int32, device=device) @@ -320,7 +184,7 @@ def _native_attention_indices( n_visible = min((pos_in_seq + 1) // int(ratio), int(compressed_width)) comp_id = j if j < n_visible else -1 if 0 <= comp_id < seq_comp_len: - rank_id = int(rank_by_seq_major[cu_comp[seq_id] + comp_id]) + rank_id = int(seq_to_rank_row[cu_comp[seq_id] + comp_id]) if rank_id >= 0: topk[row, write_col] = int(compressed_base) + rank_id write_col += 1 @@ -328,32 +192,6 @@ def _native_attention_indices( return topk, lengths -def _native_indexer_loss_indices_one_seq( - l_local: int, - d_window: int, - window_size: int, - ratio: int, - logical_ids: torch.Tensor, - rank_by_seq_major: torch.Tensor, -): - device = logical_ids.device - compressed_width = logical_ids.shape[1] - attention_topk, _ = _native_attention_indices_one_seq( - l_local, d_window, window_size, ratio, 0, device - ) - compressed = torch.full((l_local, compressed_width), -1, dtype=torch.int32, device=device) - rank_major = torch.full_like(compressed, -1) - valid = (logical_ids >= 0) & (logical_ids < rank_by_seq_major.shape[0]) - if valid.any(): - lowered = int(d_window) + int(l_local) + logical_ids.clamp(min=0) - rank_ids = rank_by_seq_major.index_select(0, logical_ids.clamp(min=0).reshape(-1)).view_as( - logical_ids - ) - compressed = torch.where(valid, lowered, compressed) - rank_major = torch.where(valid, rank_ids, rank_major) - return torch.cat((compressed, attention_topk), dim=1), rank_major - - def _native_indexer_loss_indices( cu_seqlens: torch.Tensor, cu_seqlens_compressed: torch.Tensor, @@ -363,10 +201,10 @@ def _native_indexer_loss_indices( window_size: int, ratio: int, logical_ids: torch.Tensor, - rank_by_seq_major: torch.Tensor, + seq_to_rank_row: torch.Tensor, compressed_base: int, ): - """Reference contiguous compressed-first indexer-loss lowering for ragged THD rows.""" + """Reference compressed-first indexer-loss lowering for ragged THD rows.""" device = logical_ids.device compressed_width = logical_ids.shape[1] total_width = compressed_width + int(window_size) @@ -393,7 +231,7 @@ def _native_indexer_loss_indices( comp_id = int(logical_ids[row, col]) if 0 <= comp_id < comp_len: seq_major = cu_comp[seq_id] + comp_id - rank_id = int(rank_by_seq_major[seq_major]) + rank_id = int(seq_to_rank_row[seq_major]) if rank_id >= 0: topk[row, col] = int(compressed_base) + rank_id rank_major[row, col] = rank_id @@ -408,73 +246,10 @@ def _native_indexer_loss_indices( return topk, rank_major -def test_thd_local_rope_matches_native_forward_backward_and_reports_bandwidth(): +def test_compressor_input_compact_matches_native_forward_backward(): _require_cute_cuda() - torch.manual_seed(11) cu = _make_e2e_like_cu_seqlens() - chunk_ranges = _e2e_like_two_chunk_ranges() - rows = sum(end - start for start, end in chunk_ranges) - nope_dim = 4 - pos_dim = 4 - x = torch.randn(rows, 2, nope_dim + pos_dim, device="cuda", requires_grad=True) - cos = torch.randn(max(_E2E_RAGGED_PADDED_SEG_LENS), 1, 1, pos_dim, device="cuda") - sin = torch.randn(max(_E2E_RAGGED_PADDED_SEG_LENS), 1, 1, pos_dim, device="cuda") - global_rows = torch.cat( - [torch.arange(start, end, device="cuda") for start, end in chunk_ranges] - ) - positions = _seq_positions_from_rows(cu, global_rows) - ref_x = x.detach().clone().requires_grad_(True) - ref = _rope_reference(ref_x, cos, sin, positions, nope_dim, pos_dim) - grad = torch.randn_like(ref) - ref.backward(grad) - - fused_x = x.detach().clone().requires_grad_(True) - fused = csa_cp_layout_kernels.ThdLocalRope.apply( - fused_x, - cos, - sin, - cu, - chunk_ranges[0][0], - chunk_ranges[0][1] - chunk_ranges[0][0], - nope_dim, - pos_dim, - chunk_ranges[1][0], - False, - False, - ) - fused.backward(grad) - _assert_rope_close(fused, ref, "thd_local_rope.forward") - _assert_rope_close(fused_x.grad, ref_x.grad, "thd_local_rope.backward") - - bench_x = torch.randn(512, 8, nope_dim + pos_dim, device="cuda") - bench_grad = torch.randn_like(bench_x) - bench_cu = torch.tensor([0, 2048], dtype=torch.int32, device="cuda") - bench_cos = torch.randn(2048, 1, 1, pos_dim, device="cuda") - bench_sin = torch.randn(2048, 1, 1, pos_dim, device="cuda") - bench_pos = torch.arange(512, device="cuda") - - def run_cute(): - bx = bench_x.detach().clone().requires_grad_(True) - y = csa_cp_layout_kernels.ThdLocalRope.apply( - bx, bench_cos, bench_sin, bench_cu, 0, 512, nope_dim, pos_dim, 0, False, False - ) - y.backward(bench_grad) - - def run_torch(): - bx = bench_x.detach().clone().requires_grad_(True) - y = _rope_reference(bx, bench_cos, bench_sin, bench_pos, nope_dim, pos_dim) - y.backward(bench_grad) - - effective_bytes = 4 * _num_bytes(bench_x) - _print_perf( - "thd_local_rope_fwd_bwd", _time_cuda(run_cute), _time_cuda(run_torch), effective_bytes - ) - - -def test_compressor_input_compact_matches_native_forward_backward_and_reports_bandwidth(): - _require_cute_cuda() - cu = _make_e2e_like_cu_seqlens() - global_start, l_local = _e2e_like_contiguous_range() + global_start, l_local = _e2e_like_local_range() ratio = 128 d_comp = 128 d_window = 128 @@ -507,7 +282,7 @@ def test_compressor_input_compact_matches_native_forward_backward_and_reports_ba boundary.grad.zero_() fused = csa_cp_layout_kernels.CompressorInputCompact.apply( - hidden, boundary, cu, global_start, l_local, ratio, d_comp, d_window, c_cap + hidden, boundary, cu, global_start, ratio, d_comp, c_cap ) fused[0].backward(grad) assert torch.equal(fused[0], ref[0]) @@ -515,68 +290,17 @@ def test_compressor_input_compact_matches_native_forward_backward_and_reports_ba assert torch.equal(hidden.grad, ref_hidden_grad) assert torch.equal(boundary.grad, ref_boundary_grad) - bench_cu = torch.tensor([0, 16384], dtype=torch.int32, device="cuda") - bench_global_start = 4096 - bench_l_local = 8192 - bench_ratio = 4 - bench_d_comp = 128 - bench_d_window = 128 - bench_c_cap = (bench_l_local + bench_d_comp) // bench_ratio - bench_hidden = torch.randn(bench_l_local, 64, dtype=torch.bfloat16, device="cuda") - bench_boundary = torch.randn(bench_d_window, 64, dtype=torch.bfloat16, device="cuda") - bench_grad = torch.randn(bench_c_cap * bench_ratio, 64, dtype=torch.bfloat16, device="cuda") - - def run_cute(): - h = bench_hidden.detach().clone().requires_grad_(True) - b = bench_boundary.detach().clone().requires_grad_(True) - y = csa_cp_layout_kernels.CompressorInputCompact.apply( - h, - b, - bench_cu, - bench_global_start, - bench_l_local, - bench_ratio, - bench_d_comp, - bench_d_window, - bench_c_cap, - )[0] - y.backward(bench_grad) - - def run_torch(): - h = bench_hidden.detach().clone().requires_grad_(True) - b = bench_boundary.detach().clone().requires_grad_(True) - y = _native_compressor_input_compact( - h, - b, - bench_cu, - bench_global_start, - bench_l_local, - bench_ratio, - bench_d_comp, - bench_d_window, - bench_c_cap, - )[0] - y.backward(bench_grad) - - effective_bytes = 2 * (_num_bytes(bench_hidden, bench_boundary, bench_grad)) - _print_perf( - "compressor_input_compact_fwd_bwd", - _time_cuda(run_cute), - _time_cuda(run_torch), - effective_bytes, - ) - -def test_build_attention_indices_matches_native_and_reports_bandwidth(): +def test_build_attention_indices_matches_native(): _require_cute_cuda() cu = _make_e2e_like_cu_seqlens() - global_start, l_local = _e2e_like_contiguous_range() + global_start, l_local = _e2e_like_local_range() d_window = 128 window = 16 ratio = 4 compressed_width = 16 cu_comp = _compressed_cu_seqlens(cu, ratio) - rank_by_seq = torch.arange(int(cu_comp[-1]), dtype=torch.int32, device="cuda").flip(0) + seq_to_rank_row = torch.arange(int(cu_comp[-1]), dtype=torch.int32, device="cuda").flip(0) compressed_base = d_window + l_local fused = csa_cp_layout_kernels.build_attention_indices( cu, @@ -587,8 +311,7 @@ def test_build_attention_indices_matches_native_and_reports_bandwidth(): ratio, compressed_width, cu_seqlens_compressed=cu_comp, - rank_row_for_seq_row=rank_by_seq, - compressed_base=compressed_base, + seq_to_rank_row=seq_to_rank_row, ) expected = _native_attention_indices( cu, @@ -599,7 +322,7 @@ def test_build_attention_indices_matches_native_and_reports_bandwidth(): window, ratio, compressed_width, - rank_by_seq, + seq_to_rank_row, compressed_base, ) assert torch.equal(fused[0], expected[0]) @@ -616,8 +339,7 @@ def test_build_attention_indices_matches_native_and_reports_bandwidth(): logical_ids.shape[1], logical_ids, cu_seqlens_compressed=cu_comp, - rank_row_for_seq_row=rank_by_seq, - compressed_base=compressed_base, + seq_to_rank_row=seq_to_rank_row, ) expected_indexer = _native_attention_indices( cu, @@ -628,7 +350,7 @@ def test_build_attention_indices_matches_native_and_reports_bandwidth(): window, ratio, logical_ids.shape[1], - rank_by_seq, + seq_to_rank_row, compressed_base, logical_ids, ) @@ -644,58 +366,13 @@ def test_build_attention_indices_matches_native_and_reports_bandwidth(): ) assert torch.equal(padded[1][8:10], torch.tensor([1, 1], dtype=torch.int32, device="cuda")) - bench_l_local = 4096 - bench_window = 128 - bench_ratio = 4 - bench_compressed_width = bench_l_local // bench_ratio - bench_cu = torch.tensor([0, bench_l_local], dtype=torch.int32, device="cuda") - bench_cu_comp = torch.tensor([0, bench_compressed_width], dtype=torch.int32, device="cuda") - bench_rank_by_seq = torch.arange(bench_compressed_width, dtype=torch.int32, device="cuda") - bench_ref = _native_attention_indices_one_seq( - bench_l_local, - bench_window, - bench_window, - bench_ratio, - bench_compressed_width, - torch.device("cuda"), - ) - - def run_cute(): - csa_cp_layout_kernels.build_attention_indices( - bench_cu, - 0, - bench_l_local, - bench_window, - bench_window, - bench_ratio, - bench_compressed_width, - cu_seqlens_compressed=bench_cu_comp, - rank_row_for_seq_row=bench_rank_by_seq, - compressed_base=bench_window + bench_l_local, - ) - - def run_torch(): - _native_attention_indices_one_seq( - bench_l_local, - bench_window, - bench_window, - bench_ratio, - bench_compressed_width, - torch.device("cuda"), - ) - - effective_bytes = _num_bytes(*bench_ref) - _print_perf( - "build_attention_indices", _time_cuda(run_cute), _time_cuda(run_torch), effective_bytes - ) - -def test_build_indexer_loss_indices_matches_native_and_reports_bandwidth(): +def test_build_attention_indices_indexer_loss_mode_matches_native(): _require_cute_cuda() cu = _make_e2e_like_cu_seqlens() ratio = 4 cu_comp = _compressed_cu_seqlens(cu, ratio) - global_start, l_local = _e2e_like_contiguous_range() + global_start, l_local = _e2e_like_local_range() d_window = 128 window = 16 compressed_width = 8 @@ -705,19 +382,20 @@ def test_build_indexer_loss_indices_matches_native_and_reports_bandwidth(): .repeat(l_local, 1) ) logical_ids[1::5, -1] = -1 - rank_by_seq = torch.arange(int(cu_comp[-1]), dtype=torch.int32, device="cuda").flip(0) + seq_to_rank_row = torch.arange(int(cu_comp[-1]), dtype=torch.int32, device="cuda").flip(0) compressed_base = d_window + l_local - fused = csa_cp_layout_kernels.build_indexer_loss_indices( + fused = csa_cp_layout_kernels.build_attention_indices( cu, - cu_comp, global_start, l_local, d_window, window, ratio, + compressed_width, logical_ids, - rank_by_seq, - compressed_base=compressed_base, + cu_seqlens_compressed=cu_comp, + seq_to_rank_row=seq_to_rank_row, + for_indexer_loss=True, ) expected = _native_indexer_loss_indices( cu, @@ -728,48 +406,8 @@ def test_build_indexer_loss_indices_matches_native_and_reports_bandwidth(): window, ratio, logical_ids, - rank_by_seq, + seq_to_rank_row, compressed_base, ) assert torch.equal(fused[0], expected[0]) - assert torch.equal(fused[1], expected[1]) - - bench_l_local = 4096 - bench_window = 128 - bench_ratio = 4 - bench_compressed_width = bench_l_local // bench_ratio - bench_cu = torch.tensor([0, bench_l_local], dtype=torch.int32, device="cuda") - bench_cu_comp = torch.tensor([0, bench_compressed_width], dtype=torch.int32, device="cuda") - bench_logical = ( - torch.arange(bench_compressed_width, dtype=torch.int32, device="cuda") - .unsqueeze(0) - .expand(bench_l_local, -1) - ) - bench_rank_by_seq = torch.arange(bench_compressed_width, dtype=torch.int32, device="cuda") - bench_ref = _native_indexer_loss_indices_one_seq( - bench_l_local, bench_window, bench_window, bench_ratio, bench_logical, bench_rank_by_seq - ) - - def run_cute(): - csa_cp_layout_kernels.build_indexer_loss_indices( - bench_cu, - bench_cu_comp, - 0, - bench_l_local, - bench_window, - bench_window, - bench_ratio, - bench_logical, - bench_rank_by_seq, - compressed_base=bench_window + bench_l_local, - ) - - def run_torch(): - _native_indexer_loss_indices_one_seq( - bench_l_local, bench_window, bench_window, bench_ratio, bench_logical, bench_rank_by_seq - ) - - effective_bytes = _num_bytes(bench_logical, bench_rank_by_seq, *bench_ref) - _print_perf( - "build_indexer_loss_indices", _time_cuda(run_cute), _time_cuda(run_torch), effective_bytes - ) + assert torch.equal(fused[2], expected[1]) diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py index 698422265bf..519481dcd36 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py @@ -7,15 +7,10 @@ from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.transformer.experimental_attention_variant import csa_cp_utils from megatron.core.transformer.experimental_attention_variant.csa_cp_utils import ( - DSV4_CP_PARTITION_CONTIGUOUS, - DSV4_CP_PARTITION_TWO_CHUNK, apply_thd_cp_local_rope_fused, - build_cp_compressor_prep_compact_fused, - compute_cp_indexer_topk_logical_fused, + compute_cp_indexer_topk, exchange_cp_boundary_hidden, - local_kv_cp_chunk_ranges, - local_q_cp_chunk_ranges, - thd_cp_local_row_indices, + prepare_cp_compressor_input, ) from tests.unit_tests.test_utilities import Utils @@ -25,117 +20,29 @@ def _require_cuda(): pytest.skip("DSv4 CP CUDA utility tests require CUDA.") -def test_csa_cp_partition_mode_contract(): - """Validate the public CSA CP partition-mode names. - - Expected: omitted mode defaults to the original contiguous partition, - supported strings round-trip exactly, and unknown values fail before any CP - layout helper can silently choose the wrong row order. - """ - assert local_q_cp_chunk_ranges(None, 4, 1, 0) == ((0, 4),) - assert local_q_cp_chunk_ranges(DSV4_CP_PARTITION_CONTIGUOUS, 4, 1, 0) == ((0, 4),) - assert local_q_cp_chunk_ranges(DSV4_CP_PARTITION_TWO_CHUNK, 4, 1, 0) == ((0, 4),) - - with pytest.raises(RuntimeError, match="Unsupported CSA CP partition mode"): - local_q_cp_chunk_ranges("invalid_mode", 4, 1, 0) - - -@pytest.mark.parametrize( - "partition_mode, uses_two_chunk", - [(DSV4_CP_PARTITION_CONTIGUOUS, False), (DSV4_CP_PARTITION_TWO_CHUNK, True)], - ids=["contiguous", "two_chunk"], -) -def test_thd_cp_partition_mode_selects_expected_row_order(partition_mode, uses_two_chunk): - """Validate the CP row order selected by each partition mode. - - Expected: ``contiguous`` returns this rank's single local range, while - ``two_chunk`` returns chunk ``rank`` followed by chunk - ``2*cp_size-1-rank``. A failure means CP data slicing and layer-side layout - helpers could disagree about local packed-token order. - """ - local_rows = 16 - cp_size = 4 - - for cp_rank in range(cp_size): - ranges = local_q_cp_chunk_ranges(partition_mode, local_rows, cp_size, cp_rank) - if uses_two_chunk: - chunk_len = local_rows // 2 - total_chunks = 2 * cp_size - chunk_ids = (cp_rank, total_chunks - 1 - cp_rank) - expected = tuple( - (chunk_id * chunk_len, (chunk_id + 1) * chunk_len) for chunk_id in chunk_ids - ) - else: - expected = ((cp_rank * local_rows, (cp_rank + 1) * local_rows),) - assert ranges == expected - - -def test_two_chunk_cp_ranges_match_expected_order(): - """Validate the two-chunk CP partition contract. - - Expected: rank r owns global chunk r followed by chunk 2*cp_size-1-r. - A failure here means DSv4 chunk-aware CP helpers would disagree with the - two-chunk local row order. - """ - assert local_q_cp_chunk_ranges(DSV4_CP_PARTITION_TWO_CHUNK, 16, 1, 0) == ((0, 16),) - expected_by_rank = { - 0: ((0, 2), (14, 16)), - 1: ((2, 4), (12, 14)), - 2: ((4, 6), (10, 12)), - 3: ((6, 8), (8, 10)), - } - for rank, expected in expected_by_rank.items(): - assert local_q_cp_chunk_ranges(DSV4_CP_PARTITION_TWO_CHUNK, 4, 4, rank) == expected - - with pytest.raises(RuntimeError, match="even local_rows"): - local_q_cp_chunk_ranges(DSV4_CP_PARTITION_TWO_CHUNK, 5, 4, 0) +def _rope_reference(x, cos, sin, positions, nope_dim, pos_dim, inverse=False): + x_nope, x_pos = torch.split(x, [nope_dim, pos_dim], dim=-1) + cos_pos = cos.index_select(0, positions).view(x.shape[0], 1, pos_dim) + sin_pos = sin.index_select(0, positions).view(x.shape[0], 1, pos_dim) + if inverse: + sin_pos = -sin_pos + half = pos_dim // 2 + x1, x2 = x_pos[..., 0::2], x_pos[..., 1::2] + rotated = torch.stack( + ( + x1 * cos_pos[..., :half] - x2 * sin_pos[..., :half], + x2 * cos_pos[..., half:] + x1 * sin_pos[..., half:], + ), + dim=-1, + ).flatten(-2) + return torch.cat((x_nope, rotated), dim=-1) -def test_cp_chunk_ranges_match_partition_mode(): - """Validate local and left-boundary ranges for each CP partition mode. - - Expected: each mode returns the current rank's local rows and matching - left-boundary rows in global packed-token coordinates. - """ - assert local_q_cp_chunk_ranges(None, local_rows=4, cp_size=4, cp_rank=2) == ((8, 12),) - assert local_kv_cp_chunk_ranges( - DSV4_CP_PARTITION_CONTIGUOUS, local_rows=4, boundary_rows=2, cp_size=4, cp_rank=2 - ) == ((6, 8), (8, 12)) - - assert local_q_cp_chunk_ranges( - DSV4_CP_PARTITION_TWO_CHUNK, local_rows=4, cp_size=4, cp_rank=0 - ) == ((0, 2), (14, 16)) - assert local_kv_cp_chunk_ranges( - DSV4_CP_PARTITION_TWO_CHUNK, local_rows=4, boundary_rows=4, cp_size=4, cp_rank=0 - ) == ((-2, 0), (12, 14), (0, 2), (14, 16)) - - with pytest.raises(RuntimeError, match="boundary rows must be divisible"): - local_kv_cp_chunk_ranges( - DSV4_CP_PARTITION_TWO_CHUNK, local_rows=4, boundary_rows=3, cp_size=4, cp_rank=0 - ) - - -def test_thd_cp_local_row_indices_match_chunk_ranges(): - """Validate THD batch row selection for each CSA CP partition mode. - - Expected: data loading selects the same global packed-token rows that the - attention layer later uses for RoPE, boundary exchange, and index lowering. - """ - assert torch.equal( - thd_cp_local_row_indices(DSV4_CP_PARTITION_CONTIGUOUS, 16, 4, 2, torch.device("cpu")), - torch.tensor([8, 9, 10, 11]), - ) - assert torch.equal( - thd_cp_local_row_indices(DSV4_CP_PARTITION_TWO_CHUNK, 16, 4, 0, torch.device("cpu")), - torch.tensor([0, 1, 14, 15]), - ) - assert torch.equal( - thd_cp_local_row_indices(DSV4_CP_PARTITION_TWO_CHUNK, 16, 4, 3, torch.device("cpu")), - torch.tensor([6, 7, 8, 9]), - ) - - with pytest.raises(RuntimeError, match="divisible by cp_size"): - thd_cp_local_row_indices(DSV4_CP_PARTITION_CONTIGUOUS, 15, 4, 0, torch.device("cpu")) +def _sequence_positions(cu_seqlens, rows, clamp=False): + if clamp: + rows = rows.clamp(0, cu_seqlens[-1] - 1) + seq_ids = torch.bucketize(rows, cu_seqlens[1:], right=True).clamp_max(cu_seqlens.shape[0] - 2) + return rows - cu_seqlens[seq_ids] def test_boundary_exchange_single_rank_returns_zero_boundary_and_zero_grad(): @@ -147,23 +54,7 @@ def test_boundary_exchange_single_rank_returns_zero_boundary_and_zero_grad(): """ local = torch.arange(12, dtype=torch.float32).reshape(4, 3).requires_grad_(True) - boundary = exchange_cp_boundary_hidden(local, [], 2, DSV4_CP_PARTITION_CONTIGUOUS, None) - - assert torch.equal(boundary, torch.zeros(2, 3)) - boundary.sum().backward() - assert torch.equal(local.grad, torch.zeros_like(local)) - - -def test_two_chunk_boundary_exchange_single_rank_returns_zero_boundary_and_zero_grad(): - """Validate the no-CP two-chunk boundary exchange contract. - - Expected: with cp_group=None, the helper degenerates to one zero-filled - boundary and contributes no local gradient. A failure here means chunk-aware - CP boundary plumbing could perturb disabled-CP paths. - """ - local = torch.arange(12, dtype=torch.float32).reshape(4, 3).requires_grad_(True) - - boundary = exchange_cp_boundary_hidden(local, [], 2, DSV4_CP_PARTITION_TWO_CHUNK, None) + boundary = exchange_cp_boundary_hidden(local, 0, 2, None) assert torch.equal(boundary, torch.zeros(2, 3)) boundary.sum().backward() @@ -200,9 +91,7 @@ def test_thd_cp_left_boundary_exchange_forward_backward(): ).reshape(local_len, width) local = values.detach().clone().requires_grad_(True) - boundary = exchange_cp_boundary_hidden( - local, [], d_window, DSV4_CP_PARTITION_CONTIGUOUS, cp_group - ) + boundary = exchange_cp_boundary_hidden(local, 0, d_window, cp_group) if cp_rank == 0: expected_boundary = torch.zeros_like(boundary) else: @@ -224,78 +113,56 @@ def test_thd_cp_left_boundary_exchange_forward_backward(): Utils.destroy_model_parallel() -def test_apply_thd_cp_local_rope_fused_dispatches_equal_and_ragged_chunks(monkeypatch): - """Validate the public local-RoPE wrapper dispatch contract. - - Expected: equal one/two-chunk layouts call the fused kernel once with both - chunk starts, while ragged multi-chunk layouts call once per chunk and - concatenate the chunk results in local row order. - """ - calls = [] - - def fake_apply( - x, - cos, - sin, - cu, - chunk0_start, - chunk_len, - nope_dim, - pos_dim, - chunk1_start, - inverse, - clamp_to_valid_token, - ): - calls.append( - ( - tuple(x.shape), - int(chunk0_start), - int(chunk_len), - int(nope_dim), - int(pos_dim), - int(chunk1_start), - bool(inverse), - bool(clamp_to_valid_token), - ) - ) - return x + len(calls) * 10 - - monkeypatch.setattr( - csa_cp_utils.csa_cp_layout_kernels.ThdLocalRope, "apply", staticmethod(fake_apply) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("inverse", [False, True]) +def test_apply_thd_cp_local_rope_matches_reference_forward_backward(dtype, inverse): + _require_cuda() + torch.manual_seed(11) + cu = torch.tensor([0, 4, 12], dtype=torch.int32, device="cuda") + global_start = 2 + global_rows = torch.arange(2, 6, dtype=torch.int32, device="cuda") + positions = _sequence_positions(cu, global_rows) + nope_dim = pos_dim = 4 + x = torch.randn(4, 2, 8, dtype=dtype, device="cuda") + cos = torch.randn(8, pos_dim, dtype=dtype, device="cuda") + sin = torch.randn(8, pos_dim, dtype=dtype, device="cuda") + + ref_x = x.detach().clone().requires_grad_(True) + expected = _rope_reference(ref_x, cos, sin, positions, nope_dim, pos_dim, inverse) + grad = torch.randn_like(expected) + expected.backward(grad) + + actual_x = x.detach().clone().requires_grad_(True) + actual = apply_thd_cp_local_rope_fused( + actual_x, cos, sin, nope_dim, pos_dim, cu, global_start, inverse=inverse ) - cos = sin = torch.empty(8, 2) - cu = torch.tensor([0, 8], dtype=torch.int32) + actual.backward(grad) + rtol, atol = (1e-5, 1e-5) if dtype == torch.float32 else (2e-2, 5e-2) + torch.testing.assert_close(actual, expected, rtol=rtol, atol=atol) + torch.testing.assert_close(actual_x.grad, ref_x.grad, rtol=rtol, atol=atol) - equal = torch.arange(8, dtype=torch.float32).reshape(4, 2) - out = apply_thd_cp_local_rope_fused(equal, cos, sin, 1, 1, cu, ((2, 4), (10, 12)), inverse=True) - assert torch.equal(out, equal + 10) - assert calls[-1] == ((4, 2), 2, 2, 1, 1, 10, True, False) - ragged = torch.arange(10, dtype=torch.float32).reshape(5, 2) - out = apply_thd_cp_local_rope_fused( - ragged, cos, sin, 1, 1, cu, ((2, 4), (10, 13)), clamp_to_valid_token=True +def test_apply_thd_cp_local_rope_clamps_boundary_rows(): + _require_cuda() + cu = torch.tensor([0, 4, 12], dtype=torch.int32, device="cuda") + global_start = -2 + global_rows = torch.arange(-2, 14, dtype=torch.int32, device="cuda") + positions = _sequence_positions(cu, global_rows, clamp=True) + x = torch.randn(global_rows.shape[0], 1, 8, device="cuda") + cos = torch.randn(8, 4, device="cuda") + sin = torch.randn(8, 4, device="cuda") + expected = _rope_reference(x, cos, sin, positions, 4, 4) + actual = apply_thd_cp_local_rope_fused( + x, cos, sin, 4, 4, cu, global_start, clamp_to_valid_token=True ) - assert torch.equal(out, torch.cat((ragged[:2] + 20, ragged[2:] + 30), dim=0)) - assert calls[-2:] == [ - ((2, 2), 2, 2, 1, 1, 0, False, True), - ((3, 2), 10, 3, 1, 1, 0, False, True), - ] + torch.testing.assert_close(actual, expected) - with pytest.raises(RuntimeError, match="expects x rows"): - apply_thd_cp_local_rope_fused(equal[:3], cos, sin, 1, 1, cu, ((2, 4), (10, 12))) - -def test_compressor_prep_compact_wrapper_merges_two_chunk_outputs(monkeypatch): - """Validate public compressor-prep wrapper splitting, merge, and padding. - - Expected: two chunks are compacted independently, local compact group ids - are concatenated in local chunk order, and the wrapper pads to the shared - aligned capacity used by later fixed-shape collectives. - """ +def test_prepare_cp_compressor_input_builds_rank_row_map(monkeypatch): calls = [] - def fake_apply(hidden, boundary, cu, global_start, rows, ratio, d_comp, d_window, c_cap): - calls.append((int(global_start), int(rows), int(c_cap), tuple(boundary.shape))) + def fake_apply(hidden, boundary, cu, global_start, ratio, d_comp, c_cap): + calls.append((int(global_start), hidden.shape[0], int(c_cap), tuple(boundary.shape))) marker = len(calls) hidden_compact = torch.full((int(c_cap) * int(ratio), 1), marker, dtype=hidden.dtype) comp_ids = torch.arange(int(c_cap), dtype=torch.int32) + marker * 10 @@ -305,58 +172,28 @@ def fake_apply(hidden, boundary, cu, global_start, rows, ratio, d_comp, d_window csa_cp_utils.csa_cp_layout_kernels.CompressorInputCompact, "apply", staticmethod(fake_apply) ) hidden = torch.arange(16, dtype=torch.float32).reshape(16, 1) - boundary = torch.arange(4, dtype=torch.float32).reshape(4, 1) + boundary = torch.arange(2, dtype=torch.float32).reshape(2, 1) cu = torch.tensor([0, 32], dtype=torch.int32) cu_comp = torch.tensor([0, 8], dtype=torch.int32) - hidden_compact, comp_ids, rank_rows = build_cp_compressor_prep_compact_fused( - hidden, boundary, cu, cu_comp, ((0, 8), (24, 32)), cp_size=2, ratio=4, d_comp=8, d_window=2 + hidden_compact, comp_ids, rank_rows = prepare_cp_compressor_input( + hidden, boundary, cu, cu_comp, 0, cp_size=2, ratio=4 ) - assert calls == [(0, 8, 4, (2, 1)), (24, 8, 4, (2, 1))] + assert calls == [(0, 16, 8, (2, 1))] assert hidden_compact.shape == (32, 1) - assert torch.equal(comp_ids, torch.tensor([10, 11, 12, 13, 20, 21, 22, 23])) - assert torch.equal(rank_rows, torch.tensor([0, 1, 10, 11, 14, 15, 6, 7], dtype=torch.int32)) + assert torch.equal(comp_ids, torch.arange(10, 18, dtype=torch.int32)) + assert torch.equal(rank_rows, torch.tensor([0, 1, 2, 3, 10, 11, 12, 13], dtype=torch.int32)) -def test_compute_cp_indexer_topk_logical_fused_splits_chunks_and_passes_visible_lengths( - monkeypatch, -): - """Validate public CP indexer-topk wrapper chunk splitting. - - Expected: each local chunk gets its own trapezoid metadata and indexer-topk - call, visible K lengths are passed through, and chunk outputs are - concatenated in local row order. - """ +def test_compute_cp_indexer_topk_passes_valid_lengths(monkeypatch): topk_calls = [] def fake_indexer_topk( - q, - k, - weights, - *, - topk, - ratio, - indexer_softmax_scale, - cu_seqlens_q, - cu_seqlens_kv, - max_seqlen_q, - max_seqlen_kv, - visible_k_lengths, + q, k, _weights, *, topk, cu_seqlens_q, cu_seqlens_kv, valid_k_lengths, **_ ): topk_calls.append( - ( - tuple(q.shape), - int(topk), - int(ratio), - float(indexer_softmax_scale), - int(max_seqlen_q), - int(max_seqlen_kv), - k.clone(), - cu_seqlens_q.clone(), - cu_seqlens_kv.clone(), - visible_k_lengths.clone(), - ) + (k.clone(), cu_seqlens_q.clone(), cu_seqlens_kv.clone(), valid_k_lengths.clone()) ) return torch.full((q.shape[0], int(topk)), len(topk_calls), dtype=torch.int32), None @@ -368,34 +205,23 @@ def fake_indexer_topk( cu_q = torch.tensor([0, 10], dtype=torch.int32) cu_comp = torch.tensor([0, 3], dtype=torch.int32) - out = compute_cp_indexer_topk_logical_fused( + out = compute_cp_indexer_topk( q, weights, k_seq, cu_q, cu_comp, - ((2, 4), (8, 11)), + 8, ratio=4, topk_width=2, indexer_softmax_scale=0.5, max_seqlen_q=10, - max_seqlen_kv=3, ) - assert torch.equal(out, torch.tensor([[1, 1], [1, 1], [2, 2], [2, 2], [2, 2]])) - assert torch.equal(topk_calls[0][-4][0], k_seq[0]) - assert torch.count_nonzero(topk_calls[0][-4][1:]) == 0 - assert torch.equal(topk_calls[0][-3], torch.tensor([0, 2, 2], dtype=torch.int32)) - assert torch.equal(topk_calls[0][-2], torch.tensor([0, 1, 1], dtype=torch.int32)) - assert torch.equal(topk_calls[0][-1], torch.ones(2, dtype=torch.int32)) - assert torch.equal(topk_calls[1][-4][:2], k_seq[:2]) - assert torch.count_nonzero(topk_calls[1][-4][2:]) == 0 - assert torch.equal(topk_calls[1][-3], torch.tensor([0, 2, 3], dtype=torch.int32)) - assert torch.equal(topk_calls[1][-2], torch.tensor([0, 2, 2], dtype=torch.int32)) - assert torch.equal(topk_calls[1][-1], torch.tensor([2, 2, 0], dtype=torch.int32)) - assert ( - compute_cp_indexer_topk_logical_fused( - q, weights, k_seq[:0], cu_q, cu_comp, ((2, 4),), 4, 2, 1.0, 10, 3 - ) - is None - ) + assert torch.equal(out, torch.ones(5, 2, dtype=torch.int32)) + assert torch.equal(topk_calls[0][0][:2], k_seq[:2]) + assert torch.count_nonzero(topk_calls[0][0][2:]) == 0 + assert torch.equal(topk_calls[0][1], torch.tensor([0, 2, 5], dtype=torch.int32)) + assert torch.equal(topk_calls[0][2], torch.tensor([0, 2, 2], dtype=torch.int32)) + assert torch.equal(topk_calls[0][3], torch.tensor([2, 2, 0, 0, 0], dtype=torch.int32)) + assert compute_cp_indexer_topk(q, weights, k_seq[:0], cu_q, cu_comp, 2, 4, 2, 1.0, 10) is None diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py index 096c73d1451..9a87b8826ae 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py @@ -1632,22 +1632,6 @@ def _ref_attn_full_score( return torch.where(valid, s, torch.zeros_like(s)) -def _ref_indexer_predict_sparse(q_bshd_fp32, k_bsd_fp32, w_bsh_fp32, topk_indices, sm_scale): - """Reference for ``sparse_indexer_score_recompute_wrapper.predict``. - - Compute the full-KV indexer score, gather ``topk_indices``, softmax - over the topK axis. ``-1`` entries in topk are masked to ``-inf`` - so they contribute zero probability. - """ - qk = torch.einsum('bqhd,bkd->bqhk', q_bshd_fp32, k_bsd_fp32) - s = (torch.relu(qk) * w_bsh_fp32.unsqueeze(-1)).sum(dim=2) * sm_scale # (B, Sq, Sk) - valid = topk_indices >= 0 - safe = topk_indices.clamp(min=0).long() - s_topk = torch.gather(s, dim=-1, index=safe) - s_topk = torch.where(valid, s_topk, torch.full_like(s_topk, float('-inf'))) - return torch.softmax(s_topk, dim=-1) - - def _ref_attn_target_sparse(q_bshd_fp32, k_bsd_fp32, lse_bsh_fp32, topk_indices, softmax_scale): """Reference for ``sparse_attn_score_recompute_wrapper.target``. @@ -1835,17 +1819,16 @@ def _build_real_score_inputs(s, *, with_lse: bool = True, with_topk: bool = True class TestRealKernelScoreHelpers: - """Real-kernel parity tests for the four ``_compute_*`` score helpers + """Real-kernel parity tests for three ``_compute_*`` score helpers against PyTorch reference implementations. Single parametrized test - covers all four; numeric tolerance is bf16-friendly (raw fp32 score + covers all three; numeric tolerance is bf16-friendly (raw fp32 score sums agree to ~5%, normalized distributions to ~5e-3). """ # Each case: (id, sm_min, kernel_name, runner). The runner does the # call + ref + assertion; it returns nothing on success. @pytest.mark.parametrize( - "case", - ['sparse_indexer_predict', 'sparse_attn_target', 'dense_indexer_score', 'dense_attn_score'], + "case", ['sparse_attn_target', 'dense_indexer_score', 'dense_attn_score'] ) def test_real_score_helper(self, case, reset_lazy_kernel_state): _skip_if_real_kernels_unavailable(sm_min=10) @@ -1860,29 +1843,7 @@ def test_real_score_helper(self, case, reset_lazy_kernel_state): from megatron.core.transformer.experimental_attention_variant import dsa_kernels as _dk - if case == 'sparse_indexer_predict': - # The kernel takes sm_scale=1.0; scale is applied via weights - # pre-multiplication (relu(c·x)·W trick). Reference mirrors that. - scale = s['indexer_softmax_scale'] - w_scaled = (x['w'].float() * scale).to(x['w'].dtype) - out = _dk._compute_indexer_predict( - x['q_idx'], x['k_idx'], w_scaled, x['topk'], qhead_per_kv_head=s['idx_nh'] - ) - ref = _ref_indexer_predict_sparse( - x['q_idx'].float(), x['k_idx'].float(), w_scaled.float(), x['topk'], sm_scale=1.0 - ) - # Softmax outputs in [0, 1]; bf16 element-wise noise can break - # absolute tolerance, so compare directions via cosine similarity. - assert out.shape == ref.shape == (s['b'], s['sq'], s['topk']) - cos = torch.nn.functional.cosine_similarity( - out.flatten().unsqueeze(0).float(), ref.flatten().unsqueeze(0).float() - ).item() - assert cos > 0.99, ( - f"{case}: cos sim = {cos:.4f}, " - f"max abs diff = {(out - ref).abs().max().item():.3e}" - ) - - elif case == 'sparse_attn_target': + if case == 'sparse_attn_target': out = _dk._compute_attn_target( x['q_attn'], x['k_attn'], @@ -2817,7 +2778,7 @@ def test_thd_single_segment_matches_sbhd_b1(self, sparse_loss, reset_lazy_kernel cu_comp_idx = _make_cu_seqlens([s['n_comp']], device=dev) # B=1: per-segment [kv, compressed] layout collapses to a single - # contiguous slice — same kv_offset as the SBHD case. + # Ordered slice with the same kv_offset as the SBHD case. compressed_kv_thd = kv_full_thd[kv_offset:] _, loss_thd = fused_indexer_sparse_attn( query_thd, diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention.py index 725488f6122..155cd940ffe 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention.py @@ -6,11 +6,9 @@ import torch import torch.nn.functional as F -import megatron.core.parallel_state as parallel_state from megatron.core.extensions.transformer_engine import HAVE_TE from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed -from megatron.core.transformer.enums import AttnMaskType from megatron.core.transformer.transformer_config import MLATransformerConfig from tests.unit_tests.test_utilities import Utils @@ -357,7 +355,7 @@ def test_qkv_shapes(self): seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 ).cuda() - q, k, v, q_compressed, kv_compressed, _ = attn.get_query_key_value_tensors(hidden) + q, k, v, _, _ = attn.get_query_key_value_tensors(hidden) n_heads = self.config.num_attention_heads v_dim = self.config.v_head_dim @@ -380,7 +378,7 @@ def test_key_equals_value(self): seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 ).cuda() - q, k, v, _, _, _ = attn.get_query_key_value_tensors(hidden) + q, k, v, _, _ = attn.get_query_key_value_tensors(hidden) assert torch.equal(k, v), "key and value should be identical in wkv path" diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py index 2a7468ace7e..e2ffbde272f 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py @@ -17,11 +17,6 @@ from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed -from megatron.core.transformer.experimental_attention_variant.csa_cp_utils import ( - DSV4_CP_PARTITION_CONTIGUOUS, - DSV4_CP_PARTITION_TWO_CHUNK, - thd_cp_local_row_indices, -) from tests.unit_tests.test_utilities import Utils from tests.unit_tests.transformer.experimental_attention_variant.test_dsv4_hybrid_attention import ( _SEED, @@ -270,7 +265,6 @@ def _make_dsv4_cp_config( dsa_indexer_use_sparse_loss=True, apply_dsa_kernel_fusion=True, apply_rope_fusion=True, - csa_cp_partition_mode=DSV4_CP_PARTITION_CONTIGUOUS, ): """Build the DSv4 flash attention config used by CP tests.""" shape = _DSV4_VARIANTS[_DSV4_CP_TEST_VARIANT] @@ -299,7 +293,6 @@ def _make_dsv4_cp_config( expert_model_parallel_size=1, apply_dsa_kernel_fusion=apply_dsa_kernel_fusion, apply_rope_fusion=apply_rope_fusion, - csa_cp_partition_mode=csa_cp_partition_mode, ) @@ -311,20 +304,13 @@ def _copy_module_parameters(src, dst): param.data.copy_(src_params[name].data) -def _make_cp_partition_indices(partition_mode, padded_total_tokens, cp_size, device='cuda'): - """Return local row indices for every CP rank.""" - return tuple( - thd_cp_local_row_indices(partition_mode, padded_total_tokens, cp_size, rank, device) - for rank in range(cp_size) - ) - - -def _make_ragged_cp_case(partition_mode, cp_size, cp_rank): +def _make_ragged_cp_case(cp_size, cp_rank): """Build the ragged THD packed params and local rows for one CP rank.""" padded_total_tokens = sum(_DSV4_CP_RAGGED_PADDED_SEG_LENS) packed = _make_thd_packed_seq_params(_DSV4_CP_RAGGED_SEG_LENS, _DSV4_CP_RAGGED_PADDED_SEG_LENS) - partition_indices = _make_cp_partition_indices(partition_mode, padded_total_tokens, cp_size) - return packed, padded_total_tokens, partition_indices[cp_rank] + local_rows = padded_total_tokens // cp_size + local_indices = torch.arange(cp_rank * local_rows, (cp_rank + 1) * local_rows, device='cuda') + return packed, padded_total_tokens, local_indices def _make_hidden_and_grad(padded_total_tokens, hidden_size): @@ -445,13 +431,10 @@ def _measure_cuda_graph_time(attn, static_hidden, static_grad, packed_seq_params return _max_float_across_world(statistics.median(timings)) -def _format_scale_report( - metric, layer_number, cp_size, partition_mode, baseline, measured, scale, limit -): +def _format_scale_report(metric, layer_number, cp_size, baseline, measured, scale, limit): """Format a CP-vs-CP1 scale report line.""" return ( f"DSv4 THD CP {metric} scale layer={layer_number} cp{cp_size}: " - f"partition={partition_mode}, " f"cp1={baseline:.6f}, cp{cp_size}={measured:.6f}, " f"scale={scale:.6f}, limit={limit:.6f}" ) @@ -564,23 +547,14 @@ def _measure_cp1_cuda_graph_time(self, layer_number, packed, padded_tokens): [1, 2, 3], ids=["ratio_0_window_only", "ratio_4_indexer", "ratio_128_compressor"], ) - @pytest.mark.parametrize( - "partition_mode", - [DSV4_CP_PARTITION_CONTIGUOUS, DSV4_CP_PARTITION_TWO_CHUNK], - ids=["contiguous", "two_chunk"], - ) - def test_thd_cp_matches_full_reference_forward_backward(self, layer_number, partition_mode): + def test_thd_cp_matches_full_reference_forward_backward(self, layer_number): """CP path matches the full-sequence THD reference on ragged packed inputs using the DSv4 layer configuration. Verifies local CP output, hidden grad, and reduced parameter grads - against the sliced full-reference tensors. The test body is shared by - contiguous and two-chunk CP; only the partition mode and - local row indices differ. + against the sliced full-reference tensors. """ - packed, padded_tokens, local_idx = _make_ragged_cp_case( - partition_mode, self.cp_size, self.cp_rank - ) + packed, padded_tokens, local_idx = _make_ragged_cp_case(self.cp_size, self.cp_rank) torch.manual_seed(_SEED + layer_number) model_parallel_cuda_manual_seed(_SEED + layer_number) @@ -590,7 +564,6 @@ def test_thd_cp_matches_full_reference_forward_backward(self, layer_number, part dsa_indexer_use_sparse_loss=True, apply_dsa_kernel_fusion=self.fused_kernels_available, apply_rope_fusion=True, - csa_cp_partition_mode=partition_mode, ) config_ref = _make_dsv4_cp_config( context_parallel_size=1, @@ -624,7 +597,7 @@ def test_thd_cp_matches_full_reference_forward_backward(self, layer_number, part _assert_cp_tensor_match( local_out.detach(), ref_out.detach().index_select(0, local_idx), - f"layer={layer_number}:{partition_mode}:output", + f"layer={layer_number}:output", ) grad = torch.randn_like(ref_out) @@ -633,7 +606,7 @@ def test_thd_cp_matches_full_reference_forward_backward(self, layer_number, part _assert_cp_tensor_match( local_hidden.grad.detach(), ref_hidden.grad.index_select(0, local_idx), - f"layer={layer_number}:{partition_mode}:hidden_grad", + f"layer={layer_number}:hidden_grad", ) ref_params = dict(ref_attn.named_parameters()) @@ -643,19 +616,14 @@ def test_thd_cp_matches_full_reference_forward_backward(self, layer_number, part assert ref_grad is not None, f"Missing reference grad for {name}" grad_sum = param.grad.detach().clone() dist.all_reduce(grad_sum, group=self.pg.cp) - _assert_cp_tensor_match( - grad_sum, ref_grad, f"layer={layer_number}:{partition_mode}:param_grad:{name}" - ) + _assert_cp_tensor_match(grad_sum, ref_grad, f"layer={layer_number}:param_grad:{name}") del cp_attn, ref_attn, full_hidden, local_hidden, ref_hidden, local_out, ref_out, grad _clear_cuda_test_state() def test_thd_cp_ratio4_eval_matches_full_reference(self): - """Ratio-4 two-chunk inference lowers logical indexer top-k rows correctly.""" - partition_mode = DSV4_CP_PARTITION_TWO_CHUNK - packed, padded_tokens, local_idx = _make_ragged_cp_case( - partition_mode, self.cp_size, self.cp_rank - ) + """Ratio-4 inference lowers logical indexer top-k rows correctly.""" + packed, padded_tokens, local_idx = _make_ragged_cp_case(self.cp_size, self.cp_rank) torch.manual_seed(_SEED + 1202) model_parallel_cuda_manual_seed(_SEED + 1202) @@ -663,7 +631,6 @@ def test_thd_cp_ratio4_eval_matches_full_reference(self): context_parallel_size=self.cp_size, apply_dsa_kernel_fusion=self.fused_kernels_available, apply_rope_fusion=True, - csa_cp_partition_mode=partition_mode, ) config_ref = _make_dsv4_cp_config( context_parallel_size=1, @@ -688,7 +655,7 @@ def test_thd_cp_ratio4_eval_matches_full_reference(self): hidden_states=full_hidden, attention_mask=None, packed_seq_params=packed ) _assert_cp_tensor_match( - local_out, ref_out.index_select(0, local_idx), f"layer=2:eval:{partition_mode}:output" + local_out, ref_out.index_select(0, local_idx), "layer=2:eval:output" ) del cp_attn, ref_attn, full_hidden, local_hidden, local_out, ref_out @@ -700,14 +667,7 @@ def test_thd_cp_ratio4_eval_matches_full_reference(self): ids=["ratio_0_window_only", "ratio_4_indexer", "ratio_128_compressor"], ) @pytest.mark.parametrize("fused", [True, False]) - @pytest.mark.parametrize( - "partition_mode", - [DSV4_CP_PARTITION_CONTIGUOUS, DSV4_CP_PARTITION_TWO_CHUNK], - ids=["contiguous", "two_chunk"], - ) - def test_thd_cp_cuda_graph_matches_eager_forward_backward( - self, layer_number, fused, partition_mode - ): + def test_thd_cp_cuda_graph_matches_eager_forward_backward(self, layer_number, fused): """CUDA graph replay matches eager THD CP forward/backward. Captures the DSv4 attention layer's CP-local forward and backward @@ -725,9 +685,7 @@ def test_thd_cp_cuda_graph_matches_eager_forward_backward( context = nullcontext() if fused else _deterministic_torch_algorithms() mode = "fused" if fused else "unfused" with context: - packed, padded_tokens, local_idx = _make_ragged_cp_case( - partition_mode, self.cp_size, self.cp_rank - ) + packed, padded_tokens, local_idx = _make_ragged_cp_case(self.cp_size, self.cp_rank) torch.manual_seed(_SEED + 700 + layer_number) model_parallel_cuda_manual_seed(_SEED + 700 + layer_number) @@ -737,7 +695,6 @@ def test_thd_cp_cuda_graph_matches_eager_forward_backward( dsa_indexer_use_sparse_loss=True, apply_dsa_kernel_fusion=fused, apply_rope_fusion=True, - csa_cp_partition_mode=partition_mode, ) graph_attn = _build_attention( config, layer_number=layer_number, pg_collection=self.pg @@ -784,21 +741,19 @@ def test_thd_cp_cuda_graph_matches_eager_forward_backward( ) torch.cuda.synchronize() assert graph_param_grads.keys() == eager_param_grads.keys() - output_label = f"layer={layer_number}:{mode}:{partition_mode}:output" + output_label = f"layer={layer_number}:{mode}:output" _assert_cp_graph_bitwise_match(graph_out, eager_out, output_label) bwd_match_fn = ( _assert_cp_graph_fused_grad_match if fused else _assert_cp_graph_bitwise_match ) bwd_match_fn( - graph_hidden_grad, - eager_hidden_grad, - f"layer={layer_number}:{mode}:{partition_mode}:hidden_grad", + graph_hidden_grad, eager_hidden_grad, f"layer={layer_number}:{mode}:hidden_grad" ) for name, graph_grad in graph_param_grads.items(): bwd_match_fn( graph_grad, eager_param_grads[name], - f"layer={layer_number}:{mode}:{partition_mode}:param_grad:{name}", + f"layer={layer_number}:{mode}:param_grad:{name}", ) del graph, graph_output @@ -809,14 +764,7 @@ def test_thd_cp_cuda_graph_matches_eager_forward_backward( @pytest.mark.parametrize( "layer_number", [2, 3], ids=["ratio_4_indexer", "ratio_128_compressor"] ) - @pytest.mark.parametrize( - "partition_mode", - [DSV4_CP_PARTITION_CONTIGUOUS, DSV4_CP_PARTITION_TWO_CHUNK], - ids=["contiguous", "two_chunk"], - ) - def test_thd_cp_cuda_graph_replay_accepts_changed_padded_boundaries( - self, layer_number, partition_mode - ): + def test_thd_cp_cuda_graph_replay_accepts_changed_padded_boundaries(self, layer_number): """CUDA graph replay uses updated device cu_seqlens_padded values. The graph is captured with one padded THD layout, then replayed after @@ -844,9 +792,10 @@ def test_thd_cp_cuda_graph_replay_accepts_changed_padded_boundaries( assert not torch.equal( capture_packed.cu_seqlens_q_padded, replay_packed.cu_seqlens_q_padded ) - local_idx = _make_cp_partition_indices(partition_mode, padded_tokens, self.cp_size)[ - self.cp_rank - ] + local_rows = padded_tokens // self.cp_size + local_idx = torch.arange( + self.cp_rank * local_rows, (self.cp_rank + 1) * local_rows, device='cuda' + ) torch.manual_seed(_SEED + 1100 + layer_number) model_parallel_cuda_manual_seed(_SEED + 1100 + layer_number) @@ -856,7 +805,6 @@ def test_thd_cp_cuda_graph_replay_accepts_changed_padded_boundaries( dsa_indexer_use_sparse_loss=True, apply_dsa_kernel_fusion=True, apply_rope_fusion=True, - csa_cp_partition_mode=partition_mode, ) graph_attn = _build_attention( config, layer_number=layer_number, pg_collection=self.pg @@ -911,18 +859,18 @@ def test_thd_cp_cuda_graph_replay_accepts_changed_padded_boundaries( torch.cuda.synchronize() assert graph_param_grads.keys() == eager_param_grads.keys() _assert_cp_graph_bitwise_match( - graph_out, eager_out, f"layer={layer_number}:metadata_replay:{partition_mode}:output" + graph_out, eager_out, f"layer={layer_number}:metadata_replay:output" ) _assert_cp_graph_fused_grad_match( graph_hidden_grad, eager_hidden_grad, - f"layer={layer_number}:metadata_replay:{partition_mode}:hidden_grad", + f"layer={layer_number}:metadata_replay:hidden_grad", ) for name, graph_grad in graph_param_grads.items(): _assert_cp_graph_fused_grad_match( graph_grad, eager_param_grads[name], - f"layer={layer_number}:metadata_replay:{partition_mode}:param_grad:{name}", + f"layer={layer_number}:metadata_replay:param_grad:{name}", ) del graph, graph_output, graph_attn, eager_attn, capture_full_hidden, replay_full_hidden @@ -935,12 +883,7 @@ def test_thd_cp_cuda_graph_replay_accepts_changed_padded_boundaries( [1, 2, 3], ids=["ratio_0_window_only", "ratio_4_indexer", "ratio_128_compressor"], ) - @pytest.mark.parametrize( - "partition_mode", - [DSV4_CP_PARTITION_CONTIGUOUS, DSV4_CP_PARTITION_TWO_CHUNK], - ids=["contiguous", "two_chunk"], - ) - def test_thd_cp_peak_allocated_delta_scales_vs_cp1(self, layer_number, partition_mode): + def test_thd_cp_peak_allocated_delta_scales_vs_cp1(self, layer_number): """CP-local THD forward/backward peak allocated delta scales down vs CP1. This is a black-box guard against implementations that pass parity but @@ -950,9 +893,7 @@ def test_thd_cp_peak_allocated_delta_scales_vs_cp1(self, layer_number, partition if not self.fused_kernels_available: pytest.skip(_DSV4_CP_FUSED_KERNELS_UNAVAILABLE_REASON) - packed, padded_tokens, local_idx = _make_ragged_cp_case( - partition_mode, self.cp_size, self.cp_rank - ) + packed, padded_tokens, local_idx = _make_ragged_cp_case(self.cp_size, self.cp_rank) ref_delta = self._measure_cp1_peak_allocated_delta(layer_number, packed, padded_tokens) @@ -963,7 +904,6 @@ def test_thd_cp_peak_allocated_delta_scales_vs_cp1(self, layer_number, partition dsa_indexer_loss_coeff=1.0, dsa_indexer_use_sparse_loss=True, apply_rope_fusion=True, - csa_cp_partition_mode=partition_mode, ) cp_attn = _build_attention( config_cp, layer_number=layer_number, pg_collection=self.pg @@ -984,7 +924,6 @@ def test_thd_cp_peak_allocated_delta_scales_vs_cp1(self, layer_number, partition "memory_delta_mib", layer_number, self.cp_size, - partition_mode, ref_delta / (1024.0**2), cp_delta / (1024.0**2), scale, @@ -1001,12 +940,7 @@ def test_thd_cp_peak_allocated_delta_scales_vs_cp1(self, layer_number, partition [1, 2, 3], ids=["ratio_0_window_only", "ratio_4_indexer", "ratio_128_compressor"], ) - @pytest.mark.parametrize( - "partition_mode", - [DSV4_CP_PARTITION_CONTIGUOUS, DSV4_CP_PARTITION_TWO_CHUNK], - ids=["contiguous", "two_chunk"], - ) - def test_thd_cp_cuda_graph_time_scales_vs_cp1(self, layer_number, partition_mode): + def test_thd_cp_cuda_graph_time_scales_vs_cp1(self, layer_number): """CP-local CUDA graph replay time scales down vs CP1 on the same machine. This is a black-box performance guard: correctness and memory scaling @@ -1016,9 +950,7 @@ def test_thd_cp_cuda_graph_time_scales_vs_cp1(self, layer_number, partition_mode if not self.fused_kernels_available: pytest.skip(_DSV4_CP_FUSED_KERNELS_UNAVAILABLE_REASON) - packed, padded_tokens, local_idx = _make_ragged_cp_case( - partition_mode, self.cp_size, self.cp_rank - ) + packed, padded_tokens, local_idx = _make_ragged_cp_case(self.cp_size, self.cp_rank) ref_time_ms = self._measure_cp1_cuda_graph_time(layer_number, packed, padded_tokens) @@ -1029,7 +961,6 @@ def test_thd_cp_cuda_graph_time_scales_vs_cp1(self, layer_number, partition_mode dsa_indexer_loss_coeff=1.0, dsa_indexer_use_sparse_loss=True, apply_rope_fusion=True, - csa_cp_partition_mode=partition_mode, ) cp_attn = _build_attention( config_cp, layer_number=layer_number, pg_collection=self.pg @@ -1047,14 +978,7 @@ def test_thd_cp_cuda_graph_time_scales_vs_cp1(self, layer_number, partition_mode limit = _DSV4_CP_GRAPH_TIME_RATIO_LIMITS[self.cp_size] scale = cp_time_ms / ref_time_ms report = _format_scale_report( - "graph_time_ms", - layer_number, - self.cp_size, - partition_mode, - ref_time_ms, - cp_time_ms, - scale, - limit, + "graph_time_ms", layer_number, self.cp_size, ref_time_ms, cp_time_ms, scale, limit ) print(report) assert scale <= limit, report From cea7032bc35737b5e6a68a9841b33c8652149db3 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Sat, 27 Jun 2026 10:36:52 -0700 Subject: [PATCH 07/23] Keep DSv4 CP changes scoped Signed-off-by: kunlunl --- megatron/core/datasets/data_schedule.py | 1 + megatron/core/datasets/data_schedule_utils.py | 4 +- .../core/fusions/fused_mla_yarn_rope_apply.py | 3 +- .../experimental_attention_variant/csa.py | 111 ++++++++++++++---- .../csa_cp_utils.py | 12 +- .../deepseek_v4_hybrid_attention.py | 86 +++++++++----- .../dsa_kernels.py | 66 ++++++++++- tests/unit_tests/test_sequence_packing.py | 2 +- .../test_csa_cp_utils.py | 16 --- .../test_dsa_kernels.py | 49 +++++++- .../test_dsv4_hybrid_attention.py | 4 +- 11 files changed, 262 insertions(+), 92 deletions(-) diff --git a/megatron/core/datasets/data_schedule.py b/megatron/core/datasets/data_schedule.py index 00ca7e59cc2..7e6c0c626e1 100644 --- a/megatron/core/datasets/data_schedule.py +++ b/megatron/core/datasets/data_schedule.py @@ -647,6 +647,7 @@ def get_batch_on_this_rank_for_sequence_packing( use_dsv4_cp_slice = ( config is not None and getattr(config, 'experimental_attention_variant', None) == "dsv4_hybrid" + and cp_group.size() > 1 ) dsv4_cp_local_target_len = None pad_alignment = ( diff --git a/megatron/core/datasets/data_schedule_utils.py b/megatron/core/datasets/data_schedule_utils.py index 88e4b78394d..87374e470d5 100644 --- a/megatron/core/datasets/data_schedule_utils.py +++ b/megatron/core/datasets/data_schedule_utils.py @@ -32,7 +32,7 @@ def get_cp_slice_for_thd( before slicing. Existing cu_seqlens metadata is left unchanged. """ cp_size = cp_group.size() - if cp_size <= 1 and partition_total_tokens is None: + if cp_size <= 1: return cp_rank = cp_group.rank() # Partition with padded cumulative lengths so CP slices match the THD @@ -61,8 +61,6 @@ def get_cp_slice_for_thd( if pad_len > 0: pad_value = True if key == 'padding_mask' else 0 batch[key] = torch.cat([batch[key], batch[key].new_full((pad_len,), pad_value)]) - if cp_size <= 1: - return if use_dsv4_cp_slice: if total_tokens % cp_size != 0: raise RuntimeError( diff --git a/megatron/core/fusions/fused_mla_yarn_rope_apply.py b/megatron/core/fusions/fused_mla_yarn_rope_apply.py index b9e13fcf7cb..74636c05eb6 100644 --- a/megatron/core/fusions/fused_mla_yarn_rope_apply.py +++ b/megatron/core/fusions/fused_mla_yarn_rope_apply.py @@ -369,7 +369,8 @@ def backward(ctx, grad): total_seqlen = grad.shape[0] else: seq_num = len(ctx.cu_seqlens_q) - 1 - grad = grad.contiguous() + if ctx.has_position_ids: + grad = grad.contiguous() total_seqlen, nheads, headdim = grad.shape assert grad.stride(-1) == 1 diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index 3482bcb3896..75f82a16750 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -88,6 +88,32 @@ def get_compress_topk_idxs( return matrix.unsqueeze(0).expand(batch_size, -1, -1) +def _get_csa_compressed_capacity( + packed_seq_params: Optional[PackedSeqParams], ratio: int, total_tokens: int +) -> Optional[int]: + """Return a host-known compressed capacity for THD CUDA graph capture. + + The exact ``sum(seq_len // ratio)`` lives in ``cu_seqlens`` on device. + Reading it on the host would break CUDA graph capture, and + ``PackedSeqParams`` must stay a generic MCore contract rather than + carrying CSA-specific metadata. Use a static upper bound instead; + device-side ``cu_seqlens_compressed`` keeps the true valid rows and + downstream kernels leave the extra rows as tail padding. + """ + if packed_seq_params is None or ratio <= 1: + return None + max_seqlen = packed_seq_params.max_seqlen_q + cu_seqlens = ( + packed_seq_params.cu_seqlens_q_padded + if packed_seq_params.cu_seqlens_q_padded is not None + else packed_seq_params.cu_seqlens_q + ) + if max_seqlen is None or cu_seqlens is None: + return None + num_sequences = max(int(cu_seqlens.shape[0]) - 1, 0) + return min(int(total_tokens) // ratio, num_sequences * (int(max_seqlen) // ratio)) + + # --------------------------------------------------------------------------- # THD (packed) variants of the index helpers above. # @@ -873,9 +899,10 @@ def _forward_thd( softmax + reduce is batched across ALL compressed entries from all segments via a ``(total_compressed, ratio)`` gather index. - Non-CP callers may supply ``fixed_total_comp`` to pad the output to a - host-known CUDA-graph capacity. Rows beyond the true - ``cu_seqlens_compressed[-1]`` gather from position 0 and remain tail padding. + When ``fixed_total_comp`` is supplied the output is padded to that + static capacity so tensor shapes are host-known and CUDA-graph + capturable. Rows beyond the true ``cu_seqlens_compressed[-1]`` + gather from position 0 and are left as tail padding. Args: x: ``(total, 1, hidden_size)`` packed bf16. @@ -885,14 +912,14 @@ def _forward_thd( sync when building the rotary table for ``ratio > 1``). compressed_group_ids: CP compressor-prep group ids. When supplied, ``x`` is already packed into ``ratio``-sized groups. - fixed_total_comp: optional non-CP output row capacity. Must be at - least the true compressed count. + fixed_total_comp: when set, overrides ``cu_seqlens_compressed[-1]`` + as the output row count. Must be >= the true compressed count. Returns: ``(compressed_thd, cu_seqlens_compressed)`` where ``compressed_thd`` is ``(total_compressed, 1, head_dim)`` bf16 - (or ``None`` when no sequence has ``seg_len >= ratio`` and no - fixed capacity is set) and + (or ``None`` when no sequence has ``seg_len >= ratio`` and + ``fixed_total_comp`` is not set) and ``cu_seqlens_compressed`` is ``(B+1,)`` int32 with ``cu_seqlens_compressed[b+1] - cu_seqlens_compressed[b] = seqlen_b // ratio``. Pre-grouped CP inputs return ``None`` for this unused second value. @@ -906,7 +933,9 @@ def _forward_thd( cu_seqlens_compressed = None total_comp = compressed_group_ids.shape[0] else: - seg_compressed_lens = (cu_seqlens[1:] - cu_seqlens[:-1]) // ratio + # Per-segment compressed lengths (vectorized). + seq_lens = cu_seqlens[1:] - cu_seqlens[:-1] + seg_compressed_lens = seq_lens // ratio cu_seqlens_compressed = torch.cat( [ torch.zeros(1, dtype=cu_seqlens.dtype, device=device), @@ -947,9 +976,8 @@ def _forward_thd( offsets = torch.arange(ratio, device=device, dtype=base.dtype).unsqueeze(0) gather_idx = base + offsets # (total_comp, ratio) - # Index flat kv/score into ``(total_comp, ratio, 1, coff * d)`` groups. kv_grouped = kv[gather_idx] # (total_comp, ratio, 1, coff * d) - score_grouped = score[gather_idx] # same + score_grouped = score[gather_idx] # APE: (ratio, coff * d) → broadcast (1, ratio, 1, coff * d). score_grouped = score_grouped + self.ape.view(1, ratio, 1, -1) @@ -1047,14 +1075,13 @@ def forward( "to avoid a GPU→CPU sync that breaks CUDA graph capture." ) max_seqlen_q = int(packed_seq_params.max_seqlen_q) - fixed_total_comp = None - if self.compress_ratio > 1: - fixed_total_comp = min( - x.shape[0] // self.compress_ratio, - (cu_seqlens.shape[0] - 1) * (max_seqlen_q // self.compress_ratio), - ) result = self._forward_thd( - x, cu_seqlens, max_seqlen_q=max_seqlen_q, fixed_total_comp=fixed_total_comp + x, + cu_seqlens, + max_seqlen_q=max_seqlen_q, + fixed_total_comp=_get_csa_compressed_capacity( + packed_seq_params, self.compress_ratio, x.shape[0] + ), ) else: result = self._forward_sbhd(x) @@ -1406,6 +1433,31 @@ def __init__( # Private helpers – each owns one logical slice of the forward pass. # ------------------------------------------------------------------ + def _build_kv_full( + self, kv: torch.Tensor, x: torch.Tensor + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], int]: + """Concatenate original KV with compressed KV (if applicable). + + Returns: + kv_full: [n_kv, b, v_head_dim] original + compressed KV. + compressed_kv: [n_compressed, b, v_head_dim] or None. + n_compressed: number of compressed positions (0 when unused). + """ + if self.compressor is not None and self.compress_ratio > 1: + compressed_kv = self.compressor(x) + if compressed_kv is not None: + kv_full = torch.cat([kv, compressed_kv], dim=0) + n_compressed = compressed_kv.size(0) + else: + kv_full = kv + compressed_kv = None + n_compressed = 0 + else: + kv_full = kv + compressed_kv = None + n_compressed = 0 + return kv_full, compressed_kv, n_compressed + def _forward_unfused_csa( self, query: torch.Tensor, @@ -1535,6 +1587,7 @@ def _forward_fused_indexer_inference( x: torch.Tensor, qr: torch.Tensor, kv_full: torch.Tensor, + n_compressed: int, offset: int, window_idxs: torch.Tensor, packed_seq_params: Optional[PackedSeqParams], @@ -1580,6 +1633,7 @@ def _forward_fused_indexer_training( x: torch.Tensor, qr: torch.Tensor, kv_full: torch.Tensor, + n_compressed: int, offset: int, window_idxs: torch.Tensor, packed_seq_params: Optional[PackedSeqParams], @@ -1672,11 +1726,7 @@ def forward( sq, b, np, hn = query.size() kv = key.squeeze(-2) # [sq, b, 1, v_head_dim] -> [sq, b, v_head_dim] - compressed_kv = ( - self.compressor(x) if self.compressor is not None and self.compress_ratio > 1 else None - ) - n_compressed = 0 if compressed_kv is None else compressed_kv.shape[0] - kv_full = kv if compressed_kv is None else torch.cat((kv, compressed_kv), dim=0) + kv_full, compressed_kv, n_compressed = self._build_kv_full(kv, x) offset = sq # compressed indices start after original positions window_idxs = get_window_topk_idxs(self.window_size, b, sq, query.device) @@ -1700,11 +1750,11 @@ def forward( ) elif has_indexer_compressed and self.training and torch.is_grad_enabled(): output, indexer_loss = self._forward_fused_indexer_training( - query, x, qr, kv_full, offset, window_idxs, packed_seq_params + query, x, qr, kv_full, n_compressed, offset, window_idxs, packed_seq_params ) elif has_indexer_compressed: output = self._forward_fused_indexer_inference( - query, x, qr, kv_full, offset, window_idxs, packed_seq_params + query, x, qr, kv_full, n_compressed, offset, window_idxs, packed_seq_params ) else: output = self._forward_fused_no_indexer( @@ -1736,6 +1786,7 @@ def _forward_unfused_csa_thd( cu_seqlens_kv_full: torch.Tensor, cu_seqlens_compressed: torch.Tensor, window_idxs: torch.Tensor, + max_seqlen_q: int, max_seqlen_compressed_idx: int, packed_seq_params: PackedSeqParams, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: @@ -1953,6 +2004,7 @@ def _forward_fused_indexer_inference_thd( window_idxs: torch.Tensor, max_seqlen_q: int, max_seqlen_compressed_idx: int, + max_seqlen_kv: int, ) -> torch.Tensor: """Path C (THD): separate indexer forward (no loss) + fused sparse attn (compact). @@ -2021,6 +2073,8 @@ def _forward_fused_indexer_training_thd( x: torch.Tensor, qr: torch.Tensor, packed_seq_params: PackedSeqParams, + total_q: int, + np_: int, cu_seqlens_q: torch.Tensor, cu_seqlens_kv: torch.Tensor, cu_seqlens_kv_full: torch.Tensor, @@ -2386,7 +2440,7 @@ def _forward_thd( # DSv4 hybrid adds a dummy batch dim to keep the MQA-head # unsqueeze symmetric with SBHD) # x, qr : (total_q, 1, *) - total_q, np_, _ = query.shape + total_q, _np, _ = query.shape cu_seqlens_q = ( packed_seq_params.cu_seqlens_q_padded @@ -2399,6 +2453,7 @@ def _forward_thd( else packed_seq_params.cu_seqlens_kv ) max_seqlen_q = int(packed_seq_params.max_seqlen_q) + max_seqlen_kv = int(packed_seq_params.max_seqlen_kv) # Squeeze the dummy b=1 and MQA head-dim to get the KV-flat layout. # (key arrives as (total_kv, 1, 1, hn) for MQA.) @@ -2454,13 +2509,14 @@ def _forward_thd( kv_full_thd, compressed_kv, n_compressed_total, - np_, + _np, total_q, cu_seqlens_q, cu_seqlens_kv, cu_seqlens_kv_full, cu_seqlens_compressed, window_idxs, + max_seqlen_q, max_seqlen_compressed_idx, packed_seq_params, ) @@ -2470,6 +2526,8 @@ def _forward_thd( x, qr, packed_seq_params, + total_q, + _np, cu_seqlens_q, cu_seqlens_kv, cu_seqlens_kv_full, @@ -2493,6 +2551,7 @@ def _forward_thd( window_idxs, max_seqlen_q, max_seqlen_compressed_idx, + max_seqlen_kv, ) else: output = self._forward_fused_no_indexer_thd( diff --git a/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py b/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py index 2e1703a0470..09b013ca9f0 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py +++ b/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py @@ -84,8 +84,8 @@ class _LeftBoundaryExchange(torch.autograd.Function): @staticmethod def forward(ctx, tensor: torch.Tensor, d_window: int, cp_group: torch.distributed.ProcessGroup): """Receive fixed left-boundary hidden rows needed by this CP rank.""" - cp_size = 1 if cp_group is None else cp_group.size() - cp_rank = 0 if cp_group is None else cp_group.rank() + cp_size = cp_group.size() + cp_rank = cp_group.rank() ctx.cp_group = cp_group ctx.d_window = d_window ctx.input_shape = tensor.shape @@ -95,8 +95,6 @@ def forward(ctx, tensor: torch.Tensor, d_window: int, cp_group: torch.distribute f"local_rows={tensor.shape[0]}, D_window={d_window}." ) boundary = tensor.new_zeros((d_window,) + tuple(tensor.shape[1:])) - if cp_size <= 1: - return boundary ops = [] if cp_rank > 0: @@ -120,12 +118,10 @@ def forward(ctx, tensor: torch.Tensor, d_window: int, cp_group: torch.distribute def backward(ctx, grad_boundary: torch.Tensor): """Send boundary gradients back to ranks that own those hidden rows.""" cp_group = ctx.cp_group - cp_size = 1 if cp_group is None else cp_group.size() - cp_rank = 0 if cp_group is None else cp_group.rank() + cp_size = cp_group.size() + cp_rank = cp_group.rank() d_window = ctx.d_window grad_input = grad_boundary.new_zeros(ctx.input_shape) - if cp_size <= 1: - return grad_input, None, None ops = [] if cp_rank > 0: diff --git a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py index daccda6198b..3cadb969c3f 100644 --- a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py +++ b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py @@ -128,6 +128,7 @@ def __init__( self.config.csa_compress_rotary_base if use_compressed_yarn else self.config.rotary_base ) self._dsv4_compress_ratio = compress_ratio + self._dsv4_rope_base = rope_base self._dsv4_uses_yarn_rope = use_compressed_yarn if not use_compressed_yarn: self.rotary_pos_emb = RotaryEmbedding( @@ -264,14 +265,10 @@ def forward( cp_size = self.pg_collection.cp.size() qkv_format = packed_seq_params.qkv_format if packed_seq_params is not None else None - if cp_size > 1 and qkv_format != 'thd': - raise RuntimeError( - "DSv4 context parallel support requires THD packed sequence input; " - f"got cp_size={cp_size}, qkv_format={qkv_format!r}." - ) + use_thd_cp = cp_size > 1 and qkv_format == 'thd' boundary_hidden = None - if cp_size > 1: + if use_thd_cp: boundary_hidden = cp_utils.exchange_cp_boundary_hidden( hidden_states, self._dsv4_compress_ratio, @@ -284,7 +281,7 @@ def forward( # ===================== # Get the query, key and value tensors based on the type of attention - # self or cross attn. - query, key, value, q_compressed, boundary_kv = self.get_query_key_value_tensors( + qkv = self.get_query_key_value_tensors( hidden_states, key_value_states, position_ids, @@ -292,6 +289,11 @@ def forward( inference_context=inference_context, boundary_hidden=boundary_hidden, ) + if use_thd_cp: + query, key, value, q_compressed, kv_compressed, boundary_kv = qkv + else: + query, key, value, q_compressed, kv_compressed = qkv + boundary_kv = None # TODO: Currently, TE can only accept contiguous tensors for MLA query = query.contiguous() @@ -378,7 +380,7 @@ def forward( else: rotary_pos_emb = self.rotary_pos_emb(rope_seqlen, packed_seq=packed_seq) if self.config.apply_rope_fusion: - if cp_size > 1: + if use_thd_cp: if rope_seqlen is None: raise RuntimeError("DSv4 THD CP inverse RoPE requires max_seqlen_kv.") global_start = self.pg_collection.cp.rank() * core_attn_out.shape[0] @@ -576,8 +578,8 @@ def get_query_key_value_tensors( Derives `query`, `key` and `value` tensors from `hidden_states`. Returns: - Tuple of ``(query, key, value, q_compressed, boundary_kv)``. - ``boundary_kv`` carries CP boundary rows when the DSv4 CP path is active. + Tuple of ``(query, key, value, q_compressed, kv_compressed)``. The THD CP + path appends ``boundary_kv`` carrying the projected left-boundary rows. """ # s = sequence length, b = batch size, h = hidden size, n = num attention heads # Attention heads [s, b, n*h] @@ -643,6 +645,7 @@ def get_query_key_value_tensors( q_compressed, _ = self.linear_q_down_proj(hidden_states) kv_compressed = hidden_states + k_pos_emb = None boundary_kv_compressed = boundary_hidden if packed_seq_params is not None: @@ -667,7 +670,7 @@ def get_query_key_value_tensors( # ========================================= def qkv_up_proj_and_rope_apply( - q_compressed, kv_compressed, rotary_pos_emb, boundary_kv_compressed + q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb, boundary_kv_compressed=None ): """ Apply the up projection and RoPE to the query and key. @@ -694,9 +697,13 @@ def qkv_up_proj_and_rope_apply( kv = self.kv_layernorm(kv) boundary_kv = None + # [num_tokens, qk_pos_emb_head_dim] -> [num_tokens, 1, qk_pos_emb_head_dim] + if k_pos_emb is not None: + k_pos_emb = torch.unsqueeze(k_pos_emb, -2) + cp_size = self.pg_collection.cp.size() if self.config.apply_rope_fusion: - if cp_size > 1: + if cp_size > 1 and packed_seq: cp_rank = self.pg_collection.cp.rank() # Rank r owns global rows [r * local_rows, (r + 1) * local_rows). assert packed_seq, "DSv4 CP fused RoPE expects THD packed sequence input." @@ -721,10 +728,11 @@ def qkv_up_proj_and_rope_apply( global_start - boundary_rows, clamp_to_valid_token=bool(boundary_rows), ) - if boundary_rows: + if boundary_kv_compressed is not None: boundary_kv = kv[:boundary_rows] kv = kv[boundary_rows:] else: + cp_rank = self.pg_collection.cp.rank() query = fused_mla_rope_inplace( q, rotary_pos_cos, @@ -732,8 +740,8 @@ def qkv_up_proj_and_rope_apply( self.config.qk_head_dim, self.config.qk_pos_emb_head_dim, cu_seqlens_q, - self.pg_collection.cp.rank(), - self.pg_collection.cp.size(), + cp_rank, + cp_size, remove_interleaving=True, ) kv = kv.unsqueeze(-2) @@ -744,8 +752,8 @@ def qkv_up_proj_and_rope_apply( self.config.qk_head_dim, self.config.qk_pos_emb_head_dim, cu_seqlens_q, - self.pg_collection.cp.rank(), - self.pg_collection.cp.size(), + cp_rank, + cp_size, remove_interleaving=True, ) key = kv @@ -818,24 +826,46 @@ def qkv_up_proj_and_rope_apply( if boundary_kv is not None: boundary_kv = boundary_kv.contiguous() + if boundary_kv is None: + return query, key, value return query, key, value, boundary_kv if self.recompute_up_proj: quantization = self.config.fp8 or self.config.fp4 self.qkv_up_checkpoint = tensor_parallel.CheckpointWithoutOutput(fp8=quantization) - query, key, value, boundary_kv = self.qkv_up_checkpoint.checkpoint( - qkv_up_proj_and_rope_apply, - q_compressed, - kv_compressed, - rotary_pos_emb, - boundary_kv_compressed, - ) + if boundary_kv_compressed is None: + query, key, value = self.qkv_up_checkpoint.checkpoint( + qkv_up_proj_and_rope_apply, + q_compressed, + kv_compressed, + k_pos_emb, + rotary_pos_emb, + ) + boundary_kv = None + else: + query, key, value, boundary_kv = self.qkv_up_checkpoint.checkpoint( + qkv_up_proj_and_rope_apply, + q_compressed, + kv_compressed, + k_pos_emb, + rotary_pos_emb, + boundary_kv_compressed, + ) else: - query, key, value, boundary_kv = qkv_up_proj_and_rope_apply( - q_compressed, kv_compressed, rotary_pos_emb, boundary_kv_compressed - ) + if boundary_kv_compressed is None: + query, key, value = qkv_up_proj_and_rope_apply( + q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb + ) + boundary_kv = None + else: + query, key, value, boundary_kv = qkv_up_proj_and_rope_apply( + q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb, boundary_kv_compressed + ) - return query, key, value, q_compressed, boundary_kv + result = (query, key, value, q_compressed, kv_compressed) + if boundary_kv is not None: + return result + (boundary_kv,) + return result def backward_dw(self) -> NoReturn: """Execute weight gradient computation""" diff --git a/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py b/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py index 69dbf53d542..3bc19b79f1e 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py +++ b/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py @@ -717,6 +717,63 @@ def _thd_to_fake_bshd(*tensors: Tensor) -> Tuple[Tensor, ...]: return tuple(t.unsqueeze(0) for t in tensors) +def _compute_indexer_predict( + q_indexer: Tensor, + k_indexer: Tensor, + weights: Tensor, + topk_indices: Tensor, + qhead_per_kv_head: int, + *, + topk_indices_global: bool = False, +) -> Tensor: + """Compute ``predict`` distribution (softmax over top-K of indexer scores). + + Wraps `cudnn.DSA.sparse_indexer_score_recompute_wrapper`. + This function is not used now, but it is kept for potential future use. + + Two layouts: + + * **BSHD** (default; 4-D q): ``q (B, S_q, H, D)``, ``k (B, S_k, D)``, + ``w (B, S_q, H)``, ``topk (B, S_q, topk)``. + * **THD packed** (3-D q): ``q (total_q, H, D)``, ``k (total_k, D)``, + ``w (total_q, H)``, ``topk (total_q, topk)``. Internally + fake-BSHD'd with ``B=1`` so the wrapper's 4-D-Q shape check + passes; ``topk_indices_global=True`` is required (and enforced) so + the kernel decodes the flat ids directly as positions into the + ``(1*total_k, D)`` view. + + Output shape matches the layout: BSHD ``(B, S_q, topk)`` or + THD ``(total_q, topk)``, fp32 softmax over the top-K axis. + """ + _ensure_dsa_namespace() + is_thd = q_indexer.ndim == 3 + if is_thd: + if not topk_indices_global: + raise ValueError( + "THD ``_compute_indexer_predict`` requires " + "``topk_indices_global=True`` so the kernel addresses K " + "by flat ids over the packed ``(total_k, D)`` buffer." + ) + q_bshd, k_bsd, w_bsh, topk_bst = _thd_to_fake_bshd( + q_indexer, k_indexer, weights, topk_indices + ) + else: + q_bshd, k_bsd, w_bsh, topk_bst = q_indexer, k_indexer, weights, topk_indices + + result = _DSA.sparse_indexer_score_recompute_wrapper( + q_bshd, + k_bsd, + w_bsh, + topk_bst, + qhead_per_kv_head=qhead_per_kv_head, + topk_indices_global=topk_indices_global, + ) + predict = result["predict"] + if is_thd: + predict = predict.squeeze(0) + return predict + + def _compute_attn_target( q_attn: Tensor, k_attn: Tensor, @@ -729,9 +786,10 @@ def _compute_attn_target( ) -> Tensor: """Compute ``target`` distribution (L1-normalised head-sum softmax). - Wraps :attr:`cudnn.DSA.sparse_attn_score_recompute_wrapper`. A 4-D q - is BSHD; a 3-D q is THD and gets fake-BSHD'd with ``B=1`` before the - wrapper call so the 4-D-Q shape check passes. + Wraps :attr:`cudnn.DSA.sparse_attn_score_recompute_wrapper`. Same + layout convention as :func:`_compute_indexer_predict`: 4-D q is + BSHD; 3-D q is THD and gets fake-BSHD'd with ``B=1`` before the + wrapper call (so the 4-D-Q shape check passes). """ _ensure_dsa_namespace() is_thd = q_attn.ndim == 3 @@ -1010,6 +1068,7 @@ def forward( # THD: skip the permute; tensors are already flat. if is_thd: total_q = q_indexer.shape[0] + idx_nh = q_indexer.shape[1] np_, d = query.shape[1], query.shape[2] q_indexer_flat = q_indexer @@ -1018,6 +1077,7 @@ def forward( else: sq, b, np_, d = query.shape skv = kv_full.shape[0] + idx_nh = q_indexer.shape[2] q_indexer_flat = q_indexer.permute(1, 0, 2, 3).contiguous() k_indexer_flat = k_indexer.permute(1, 0, 2).contiguous() diff --git a/tests/unit_tests/test_sequence_packing.py b/tests/unit_tests/test_sequence_packing.py index 3e123f5a49f..3b5bfeffad4 100644 --- a/tests/unit_tests/test_sequence_packing.py +++ b/tests/unit_tests/test_sequence_packing.py @@ -209,7 +209,7 @@ def test_dsv4_thd_cp_slice_uses_static_partition_total(): @pytest.mark.parametrize( ("alignment", "cuda_graph_impl", "local_cp_size", "total_tokens", "local_target"), - [("max", "none", 1, 6, 8), (4, "none", 2, 10, 8), (8, "transformer_engine", 2, 10, 8)], + [(4, "none", 2, 10, 8), (8, "transformer_engine", 2, 10, 8)], ) def test_dsv4_thd_dynamic_cp_pads_before_slicing( monkeypatch, alignment, cuda_graph_impl, local_cp_size, total_tokens, local_target diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py index 519481dcd36..e0e31e576d6 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py @@ -45,22 +45,6 @@ def _sequence_positions(cu_seqlens, rows, clamp=False): return rows - cu_seqlens[seq_ids] -def test_boundary_exchange_single_rank_returns_zero_boundary_and_zero_grad(): - """Validate the no-CP boundary exchange contract. - - Expected: with cp_group=None, the fixed left boundary is zero-filled and its - backward path contributes no gradient to the local tensor. A failure here - means the CP path could invent boundary tokens or bogus local gradients. - """ - local = torch.arange(12, dtype=torch.float32).reshape(4, 3).requires_grad_(True) - - boundary = exchange_cp_boundary_hidden(local, 0, 2, None) - - assert torch.equal(boundary, torch.zeros(2, 3)) - boundary.sum().backward() - assert torch.equal(local.grad, torch.zeros_like(local)) - - def test_thd_cp_left_boundary_exchange_forward_backward(): """Validate distributed CP boundary exchange forward/backward. diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py index 9a87b8826ae..096c73d1451 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py @@ -1632,6 +1632,22 @@ def _ref_attn_full_score( return torch.where(valid, s, torch.zeros_like(s)) +def _ref_indexer_predict_sparse(q_bshd_fp32, k_bsd_fp32, w_bsh_fp32, topk_indices, sm_scale): + """Reference for ``sparse_indexer_score_recompute_wrapper.predict``. + + Compute the full-KV indexer score, gather ``topk_indices``, softmax + over the topK axis. ``-1`` entries in topk are masked to ``-inf`` + so they contribute zero probability. + """ + qk = torch.einsum('bqhd,bkd->bqhk', q_bshd_fp32, k_bsd_fp32) + s = (torch.relu(qk) * w_bsh_fp32.unsqueeze(-1)).sum(dim=2) * sm_scale # (B, Sq, Sk) + valid = topk_indices >= 0 + safe = topk_indices.clamp(min=0).long() + s_topk = torch.gather(s, dim=-1, index=safe) + s_topk = torch.where(valid, s_topk, torch.full_like(s_topk, float('-inf'))) + return torch.softmax(s_topk, dim=-1) + + def _ref_attn_target_sparse(q_bshd_fp32, k_bsd_fp32, lse_bsh_fp32, topk_indices, softmax_scale): """Reference for ``sparse_attn_score_recompute_wrapper.target``. @@ -1819,16 +1835,17 @@ def _build_real_score_inputs(s, *, with_lse: bool = True, with_topk: bool = True class TestRealKernelScoreHelpers: - """Real-kernel parity tests for three ``_compute_*`` score helpers + """Real-kernel parity tests for the four ``_compute_*`` score helpers against PyTorch reference implementations. Single parametrized test - covers all three; numeric tolerance is bf16-friendly (raw fp32 score + covers all four; numeric tolerance is bf16-friendly (raw fp32 score sums agree to ~5%, normalized distributions to ~5e-3). """ # Each case: (id, sm_min, kernel_name, runner). The runner does the # call + ref + assertion; it returns nothing on success. @pytest.mark.parametrize( - "case", ['sparse_attn_target', 'dense_indexer_score', 'dense_attn_score'] + "case", + ['sparse_indexer_predict', 'sparse_attn_target', 'dense_indexer_score', 'dense_attn_score'], ) def test_real_score_helper(self, case, reset_lazy_kernel_state): _skip_if_real_kernels_unavailable(sm_min=10) @@ -1843,7 +1860,29 @@ def test_real_score_helper(self, case, reset_lazy_kernel_state): from megatron.core.transformer.experimental_attention_variant import dsa_kernels as _dk - if case == 'sparse_attn_target': + if case == 'sparse_indexer_predict': + # The kernel takes sm_scale=1.0; scale is applied via weights + # pre-multiplication (relu(c·x)·W trick). Reference mirrors that. + scale = s['indexer_softmax_scale'] + w_scaled = (x['w'].float() * scale).to(x['w'].dtype) + out = _dk._compute_indexer_predict( + x['q_idx'], x['k_idx'], w_scaled, x['topk'], qhead_per_kv_head=s['idx_nh'] + ) + ref = _ref_indexer_predict_sparse( + x['q_idx'].float(), x['k_idx'].float(), w_scaled.float(), x['topk'], sm_scale=1.0 + ) + # Softmax outputs in [0, 1]; bf16 element-wise noise can break + # absolute tolerance, so compare directions via cosine similarity. + assert out.shape == ref.shape == (s['b'], s['sq'], s['topk']) + cos = torch.nn.functional.cosine_similarity( + out.flatten().unsqueeze(0).float(), ref.flatten().unsqueeze(0).float() + ).item() + assert cos > 0.99, ( + f"{case}: cos sim = {cos:.4f}, " + f"max abs diff = {(out - ref).abs().max().item():.3e}" + ) + + elif case == 'sparse_attn_target': out = _dk._compute_attn_target( x['q_attn'], x['k_attn'], @@ -2778,7 +2817,7 @@ def test_thd_single_segment_matches_sbhd_b1(self, sparse_loss, reset_lazy_kernel cu_comp_idx = _make_cu_seqlens([s['n_comp']], device=dev) # B=1: per-segment [kv, compressed] layout collapses to a single - # Ordered slice with the same kv_offset as the SBHD case. + # contiguous slice — same kv_offset as the SBHD case. compressed_kv_thd = kv_full_thd[kv_offset:] _, loss_thd = fused_indexer_sparse_attn( query_thd, diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention.py index 155cd940ffe..6a8ec00f624 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention.py @@ -6,9 +6,11 @@ import torch import torch.nn.functional as F +import megatron.core.parallel_state as parallel_state from megatron.core.extensions.transformer_engine import HAVE_TE from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.enums import AttnMaskType from megatron.core.transformer.transformer_config import MLATransformerConfig from tests.unit_tests.test_utilities import Utils @@ -355,7 +357,7 @@ def test_qkv_shapes(self): seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 ).cuda() - q, k, v, _, _ = attn.get_query_key_value_tensors(hidden) + q, k, v, q_compressed, kv_compressed = attn.get_query_key_value_tensors(hidden) n_heads = self.config.num_attention_heads v_dim = self.config.v_head_dim From db75531da3965684e400c9b3f1b3d9aff3b073f6 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Sat, 27 Jun 2026 22:05:15 -0700 Subject: [PATCH 08/23] Address DSv4 CP review feedback Signed-off-by: kunlunl --- megatron/core/datasets/data_schedule.py | 22 ++--- megatron/core/datasets/data_schedule_utils.py | 14 ++-- .../experimental_attention_variant/csa.py | 81 ++++++++++--------- .../csa_cp_utils.py | 5 +- .../deepseek_v4_hybrid_attention.py | 3 - tests/unit_tests/test_sequence_packing.py | 2 +- 6 files changed, 66 insertions(+), 61 deletions(-) diff --git a/megatron/core/datasets/data_schedule.py b/megatron/core/datasets/data_schedule.py index 7e6c0c626e1..30dd1e12b92 100644 --- a/megatron/core/datasets/data_schedule.py +++ b/megatron/core/datasets/data_schedule.py @@ -644,12 +644,12 @@ def get_batch_on_this_rank_for_sequence_packing( group_size=local_cp_size_val ) - use_dsv4_cp_slice = ( + use_contiguous_cp_slice = ( config is not None and getattr(config, 'experimental_attention_variant', None) == "dsv4_hybrid" and cp_group.size() > 1 ) - dsv4_cp_local_target_len = None + contiguous_cp_local_target_len = None pad_alignment = ( getattr(config, 'pad_packed_seq_alignment', None) if config is not None else None ) @@ -661,15 +661,15 @@ def get_batch_on_this_rank_for_sequence_packing( getattr(config, 'thd_max_packed_sequences', None), getattr(config, 'cuda_graph_impl', 'none') != 'none', ) - if is_tp_rank_0 and use_dsv4_cp_slice and pad_alignment is not None: + if is_tp_rank_0 and use_contiguous_cp_slice and pad_alignment is not None: if target_len is not None: - dsv4_cp_local_target_len = target_len + contiguous_cp_local_target_len = target_len else: # Fix the local width before slicing so later padding cannot shift rank origins. assert alignment is not None total_rows = int(batch['cu_seqlens_padded'][-1].item()) local_rows = (total_rows + cp_group.size() - 1) // cp_group.size() - dsv4_cp_local_target_len = ((local_rows + alignment - 1) // alignment) * alignment + contiguous_cp_local_target_len = ((local_rows + alignment - 1) // alignment) * alignment # Build padding_mask before CP slicing while tensors still have the full # packed length represented by cu_seqlens_padded[-1]. @@ -686,15 +686,15 @@ def get_batch_on_this_rank_for_sequence_packing( if is_first_or_last_stage or mtp_on_this_rank: cp_slice_keys.extend(['tokens', 'position_ids', 'labels', 'loss_mask']) partition_total_tokens = ( - dsv4_cp_local_target_len * cp_group.size() - if dsv4_cp_local_target_len is not None + contiguous_cp_local_target_len * cp_group.size() + if contiguous_cp_local_target_len is not None else None ) get_cp_slice_for_thd( batch, cp_group, keys=cp_slice_keys, - use_dsv4_cp_slice=use_dsv4_cp_slice, + use_contiguous_cp_slice=use_contiguous_cp_slice, partition_total_tokens=partition_total_tokens, ) @@ -710,8 +710,10 @@ def get_batch_on_this_rank_for_sequence_packing( # Broadcast total_tokens because padding_mask is prepared on every PP stage. # Tokens/labels/loss_mask/position_ids use the same length on stages that own them. if is_tp_rank_0: - if dsv4_cp_local_target_len is not None: - total_tokens = torch.tensor([dsv4_cp_local_target_len], dtype=torch.int32, device=dev) + if contiguous_cp_local_target_len is not None: + total_tokens = torch.tensor( + [contiguous_cp_local_target_len], dtype=torch.int32, device=dev + ) else: # Under VPP, the last PP stage has labels but no tokens, so derive # total_tokens from cu_seqlens_padded, which is present on every diff --git a/megatron/core/datasets/data_schedule_utils.py b/megatron/core/datasets/data_schedule_utils.py index 87374e470d5..7b47565e45c 100644 --- a/megatron/core/datasets/data_schedule_utils.py +++ b/megatron/core/datasets/data_schedule_utils.py @@ -15,19 +15,21 @@ def get_cp_slice_for_thd( batch, cp_group, keys: Optional[Sequence[str]] = None, - use_dsv4_cp_slice: bool = False, + use_contiguous_cp_slice: bool = False, partition_total_tokens: Optional[int] = None, ): """Partition sequence data for context parallelism in THD format. - Uses TE's THD partitioned indices by default. DSv4 instead assigns each rank one - ordered range of the global rows. Only keys present in the batch are sliced. + Uses TE's THD partitioned indices to split the packed sequence across CP ranks. + With ``use_contiguous_cp_slice=True``, the flattened rows are split into ``cp_size`` equal + slices, and each rank receives the slice at its rank index. + Only keys present in the batch are sliced. Args: batch: Dict with packed sequence data. cp_group: Context parallel process group. keys: Sequence data keys to slice. Defaults to the original THD data tensors. - use_dsv4_cp_slice: Whether to use the DSv4 row assignment. + use_contiguous_cp_slice: Whether to assign one contiguous row range per CP rank. partition_total_tokens: Optional total used to tail-pad tensors selected by ``keys`` before slicing. Existing cu_seqlens metadata is left unchanged. """ @@ -61,10 +63,10 @@ def get_cp_slice_for_thd( if pad_len > 0: pad_value = True if key == 'padding_mask' else 0 batch[key] = torch.cat([batch[key], batch[key].new_full((pad_len,), pad_value)]) - if use_dsv4_cp_slice: + if use_contiguous_cp_slice: if total_tokens % cp_size != 0: raise RuntimeError( - f"DSv4 CP slicing requires total_tokens={total_tokens} to be divisible by " + f"Contiguous CP slicing requires total_tokens={total_tokens} to be divisible by " f"cp_size={cp_size}." ) local_rows = total_tokens // cp_size diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index 75f82a16750..497f8e8f125 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -655,13 +655,21 @@ def _unfused_indexer_sparse_attn_from_topk( row_valid = valid.any(dim=-1, keepdim=True) safe_indices = indexer_topk_indices.clamp(min=0).long() - index_scores_full = torch.einsum("rhd,kd->rhk", q_indexer.float(), k_indexer.float()) - index_scores_full = torch.relu(index_scores_full) weights_scaled = weights.float() * float(indexer_softmax_scale) - index_scores_full = index_scores_full * weights_scaled.unsqueeze(-1) - index_scores_full = index_scores_full.sum(dim=1) - - predict_logits = torch.gather(index_scores_full, dim=-1, index=safe_indices) + predict_chunks = [] + # Avoid materializing the full [local_q, index_heads, global_k] score tensor. + for start in range(0, total_q, 512): + end = start + 512 + chunk_indices = safe_indices[start:end] + selected_k_indexer = k_indexer.index_select(0, chunk_indices.reshape(-1)).reshape( + chunk_indices.shape[0], indexer_topk, -1 + ) + chunk_scores = torch.einsum( + "rhd,rkd->rhk", q_indexer[start:end].float(), selected_k_indexer.float() + ) + chunk_scores = torch.relu(chunk_scores) * weights_scaled[start:end].unsqueeze(-1) + predict_chunks.append(chunk_scores.sum(dim=1)) + predict_logits = torch.cat(predict_chunks) predict_logits = predict_logits.masked_fill(~valid, float("-inf")) predict_logits = predict_logits.masked_fill(~row_valid, 0.0) predict = torch.softmax(predict_logits, dim=-1, dtype=torch.float32) @@ -681,8 +689,9 @@ def _unfused_indexer_sparse_attn_from_topk( exp_scores = torch.exp(attn_scores - scores_max) exp_sink = torch.exp(sink - scores_max) attn_probs = exp_scores / (exp_scores.sum(dim=-1, keepdim=True) + exp_sink) - attn_probs = attn_probs * row_valid.unsqueeze(1).float() target = attn_probs.sum(dim=1) + target = target / target.sum(dim=-1, keepdim=True).clamp(min=1e-10) + target = target * row_valid.float() eps = torch.finfo(torch.float32).tiny target = target.clamp(min=eps) @@ -999,7 +1008,21 @@ def _forward_thd( compressed_thd = self.norm(compressed_thd.to(dtype)) - if not pre_grouped: + if pre_grouped: + rotary_pos_cos, rotary_pos_sin = self.rotary_pos_emb.get_cached_cos_sin( + int(max_seqlen_q), dtype=compressed_thd.dtype, packed_seq=True, mscale=1.0 + ) + compressed_thd = fused_mla_rope_inplace( + compressed_thd, + rotary_pos_cos, + rotary_pos_sin, + self.head_dim - self.qk_pos_emb_head_dim, + self.qk_pos_emb_head_dim, + cu_seqlens_q=cu_seqlens, + remove_interleaving=True, + position_ids=compressed_group_ids[:total_comp].clamp_min(0) * ratio, + ) + else: # RoPE: applied in a single vectorized THD call. max_seqlen_rope = (max_seqlen_q // ratio) * ratio if max_seqlen_q is not None else None compressed_thd = _apply_rope( @@ -1014,32 +1037,6 @@ def _forward_thd( cu_seqlens=cu_seqlens_compressed, max_seqlen_rope=max_seqlen_rope, ) - else: - rotary_pos_cos, rotary_pos_sin = self.rotary_pos_emb.get_cached_cos_sin( - int(max_seqlen_q), dtype=compressed_thd.dtype, packed_seq=True, mscale=1.0 - ) - positions = compressed_group_ids[:total_comp].clamp_min(0).to(torch.long) * ratio - cos = rotary_pos_cos.flatten(1).index_select(0, positions) - sin = rotary_pos_sin.flatten(1).index_select(0, positions) - pos_dim = self.qk_pos_emb_head_dim - table_shape = (total_comp,) + (1,) * (compressed_thd.ndim - 2) + (pos_dim,) - cos = cos[:, :pos_dim].view(table_shape) - sin = sin[:, :pos_dim].view(table_shape) - x_nope, x_pos = torch.split(compressed_thd, [self.head_dim - pos_dim, pos_dim], dim=-1) - x1, x2 = x_pos[..., 0::2], x_pos[..., 1::2] - half = pos_dim // 2 - dtype = compressed_thd.dtype - left = ( - (x1.float() * cos[..., :half].float()).to(dtype).float() - - (x2.float() * sin[..., :half].float()).to(dtype).float() - ).to(dtype) - right = ( - (x2.float() * cos[..., half:].float()).to(dtype).float() - + (x1.float() * sin[..., half:].float()).to(dtype).float() - ).to(dtype) - compressed_thd = torch.cat( - (x_nope, torch.stack((left, right), dim=-1).flatten(-2)), dim=-1 - ) if self.rotate: compressed_thd = rotate_activation(compressed_thd) @@ -2200,8 +2197,9 @@ def _forward_thd_cp( compressed_kv_rank_major = kv_local.new_empty((0, kv_local.shape[-1])) cu_seqlens_compressed = None - # Logical top-k uses per-sequence comp ids; the map lowers them to - # rank-major gathered rows when needed. + # ``compressed_topk`` records which compressed blocks each query selects + # within its sequence. ``seq_to_rank_row`` maps each compressed block to + # the row where its K and KV are stored in the all-gathered buffer. compressed_topk = seq_to_rank_row = None ratio = self.compress_ratio indexer = self.indexer @@ -2221,9 +2219,10 @@ def _forward_thd_cp( torch.cumsum(compressed_lens, dim=0, dtype=torch.int32), ) ) - # Compact local+boundary tokens for Compressor. ``compressed_group_ids`` - # drives compressed RoPE; the row map directly lowers global - # sequence-major compressed ids into rank-major all-gather rows. + # ``hidden_compact`` packs the local and boundary tokens needed by the + # Compressor. ``compressed_group_ids`` gives each compressed block's + # position within its sequence for RoPE. ``seq_to_rank_row`` maps each + # block to its row in the all-gathered K and KV buffers. hidden_compact, compressed_group_ids, seq_to_rank_row = ( cp_utils.prepare_cp_compressor_input( x, @@ -2320,7 +2319,9 @@ def _forward_thd_cp( # Final indices address these source rows directly. This avoids a # second per-sequence packed KV layout and lets torch.cat own backward. kv_full_thd = torch.cat((boundary_kv, kv_local, compressed_kv_rank_major), dim=0) - use_indexer_loss = training_with_grad and compressed_topk is not None + use_indexer_loss = ( + training_with_grad and indexer_loss_coeff > 0 and compressed_topk is not None + ) compressed_width = ( compressed_topk.shape[-1] if compressed_topk is not None diff --git a/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py b/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py index 09b013ca9f0..457d652cc40 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py +++ b/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py @@ -55,8 +55,11 @@ def apply_thd_cp_local_rope_fused( squeezed_head = x.ndim == 2 rope_input = x.squeeze(1) if squeezed_batch else x rope_input = rope_input.unsqueeze(1) if squeezed_head else rope_input + if inverse: + # Sparse attention saves its unrotated output for backward. + rope_input = rope_input.clone() output = fused_mla_rope_inplace( - rope_input.clone(), + rope_input, cos, sin, nope_dim, diff --git a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py index 3cadb969c3f..928dce2d074 100644 --- a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py +++ b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py @@ -381,8 +381,6 @@ def forward( rotary_pos_emb = self.rotary_pos_emb(rope_seqlen, packed_seq=packed_seq) if self.config.apply_rope_fusion: if use_thd_cp: - if rope_seqlen is None: - raise RuntimeError("DSv4 THD CP inverse RoPE requires max_seqlen_kv.") global_start = self.pg_collection.cp.rank() * core_attn_out.shape[0] core_attn_out = cp_utils.apply_thd_cp_local_rope_fused( core_attn_out, @@ -706,7 +704,6 @@ def qkv_up_proj_and_rope_apply( if cp_size > 1 and packed_seq: cp_rank = self.pg_collection.cp.rank() # Rank r owns global rows [r * local_rows, (r + 1) * local_rows). - assert packed_seq, "DSv4 CP fused RoPE expects THD packed sequence input." global_start = cp_rank * q.shape[0] query = cp_utils.apply_thd_cp_local_rope_fused( q, diff --git a/tests/unit_tests/test_sequence_packing.py b/tests/unit_tests/test_sequence_packing.py index 3b5bfeffad4..b4273f4c696 100644 --- a/tests/unit_tests/test_sequence_packing.py +++ b/tests/unit_tests/test_sequence_packing.py @@ -196,7 +196,7 @@ def test_dsv4_thd_cp_slice_uses_static_partition_total(): } get_cp_slice_for_thd( - batch, _MockCPGroup(size=4, rank=2), use_dsv4_cp_slice=True, partition_total_tokens=16 + batch, _MockCPGroup(size=4, rank=2), use_contiguous_cp_slice=True, partition_total_tokens=16 ) assert torch.equal(batch["tokens"], torch.arange(8, 12, dtype=torch.int64)) From 2acda35b90739f1d9dd4aa59a69bcf4524593bb6 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Sun, 28 Jun 2026 00:04:23 -0700 Subject: [PATCH 09/23] Complete DSv4 THD CP review fixes Signed-off-by: kunlunl --- .../experimental_attention_variant/csa.py | 70 +++++-- .../csa_cp_utils.py | 154 ++++++++++---- .../deepseek_v4_hybrid_attention.py | 130 +++++++----- .../test_csa_cp_layout_kernels.py | 198 ++++++++++++++++++ .../test_csa_cp_utils.py | 101 ++++++++- .../test_dsv4_hybrid_attention_cp.py | 43 ++-- 6 files changed, 562 insertions(+), 134 deletions(-) diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index 497f8e8f125..dc1f0cd800a 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -1009,19 +1009,33 @@ def _forward_thd( compressed_thd = self.norm(compressed_thd.to(dtype)) if pre_grouped: - rotary_pos_cos, rotary_pos_sin = self.rotary_pos_emb.get_cached_cos_sin( - int(max_seqlen_q), dtype=compressed_thd.dtype, packed_seq=True, mscale=1.0 - ) - compressed_thd = fused_mla_rope_inplace( - compressed_thd, - rotary_pos_cos, - rotary_pos_sin, - self.head_dim - self.qk_pos_emb_head_dim, - self.qk_pos_emb_head_dim, - cu_seqlens_q=cu_seqlens, - remove_interleaving=True, - position_ids=compressed_group_ids[:total_comp].clamp_min(0) * ratio, - ) + position_ids = compressed_group_ids[:total_comp].clamp_min(0) * ratio + if self.config.apply_rope_fusion: + rotary_pos_cos, rotary_pos_sin = self.rotary_pos_emb.get_cached_cos_sin( + int(max_seqlen_q), dtype=compressed_thd.dtype, packed_seq=True, mscale=1.0 + ) + compressed_thd = fused_mla_rope_inplace( + compressed_thd, + rotary_pos_cos, + rotary_pos_sin, + self.head_dim - self.qk_pos_emb_head_dim, + self.qk_pos_emb_head_dim, + cu_seqlens_q=cu_seqlens, + remove_interleaving=True, + position_ids=position_ids, + ) + else: + rope_result = self.rotary_pos_emb(int(max_seqlen_q), packed_seq=True) + rotary_pos_emb = rope_result[0] if isinstance(rope_result, tuple) else rope_result + compressed_thd = _apply_unfused_rope( + compressed_thd, + torch.index_select(rotary_pos_emb, 0, position_ids.long()), + self.head_dim - self.qk_pos_emb_head_dim, + self.qk_pos_emb_head_dim, + self.config, + None, + self.pg_collection.cp, + ) else: # RoPE: applied in a single vectorized THD call. max_seqlen_rope = (max_seqlen_q // ratio) * ratio if max_seqlen_q is not None else None @@ -2244,7 +2258,8 @@ def _forward_thd_cp( ): raise RuntimeError( "DSv4 THD CP path currently supports sparse indexer loss only. " - "Dense CP-aware indexer loss needs a CP-aware dense score kernel." + "The current dense score and backward APIs cannot express the " + "sequence position where this rank's local queries begin." ) indexer_x, indexer_qr = x.detach(), qr.detach() @@ -2256,11 +2271,11 @@ def _forward_thd_cp( q_indexer_cp = q_indexer_cp.reshape( l_local, indexer.index_n_heads, indexer.index_head_dim ) - rotary_pos_cos, rotary_pos_sin = indexer.rotary_pos_emb.get_cached_cos_sin( - max_seqlen_q, dtype=q_indexer_cp.dtype, packed_seq=True, mscale=1.0 - ) - q_indexer_cp = rotate_activation( - cp_utils.apply_thd_cp_local_rope_fused( + if self.config.apply_rope_fusion: + rotary_pos_cos, rotary_pos_sin = indexer.rotary_pos_emb.get_cached_cos_sin( + max_seqlen_q, dtype=q_indexer_cp.dtype, packed_seq=True, mscale=1.0 + ) + q_indexer_cp = cp_utils.apply_thd_cp_local_rope_fused( q_indexer_cp, rotary_pos_cos, rotary_pos_sin, @@ -2269,7 +2284,21 @@ def _forward_thd_cp( cu_seqlens, global_start, ) - ) + else: + rope_result = indexer.rotary_pos_emb(max_seqlen_q, packed_seq=True) + rotary_pos_emb = ( + rope_result[0] if isinstance(rope_result, tuple) else rope_result + ) + q_indexer_cp = cp_utils.apply_thd_cp_local_rope_unfused( + q_indexer_cp, + rotary_pos_emb, + indexer.index_head_dim - indexer.qk_pos_emb_head_dim, + indexer.qk_pos_emb_head_dim, + cu_seqlens, + global_start, + self.config, + ) + q_indexer_cp = rotate_activation(q_indexer_cp) weights_indexer_cp, _ = indexer.linear_weights_proj(indexer_x) weights_indexer_cp = weights_indexer_cp.squeeze(1) * (indexer.index_n_heads**-0.5) @@ -2301,6 +2330,7 @@ def _forward_thd_cp( indexer.index_topk, indexer.softmax_scale, max_seqlen_q=max_seqlen_q, + use_fused=self.apply_dsa_kernel_fusion, ) # ---- Step 5: attention compressed KV path ------------------------- diff --git a/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py b/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py index 457d652cc40..88c47b85a92 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py +++ b/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py @@ -13,6 +13,7 @@ import torch.distributed as dist from megatron.core.fusions.fused_mla_yarn_rope_apply import fused_mla_rope_inplace +from megatron.core.models.common.embeddings.rope_utils import _apply_rotary_pos_emb_bshd from megatron.core.transformer.experimental_attention_variant import csa_cp_layout_kernels from megatron.core.transformer.experimental_attention_variant.dsa_kernels import indexer_topk @@ -21,42 +22,44 @@ # ============================================================================= -def apply_thd_cp_local_rope_fused( - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - nope_dim: int, - pos_dim: int, - cu_seqlens_padded: torch.Tensor, - global_start: int, - inverse: bool = False, - clamp_to_valid_token: bool = False, +def _thd_cp_position_ids( + cu_seqlens_padded: torch.Tensor, global_start: int, local_rows: int ) -> torch.Tensor: - """Apply fused non-interleaved RoPE to local THD CP rows.""" + """Map a consecutive CP row interval to positions within packed sequences.""" global_rows = torch.arange( int(global_start), - int(global_start) + x.shape[0], + int(global_start) + int(local_rows), dtype=cu_seqlens_padded.dtype, - device=x.device, + device=cu_seqlens_padded.device, ) - if clamp_to_valid_token: - global_rows = torch.minimum(global_rows.clamp_min(0), cu_seqlens_padded[-1] - 1) sequence_ids = torch.bucketize( global_rows, cu_seqlens_padded[1:], out_int32=True, right=True ).clamp_max(cu_seqlens_padded.shape[0] - 2) sequence_starts = cu_seqlens_padded[sequence_ids] sequence_ends = cu_seqlens_padded[sequence_ids + 1] valid_rows = (global_rows >= sequence_starts) & (global_rows < sequence_ends) - position_ids = torch.where( - valid_rows, global_rows - sequence_starts, torch.zeros_like(global_rows) - ) + return torch.where(valid_rows, global_rows - sequence_starts, 0) + + +def apply_thd_cp_local_rope_fused( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + nope_dim: int, + pos_dim: int, + cu_seqlens_padded: torch.Tensor, + global_start: int, + inverse: bool = False, +) -> torch.Tensor: + """Apply fused non-interleaved RoPE to local THD CP rows.""" + position_ids = _thd_cp_position_ids(cu_seqlens_padded, global_start, x.shape[0]) squeezed_batch = x.ndim == 4 and x.shape[1] == 1 squeezed_head = x.ndim == 2 rope_input = x.squeeze(1) if squeezed_batch else x rope_input = rope_input.unsqueeze(1) if squeezed_head else rope_input if inverse: - # Sparse attention saves its unrotated output for backward. + # The fused kernel is in-place, but sparse-attention backward needs its original output. rope_input = rope_input.clone() output = fused_mla_rope_inplace( rope_input, @@ -76,6 +79,42 @@ def apply_thd_cp_local_rope_fused( return output +def apply_thd_cp_local_rope_unfused( + x: torch.Tensor, + rotary_pos_emb: torch.Tensor, + nope_dim: int, + pos_dim: int, + cu_seqlens_padded: torch.Tensor, + global_start: int, + config, + inverse: bool = False, +) -> torch.Tensor: + """Apply unfused RoPE to a consecutive interval of packed CP rows.""" + position_ids = _thd_cp_position_ids(cu_seqlens_padded, global_start, x.shape[0]) + freqs = torch.index_select(rotary_pos_emb, 0, position_ids.long()) + + squeezed_batch = x.ndim == 4 and x.shape[1] == 1 + squeezed_head = x.ndim == 2 + rope_input = x.squeeze(1) if squeezed_batch else x + rope_input = rope_input.unsqueeze(1) if squeezed_head else rope_input + content, rotary = torch.split(rope_input, [nope_dim, pos_dim], dim=-1) + rotary = _apply_rotary_pos_emb_bshd( + rotary, + freqs, + rotary_interleaved=config.rotary_interleaved, + mscale=1.0, + mla_rotary_interleaved=True, + inverse=inverse, + mla_output_remove_interleaving=True, + ) + output = torch.cat((content, rotary), dim=-1) + if squeezed_batch: + return output.unsqueeze(1) + if squeezed_head: + return output.squeeze(1) + return output + + # ============================================================================= # Boundary Hidden Exchange # ============================================================================= @@ -242,24 +281,9 @@ def compute_cp_indexer_topk( topk_width: int, indexer_softmax_scale: float, max_seqlen_q: int, + use_fused: bool, ) -> Optional[torch.Tensor]: - """Run indexer top-k for this rank's query block. - - This is a workaround for the DSA fused indexer forward mask contract. A - rank may need a mask whose first local query already sees a compressed K - prefix, for example: - - ``` - 1 1 1 0 0 0 0 0 - 1 1 1 1 0 0 0 0 - 1 1 1 1 1 0 0 0 - ``` - - The fused kernel assumes the last Q row can see all K rows, then derives - earlier rows' causal masks from that endpoint. A rank's last Q row may - only see a prefix of the global K rows, so this helper copies that visible - prefix before calling the kernel. - """ + """Run indexer top-k for this rank's query block.""" topk_width = int(topk_width) if topk_width == 0 or k_indexer_seq_major.shape[0] == 0: return None @@ -272,6 +296,64 @@ def compute_cp_indexer_topk( f"{l_local}, got {weights_indexer_local.shape[0]}." ) + if not use_fused: + global_rows = torch.arange( + global_start, + global_start + l_local, + dtype=cu_seqlens_q.dtype, + device=cu_seqlens_q.device, + ) + sequence_ids = torch.bucketize( + global_rows, cu_seqlens_q[1:], out_int32=True, right=True + ).clamp_max(cu_seqlens_q.shape[0] - 2) + positions = global_rows - cu_seqlens_q[sequence_ids] + visible_k = torch.minimum( + torch.div(positions + 1, int(ratio), rounding_mode="floor"), + cu_seqlens_compressed[sequence_ids + 1] - cu_seqlens_compressed[sequence_ids], + ).clamp_min(0) + valid_q = (global_rows >= cu_seqlens_q[sequence_ids]) & ( + global_rows < cu_seqlens_q[sequence_ids + 1] + ) + + k_rows = torch.arange( + k_indexer_seq_major.shape[0], + dtype=cu_seqlens_compressed.dtype, + device=cu_seqlens_compressed.device, + ) + k_sequence_ids = torch.bucketize( + k_rows, cu_seqlens_compressed[1:], out_int32=True, right=True + ).clamp_max(cu_seqlens_compressed.shape[0] - 2) + k_positions = k_rows - cu_seqlens_compressed[k_sequence_ids] + output = torch.full( + (l_local, topk_width), -1, dtype=torch.int32, device=q_indexer_local.device + ) + selected_width = min(topk_width, k_indexer_seq_major.shape[0]) + for start in range(0, l_local, 128): + end = min(start + 128, l_local) + scores = torch.einsum( + "rhd,kd->rhk", q_indexer_local[start:end].float(), k_indexer_seq_major.float() + ) + scores = torch.relu(scores) * weights_indexer_local[start:end].float().unsqueeze(-1) + scores = scores.sum(dim=1) * float(indexer_softmax_scale) + valid_k = ( + (k_sequence_ids.unsqueeze(0) == sequence_ids[start:end].unsqueeze(1)) + & (k_positions.unsqueeze(0) < visible_k[start:end].unsqueeze(1)) + & valid_q[start:end].unsqueeze(1) + ) + scores = scores.masked_fill(~valid_k, float("-inf")) + values, rows = torch.topk(scores, selected_width, dim=-1) + local_rows = k_positions[rows].to(torch.int32) + output[start:end, :selected_width] = torch.where(torch.isfinite(values), local_rows, -1) + return output + + # The fused kernel derives a bottom-right causal mask from each Q/K + # segment. Copy only the prefix visible to this rank so that mask starts + # at the rank's true position in the global sequence. For example: + + # 1 1 1 0 0 0 0 0 + # 1 1 1 1 0 0 0 0 + # 1 1 1 1 1 0 0 0 + global_end = global_start + l_local zero = torch.zeros((1,), dtype=cu_seqlens_q.dtype, device=cu_seqlens_q.device) k_rows = torch.arange( diff --git a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py index 928dce2d074..91c14892aee 100644 --- a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py +++ b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py @@ -409,6 +409,18 @@ def forward( ) if packed_seq: core_attn_out = core_attn_out.unsqueeze(1) + elif use_thd_cp: + global_start = self.pg_collection.cp.rank() * core_attn_out.shape[0] + core_attn_out = cp_utils.apply_thd_cp_local_rope_unfused( + core_attn_out, + rotary_pos_emb, + nope_dim, + pos_dim, + cu_seqlens_kv, + global_start, + self.config, + inverse=True, + ) else: content_part, rot_part = torch.split( core_attn_out, [core_attn_out.size(-1) - pos_dim, pos_dim], dim=-1 @@ -723,7 +735,6 @@ def qkv_up_proj_and_rope_apply( self.config.qk_pos_emb_head_dim, cu_seqlens_q, global_start - boundary_rows, - clamp_to_valid_token=bool(boundary_rows), ) if boundary_kv_compressed is not None: boundary_kv = kv[:boundary_rows] @@ -757,65 +768,76 @@ def qkv_up_proj_and_rope_apply( value = kv else: if packed_seq and cp_size > 1: - raise RuntimeError( - "DSv4 THD CP requires apply_rope_fusion=True; " - "the unfused THD CP RoPE path is not implemented." + global_start = self.pg_collection.cp.rank() * q.shape[0] + query = cp_utils.apply_thd_cp_local_rope_unfused( + q, + rotary_pos_emb, + self.config.qk_head_dim, + self.config.qk_pos_emb_head_dim, + cu_seqlens_q, + global_start, + self.config, ) - - q_len = q.size()[0] - if packed_seq_params is None or self.config.context_parallel_size == 1: + kv = cp_utils.apply_thd_cp_local_rope_unfused( + kv.unsqueeze(-2), + rotary_pos_emb, + self.config.qk_head_dim, + self.config.qk_pos_emb_head_dim, + cu_seqlens_kv, + global_start - boundary_rows, + self.config, + ) + if boundary_kv_compressed is not None: + boundary_kv = kv[:boundary_rows] + kv = kv[boundary_rows:] + key = value = kv + else: + q_len = q.size()[0] # Shorten rotary_pos_emb to the sequence length when inference_params - # is not provided. This makes sure we can run forward directly with - # any sequence length. During training, the sequence length is always - # the full rotary_pos_emb length, except for sequence packing + CP. - # When sequence packing and context parallel are both enabled, the - # position embedding will not split rotary_pos_emb, so it may exceed - # the sequence length on this CP rank, but we need the full rotary_pos_emb - # to cover the full sequence, so we do not shorten it here. + # is not provided so direct forward accepts any sequence length. rotary_pos_emb = rotary_pos_emb[0:q_len] - # q_no_pe: [num_tokens, n, qk_head_dim] - # q_pos_emb: [num_tokens, n, qk_pos_emb_head_dim] - q_no_pe, q_pos_emb = torch.split( - q, [self.config.qk_head_dim, self.config.qk_pos_emb_head_dim], dim=-1 - ) - - # RoPE and query (shared for wkv and latent) - # q_pos_emb: [num_tokens, n, qk_pos_emb_head_dim] - q_pos_emb = apply_rotary_pos_emb( - q_pos_emb, - rotary_pos_emb, - config=self.config, - cu_seqlens=cu_seqlens_q, - mscale=mscale, - cp_group=self.pg_collection.cp, - mla_rotary_interleaved=True, - mla_output_remove_interleaving=True, - max_seqlen=rope_max_seqlen_q, - ) - # query: [num_tokens, n, (qk_head_dim + v_head_dim)] - query = torch.cat([q_no_pe, q_pos_emb], dim=-1) - - pos_dim = self.config.qk_pos_emb_head_dim - kv_no_pe, k_pos_emb = torch.split(kv, [kv.size(-1) - pos_dim, pos_dim], dim=-1) + # q_no_pe: [num_tokens, n, qk_head_dim] + # q_pos_emb: [num_tokens, n, qk_pos_emb_head_dim] + q_no_pe, q_pos_emb = torch.split( + q, [self.config.qk_head_dim, self.config.qk_pos_emb_head_dim], dim=-1 + ) - # k_pos_emb:[num_tokens, 1, qk_pos_emb_head_dim] - k_pos_emb = apply_rotary_pos_emb( - k_pos_emb, - rotary_pos_emb, - config=self.config, - cu_seqlens=cu_seqlens_kv, - mscale=mscale, - cp_group=self.pg_collection.cp, - mla_rotary_interleaved=True, - mla_output_remove_interleaving=True, - max_seqlen=rope_max_seqlen_kv, - ) + # RoPE and query (shared for wkv and latent) + # q_pos_emb: [num_tokens, n, qk_pos_emb_head_dim] + q_pos_emb = apply_rotary_pos_emb( + q_pos_emb, + rotary_pos_emb, + config=self.config, + cu_seqlens=cu_seqlens_q, + mscale=mscale, + cp_group=self.pg_collection.cp, + mla_rotary_interleaved=True, + mla_output_remove_interleaving=True, + max_seqlen=rope_max_seqlen_q, + ) + # query: [num_tokens, n, (qk_head_dim + v_head_dim)] + query = torch.cat([q_no_pe, q_pos_emb], dim=-1) + + pos_dim = self.config.qk_pos_emb_head_dim + kv_no_pe, k_pos_emb = torch.split(kv, [kv.size(-1) - pos_dim, pos_dim], dim=-1) + + # k_pos_emb:[num_tokens, 1, qk_pos_emb_head_dim] + k_pos_emb = apply_rotary_pos_emb( + k_pos_emb, + rotary_pos_emb, + config=self.config, + cu_seqlens=cu_seqlens_kv, + mscale=mscale, + cp_group=self.pg_collection.cp, + mla_rotary_interleaved=True, + mla_output_remove_interleaving=True, + max_seqlen=rope_max_seqlen_kv, + ) - # Single head: key = value = [num_tokens, 1, v_head_dim] - kv = torch.cat([kv_no_pe, k_pos_emb], dim=-1).unsqueeze(-2) - key = kv - value = kv + # Single head: key = value = [num_tokens, 1, v_head_dim] + kv = torch.cat([kv_no_pe, k_pos_emb], dim=-1).unsqueeze(-2) + key = value = kv query = query.contiguous() key = key.contiguous() diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_layout_kernels.py b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_layout_kernels.py index b24513ca714..0aa688d1067 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_layout_kernels.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_layout_kernels.py @@ -6,6 +6,9 @@ import torch from megatron.core.transformer.experimental_attention_variant import csa_cp_layout_kernels +from megatron.core.transformer.experimental_attention_variant.csa_cp_utils import ( + prepare_cp_compressor_input, +) # This file guards only DSv4 CP layout/metadata kernels. Layer-level CUDA graph # tests guard graph capture/replay behavior. @@ -411,3 +414,198 @@ def test_build_attention_indices_indexer_loss_mode_matches_native(): ) assert torch.equal(fused[0], expected[0]) assert torch.equal(fused[2], expected[1]) + + +@pytest.mark.parametrize( + ("ratio", "lengths"), + [(4, (3, 10, 20, 5, 33, 10, 27, 20)), (128, (3, 130, 170, 5, 260, 129, 200, 127))], + ids=["ratio4", "ratio128"], +) +@pytest.mark.parametrize("cp_size", [2, 4]) +def test_composed_cp_layout_maps_every_index_and_gradient_to_its_source(ratio, lengths, cp_size): + """Compose compaction, rank-major gather, and final index lowering.""" + _require_cute_cuda() + total = sum(lengths) + local_rows = total // cp_size + d_window = 8 if ratio == 4 else ratio + window_size = 4 + cu = torch.tensor( + [0] + list(torch.tensor(lengths).cumsum(0).tolist()), dtype=torch.int32, device="cuda" + ) + cu_compressed = _compressed_cu_seqlens(cu, ratio) + hidden = ( + torch.arange(1, total + 1, dtype=torch.float32, device="cuda") + .unsqueeze(1) + .requires_grad_(True) + ) + + boundaries = [] + locals_ = [] + compact_values = [] + compact_first_tokens = [] + compact_ids = [] + row_maps = [] + for rank in range(cp_size): + start = rank * local_rows + local = hidden[start : start + local_rows] + boundary = torch.cat( + ( + hidden.new_zeros((max(0, d_window - start), 1)), + hidden[max(0, start - d_window) : start], + ) + ) + compact, group_ids, row_map = prepare_cp_compressor_input( + local, boundary, cu, cu_compressed, start, cp_size, ratio + ) + grouped = compact.reshape(group_ids.shape[0], ratio, 1) + boundaries.append(boundary) + locals_.append(local) + compact_values.append(grouped.sum(dim=1)) + compact_first_tokens.append(grouped[:, 0, 0]) + compact_ids.append(group_ids) + row_maps.append(row_map) + + capacity = compact_ids[0].shape[0] + expected_map = torch.full((total // ratio,), -1, dtype=torch.int32, device="cuda") + physical_to_tokens = {} + logical_values = [] + logical_row = 0 + seq_start = 0 + for seq_len in lengths: + for group in range(seq_len // ratio): + first_token = seq_start + group * ratio + owner = (first_token + ratio - 1) // local_rows + matches = torch.nonzero( + compact_first_tokens[owner] == first_token + 1, as_tuple=False + ).flatten() + assert matches.numel() == 1 + slot = int(matches[0]) + assert int(compact_ids[owner][slot]) == group + physical_row = owner * capacity + slot + expected_map[logical_row] = physical_row + physical_to_tokens[physical_row] = range(first_token, first_token + ratio) + logical_values.append(sum(range(first_token + 1, first_token + ratio + 1))) + logical_row += 1 + seq_start += seq_len + + for row_map in row_maps: + assert torch.equal(row_map, expected_map) + reachable = set(int(row) for row in expected_map[expected_map >= 0].cpu().tolist()) + all_ids = torch.cat(compact_ids) + assert all(int(row) not in reachable for row in torch.nonzero(all_ids < 0).flatten().tolist()) + + compressed_rank_major = torch.cat(compact_values) + sequence_major = torch.index_select( + compressed_rank_major, 0, expected_map[:logical_row].long() + ).squeeze(1) + assert torch.equal( + sequence_major, torch.tensor(logical_values, dtype=torch.float32, device="cuda") + ) + compressed_width = 3 if ratio == 4 else max(lengths) // ratio + loss = hidden.new_zeros(()) + expected_grad = torch.zeros_like(hidden) + cu_list = [int(value) for value in cu.cpu().tolist()] + for rank in range(cp_size): + start = rank * local_rows + logical_topk = None + if ratio == 4: + logical_topk = torch.full( + (local_rows, compressed_width), -1, dtype=torch.int32, device="cuda" + ) + for row, global_row in enumerate(range(start, start + local_rows)): + seq = next(i for i in range(len(lengths)) if global_row < cu_list[i + 1]) + visible = min((global_row - cu_list[seq] + 1) // ratio, lengths[seq] // ratio) + selected = list(range(visible - 1, max(-1, visible - compressed_width - 1), -1)) + if selected: + logical_topk[row, : len(selected)] = torch.tensor( + selected, dtype=torch.int32, device="cuda" + ) + + actual = csa_cp_layout_kernels.build_attention_indices( + cu, + start, + local_rows, + d_window, + window_size, + ratio, + compressed_width, + logical_topk, + cu_seqlens_compressed=cu_compressed, + seq_to_rank_row=expected_map, + ) + expected = _native_attention_indices( + cu, + cu_compressed, + start, + local_rows, + d_window, + window_size, + ratio, + compressed_width, + expected_map, + d_window + local_rows, + logical_topk, + ) + assert torch.equal(actual[0], expected[0]) + assert torch.equal(actual[1], expected[1]) + + if logical_topk is not None: + actual = csa_cp_layout_kernels.build_attention_indices( + cu, + start, + local_rows, + d_window, + window_size, + ratio, + compressed_width, + logical_topk, + cu_seqlens_compressed=cu_compressed, + seq_to_rank_row=expected_map, + for_indexer_loss=True, + ) + expected_loss_indices = _native_indexer_loss_indices( + cu, + cu_compressed, + start, + local_rows, + d_window, + window_size, + ratio, + logical_topk, + expected_map, + d_window + local_rows, + ) + assert torch.equal(actual[0], expected_loss_indices[0]) + assert torch.equal(actual[2], expected_loss_indices[1]) + + indices = actual[0] + valid = indices >= 0 + kv = torch.cat((boundaries[rank], locals_[rank], compressed_rank_major)) + selected = torch.index_select(kv, 0, indices.clamp_min(0).long().flatten()).reshape( + indices.shape + ) + coefficients = ( + torch.arange(indices.numel(), dtype=torch.float32, device="cuda").reshape(indices.shape) + + rank * indices.numel() + + 1 + ) + loss = loss + (selected * coefficients * valid).sum() + + compressed_base = d_window + local_rows + for row, index_row in enumerate(indices.cpu().tolist()): + for column, index in enumerate(index_row): + if index < 0: + continue + coefficient = float(coefficients[row, column]) + if index < compressed_base: + token = start - d_window + index + assert 0 <= token <= start + row + expected_grad[token] += coefficient + else: + physical_row = index - compressed_base + assert physical_row in physical_to_tokens + for token in physical_to_tokens[physical_row]: + expected_grad[token] += coefficient + + loss.backward() + assert torch.equal(hidden.grad, expected_grad) diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py index e0e31e576d6..37bc61a87f1 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py @@ -1,5 +1,7 @@ # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +from types import SimpleNamespace + import pytest import torch @@ -8,6 +10,7 @@ from megatron.core.transformer.experimental_attention_variant import csa_cp_utils from megatron.core.transformer.experimental_attention_variant.csa_cp_utils import ( apply_thd_cp_local_rope_fused, + apply_thd_cp_local_rope_unfused, compute_cp_indexer_topk, exchange_cp_boundary_hidden, prepare_cp_compressor_input, @@ -38,11 +41,11 @@ def _rope_reference(x, cos, sin, positions, nope_dim, pos_dim, inverse=False): return torch.cat((x_nope, rotated), dim=-1) -def _sequence_positions(cu_seqlens, rows, clamp=False): - if clamp: - rows = rows.clamp(0, cu_seqlens[-1] - 1) +def _sequence_positions(cu_seqlens, rows): seq_ids = torch.bucketize(rows, cu_seqlens[1:], right=True).clamp_max(cu_seqlens.shape[0] - 2) - return rows - cu_seqlens[seq_ids] + starts = cu_seqlens[seq_ids] + ends = cu_seqlens[seq_ids + 1] + return torch.where((rows >= starts) & (rows < ends), rows - starts, 0) def test_thd_cp_left_boundary_exchange_forward_backward(): @@ -126,20 +129,48 @@ def test_apply_thd_cp_local_rope_matches_reference_forward_backward(dtype, inver torch.testing.assert_close(actual_x.grad, ref_x.grad, rtol=rtol, atol=atol) -def test_apply_thd_cp_local_rope_clamps_boundary_rows(): +def test_apply_thd_cp_local_rope_maps_invalid_boundary_rows_to_position_zero(): _require_cuda() cu = torch.tensor([0, 4, 12], dtype=torch.int32, device="cuda") global_start = -2 global_rows = torch.arange(-2, 14, dtype=torch.int32, device="cuda") - positions = _sequence_positions(cu, global_rows, clamp=True) + positions = _sequence_positions(cu, global_rows) x = torch.randn(global_rows.shape[0], 1, 8, device="cuda") cos = torch.randn(8, 4, device="cuda") sin = torch.randn(8, 4, device="cuda") expected = _rope_reference(x, cos, sin, positions, 4, 4) - actual = apply_thd_cp_local_rope_fused( - x, cos, sin, 4, 4, cu, global_start, clamp_to_valid_token=True + actual = apply_thd_cp_local_rope_fused(x, cos, sin, 4, 4, cu, global_start) + torch.testing.assert_close(actual, expected) + + +@pytest.mark.parametrize("inverse", [False, True]) +def test_apply_thd_cp_local_rope_unfused_matches_explicit_positions(inverse): + _require_cuda() + torch.manual_seed(17) + cu = torch.tensor([0, 3, 9], dtype=torch.int32, device="cuda") + global_start = 2 + global_rows = torch.arange(2, 7, dtype=torch.int32, device="cuda") + positions = _sequence_positions(cu, global_rows) + x = torch.randn(5, 2, 8, device="cuda", requires_grad=True) + freqs = torch.randn(6, 1, 1, 4, device="cuda") + config = SimpleNamespace( + apply_rope_fusion=False, rotary_interleaved=False, multi_latent_attention=False + ) + + expected_x = x.detach().clone().requires_grad_(True) + expected = _rope_reference( + expected_x, torch.cos(freqs[:, 0, 0]), torch.sin(freqs[:, 0, 0]), positions, 4, 4, inverse + ) + grad = torch.randn_like(expected) + expected.backward(grad) + + actual_x = x.detach().clone().requires_grad_(True) + actual = apply_thd_cp_local_rope_unfused( + actual_x, freqs, 4, 4, cu, global_start, config, inverse=inverse ) + actual.backward(grad) torch.testing.assert_close(actual, expected) + torch.testing.assert_close(actual_x.grad, expected_x.grad) def test_prepare_cp_compressor_input_builds_rank_row_map(monkeypatch): @@ -200,6 +231,7 @@ def fake_indexer_topk( topk_width=2, indexer_softmax_scale=0.5, max_seqlen_q=10, + use_fused=True, ) assert torch.equal(out, torch.ones(5, 2, dtype=torch.int32)) @@ -208,4 +240,55 @@ def fake_indexer_topk( assert torch.equal(topk_calls[0][1], torch.tensor([0, 2, 5], dtype=torch.int32)) assert torch.equal(topk_calls[0][2], torch.tensor([0, 2, 2], dtype=torch.int32)) assert torch.equal(topk_calls[0][3], torch.tensor([2, 2, 0, 0, 0], dtype=torch.int32)) - assert compute_cp_indexer_topk(q, weights, k_seq[:0], cu_q, cu_comp, 2, 4, 2, 1.0, 10) is None + assert ( + compute_cp_indexer_topk(q, weights, k_seq[:0], cu_q, cu_comp, 2, 4, 2, 1.0, 10, True) + is None + ) + + +def test_compute_cp_indexer_topk_unfused_uses_exact_global_positions(monkeypatch): + def fail_if_fused(*_args, **_kwargs): + raise AssertionError("unfused top-k must not call the fused kernel") + + monkeypatch.setattr(csa_cp_utils, "indexer_topk", fail_if_fused) + torch.manual_seed(19) + q = torch.randn(8, 2, 3) + weights = torch.randn(8, 2) + k = torch.randn(4, 3) + cu_q = torch.tensor([0, 5, 13, 20], dtype=torch.int32) + cu_k = torch.tensor([0, 1, 3, 4], dtype=torch.int32) + ratio = 4 + topk_width = 3 + scale = 0.7 + + actual = compute_cp_indexer_topk( + q, + weights, + k, + cu_q, + cu_k, + global_start=7, + ratio=ratio, + topk_width=topk_width, + indexer_softmax_scale=scale, + max_seqlen_q=8, + use_fused=False, + ) + + expected = torch.full((q.shape[0], topk_width), -1, dtype=torch.int32) + for local_row, global_row in enumerate(range(7, 15)): + seq = next(i for i in range(len(cu_q) - 1) if global_row < int(cu_q[i + 1])) + visible = min((global_row - int(cu_q[seq]) + 1) // ratio, int(cu_k[seq + 1] - cu_k[seq])) + candidates = [] + for local_k in range(visible): + k_row = int(cu_k[seq]) + local_k + score = 0.0 + for head in range(q.shape[1]): + dot = torch.dot(q[local_row, head], k[k_row]).item() + score += max(dot, 0.0) * weights[local_row, head].item() * scale + candidates.append((score, local_k)) + candidates.sort(reverse=True) + for column, (_, local_k) in enumerate(candidates[:topk_width]): + expected[local_row, column] = local_k + + assert torch.equal(actual, expected) diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py index e2ffbde272f..2f2cfb0f9a7 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py @@ -547,7 +547,8 @@ def _measure_cp1_cuda_graph_time(self, layer_number, packed, padded_tokens): [1, 2, 3], ids=["ratio_0_window_only", "ratio_4_indexer", "ratio_128_compressor"], ) - def test_thd_cp_matches_full_reference_forward_backward(self, layer_number): + @pytest.mark.parametrize("apply_rope_fusion", [True, False], ids=["fused_rope", "unfused_rope"]) + def test_thd_cp_matches_full_reference_forward_backward(self, layer_number, apply_rope_fusion): """CP path matches the full-sequence THD reference on ragged packed inputs using the DSv4 layer configuration. @@ -558,12 +559,13 @@ def test_thd_cp_matches_full_reference_forward_backward(self, layer_number): torch.manual_seed(_SEED + layer_number) model_parallel_cuda_manual_seed(_SEED + layer_number) + apply_dsa_kernel_fusion = self.fused_kernels_available config_cp = _make_dsv4_cp_config( context_parallel_size=self.cp_size, dsa_indexer_loss_coeff=1.0, dsa_indexer_use_sparse_loss=True, - apply_dsa_kernel_fusion=self.fused_kernels_available, - apply_rope_fusion=True, + apply_dsa_kernel_fusion=apply_dsa_kernel_fusion, + apply_rope_fusion=apply_rope_fusion, ) config_ref = _make_dsv4_cp_config( context_parallel_size=1, @@ -571,7 +573,7 @@ def test_thd_cp_matches_full_reference_forward_backward(self, layer_number): dsa_indexer_use_sparse_loss=True, # Numerical parity uses the unfused CP1 reference so failures point # at the CP path instead of non-CP fused RoPE behavior. - apply_dsa_kernel_fusion=self.fused_kernels_available, + apply_dsa_kernel_fusion=apply_dsa_kernel_fusion, apply_rope_fusion=False, ) cp_attn = _build_attention( @@ -597,7 +599,7 @@ def test_thd_cp_matches_full_reference_forward_backward(self, layer_number): _assert_cp_tensor_match( local_out.detach(), ref_out.detach().index_select(0, local_idx), - f"layer={layer_number}:output", + f"layer={layer_number}:dsa={apply_dsa_kernel_fusion}:rope={apply_rope_fusion}:output", ) grad = torch.randn_like(ref_out) @@ -606,7 +608,7 @@ def test_thd_cp_matches_full_reference_forward_backward(self, layer_number): _assert_cp_tensor_match( local_hidden.grad.detach(), ref_hidden.grad.index_select(0, local_idx), - f"layer={layer_number}:hidden_grad", + f"layer={layer_number}:dsa={apply_dsa_kernel_fusion}:rope={apply_rope_fusion}:hidden_grad", ) ref_params = dict(ref_attn.named_parameters()) @@ -616,7 +618,12 @@ def test_thd_cp_matches_full_reference_forward_backward(self, layer_number): assert ref_grad is not None, f"Missing reference grad for {name}" grad_sum = param.grad.detach().clone() dist.all_reduce(grad_sum, group=self.pg.cp) - _assert_cp_tensor_match(grad_sum, ref_grad, f"layer={layer_number}:param_grad:{name}") + _assert_cp_tensor_match( + grad_sum, + ref_grad, + f"layer={layer_number}:dsa={apply_dsa_kernel_fusion}:" + f"rope={apply_rope_fusion}:param_grad:{name}", + ) del cp_attn, ref_attn, full_hidden, local_hidden, ref_hidden, local_out, ref_out, grad _clear_cuda_test_state() @@ -666,8 +673,14 @@ def test_thd_cp_ratio4_eval_matches_full_reference(self): [1, 2, 3], ids=["ratio_0_window_only", "ratio_4_indexer", "ratio_128_compressor"], ) - @pytest.mark.parametrize("fused", [True, False]) - def test_thd_cp_cuda_graph_matches_eager_forward_backward(self, layer_number, fused): + @pytest.mark.parametrize( + ("dsa_fused", "rope_fused"), + [(True, True), (False, True), (False, False)], + ids=["fused", "unfused_dsa", "unfused_dsa_rope"], + ) + def test_thd_cp_cuda_graph_matches_eager_forward_backward( + self, layer_number, dsa_fused, rope_fused + ): """CUDA graph replay matches eager THD CP forward/backward. Captures the DSv4 attention layer's CP-local forward and backward @@ -679,11 +692,11 @@ def test_thd_cp_cuda_graph_matches_eager_forward_backward(self, layer_number, fu plus elementwise ``assert_close`` gates because they may be nondeterministic against eager execution. """ - if fused and not self.fused_kernels_available: + if (dsa_fused or rope_fused) and not self.fused_kernels_available: pytest.skip(_DSV4_CP_FUSED_KERNELS_UNAVAILABLE_REASON) - context = nullcontext() if fused else _deterministic_torch_algorithms() - mode = "fused" if fused else "unfused" + context = nullcontext() if dsa_fused else _deterministic_torch_algorithms() + mode = f"dsa_fused={dsa_fused}:rope_fused={rope_fused}" with context: packed, padded_tokens, local_idx = _make_ragged_cp_case(self.cp_size, self.cp_rank) @@ -693,8 +706,8 @@ def test_thd_cp_cuda_graph_matches_eager_forward_backward(self, layer_number, fu context_parallel_size=self.cp_size, dsa_indexer_loss_coeff=1.0, dsa_indexer_use_sparse_loss=True, - apply_dsa_kernel_fusion=fused, - apply_rope_fusion=True, + apply_dsa_kernel_fusion=dsa_fused, + apply_rope_fusion=rope_fused, ) graph_attn = _build_attention( config, layer_number=layer_number, pg_collection=self.pg @@ -744,7 +757,7 @@ def test_thd_cp_cuda_graph_matches_eager_forward_backward(self, layer_number, fu output_label = f"layer={layer_number}:{mode}:output" _assert_cp_graph_bitwise_match(graph_out, eager_out, output_label) bwd_match_fn = ( - _assert_cp_graph_fused_grad_match if fused else _assert_cp_graph_bitwise_match + _assert_cp_graph_fused_grad_match if dsa_fused else _assert_cp_graph_bitwise_match ) bwd_match_fn( graph_hidden_grad, eager_hidden_grad, f"layer={layer_number}:{mode}:hidden_grad" From 867d4c80458a249ccbb9626cf4efb50c8e06c38d Mon Sep 17 00:00:00 2001 From: kunlunl Date: Mon, 29 Jun 2026 01:00:55 -0700 Subject: [PATCH 10/23] Fix THD padding on metadata-only pipeline stages Signed-off-by: kunlunl --- megatron/core/packed_seq_params.py | 17 ++++----- .../transformer/test_thd_cuda_graph.py | 36 +++++++++++++++++++ 2 files changed, 45 insertions(+), 8 deletions(-) diff --git a/megatron/core/packed_seq_params.py b/megatron/core/packed_seq_params.py index 3095b1b8464..9352e8ef055 100644 --- a/megatron/core/packed_seq_params.py +++ b/megatron/core/packed_seq_params.py @@ -198,6 +198,7 @@ def _resolve_thd_padding_lengths( packed_seq_params: PackedSeqParams, target_len: Optional[int], alignment: Optional[int], + padding_mask: Optional[Tensor] = None, cp_group: Optional[dist.ProcessGroup] = None, cp_size: Optional[int] = None, cp_rank: Optional[int] = None, @@ -205,7 +206,7 @@ def _resolve_thd_padding_lengths( """Resolve local/global THD padding lengths without changing tensors. Returns: - local_actual_T: Current rank's token-like tensor length. + local_actual_T: Current rank's physical row count. global_actual_T: Global packed length represented by THD metadata. local_target_len: Current rank's padded token-like tensor length. global_target_len: Global padded endpoint represented by THD metadata. @@ -216,10 +217,10 @@ def _resolve_thd_padding_lengths( packed_seq_params, cp_group=cp_group, cp_size=cp_size, cp_rank=cp_rank ) - # Find the first token-like tensor that carries this rank's local length. + # Find a tensor that already carries this rank's local physical row count. local_tensor_T = None mask_device = None - for candidate in (tokens, labels, loss_mask, position_ids): + for candidate in (tokens, labels, loss_mask, position_ids, padding_mask): if candidate is not None: local_tensor_T = int(candidate.shape[-1]) mask_device = candidate.device @@ -234,7 +235,7 @@ def _resolve_thd_padding_lengths( else: assert has_local_tensor, ( "packed_seq_params.cu_seqlens_q must be available to derive padding_mask " - "when tokens/labels/loss_mask/position_ids are all None." + "when tokens/labels/loss_mask/position_ids/padding_mask are all None." ) global_actual_T = local_tensor_T * cp_size @@ -382,10 +383,9 @@ def pad_sequence_for_thd( cp_rank: Explicit context-parallel rank used with ``cp_size``. Notes: - - THD CP slicing is defined by Transformer Engine. On metadata-only - stages, Megatron asks TE which packed rows this CP rank would receive - and uses that row count as the local length instead of assuming equal - division by CP size. + - On stages without token-like tensors, an already CP-sliced padding + mask supplies the local physical row count. Legacy callers without a + padding mask ask TE which packed rows this CP rank would receive. - When ``pad_by_appending_dummy_seq`` is true, the padding tail is also represented as an ordinary dummy sequence in cu_seqlens metadata. - ``max_num_seqs`` pads all four cu_seqlens tensors; this is required @@ -408,6 +408,7 @@ def pad_sequence_for_thd( packed_seq_params, target_len=target_len, alignment=alignment, + padding_mask=padding_mask, cp_group=cp_group, cp_size=cp_size, cp_rank=cp_rank, diff --git a/tests/unit_tests/transformer/test_thd_cuda_graph.py b/tests/unit_tests/transformer/test_thd_cuda_graph.py index 92a81b2fcdc..665e4d6ad14 100644 --- a/tests/unit_tests/transformer/test_thd_cuda_graph.py +++ b/tests/unit_tests/transformer/test_thd_cuda_graph.py @@ -474,6 +474,42 @@ def test_padding_mask_preserves_existing_padding(self): ), ) + @pytest.mark.internal + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_cp_metadata_only_uses_local_padding_mask(self, monkeypatch): + """A middle PP stage must not repartition already CP-sliced physical rows.""" + packed_seq_params = _make_psp([15]) + cu_seqlens_padded = torch.tensor([0, 16], dtype=torch.int32, device="cuda") + packed_seq_params.cu_seqlens_q_padded = cu_seqlens_padded + packed_seq_params.cu_seqlens_kv_padded = cu_seqlens_padded.clone() + padding_mask = torch.tensor( + [[False, False, False, True]], dtype=torch.bool, device="cuda" + ) + + def fail_if_called(*_args, **_kwargs): + pytest.fail("already sliced padding_mask must avoid TE repartitioning") + + monkeypatch.setattr( + "megatron.core.extensions.transformer_engine.get_thd_partitioned_indices", + fail_if_called, + ) + + _, _, _, _, padded_params, padded_mask = pad_sequence_for_thd( + None, + None, + None, + None, + packed_seq_params, + target_len=4, + padding_mask=padding_mask, + cp_size=4, + cp_rank=3, + ) + + assert torch.equal(padded_mask, padding_mask) + assert padded_params.cu_seqlens_q.tolist() == [0, 15, 16] + assert padded_params.cu_seqlens_q_padded.tolist() == [0, 16, 16] + @pytest.mark.internal @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") def test_cu_seqlens_fill_value(self): From e75190a21ad6a2706d8539c3cd07f77281378230 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Tue, 30 Jun 2026 02:33:28 -0700 Subject: [PATCH 11/23] Fix hash MoE inputs for THD CUDA graphs Signed-off-by: kunlunl --- .../core/transformer/transformer_layer.py | 19 ++++++------------- .../transformer/test_thd_cuda_graph.py | 17 +++++++++++++++++ 2 files changed, 23 insertions(+), 13 deletions(-) diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index d3571152966..2ad6e150b53 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -1257,8 +1257,13 @@ def get_layer_static_inputs(self, seq_length, micro_batch_size): and self.config.moe_n_hash_layers > 0 and getattr(self.mlp.router, 'is_hash_layer', False) ): + input_ids_shape = ( + (1, self.config.max_seqlen_per_dp_cp_rank) + if self._is_thd_cuda_graph() + else (micro_batch_size, seq_length) + ) static_inputs["input_ids"] = torch.zeros( - (micro_batch_size, seq_length), dtype=torch.long, device=torch.cuda.current_device() + input_ids_shape, dtype=torch.long, device=torch.cuda.current_device() ) return static_inputs @@ -1897,18 +1902,6 @@ def get_layer_static_inputs(self, seq_length, micro_batch_size): device=hs.device, ) - # Add input_ids for hash-based MoE routing under CUDA graphs. - # Only add for layers that actually use hash routing, - # since other layers (e.g. on later PP stages) receive input_ids=None. - if ( - self.is_moe_layer - and self.config.moe_n_hash_layers > 0 - and getattr(self.mlp.router, 'is_hash_layer', False) - ): - static_inputs["input_ids"] = torch.zeros( - (micro_batch_size, seq_length), dtype=torch.long, device=torch.cuda.current_device() - ) - return static_inputs def _get_submodules_under_cudagraphs(self): diff --git a/tests/unit_tests/transformer/test_thd_cuda_graph.py b/tests/unit_tests/transformer/test_thd_cuda_graph.py index 665e4d6ad14..70245498cac 100644 --- a/tests/unit_tests/transformer/test_thd_cuda_graph.py +++ b/tests/unit_tests/transformer/test_thd_cuda_graph.py @@ -614,6 +614,23 @@ def test_thd_static_padding_mask_is_unmasked_for_capture(self): assert static_inputs["padding_mask"].shape == (1, 128) assert not static_inputs["padding_mask"].any() + @pytest.mark.internal + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_thd_hash_router_input_ids_use_per_rank_capacity(self): + """Hash routing input IDs must use the fixed [1, local_tokens] THD shape.""" + layer = _build_layer(256, 4, 4, 1024, 128, 8) + layer.config.context_parallel_size = 2 + layer.config.sequence_packing_scheduler = "dp_balanced" + layer.config.cuda_graph_impl = "transformer_engine" + layer.config.moe_n_hash_layers = 1 + layer.is_moe_layer = True + layer.mlp.router = torch.nn.Identity() + layer.mlp.router.is_hash_layer = True + + static_inputs = layer.get_layer_static_inputs(seq_length=256, micro_batch_size=2) + + assert static_inputs["input_ids"].shape == (1, 128) + class TestDynamicMicrobatchSlots: From 83b19386c1d17cbef1e296614b44183512c6300a Mon Sep 17 00:00:00 2001 From: kunlunl Date: Tue, 30 Jun 2026 05:42:56 -0700 Subject: [PATCH 12/23] Harden DSv4 THD context parallel support Signed-off-by: kunlunl --- megatron/core/datasets/data_schedule.py | 20 ++- megatron/core/datasets/data_schedule_utils.py | 16 +- .../core/extensions/transformer_engine.py | 3 +- megatron/core/packed_seq_params.py | 2 + .../experimental_attention_variant/csa.py | 8 +- .../csa_cp_layout_kernels.py | 168 ++++++++++-------- .../deepseek_v4_hybrid_attention.py | 45 +++-- .../dsa_kernels.py | 10 +- .../core/transformer/transformer_config.py | 31 ++++ .../core/transformer/transformer_layer.py | 1 + tests/unit_tests/test_sequence_packing.py | 8 +- .../test_csa_cp_layout_kernels.py | 71 ++++++++ .../test_dsv4_hybrid_attention_cp.py | 61 +++++++ .../transformer/test_thd_cuda_graph.py | 6 + 14 files changed, 338 insertions(+), 112 deletions(-) diff --git a/megatron/core/datasets/data_schedule.py b/megatron/core/datasets/data_schedule.py index 30dd1e12b92..e51e56dfe11 100644 --- a/megatron/core/datasets/data_schedule.py +++ b/megatron/core/datasets/data_schedule.py @@ -587,8 +587,8 @@ def get_batch_on_this_rank_for_sequence_packing( data_iterator (Iterator): The data iterator to get the batch from. mtp_on_this_rank (bool): Whether to use multi-token prediction. vp_stage (Optional[int]): The stage of the pipeline. - config: Model parallel config used for optional THD packed-sequence padding. - When None or config.pad_packed_seq_alignment is None, no padding is applied. + config: Model config used for CP partitioning and optional THD packed-sequence padding. + When None, CP partitioning defaults to zigzag and no padding is applied. Returns: tuple of (tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params, padding_mask) @@ -644,11 +644,7 @@ def get_batch_on_this_rank_for_sequence_packing( group_size=local_cp_size_val ) - use_contiguous_cp_slice = ( - config is not None - and getattr(config, 'experimental_attention_variant', None) == "dsv4_hybrid" - and cp_group.size() > 1 - ) + cp_partition_mode = config.cp_partition_mode if config is not None else "zigzag" contiguous_cp_local_target_len = None pad_alignment = ( getattr(config, 'pad_packed_seq_alignment', None) if config is not None else None @@ -661,7 +657,12 @@ def get_batch_on_this_rank_for_sequence_packing( getattr(config, 'thd_max_packed_sequences', None), getattr(config, 'cuda_graph_impl', 'none') != 'none', ) - if is_tp_rank_0 and use_contiguous_cp_slice and pad_alignment is not None: + if ( + is_tp_rank_0 + and cp_group.size() > 1 + and cp_partition_mode == "contiguous" + and pad_alignment is not None + ): if target_len is not None: contiguous_cp_local_target_len = target_len else: @@ -694,7 +695,7 @@ def get_batch_on_this_rank_for_sequence_packing( batch, cp_group, keys=cp_slice_keys, - use_contiguous_cp_slice=use_contiguous_cp_slice, + cp_partition_mode=cp_partition_mode, partition_total_tokens=partition_total_tokens, ) @@ -835,6 +836,7 @@ def get_batch_on_this_rank_for_sequence_packing( max_seqlen_kv=max_seqlen, local_cp_size=local_cp_size, cp_group=cp_group, + cp_partition_mode=cp_partition_mode, pad_between_seqs=False, ) diff --git a/megatron/core/datasets/data_schedule_utils.py b/megatron/core/datasets/data_schedule_utils.py index 7b47565e45c..25525751f4b 100644 --- a/megatron/core/datasets/data_schedule_utils.py +++ b/megatron/core/datasets/data_schedule_utils.py @@ -3,7 +3,7 @@ from collections import deque from functools import lru_cache from math import ceil, log2 -from typing import Callable, Dict, List, Optional, Sequence, Tuple +from typing import Callable, Dict, List, Literal, Optional, Sequence, Tuple import torch @@ -15,21 +15,20 @@ def get_cp_slice_for_thd( batch, cp_group, keys: Optional[Sequence[str]] = None, - use_contiguous_cp_slice: bool = False, + cp_partition_mode: Literal["zigzag", "contiguous"] = "zigzag", partition_total_tokens: Optional[int] = None, ): """Partition sequence data for context parallelism in THD format. - Uses TE's THD partitioned indices to split the packed sequence across CP ranks. - With ``use_contiguous_cp_slice=True``, the flattened rows are split into ``cp_size`` equal - slices, and each rank receives the slice at its rank index. + ``zigzag`` uses TE's THD partitioned indices. ``contiguous`` splits the flattened rows into + ``cp_size`` equal slices and assigns one slice to each rank. Only keys present in the batch are sliced. Args: batch: Dict with packed sequence data. cp_group: Context parallel process group. keys: Sequence data keys to slice. Defaults to the original THD data tensors. - use_contiguous_cp_slice: Whether to assign one contiguous row range per CP rank. + cp_partition_mode: How to assign packed rows to CP ranks. partition_total_tokens: Optional total used to tail-pad tensors selected by ``keys`` before slicing. Existing cu_seqlens metadata is left unchanged. """ @@ -63,7 +62,7 @@ def get_cp_slice_for_thd( if pad_len > 0: pad_value = True if key == 'padding_mask' else 0 batch[key] = torch.cat([batch[key], batch[key].new_full((pad_len,), pad_value)]) - if use_contiguous_cp_slice: + if cp_partition_mode == "contiguous": if total_tokens % cp_size != 0: raise RuntimeError( f"Contiguous CP slicing requires total_tokens={total_tokens} to be divisible by " @@ -77,6 +76,9 @@ def get_cp_slice_for_thd( batch[key] = batch[key][row_slice] return + if cp_partition_mode != "zigzag": + raise ValueError(f"Unsupported CP partition mode: {cp_partition_mode}") + index = get_thd_partitioned_indices(cu_seqlens, total_tokens, cp_size, cp_rank) for key in keys: if key in batch and batch[key] is not None: diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 64bce710993..208762cd2e2 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -1748,9 +1748,10 @@ def __init__( self.kept_packed_seq_params.discard("cu_seqlens_q_padded") self.kept_packed_seq_params.discard("cu_seqlens_kv_padded") - # total_tokens and seq_idx are only for Mamba and should not be forwarded to TE attention. + # These fields are MCore-only and should not be forwarded to TE attention. self.kept_packed_seq_params.discard("total_tokens") self.kept_packed_seq_params.discard("seq_idx") + self.kept_packed_seq_params.discard("cp_partition_mode") if config.qk_clip or config.log_max_attention_logit: # qk-clip is only supported in TE 2.9.0 and later diff --git a/megatron/core/packed_seq_params.py b/megatron/core/packed_seq_params.py index 9352e8ef055..244c666a25b 100644 --- a/megatron/core/packed_seq_params.py +++ b/megatron/core/packed_seq_params.py @@ -27,6 +27,7 @@ class PackedSeqParams: total_tokens: int = None seq_idx: Tensor = None pad_between_seqs: Optional[bool] = None + cp_partition_mode: Literal["zigzag", "contiguous"] = "zigzag" def __post_init__(self): """Pre-compute seq_idx for Mamba mixer CUDA graph compatibility. @@ -475,6 +476,7 @@ def pad_sequence_for_thd( ), local_cp_size=packed_seq_params.local_cp_size, cp_group=packed_seq_params.cp_group, + cp_partition_mode=packed_seq_params.cp_partition_mode, total_tokens=local_target_len if target_cu_entries is None else None, pad_between_seqs=False if has_dummy_padding_seq else packed_seq_params.pad_between_seqs, ) diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index dc1f0cd800a..82a9848b084 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -1703,11 +1703,11 @@ def forward( attention_mask: torch.Tensor, x: torch.Tensor = None, qr: torch.Tensor = None, - boundary_hidden: Optional[torch.Tensor] = None, - boundary_kv: Optional[torch.Tensor] = None, attn_mask_type: AttnMaskType = None, attention_bias: torch.Tensor = None, packed_seq_params: PackedSeqParams = None, + boundary_hidden: Optional[torch.Tensor] = None, + boundary_kv: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Forward pass for CompressedSparseAttention. @@ -2173,8 +2173,8 @@ def _forward_thd_cp( key: torch.Tensor, x: torch.Tensor, qr: torch.Tensor, - boundary_hidden: torch.Tensor, - boundary_kv: torch.Tensor, + boundary_hidden: Optional[torch.Tensor], + boundary_kv: Optional[torch.Tensor], packed_seq_params: PackedSeqParams, ) -> torch.Tensor: """THD-packed context-parallel branch. diff --git a/megatron/core/transformer/experimental_attention_variant/csa_cp_layout_kernels.py b/megatron/core/transformer/experimental_attention_variant/csa_cp_layout_kernels.py index a8ef4df4cb9..a77f72d56f0 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa_cp_layout_kernels.py +++ b/megatron/core/transformer/experimental_attention_variant/csa_cp_layout_kernels.py @@ -17,6 +17,7 @@ import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute + from cutlass import utils from cutlass.cute.runtime import from_dlpack, make_fake_stream _CUTE_AVAILABLE = True @@ -388,92 +389,111 @@ def _build_attention_indices_kernel( seq_major_rows: cutlass.Int32, compressed_base: cutlass.Int32, total_width: cutlass.Int32, - index_mode: cutlass.Int32, + index_mode: cutlass.Constexpr, ): tidx, _, _ = cute.arch.thread_idx() bidx, _, _ = cute.arch.block_idx() linear = bidx * 128 + tidx # 0: selected top-k, 1: all visible compressed rows, 2: indexer loss. - row = linear - col = cutlass.Int32(0) - if index_mode != 0: - if linear < l_local * total_width: - row = linear // total_width - col = linear - row * total_width - else: - row = l_local + if cutlass.const_expr(index_mode == 0): + row = linear + else: + row = bidx if row < l_local: global_q = global_start + row seq_start_found = cutlass.Int32(-1) seq_comp_start = cutlass.Int32(0) seq_comp_len = cutlass.Int32(0) - for seq in range(n_seq): - seq_start = cu_seqlens[seq] - seq_end = cu_seqlens[seq + 1] - if global_q >= seq_start and global_q < seq_end: - seq_start_found = seq_start - if ratio > 1 and compressed_width > 0: - seq_comp_start = cu_seqlens_compressed[seq] - seq_comp_len = cu_seqlens_compressed[seq + 1] - seq_comp_start - - if index_mode != 0: - topk_value = cutlass.Int32(-1) - rank_major_value = cutlass.Int32(-1) - length = cutlass.Int32(0) - if seq_start_found >= 0: - window_start = global_q - window_size + 1 - if window_start < seq_start_found: - window_start = seq_start_found - window_count = global_q - window_start + 1 - if index_mode == 2: - if col < compressed_width: - comp_id = indexer_topk[row, col] - if comp_id >= 0 and comp_id < seq_comp_len: - seq_major_id = seq_comp_start + comp_id - if seq_major_id < seq_major_rows: - rank_major_value = seq_to_rank_row[seq_major_id] - if rank_major_value >= 0: - topk_value = compressed_base + rank_major_value + if cutlass.const_expr(index_mode == 0): + for seq in range(n_seq): + seq_start = cu_seqlens[seq] + seq_end = cu_seqlens[seq + 1] + if global_q >= seq_start and global_q < seq_end: + seq_start_found = seq_start + if ratio > 1 and compressed_width > 0: + seq_comp_start = cu_seqlens_compressed[seq] + seq_comp_len = cu_seqlens_compressed[seq + 1] - seq_comp_start + else: + shared = utils.SmemAllocator().allocate_tensor(cutlass.Int32, cute.make_layout(3)) + if tidx == 0: + for seq in range(n_seq): + seq_start = cu_seqlens[seq] + seq_end = cu_seqlens[seq + 1] + if global_q >= seq_start and global_q < seq_end: + seq_start_found = seq_start + if ratio > 1 and compressed_width > 0: + seq_comp_start = cu_seqlens_compressed[seq] + seq_comp_len = cu_seqlens_compressed[seq + 1] - seq_comp_start + shared[0] = seq_start_found + shared[1] = seq_comp_start + shared[2] = seq_comp_len + cute.arch.sync_threads() + seq_start_found = shared[0] + seq_comp_start = shared[1] + seq_comp_len = shared[2] + + if cutlass.const_expr(index_mode != 0): + col = tidx + while col < total_width: + topk_value = cutlass.Int32(-1) + rank_major_value = cutlass.Int32(-1) + length = cutlass.Int32(0) + if seq_start_found >= 0: + window_start = global_q - window_size + 1 + if window_start < seq_start_found: + window_start = seq_start_found + window_count = global_q - window_start + 1 + if cutlass.const_expr(index_mode == 2): + if col < compressed_width: + comp_id = indexer_topk[row, col] + if comp_id >= 0 and comp_id < seq_comp_len: + seq_major_id = seq_comp_start + comp_id + if seq_major_id < seq_major_rows: + rank_major_value = seq_to_rank_row[seq_major_id] + if rank_major_value >= 0: + topk_value = compressed_base + rank_major_value + else: + window_col = col - compressed_width + if window_col < window_count: + pos = window_start + window_col + if pos < global_start: + topk_value = pos - (global_start - d_window) + else: + topk_value = d_window + pos - global_start else: - window_col = col - compressed_width - if window_col < window_count: - pos = window_start + window_col + comp_count = cutlass.Int32(0) + if ratio > 1 and compressed_width > 0: + comp_count = (global_q - seq_start_found + 1) // ratio + if comp_count > compressed_width: + comp_count = compressed_width + if comp_count > seq_comp_len: + comp_count = seq_comp_len + length = window_count + comp_count + if col < window_count: + pos = window_start + col if pos < global_start: topk_value = pos - (global_start - d_window) else: topk_value = d_window + pos - global_start - else: - comp_count = cutlass.Int32(0) - if ratio > 1 and compressed_width > 0: - comp_count = (global_q - seq_start_found + 1) // ratio - if comp_count > compressed_width: - comp_count = compressed_width - if comp_count > seq_comp_len: - comp_count = seq_comp_len - length = window_count + comp_count - if col < window_count: - pos = window_start + col - if pos < global_start: - topk_value = pos - (global_start - d_window) - else: - topk_value = d_window + pos - global_start - elif col < length: - seq_major_id = seq_comp_start + col - window_count - if seq_major_id < seq_major_rows: - rank_major_id = seq_to_rank_row[seq_major_id] - if rank_major_id >= 0: - topk_value = compressed_base + rank_major_id - elif total_width > 0 and index_mode != 2: - length = 1 - if col == 0: - topk_value = 0 - topk_idxs[row, col] = topk_value - if index_mode == 2 and col < compressed_width: - indexer_rank_major[row, col] = rank_major_value - elif col == 0: - topk_length[row] = length + elif col < length: + seq_major_id = seq_comp_start + col - window_count + if seq_major_id < seq_major_rows: + rank_major_id = seq_to_rank_row[seq_major_id] + if rank_major_id >= 0: + topk_value = compressed_base + rank_major_id + elif total_width > 0 and cutlass.const_expr(index_mode != 2): + length = 1 + if col == 0: + topk_value = 0 + topk_idxs[row, col] = topk_value + if cutlass.const_expr(index_mode == 2): + if col < compressed_width: + indexer_rank_major[row, col] = rank_major_value + elif col == 0: + topk_length[row] = length + col = col + 128 else: for out_col in range(total_width): topk_idxs[row, out_col] = -1 @@ -528,7 +548,7 @@ def _build_attention_indices_launch( seq_major_rows: cutlass.Int32, compressed_base: cutlass.Int32, total_width: cutlass.Int32, - index_mode: cutlass.Int32, + index_mode: cutlass.Constexpr, launch_work: cutlass.Int32, stream: cuda.CUstream, ): @@ -602,7 +622,10 @@ def _run_compiled_launch( cap = torch.cuda.get_device_capability() arch = {(9, 0): "sm_90a", (10, 0): "sm_100a", (10, 3): "sm_103a"}.get(cap) if arch is None: - raise RuntimeError(f"Unsupported GPU compute capability {cap} for CSA CP CuTe kernels.") + raise RuntimeError( + f"Unsupported GPU compute capability {cap} for CSA CP CuTe kernels; " + "supported architectures: sm_90a, sm_100a, sm_103a." + ) cute_tensor_args = [] for tensor in tensor_args: cute_tensor = from_dlpack(tensor.detach(), assumed_align=16, enable_tvm_ffi=True) @@ -788,7 +811,7 @@ def build_attention_indices( else: topk_length_kernel = torch.empty((l_local,), dtype=torch.int32, device=cu_seqlens.device) indexer_rank_major = torch.empty((1, 1), dtype=torch.int32, device=cu_seqlens.device) - launch_work = l_local * total_width if index_mode else l_local + launch_work = l_local * 128 if index_mode else l_local compressed_base = int(d_window) + l_local _run_compiled_launch( _build_attention_indices_launch, @@ -815,6 +838,7 @@ def build_attention_indices( index_mode, launch_work, ), + static_arg_indices=(10,), ) if for_indexer_loss: return topk_idxs, None, indexer_rank_major diff --git a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py index 91c14892aee..ec4c76618ff 100644 --- a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py +++ b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py @@ -254,18 +254,21 @@ def forward( inference_context is None and inference_params is None ), "Inference is not supported for DSv4HybridAttention." - # Set the right cp group for dynamic-cp. Mirrors Attention.forward: - # both QKV RoPE and the post-attention inverse RoPE use - # self.pg_collection.cp, which must point at this microbatch's dynamic - # CP group. Restored before every return. + # Select this microbatch's dynamic CP group. QKV captures it explicitly + # for recompute; the rest of this forward reads it from pg_collection. + # Restore the static group before returning. _orig_cp_group = self.pg_collection.cp + cp_group = _orig_cp_group if packed_seq_params is not None and packed_seq_params.local_cp_size is not None: assert packed_seq_params.cp_group is not None, "cp_group must be set in dynamic-cp mode" - self.pg_collection.cp = packed_seq_params.cp_group + cp_group = packed_seq_params.cp_group - cp_size = self.pg_collection.cp.size() + cp_size = cp_group.size() qkv_format = packed_seq_params.qkv_format if packed_seq_params is not None else None use_thd_cp = cp_size > 1 and qkv_format == 'thd' + if use_thd_cp and packed_seq_params.cp_partition_mode != "contiguous": + raise ValueError("DSv4 THD CP requires a contiguous CP partition.") + self.pg_collection.cp = cp_group boundary_hidden = None if use_thd_cp: @@ -680,7 +683,12 @@ def get_query_key_value_tensors( # ========================================= def qkv_up_proj_and_rope_apply( - q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb, boundary_kv_compressed=None + q_compressed, + kv_compressed, + k_pos_emb, + rotary_pos_emb, + cp_group, + boundary_kv_compressed=None, ): """ Apply the up projection and RoPE to the query and key. @@ -711,10 +719,10 @@ def qkv_up_proj_and_rope_apply( if k_pos_emb is not None: k_pos_emb = torch.unsqueeze(k_pos_emb, -2) - cp_size = self.pg_collection.cp.size() + cp_size = cp_group.size() if self.config.apply_rope_fusion: if cp_size > 1 and packed_seq: - cp_rank = self.pg_collection.cp.rank() + cp_rank = cp_group.rank() # Rank r owns global rows [r * local_rows, (r + 1) * local_rows). global_start = cp_rank * q.shape[0] query = cp_utils.apply_thd_cp_local_rope_fused( @@ -740,7 +748,7 @@ def qkv_up_proj_and_rope_apply( boundary_kv = kv[:boundary_rows] kv = kv[boundary_rows:] else: - cp_rank = self.pg_collection.cp.rank() + cp_rank = cp_group.rank() query = fused_mla_rope_inplace( q, rotary_pos_cos, @@ -768,7 +776,7 @@ def qkv_up_proj_and_rope_apply( value = kv else: if packed_seq and cp_size > 1: - global_start = self.pg_collection.cp.rank() * q.shape[0] + global_start = cp_group.rank() * q.shape[0] query = cp_utils.apply_thd_cp_local_rope_unfused( q, rotary_pos_emb, @@ -811,7 +819,7 @@ def qkv_up_proj_and_rope_apply( config=self.config, cu_seqlens=cu_seqlens_q, mscale=mscale, - cp_group=self.pg_collection.cp, + cp_group=cp_group, mla_rotary_interleaved=True, mla_output_remove_interleaving=True, max_seqlen=rope_max_seqlen_q, @@ -829,7 +837,7 @@ def qkv_up_proj_and_rope_apply( config=self.config, cu_seqlens=cu_seqlens_kv, mscale=mscale, - cp_group=self.pg_collection.cp, + cp_group=cp_group, mla_rotary_interleaved=True, mla_output_remove_interleaving=True, max_seqlen=rope_max_seqlen_kv, @@ -859,6 +867,7 @@ def qkv_up_proj_and_rope_apply( kv_compressed, k_pos_emb, rotary_pos_emb, + self.pg_collection.cp, ) boundary_kv = None else: @@ -868,17 +877,23 @@ def qkv_up_proj_and_rope_apply( kv_compressed, k_pos_emb, rotary_pos_emb, + self.pg_collection.cp, boundary_kv_compressed, ) else: if boundary_kv_compressed is None: query, key, value = qkv_up_proj_and_rope_apply( - q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb + q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb, self.pg_collection.cp ) boundary_kv = None else: query, key, value, boundary_kv = qkv_up_proj_and_rope_apply( - q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb, boundary_kv_compressed + q_compressed, + kv_compressed, + k_pos_emb, + rotary_pos_emb, + self.pg_collection.cp, + boundary_kv_compressed, ) result = (query, key, value, q_compressed, kv_compressed) diff --git a/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py b/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py index 3bc19b79f1e..53f67f9b314 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py +++ b/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py @@ -514,7 +514,6 @@ def _indexer_topk_core( so the SBHD→BSHD permute can be performed once and reused across the indexer forward and the score-recompute backward kernels. """ - _ensure_dsa_namespace() is_thd = cu_seqlens_q is not None device = q.device @@ -526,7 +525,10 @@ def _indexer_topk_core( raise ValueError(f"THD k must be (total_k, idx_hd), got {k.shape}") if w.ndim != 2: raise ValueError(f"THD w must be (total_q, idx_nh), got {w.shape}") + if max_seqlen_kv == 0 or k.shape[0] == 0: + raise ValueError("indexer_topk requires at least one K row.") + _ensure_dsa_namespace() # Kernel wants k as 3-D ``(total_k, h_kv, idx_hd)``. scores = _DSA.indexer_forward_wrapper( q, @@ -576,6 +578,10 @@ def _indexer_topk_core( ) seq_lens = torch.where(row_valid, seq_lens, torch.zeros_like(seq_lens)) else: + if k.shape[1] == 0: + raise ValueError("indexer_topk requires at least one K row.") + + _ensure_dsa_namespace() # Kernel wants k as 4-D ``(b, sk, h_kv, idx_hd)``. scores = _DSA.indexer_forward_wrapper(q, k.unsqueeze(2), w, ratio=ratio)[ "scores" @@ -605,7 +611,7 @@ def _indexer_topk_core( if is_thd: row_valid = (topk_indices >= 0) & (topk_indices < seq_lens.unsqueeze(1)) topk_indices = topk_indices.masked_fill(~row_valid, -1) - safe_topk = topk_indices.clamp(min=0, max=max(sk - 1, 0)).to(torch.long) + safe_topk = topk_indices.clamp(min=0, max=sk - 1).to(torch.long) selected_scores = torch.gather(scores_flat, dim=-1, index=safe_topk) selected_valid = (topk_indices >= 0) & (topk_indices < sk) & torch.isfinite(selected_scores) topk_indices = topk_indices.masked_fill(~selected_valid, -1) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 1a44adcb0d7..12eecbbd5c4 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -307,6 +307,9 @@ class TransformerConfig(ModelParallelConfig): ) """Type of attention variant to use. Currently support gated_delta_net, dsa, and dsv4_hybrid.""" + cp_partition_mode: Literal["zigzag", "contiguous"] = "zigzag" + """How THD sequence rows are partitioned across context-parallel ranks.""" + #################### # DSA #################### @@ -1454,6 +1457,21 @@ def __post_init__(self): self.experimental_attention_variant = self.linear_attention_type self.linear_attention_type = None + if self.cp_partition_mode not in ("zigzag", "contiguous"): + raise ValueError(f"Unsupported cp_partition_mode: {self.cp_partition_mode}") + + if self.context_parallel_size > 1: + if ( + self.experimental_attention_variant == "dsv4_hybrid" + and self.cp_partition_mode != "contiguous" + ): + raise ValueError("DSv4 Hybrid with CP requires cp_partition_mode='contiguous'.") + if ( + self.experimental_attention_variant != "dsv4_hybrid" + and self.cp_partition_mode != "zigzag" + ): + raise ValueError("cp_partition_mode='contiguous' is only supported with DSv4.") + if self.experimental_attention_variant in ["gated_delta_net"]: assert ( self.linear_attention_freq is not None @@ -1516,6 +1534,19 @@ def __post_init__(self): assert all( ratio in [0, 4, 128] for ratio in self.csa_compress_ratios ), "csa_compress_ratios must be 0, 4, or 128" + if ( + self.context_parallel_size > 1 + and 4 in self.csa_compress_ratios + and not self.csa_dense_mode + and (self.dsa_indexer_loss_coeff or 0.0) > 0 + and not self.dsa_indexer_use_sparse_loss + ): + # The current cuDNN fused indexer-loss kernel cannot express the + # CP-local query-offset semantics required by dense loss. + raise ValueError( + "DSv4 CP dense indexer loss is unsupported because the current cuDNN " + "fused kernel lacks CP query-offset semantics." + ) assert ( self.tensor_model_parallel_size == 1 ), "DSv4 Hybrid Attention only supports TP size 1." diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index 2ad6e150b53..b283313b963 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -1326,6 +1326,7 @@ def _reconstruct_packed_seq_params_from_kwargs(self, kwargs): max_seqlen = self.config.max_seqlen_per_dp_cp_rank * self.config.context_parallel_size packed_seq_params = PackedSeqParams( qkv_format='thd', + cp_partition_mode=self.config.cp_partition_mode, cu_seqlens_q=kwargs.pop('cu_seqlens_q'), cu_seqlens_kv=kwargs.pop('cu_seqlens_kv'), cu_seqlens_q_padded=kwargs.pop('cu_seqlens_q_padded'), diff --git a/tests/unit_tests/test_sequence_packing.py b/tests/unit_tests/test_sequence_packing.py index b4273f4c696..ac4eff431b3 100644 --- a/tests/unit_tests/test_sequence_packing.py +++ b/tests/unit_tests/test_sequence_packing.py @@ -196,7 +196,10 @@ def test_dsv4_thd_cp_slice_uses_static_partition_total(): } get_cp_slice_for_thd( - batch, _MockCPGroup(size=4, rank=2), use_contiguous_cp_slice=True, partition_total_tokens=16 + batch, + _MockCPGroup(size=4, rank=2), + cp_partition_mode="contiguous", + partition_total_tokens=16, ) assert torch.equal(batch["tokens"], torch.arange(8, 12, dtype=torch.int64)) @@ -223,7 +226,7 @@ def test_dsv4_thd_dynamic_cp_pads_before_slicing( cp=_MockCPGroup(size=4, rank=cp_rank), ) config = SimpleNamespace( - experimental_attention_variant="dsv4_hybrid", + cp_partition_mode="contiguous", pad_packed_seq_alignment=alignment, max_seqlen_per_dp_cp_rank=8, thd_max_packed_sequences=None, @@ -300,6 +303,7 @@ def record_padding( torch.arange(global_start, global_start + local_target) >= total_tokens, ) assert packed_seq_params.local_cp_size == local_cp_size + assert packed_seq_params.cp_partition_mode == "contiguous" assert pad_input_lengths == [local_target] diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_layout_kernels.py b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_layout_kernels.py index 0aa688d1067..2b624a1236f 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_layout_kernels.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_layout_kernels.py @@ -416,6 +416,77 @@ def test_build_attention_indices_indexer_loss_mode_matches_native(): assert torch.equal(fused[2], expected[1]) +def test_build_attention_indices_many_short_sequences_match_native(): + _require_cute_cuda() + cu = torch.arange(0, 4097, 32, dtype=torch.int32, device="cuda") + ratio = 4 + cu_comp = _compressed_cu_seqlens(cu, ratio) + global_start = 1024 + l_local = 1024 + d_window = 128 + window = 16 + compressed_width = 128 + seq_to_rank_row = torch.arange(int(cu_comp[-1]), dtype=torch.int32, device="cuda").flip(0) + compressed_base = d_window + l_local + + actual = csa_cp_layout_kernels.build_attention_indices( + cu, + global_start, + l_local, + d_window, + window, + ratio, + compressed_width, + cu_seqlens_compressed=cu_comp, + seq_to_rank_row=seq_to_rank_row, + ) + expected = _native_attention_indices( + cu, + cu_comp, + global_start, + l_local, + d_window, + window, + ratio, + compressed_width, + seq_to_rank_row, + compressed_base, + ) + assert torch.equal(actual[0], expected[0]) + assert torch.equal(actual[1], expected[1]) + + logical_ids = torch.arange(compressed_width, dtype=torch.int32, device="cuda").expand( + l_local, -1 + ) + actual = csa_cp_layout_kernels.build_attention_indices( + cu, + global_start, + l_local, + d_window, + window, + ratio, + compressed_width, + logical_ids, + cu_seqlens_compressed=cu_comp, + seq_to_rank_row=seq_to_rank_row, + for_indexer_loss=True, + ) + expected = _native_indexer_loss_indices( + cu, + cu_comp, + global_start, + l_local, + d_window, + window, + ratio, + logical_ids, + seq_to_rank_row, + compressed_base, + ) + assert torch.equal(actual[0], expected[0]) + assert torch.equal(actual[2], expected[1]) + + @pytest.mark.parametrize( ("ratio", "lengths"), [(4, (3, 10, 20, 5, 33, 10, 27, 20)), (128, (3, 130, 170, 5, 260, 129, 200, 127))], diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py index 2f2cfb0f9a7..d260a044053 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py @@ -123,6 +123,7 @@ def _make_thd_packed_seq_params(seg_lens, padded_seg_lens=None, device='cuda'): max_seqlen_q=max_len, max_seqlen_kv=max_len, qkv_format='thd', + cp_partition_mode='contiguous', ) @@ -284,6 +285,7 @@ def _make_dsv4_cp_config( dsa_indexer_loss_coeff=dsa_indexer_loss_coeff, dsa_indexer_use_sparse_loss=dsa_indexer_use_sparse_loss, context_parallel_size=context_parallel_size, + cp_partition_mode="contiguous" if context_parallel_size > 1 else "zigzag", csa_dense_mode=False, csa_compress_rotary_base=shape["csa_compress_rotary_base"], layernorm_epsilon=1e-6, @@ -628,6 +630,65 @@ def test_thd_cp_matches_full_reference_forward_backward(self, layer_number, appl del cp_attn, ref_attn, full_hidden, local_hidden, ref_hidden, local_out, ref_out, grad _clear_cuda_test_state() + @pytest.mark.parametrize("apply_rope_fusion", [True, False], ids=["fused", "unfused"]) + def test_dynamic_cp_mla_up_proj_recompute_matches_eager(self, apply_rope_fusion): + """Selective MLA recompute must retain the microbatch's dynamic CP group.""" + if apply_rope_fusion and not self.fused_kernels_available: + pytest.skip(_DSV4_CP_FUSED_KERNELS_UNAVAILABLE_REASON) + + packed, padded_tokens, local_idx = _make_ragged_cp_case(self.cp_size, self.cp_rank) + packed.local_cp_size = self.cp_size + packed.cp_group = self.pg.cp + + torch.manual_seed(_SEED + 1300) + model_parallel_cuda_manual_seed(_SEED + 1300) + eager_config = _make_dsv4_cp_config( + context_parallel_size=self.cp_size, + dsa_indexer_loss_coeff=0.0, + apply_dsa_kernel_fusion=False, + apply_rope_fusion=apply_rope_fusion, + ) + recompute_config = _make_dsv4_cp_config( + context_parallel_size=self.cp_size, + dsa_indexer_loss_coeff=0.0, + apply_dsa_kernel_fusion=False, + apply_rope_fusion=apply_rope_fusion, + ) + recompute_config.recompute_granularity = "selective" + recompute_config.recompute_modules = ["mla_up_proj"] + + eager_attn = _build_attention( + eager_config, layer_number=1, pg_collection=self.ref_pg + ).cuda() + recompute_attn = _build_attention( + recompute_config, layer_number=1, pg_collection=self.ref_pg + ).cuda() + eager_attn.train() + recompute_attn.train() + _copy_module_parameters(eager_attn, recompute_attn) + + full_hidden, full_grad = _make_hidden_and_grad(padded_tokens, eager_config.hidden_size) + hidden = full_hidden.index_select(0, local_idx) + grad = full_grad.index_select(0, local_idx) + eager_result = _run_dsv4_attention_forward_backward( + eager_attn, hidden.detach().clone().requires_grad_(True), grad, packed + ) + recompute_result = _run_dsv4_attention_forward_backward( + recompute_attn, hidden.detach().clone().requires_grad_(True), grad, packed + ) + + mode = f"dynamic_cp:rope_fused={apply_rope_fusion}" + _assert_cp_tensor_match(recompute_result[0], eager_result[0], f"{mode}:output") + _assert_cp_tensor_match(recompute_result[1], eager_result[1], f"{mode}:hidden_grad") + assert recompute_result[2].keys() == eager_result[2].keys() + for name, recompute_grad in recompute_result[2].items(): + _assert_cp_tensor_match( + recompute_grad, eager_result[2][name], f"{mode}:param_grad:{name}" + ) + + del eager_attn, recompute_attn, full_hidden, full_grad, hidden, grad + _clear_cuda_test_state() + def test_thd_cp_ratio4_eval_matches_full_reference(self): """Ratio-4 inference lowers logical indexer top-k rows correctly.""" packed, padded_tokens, local_idx = _make_ragged_cp_case(self.cp_size, self.cp_rank) diff --git a/tests/unit_tests/transformer/test_thd_cuda_graph.py b/tests/unit_tests/transformer/test_thd_cuda_graph.py index 70245498cac..0d7298524e1 100644 --- a/tests/unit_tests/transformer/test_thd_cuda_graph.py +++ b/tests/unit_tests/transformer/test_thd_cuda_graph.py @@ -271,6 +271,8 @@ def test_generic_alignment_appends_dummy_padding_sequence(self): """Generic THD padding covers tail slots with an independent dummy sequence.""" seqlens, total_T = [50, 30], 80 psp = _make_psp(seqlens) + # Use the non-default mode so losing it during rebuild cannot silently pass as zigzag. + psp.cp_partition_mode = "contiguous" orig = psp.cu_seqlens_q.clone() p_tok, _, _, _, p, mask = pad_sequence_for_thd( torch.ones(1, total_T, device="cuda"), None, None, None, psp, alignment=64 @@ -280,6 +282,7 @@ def test_generic_alignment_appends_dummy_padding_sequence(self): assert torch.equal(p.cu_seqlens_q, expected) assert torch.equal(p.cu_seqlens_q_padded, expected) assert p.pad_between_seqs is False + assert p.cp_partition_mode == "contiguous" assert mask.shape == (1, 128) assert not mask[0, :total_T].any() and mask[0, total_T:].all() @@ -570,6 +573,8 @@ def test_round_trip(self): ) } layer = _build_layer(256, 4, 4, 1024, 128, 8) + # Use the non-default mode so losing it during reconstruction is observable. + layer.config.cp_partition_mode = "contiguous" kw = {'packed_seq_params': psp, 'other': 'kept'} TransformerLayer._decompose_packed_seq_params_to_kwargs(kw) assert 'packed_seq_params' not in kw and 'cu_seqlens_q' in kw @@ -577,6 +582,7 @@ def test_round_trip(self): r = kw['packed_seq_params'] assert r.qkv_format == 'thd' and r.max_seqlen_q == 128 assert r.pad_between_seqs is False + assert r.cp_partition_mode == "contiguous" for k, v in orig.items(): assert torch.equal(getattr(r, k), v) From 1df04cbd64164d62dd22a9167f0a8b8967c633df Mon Sep 17 00:00:00 2001 From: kunlunl Date: Tue, 30 Jun 2026 05:52:46 -0700 Subject: [PATCH 13/23] Fix hybrid_moe_model golden config test Signed-off-by: kunlunl --- tests/unit_tests/models/test_hybrid_moe_model.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index a86d7fa2144..bdfd9888f55 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -62,6 +62,7 @@ "config_logger_dir": "", "context_parallel_size": 1, "cp_comm_type": "p2p", + "cp_partition_mode": "zigzag", "cpu_offloading": False, "cpu_offloading_activations": True, "cpu_offloading_double_buffering": False, From 50e63e7ae63434982edc3d2d7cc5b9cc207a76e0 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Tue, 30 Jun 2026 09:16:00 -0700 Subject: [PATCH 14/23] fix: address DSv4 CP review feedback Signed-off-by: kunlunl --- .../core/transformer/experimental_attention_variant/csa.py | 6 +++++- .../experimental_attention_variant/csa_cp_utils.py | 3 +++ .../deepseek_v4_hybrid_attention.py | 4 ++++ megatron/core/transformer/transformer_config.py | 7 +++++++ .../test_dsv4_hybrid_attention_cp.py | 5 +++-- 5 files changed, 22 insertions(+), 3 deletions(-) diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index 82a9848b084..4933154cfb2 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -2318,7 +2318,8 @@ def _forward_thd_cp( k_indexer_seq_major = torch.index_select( k_indexer_rank_major, 0, seq_to_rank_row.clamp_min(0) ) - # Top-k is still in logical compressed coordinates here. + # Each top-k entry is still a logical compressed id within that + # query's sequence here. compressed_topk = cp_utils.compute_cp_indexer_topk( q_indexer_cp, weights_indexer_cp, @@ -2357,6 +2358,9 @@ def _forward_thd_cp( if compressed_topk is not None else (max_seqlen_q // ratio if ratio > 1 else 0) ) + # Lower the logical ids into two physical spaces: indexer_topk_rank_major + # addresses the rank-major compressed buffers without kv_full_thd's + # compressed base, while topk_idxs addresses final rows in kv_full_thd. topk_idxs, topk_length, indexer_topk_rank_major = ( csa_cp_layout_kernels.build_attention_indices( cu_seqlens, diff --git a/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py b/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py index 88c47b85a92..f1a84ed3be4 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py +++ b/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py @@ -225,6 +225,9 @@ def prepare_cp_compressor_input( with ``ratio=4``, ``comp_id=3`` maps to RoPE position ``12``. ``seq_to_rank_row``: map from global sequence-major compressed rows to their canonical rank-major all-gather rows. + If rank 0 owns ``A0, A1`` and rank 1 owns ``B0, B1``, with four + slots per rank, logical rows ``[A0, A1, B0, B1]`` are stored as + ``[A0, A1, pad, pad | B0, B1, pad, pad]`` and map to ``[0, 1, 4, 5]``. """ cp_size = int(cp_size) ratio = int(ratio) diff --git a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py index ec4c76618ff..11d06402469 100644 --- a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py +++ b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py @@ -265,6 +265,8 @@ def forward( cp_size = cp_group.size() qkv_format = packed_seq_params.qkv_format if packed_seq_params is not None else None + if cp_size > 1 and qkv_format != 'thd': + raise ValueError("DSv4 Hybrid with CP requires qkv_format='thd'.") use_thd_cp = cp_size > 1 and qkv_format == 'thd' if use_thd_cp and packed_seq_params.cp_partition_mode != "contiguous": raise ValueError("DSv4 THD CP requires a contiguous CP partition.") @@ -657,6 +659,8 @@ def get_query_key_value_tensors( # q_compressed: [s, b, q_lora_rank] q_compressed, _ = self.linear_q_down_proj(hidden_states) + # Despite their legacy names, these are hidden-state inputs to linear_kv_proj; + # DSv4's actual compressed KV is produced later by the CSA compressor. kv_compressed = hidden_states k_pos_emb = None boundary_kv_compressed = boundary_hidden diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 12eecbbd5c4..44a323005d9 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1460,6 +1460,13 @@ def __post_init__(self): if self.cp_partition_mode not in ("zigzag", "contiguous"): raise ValueError(f"Unsupported cp_partition_mode: {self.cp_partition_mode}") + if self.experimental_attention_variant == "dsv4_hybrid" and ( + self.context_parallel_size > 1 or self.dynamic_context_parallel + ): + assert ( + self.sequence_packing_scheduler is not None + ), "DSv4 Hybrid with CP requires a sequence_packing_scheduler for THD inputs." + if self.context_parallel_size > 1: if ( self.experimental_attention_variant == "dsv4_hybrid" diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py index d260a044053..9bb0dfafc84 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py @@ -286,6 +286,7 @@ def _make_dsv4_cp_config( dsa_indexer_use_sparse_loss=dsa_indexer_use_sparse_loss, context_parallel_size=context_parallel_size, cp_partition_mode="contiguous" if context_parallel_size > 1 else "zigzag", + sequence_packing_scheduler="dp_balanced" if context_parallel_size > 1 else None, csa_dense_mode=False, csa_compress_rotary_base=shape["csa_compress_rotary_base"], layernorm_epsilon=1e-6, @@ -658,10 +659,10 @@ def test_dynamic_cp_mla_up_proj_recompute_matches_eager(self, apply_rope_fusion) recompute_config.recompute_modules = ["mla_up_proj"] eager_attn = _build_attention( - eager_config, layer_number=1, pg_collection=self.ref_pg + eager_config, layer_number=2, pg_collection=self.ref_pg ).cuda() recompute_attn = _build_attention( - recompute_config, layer_number=1, pg_collection=self.ref_pg + recompute_config, layer_number=2, pg_collection=self.ref_pg ).cuda() eager_attn.train() recompute_attn.train() From 46ec4cf1f271b6bc2a2444bd76bae0cd144184c5 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Tue, 30 Jun 2026 11:14:06 -0700 Subject: [PATCH 15/23] Simplify fused MLA RoPE imports Signed-off-by: kunlunl --- .../deepseek_v4_hybrid_attention.py | 15 ++++----------- 1 file changed, 4 insertions(+), 11 deletions(-) diff --git a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py index f84647be157..3ad6092c310 100644 --- a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py +++ b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py @@ -8,6 +8,10 @@ from megatron.core import tensor_parallel from megatron.core.extensions.transformer_engine import HAVE_TE +from megatron.core.fusions.fused_mla_yarn_rope_apply import ( + fused_mla_rope_inplace, + fused_mla_rope_out_of_place, +) from megatron.core.models.common.embeddings import ( RotaryEmbedding, YarnRotaryEmbedding, @@ -26,16 +30,6 @@ from megatron.core.typed_torch import apply_module from megatron.core.utils import get_pg_size, is_te_min_version -try: - from megatron.core.fusions.fused_mla_yarn_rope_apply import ( - fused_mla_rope_inplace, - fused_mla_rope_out_of_place, - ) -except Exception: - fused_mla_rope_inplace = None - fused_mla_rope_out_of_place = None - - if HAVE_TE: from megatron.core.extensions.transformer_engine import TELinear, set_save_original_input else: @@ -407,7 +401,6 @@ def forward( # Fused DSA backward retains the raw attention output O. Applying # inverse RoPE to its view in-place corrupts the retained O used by # the softmax backward, so this call needs private storage. - assert fused_mla_rope_out_of_place is not None core_attn_out = fused_mla_rope_out_of_place( core_attn_out, rotary_pos_cos, From cad0ac5dd1a0212a319f38e385c8530743b5da38 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Tue, 30 Jun 2026 11:14:45 -0700 Subject: [PATCH 16/23] Format DSv4 CP tests Signed-off-by: kunlunl --- tests/unit_tests/transformer/test_thd_cuda_graph.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/unit_tests/transformer/test_thd_cuda_graph.py b/tests/unit_tests/transformer/test_thd_cuda_graph.py index 0d7298524e1..b8fe2de5896 100644 --- a/tests/unit_tests/transformer/test_thd_cuda_graph.py +++ b/tests/unit_tests/transformer/test_thd_cuda_graph.py @@ -485,9 +485,7 @@ def test_cp_metadata_only_uses_local_padding_mask(self, monkeypatch): cu_seqlens_padded = torch.tensor([0, 16], dtype=torch.int32, device="cuda") packed_seq_params.cu_seqlens_q_padded = cu_seqlens_padded packed_seq_params.cu_seqlens_kv_padded = cu_seqlens_padded.clone() - padding_mask = torch.tensor( - [[False, False, False, True]], dtype=torch.bool, device="cuda" - ) + padding_mask = torch.tensor([[False, False, False, True]], dtype=torch.bool, device="cuda") def fail_if_called(*_args, **_kwargs): pytest.fail("already sliced padding_mask must avoid TE repartitioning") From 9fb43fed8f58ae26b8bfad91dade2f118bdec069 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Wed, 1 Jul 2026 05:54:56 -0700 Subject: [PATCH 17/23] Support dense indexer loss for DSv4 CP Signed-off-by: kunlunl --- .../experimental_attention_variant/csa.py | 108 +++++-- .../csa_cp_utils.py | 84 +++--- .../dsa_kernels.py | 277 +++++++++++------- .../core/transformer/transformer_config.py | 13 - .../test_csa_cp_layout_kernels.py | 2 +- .../test_csa_cp_utils.py | 44 +-- .../test_dsa_kernels.py | 54 ++-- .../test_dsv4_hybrid_attention_cp.py | 33 ++- 8 files changed, 375 insertions(+), 240 deletions(-) diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index 4933154cfb2..0c3547d0b7d 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -635,17 +635,91 @@ def _unfused_indexer_sparse_attn_from_topk( indexer_softmax_scale: float, loss_coeff: float, loss_divisor: float, + sparse_loss: bool, + ratio: int, + _max_seqlen_q: int, + indexer_layout: Tuple[torch.Tensor, torch.Tensor, torch.Tensor], q_padding_mask: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: - """PyTorch sparse attention plus caller-supplied top-k indexer loss for THD CP. + """PyTorch sparse attention plus caller-supplied indexer loss for THD CP. This mirrors the fused CP top-k path at the tensor-contract level: ``topk_indices`` indexes ``kv_full`` for sparse attention, and - ``indexer_topk_indices`` indexes rank-major compressed K for indexer loss. + ``indexer_topk_indices`` indexes compressed K for sparse indexer loss. + ``_max_seqlen_q`` is unused here; it is retained to match the fused + callable's signature, with the leading underscore marking it as unused. """ output = unfused_compressed_sparse_attn(query, kv_full, attn_sink, topk_indices, softmax_scale) total_q, np_, hn = query.shape + cu_seqlens_q, cu_seqlens_k, q_causal_offsets = indexer_layout + + if not sparse_loss: + q_rows = torch.arange(total_q, dtype=cu_seqlens_q.dtype, device=query.device) + q_sequence_ids = batch_of_row(cu_seqlens_q, total_q=total_q) + q_positions = q_rows - cu_seqlens_q[q_sequence_ids] + q_causal_offsets[q_sequence_ids] + visible_k = torch.minimum( + torch.div(q_positions + 1, ratio, rounding_mode="floor"), + (cu_seqlens_k[1:] - cu_seqlens_k[:-1])[q_sequence_ids], + ) + q_valid = q_rows < cu_seqlens_q[-1] + if q_padding_mask is not None: + q_valid = q_valid & ~q_padding_mask + + k_rows = torch.arange(k_indexer.shape[0], dtype=cu_seqlens_k.dtype, device=k_indexer.device) + k_sequence_ids = torch.bucketize( + k_rows, cu_seqlens_k[1:], out_int32=True, right=True + ).clamp_max(cu_seqlens_k.shape[0] - 2) + k_positions = k_rows - cu_seqlens_k[k_sequence_ids] + + k_indexer_float = k_indexer.float() + compressed_kv_float = compressed_kv.detach().float() + weights_scaled = weights.float() * float(indexer_softmax_scale) + raw_local_loss = query.new_zeros((), dtype=torch.float32) + for start in range(0, total_q, 512): + end = min(start + 512, total_q) + valid = ( + (q_sequence_ids[start:end].unsqueeze(1) == k_sequence_ids.unsqueeze(0)) + & (k_positions.unsqueeze(0) < visible_k[start:end].unsqueeze(1)) + & q_valid[start:end].unsqueeze(1) + ) + row_valid = valid.any(dim=-1, keepdim=True) + + predict_logits = torch.einsum( + "rhd,kd->rhk", q_indexer[start:end].float(), k_indexer_float + ) + predict_logits = torch.relu(predict_logits) * weights_scaled[start:end].unsqueeze(-1) + predict_logits = predict_logits.sum(dim=1).masked_fill(~valid, float("-inf")) + predict_logits = torch.where(row_valid, predict_logits, 0.0) + predict = torch.softmax(predict_logits, dim=-1, dtype=torch.float32) + predict = predict * row_valid.float() + + target_logits = torch.einsum( + "rhd,kd->rhk", query[start:end].detach().float(), compressed_kv_float + ) + target_logits = (target_logits * softmax_scale).masked_fill( + ~valid.unsqueeze(1), float("-inf") + ) + target_logits = torch.where(row_valid.unsqueeze(1), target_logits, 0.0) + sink = attn_sink.view(1, np_, 1).float() + scores_max = torch.maximum(target_logits.max(dim=-1, keepdim=True).values, sink) + exp_scores = torch.exp(target_logits - scores_max) + exp_sink = torch.exp(sink - scores_max) + target = exp_scores / (exp_scores.sum(dim=-1, keepdim=True) + exp_sink) + target = (target * row_valid.unsqueeze(1).float()).sum(dim=1) + eps = torch.finfo(torch.float32).tiny + target = target / target.sum(dim=-1, keepdim=True).clamp(min=eps) + target = target.clamp(min=eps) + predict = predict.clamp(min=eps) + + kl_per_row = (target * (torch.log(target) - torch.log(predict))).sum(dim=-1) + kl_per_row = torch.where( + row_valid.squeeze(-1), kl_per_row, torch.zeros_like(kl_per_row) + ) + raw_local_loss = raw_local_loss + kl_per_row.sum() + + return output, raw_local_loss * float(loss_coeff) / float(loss_divisor) + indexer_topk = indexer_topk_indices.shape[-1] if q_padding_mask is not None: @@ -2219,7 +2293,7 @@ def _forward_thd_cp( indexer = self.indexer indexer_loss_coeff = self.config.dsa_indexer_loss_coeff or 0.0 training_with_grad = self.training and torch.is_grad_enabled() - + sparse_indexer_loss = self.config.dsa_indexer_use_sparse_loss if self.compressor is not None and ratio > 1: # ---- Step 3: build fixed-capacity compressor input ---------------- @@ -2251,17 +2325,6 @@ def _forward_thd_cp( if indexer is not None: # ---- Step 4: optional indexer compressed path ----------------- - if ( - training_with_grad - and indexer_loss_coeff > 0 - and not self.config.dsa_indexer_use_sparse_loss - ): - raise RuntimeError( - "DSv4 THD CP path currently supports sparse indexer loss only. " - "The current dense score and backward APIs cannot express the " - "sequence position where this rank's local queries begin." - ) - indexer_x, indexer_qr = x.detach(), qr.detach() if indexer_x.shape[1] != 1: raise RuntimeError( @@ -2320,7 +2383,7 @@ def _forward_thd_cp( ) # Each top-k entry is still a logical compressed id within that # query's sequence here. - compressed_topk = cp_utils.compute_cp_indexer_topk( + compressed_topk, indexer_layout = cp_utils.compute_cp_indexer_topk( q_indexer_cp, weights_indexer_cp, k_indexer_seq_major, @@ -2378,6 +2441,13 @@ def _forward_thd_cp( ) if use_indexer_loss: # ---- Step 7a: indexer-loss path ---------------------------------- + k_indexer_for_loss = k_indexer_rank_major + compressed_kv_for_loss = compressed_kv_rank_major + if not sparse_indexer_loss: + k_indexer_for_loss = k_indexer_seq_major + compressed_kv_for_loss = torch.index_select( + compressed_kv_rank_major, 0, seq_to_rank_row.clamp_min(0) + ) cu_seqlens_q_unpadded = None if ( packed_seq_params.cu_seqlens_q is not None @@ -2410,14 +2480,18 @@ def _forward_thd_cp( self.attn_sink.float(), topk_idxs, q_indexer_cp, - k_indexer_rank_major, + k_indexer_for_loss, weights_indexer_cp, indexer_topk_rank_major, - compressed_kv_rank_major, + compressed_kv_for_loss, self.softmax_scale, indexer.softmax_scale, indexer_loss_coeff, 1 if self.config.calculate_per_token_loss else l_local * cp_size, + sparse_indexer_loss, + ratio, + max_seqlen_q, + indexer_layout, q_padding_mask, ) if indexer_loss_coeff > 0: diff --git a/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py b/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py index f1a84ed3be4..dabe4f2d02b 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py +++ b/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py @@ -273,6 +273,33 @@ def prepare_cp_compressor_input( return hidden_compact, compressed_group_ids, seq_to_rank_row +@torch.compile +def _build_cp_indexer_layout( + cu_seqlens_q: torch.Tensor, + cu_seqlens_compressed: torch.Tensor, + global_start: int, + local_rows: int, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build the indexer's packed local-Q/full-K metadata.""" + # Each real Q segment intersects its sequence with this rank's row interval, + # while K keeps the sequence's full compressed segment. The final synthetic + # segment holds CP capacity padding and has zero K rows. Causal offsets + # restore each non-empty local Q segment's position in the original sequence. + global_end = global_start + local_rows + zero = torch.zeros((1,), dtype=cu_seqlens_q.dtype, device=cu_seqlens_q.device) + local_starts = cu_seqlens_q[:-1].clamp_min(global_start) + local_ends = cu_seqlens_q[1:].clamp_max(global_end) + q_lens = (local_ends - local_starts).clamp_min(0) + q_prefix = torch.cumsum(q_lens, dim=0, dtype=torch.int32) + padding_q = (global_end - cu_seqlens_q[-1].clamp_min(global_start)).clamp_min(0) + cu_q_topk = torch.cat((zero, q_prefix, (q_prefix[-1] + padding_q).view(1))) + cu_k_topk = torch.cat((cu_seqlens_compressed, cu_seqlens_compressed[-1:])) + q_causal_offsets = torch.cat( + (torch.where(q_lens > 0, local_starts - cu_seqlens_q[:-1], 0), zero) + ) + return cu_q_topk, cu_k_topk, q_causal_offsets + + def compute_cp_indexer_topk( q_indexer_local: torch.Tensor, weights_indexer_local: torch.Tensor, @@ -285,11 +312,11 @@ def compute_cp_indexer_topk( indexer_softmax_scale: float, max_seqlen_q: int, use_fused: bool, -) -> Optional[torch.Tensor]: - """Run indexer top-k for this rank's query block.""" +) -> Tuple[Optional[torch.Tensor], Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]]: + """Return local top-k and its local-Q/full-K packed layout.""" topk_width = int(topk_width) if topk_width == 0 or k_indexer_seq_major.shape[0] == 0: - return None + return None, None global_start = int(global_start) l_local = q_indexer_local.shape[0] @@ -299,6 +326,10 @@ def compute_cp_indexer_topk( f"{l_local}, got {weights_indexer_local.shape[0]}." ) + cu_q_topk, cu_k_topk, q_causal_offsets = _build_cp_indexer_layout( + cu_seqlens_q, cu_seqlens_compressed, global_start, l_local + ) + if not use_fused: global_rows = torch.arange( global_start, @@ -347,50 +378,11 @@ def compute_cp_indexer_topk( values, rows = torch.topk(scores, selected_width, dim=-1) local_rows = k_positions[rows].to(torch.int32) output[start:end, :selected_width] = torch.where(torch.isfinite(values), local_rows, -1) - return output - - # The fused kernel derives a bottom-right causal mask from each Q/K - # segment. Copy only the prefix visible to this rank so that mask starts - # at the rank's true position in the global sequence. For example: + return output, (cu_q_topk, cu_k_topk, q_causal_offsets) - # 1 1 1 0 0 0 0 0 - # 1 1 1 1 0 0 0 0 - # 1 1 1 1 1 0 0 0 - - global_end = global_start + l_local - zero = torch.zeros((1,), dtype=cu_seqlens_q.dtype, device=cu_seqlens_q.device) - k_rows = torch.arange( - k_indexer_seq_major.shape[0], - dtype=cu_seqlens_compressed.dtype, - device=cu_seqlens_compressed.device, - ) - local_starts = cu_seqlens_q[:-1].clamp_min(global_start) - local_ends = cu_seqlens_q[1:].clamp_max(global_end) - q_lens = (local_ends - local_starts).clamp_min(0) - k_lens = torch.minimum( - torch.div((local_ends - cu_seqlens_q[:-1]).clamp_min(0), int(ratio), rounding_mode="floor"), - cu_seqlens_compressed[1:] - cu_seqlens_compressed[:-1], - ) - k_lens = torch.where(q_lens > 0, k_lens, 0) - q_prefix = torch.cumsum(q_lens, dim=0, dtype=torch.int32) - k_prefix = torch.cumsum(k_lens, dim=0, dtype=torch.int32) - padding_q = (global_end - cu_seqlens_q[-1].clamp_min(global_start)).clamp_min(0) - cu_q_topk = torch.cat((zero, q_prefix, (q_prefix[-1] + padding_q).view(1))) - cu_k_topk = torch.cat((zero, k_prefix, k_prefix[-1:])) - - k_seq_ids = torch.bucketize(k_rows, cu_k_topk[1:-1], out_int32=True, right=True).clamp_max( - cu_seqlens_q.shape[0] - 2 - ) - valid_k = k_rows < k_prefix[-1] - source_rows = cu_seqlens_compressed[k_seq_ids] + k_rows - cu_k_topk[k_seq_ids] - k_for_topk = torch.index_select(k_indexer_seq_major, 0, torch.where(valid_k, source_rows, 0)) - k_for_topk = k_for_topk * valid_k.view((-1,) + (1,) * (k_for_topk.ndim - 1)) - valid_k_lengths = torch.repeat_interleave( - torch.cat((k_lens, zero)), torch.cat((q_lens, padding_q.view(1))), output_size=l_local - ) topk, _ = indexer_topk( q_indexer_local, - k_for_topk, + k_indexer_seq_major, weights_indexer_local, topk=topk_width, ratio=ratio, @@ -399,6 +391,6 @@ def compute_cp_indexer_topk( cu_seqlens_kv=cu_k_topk, max_seqlen_q=int(max_seqlen_q), max_seqlen_kv=int(max_seqlen_q) // int(ratio), - valid_k_lengths=valid_k_lengths, + q_causal_offsets=q_causal_offsets, ) - return topk + return topk, (cu_q_topk, cu_k_topk, q_causal_offsets) diff --git a/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py b/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py index 53f67f9b314..e3656c3971c 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py +++ b/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py @@ -479,7 +479,7 @@ def _indexer_topk_core( cu_seqlens_kv: Optional[Tensor] = None, max_seqlen_q: Optional[int] = None, max_seqlen_kv: Optional[int] = None, - valid_k_lengths: Optional[Tensor] = None, + q_causal_offsets: Optional[Tensor] = None, ) -> Tuple[Tensor, Tensor, Tensor]: """Layout-agnostic core for :func:`indexer_topk`. @@ -530,16 +530,15 @@ def _indexer_topk_core( _ensure_dsa_namespace() # Kernel wants k as 3-D ``(total_k, h_kv, idx_hd)``. - scores = _DSA.indexer_forward_wrapper( - q, - k.unsqueeze(1), - w, - ratio=ratio, + forward_kwargs = dict( cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_kv, max_seqlen_q=int(max_seqlen_q), max_seqlen_k=int(max_seqlen_kv), - )[ + ) + if q_causal_offsets is not None: + forward_kwargs["q_causal_offsets"] = q_causal_offsets + scores = _DSA.indexer_forward_wrapper(q, k.unsqueeze(1), w, ratio=ratio, **forward_kwargs)[ "scores" ] # (total_q, max_seqlen_kv) fp32, -inf on masked positions # Defensive contiguify (wrapper may return a stride-padded slice). @@ -547,36 +546,18 @@ def _indexer_topk_core( sk = int(max_seqlen_kv) total_q = q.shape[0] - if valid_k_lengths is not None: - if valid_k_lengths.shape[0] != total_q: - raise ValueError( - "valid_k_lengths must have one entry per THD query row: " - f"got={valid_k_lengths.shape[0]}, expected={total_q}." - ) - if valid_k_lengths.dtype != torch.int32: - raise ValueError(f"valid_k_lengths must be int32, got {valid_k_lengths.dtype}.") - if valid_k_lengths.device != device: - raise ValueError("valid_k_lengths must be on the same device as q_indexer.") - seq_lens = valid_k_lengths - else: - # Per-row valid KV length: for token ``i`` in batch ``b``, - # pos_in_seq = i - cu_seqlens_q[b] - # valid = min((pos_in_seq + 1) // ratio, seqlen_kv[b]) - # Callers with nonstandard Q/K alignment pass ``valid_k_lengths`` - # when this default bottom-right mask formula does not apply. - row_idx = torch.arange(total_q, device=device, dtype=torch.int32) - row_batch_ids = batch_of_row(cu_seqlens_q, total_q=total_q) - row_valid = row_idx < cu_seqlens_q[-1] - pos_in_seq = row_idx - cu_seqlens_q[row_batch_ids] - pos_in_seq = torch.where(row_valid, pos_in_seq, torch.zeros_like(pos_in_seq)) - seqlen_kv_per_row = (cu_seqlens_kv[1:] - cu_seqlens_kv[:-1])[row_batch_ids] - seq_lens = ( - ((pos_in_seq + 1) // ratio) - .clamp(max=seqlen_kv_per_row) - .to(torch.int32) - .contiguous() - ) - seq_lens = torch.where(row_valid, seq_lens, torch.zeros_like(seq_lens)) + row_idx = torch.arange(total_q, device=device, dtype=torch.int32) + row_batch_ids = batch_of_row(cu_seqlens_q, total_q=total_q) + row_valid = row_idx < cu_seqlens_q[-1] + pos_in_seq = row_idx - cu_seqlens_q[row_batch_ids] + if q_causal_offsets is not None: + pos_in_seq = pos_in_seq + q_causal_offsets[row_batch_ids] + pos_in_seq = torch.where(row_valid, pos_in_seq, torch.zeros_like(pos_in_seq)) + seqlen_kv_per_row = (cu_seqlens_kv[1:] - cu_seqlens_kv[:-1])[row_batch_ids] + seq_lens = ( + ((pos_in_seq + 1) // ratio).clamp(max=seqlen_kv_per_row).to(torch.int32).contiguous() + ) + seq_lens = torch.where(row_valid, seq_lens, torch.zeros_like(seq_lens)) else: if k.shape[1] == 0: raise ValueError("indexer_topk requires at least one K row.") @@ -637,7 +618,7 @@ def indexer_topk( cu_seqlens_kv: Optional[Tensor] = None, max_seqlen_q: Optional[int] = None, max_seqlen_kv: Optional[int] = None, - valid_k_lengths: Optional[Tensor] = None, + q_causal_offsets: Optional[Tensor] = None, ) -> Tuple[Tensor, Tensor]: """Score + top-K selection for inference (no KL loss, no backward). @@ -659,10 +640,8 @@ def indexer_topk( cu_seqlens_kv: THD only — ``(B+1,)`` int32 CUDA cumulative KV lens. max_seqlen_q: THD only — per-batch max Q length. max_seqlen_kv: THD only — per-batch max KV length. - valid_k_lengths: THD only — optional ``(total_q,)`` int32 CUDA tensor - overriding the per-row valid K count passed to the radix top-K - wrapper when the default bottom-right mask does not describe the - Q/K alignment. + q_causal_offsets: THD only — optional ``(B,)`` int32 CUDA tensor. Entry + ``b`` is the sequence-relative position of that segment's first Q. Returns: SBHD: ``(topk_indices (b, sq, topk), topk_length (b, sq))`` int32 @@ -676,8 +655,8 @@ def indexer_topk( "indexer_topk THD mode requires cu_seqlens_q, cu_seqlens_kv, " "max_seqlen_q, and max_seqlen_kv to all be supplied." ) - if not is_thd and valid_k_lengths is not None: - raise ValueError("valid_k_lengths is only supported in THD mode.") + if not is_thd and q_causal_offsets is not None: + raise ValueError("q_causal_offsets is only supported in THD mode.") # ``indexer_softmax_scale`` is applied via the # ``relu(c·x) = c·relu(x)`` trick (the cudnn kernel does the relu), @@ -705,7 +684,7 @@ def indexer_topk( cu_seqlens_kv=cu_seqlens_kv, max_seqlen_q=int(max_seqlen_q) if max_seqlen_q is not None else None, max_seqlen_kv=int(max_seqlen_kv) if max_seqlen_kv is not None else None, - valid_k_lengths=valid_k_lengths, + q_causal_offsets=q_causal_offsets, ) return topk_indices, topk_length @@ -869,11 +848,11 @@ def _compute_dense_indexer_score( cu_seqlens_kv: Optional[Tensor] = None, max_seqlen_q: Optional[int] = None, max_seqlen_kv: Optional[int] = None, + q_causal_offsets: Optional[Tensor] = None, ) -> Tuple[Tensor, Tensor]: """Dense indexer score forward over the full ``S_k`` axis (BSHD or THD). Wraps :attr:`cudnn.DSA.dense_indexer_score_recompute_wrapper`. - This function is not used now, but it is kept for potential future use. Layout is selected by ``cu_seqlens_*`` kwargs: * **BSHD** (``cu_seqlens_*=None``): inputs are 4-D q ``(B, S_q, H, D)``, @@ -884,10 +863,18 @@ def _compute_dense_indexer_score( ``(total_q, H)``. Outputs are ``out (total_q, max_seqlen_kv)`` + ``denom (total_q,)``. - The kernel applies the bottom-right ratio causal mask - ``col_limit = min(S_k, (q+1) // ratio)`` regardless of layout. + The ratio-causal limit is + ``min(S_k, (q_causal_offset + q + 1) // ratio)``; omitted offsets are zero. """ _ensure_dsa_namespace() + kwargs = dict( + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_kv, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_kv, + ) + if q_causal_offsets is not None: + kwargs["q_causal_offsets"] = q_causal_offsets result = _DSA.dense_indexer_score_recompute_wrapper( q_indexer, k_indexer, @@ -895,10 +882,7 @@ def _compute_dense_indexer_score( qhead_per_kv_head=qhead_per_kv_head, sm_scale=indexer_softmax_scale, ratio=ratio, - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_k=cu_seqlens_kv, - max_seqlen_q=max_seqlen_q, - max_seqlen_k=max_seqlen_kv, + **kwargs, ) return result["out"], result["denom"] @@ -915,6 +899,7 @@ def _compute_dense_attn_score( cu_seqlens_kv: Optional[Tensor] = None, max_seqlen_q: Optional[int] = None, max_seqlen_kv: Optional[int] = None, + q_causal_offsets: Optional[Tensor] = None, ) -> Tuple[Tensor, Tensor]: """Dense attention score forward over the full ``S_k`` axis (BSHD or THD). @@ -922,6 +907,14 @@ def _compute_dense_attn_score( BSHD/THD layout convention as :func:`_compute_dense_indexer_score`. """ _ensure_dsa_namespace() + kwargs = dict( + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_kv, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_kv, + ) + if q_causal_offsets is not None: + kwargs["q_causal_offsets"] = q_causal_offsets result = _DSA.dense_attn_score_recompute_wrapper( q_attn, k_attn, @@ -929,10 +922,7 @@ def _compute_dense_attn_score( softmax_scale, qhead_per_kv_head=qhead_per_kv_head, ratio=ratio, - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_k=cu_seqlens_kv, - max_seqlen_q=max_seqlen_q, - max_seqlen_k=max_seqlen_kv, + **kwargs, ) return result["out"], result["denom"] @@ -1265,6 +1255,9 @@ def forward( ratio=ratio, **dense_attn_kwargs, ) + if padding_row_mask is not None: + attn_score = attn_score.masked_fill(padding_row_mask.unsqueeze(-1), 0) + attn_l1norm = attn_l1norm.masked_fill(padding_row_mask, 0) if loss_coeff > 0: indexer_loss = _kl_loss_from_dense_scores( @@ -1343,6 +1336,10 @@ def forward( else: attn_score_for_bwd = attn_score.clone() index_score_for_bwd = index_score.clone() + index_lse_for_bwd = index_lse + if padding_row_mask is not None: + index_score_for_bwd[padding_row_mask] = 0 + index_lse_for_bwd = index_lse.masked_fill(padding_row_mask, 0) dense_bwd_kwargs = {} if is_thd: dense_bwd_kwargs = dict( @@ -1358,7 +1355,7 @@ def forward( attn_score_for_bwd, attn_l1norm, index_score_for_bwd, - index_lse, + index_lse_for_bwd, sm_scale=indexer_softmax_scale, loss_coeff=indexer_loss_coeff, grad_loss=unit_grad_loss, @@ -1523,6 +1520,10 @@ def forward( indexer_softmax_scale: float, loss_coeff: float, loss_divisor: float, + sparse_loss: bool, + ratio: int, + max_seqlen_q: int, + indexer_layout: Tuple[Tensor, Tensor, Tensor], q_padding_mask: Optional[Tensor] = None, ) -> Tuple[Tensor, Tensor]: """Run fused sparse attention using caller-supplied top-k indices.""" @@ -1533,12 +1534,6 @@ def forward( total_comp = k_indexer.shape[0] indexer_topk = indexer_topk_idxs.shape[-1] - indexer_topk_idxs_for_loss = indexer_topk_idxs - if q_padding_mask is not None: - indexer_topk_idxs_for_loss = indexer_topk_idxs.masked_fill( - q_padding_mask.unsqueeze(-1), -1 - ) - out_flat, lse, lse_indexer = _dsa_fwd_flash_mla( query, kv_full, @@ -1549,45 +1544,125 @@ def forward( indexer_topk=indexer_topk, ) - if indexer_softmax_scale != 1.0: - weights_scaled = (weights.float() * indexer_softmax_scale).to(weights.dtype) - else: - weights_scaled = weights - q_bshd, k_bsd, w_bsh, topk_bst = _thd_to_fake_bshd( - q_indexer, k_indexer, weights_scaled, indexer_topk_idxs_for_loss - ) - predict = _DSA.sparse_indexer_score_recompute_wrapper( - q_bshd, k_bsd, w_bsh, topk_bst, qhead_per_kv_head=idx_nh, topk_indices_global=True - )["predict"].squeeze(0) - target = _compute_attn_target( - query.detach(), - compressed_kv.detach(), - lse_indexer.detach(), - indexer_topk_idxs_for_loss, - softmax_scale, - qhead_per_kv_head=np_, - topk_indices_global=True, - ) - - raw_local_loss = _kl_loss_from_target_predict( - target, predict, indexer_topk_idxs_for_loss, loss_coeff, calculate_per_token_loss=True - ) - indexer_loss = raw_local_loss / loss_divisor bwd_loss_coeff = loss_coeff * total_q / loss_divisor + unit_grad_loss = torch.ones((), device=query.device, dtype=torch.float32) - if loss_coeff > 0: - ig = _DSA.indexer_backward_wrapper( - q_indexer.view(1, total_q, idx_nh, idx_hd), - weights.view(1, total_q, idx_nh), - k_indexer.view(1, total_comp, idx_hd), - target.view(1, total_q, indexer_topk), - predict.view(1, total_q, indexer_topk), - indexer_topk_idxs_for_loss.view(1, total_q, indexer_topk), - sm_scale=indexer_softmax_scale, - loss_coeff=bwd_loss_coeff, - grad_loss=torch.ones((), device=query.device, dtype=torch.float32), - block_I=128, + if sparse_loss: + indexer_topk_idxs_for_loss = indexer_topk_idxs + if q_padding_mask is not None: + indexer_topk_idxs_for_loss = indexer_topk_idxs.masked_fill( + q_padding_mask.unsqueeze(-1), -1 + ) + weights_scaled = weights + if indexer_softmax_scale != 1.0: + weights_scaled = (weights.float() * indexer_softmax_scale).to(weights.dtype) + q_bshd, k_bsd, w_bsh, topk_bst = _thd_to_fake_bshd( + q_indexer, k_indexer, weights_scaled, indexer_topk_idxs_for_loss + ) + predict = _DSA.sparse_indexer_score_recompute_wrapper( + q_bshd, k_bsd, w_bsh, topk_bst, qhead_per_kv_head=idx_nh, topk_indices_global=True + )["predict"].squeeze(0) + target = _compute_attn_target( + query.detach(), + compressed_kv.detach(), + lse_indexer.detach(), + indexer_topk_idxs_for_loss, + softmax_scale, + qhead_per_kv_head=np_, + topk_indices_global=True, + ) + raw_local_loss = _kl_loss_from_target_predict( + target, + predict, + indexer_topk_idxs_for_loss, + loss_coeff, + calculate_per_token_loss=True, + ) + indexer_loss = raw_local_loss / loss_divisor + if loss_coeff > 0: + ig = _DSA.indexer_backward_wrapper( + q_indexer.view(1, total_q, idx_nh, idx_hd), + weights.view(1, total_q, idx_nh), + k_indexer.view(1, total_comp, idx_hd), + target.view(1, total_q, indexer_topk), + predict.view(1, total_q, indexer_topk), + indexer_topk_idxs_for_loss.view(1, total_q, indexer_topk), + sm_scale=indexer_softmax_scale, + loss_coeff=bwd_loss_coeff, + grad_loss=unit_grad_loss, + block_I=128, + ) + else: + cu_seqlens_q, cu_seqlens_k, q_causal_offsets = indexer_layout + max_seqlen_k = max_seqlen_q // ratio + index_score, index_lse = _compute_dense_indexer_score( + q_indexer, + k_indexer.unsqueeze(1), + weights, + qhead_per_kv_head=idx_nh, + indexer_softmax_scale=indexer_softmax_scale, + ratio=ratio, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_k, + q_causal_offsets=q_causal_offsets, + ) + if q_padding_mask is not None: + index_score = index_score.masked_fill(q_padding_mask.unsqueeze(-1), float("-inf")) + index_lse = index_lse.masked_fill(q_padding_mask, float("-inf")) + attn_score, attn_l1norm = _compute_dense_attn_score( + query.detach(), + compressed_kv.detach().unsqueeze(1), + lse_indexer.detach(), + qhead_per_kv_head=np_, + softmax_scale=softmax_scale, + ratio=ratio, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_k, + q_causal_offsets=q_causal_offsets, + ) + if q_padding_mask is not None: + attn_score = attn_score.masked_fill(q_padding_mask.unsqueeze(-1), 0) + attn_l1norm = attn_l1norm.masked_fill(q_padding_mask, 0) + raw_local_loss = _kl_loss_from_dense_scores( + attn_score, + attn_l1norm, + index_score, + index_lse, + loss_coeff, + calculate_per_token_loss=True, ) + indexer_loss = raw_local_loss / loss_divisor + + if loss_coeff > 0: + index_score_for_bwd = index_score.clone() + index_lse_for_bwd = index_lse + if q_padding_mask is not None: + index_score_for_bwd[q_padding_mask] = 0 + index_lse_for_bwd = index_lse.masked_fill(q_padding_mask, 0) + ig = _DSA.dense_indexer_backward_wrapper( + q_indexer, + weights, + k_indexer, + attn_score, + attn_l1norm, + index_score_for_bwd, + index_lse_for_bwd, + sm_scale=indexer_softmax_scale, + loss_coeff=bwd_loss_coeff, + grad_loss=unit_grad_loss, + block_I=128, + ratio=ratio, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + q_causal_offsets=q_causal_offsets, + ) + if loss_coeff > 0: saved_grad_q_indexer = ig["d_index_q"].view(total_q, idx_nh, idx_hd) saved_grad_k_indexer = ig["d_index_k"].view(total_comp, idx_hd) saved_grad_weights = ig["d_weights"].view(total_q, idx_nh) @@ -1657,6 +1732,10 @@ def backward(ctx, grad_output, grad_loss): None, None, None, + None, + None, + None, + None, ) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 81b8744c8db..1e42a4a9438 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1544,19 +1544,6 @@ def __post_init__(self): assert all( ratio in [0, 4, 128] for ratio in self.csa_compress_ratios ), "csa_compress_ratios must be 0, 4, or 128" - if ( - self.context_parallel_size > 1 - and 4 in self.csa_compress_ratios - and not self.csa_dense_mode - and (self.dsa_indexer_loss_coeff or 0.0) > 0 - and not self.dsa_indexer_use_sparse_loss - ): - # The current cuDNN fused indexer-loss kernel cannot express the - # CP-local query-offset semantics required by dense loss. - raise ValueError( - "DSv4 CP dense indexer loss is unsupported because the current cuDNN " - "fused kernel lacks CP query-offset semantics." - ) assert ( self.tensor_model_parallel_size == 1 ), "DSv4 Hybrid Attention only supports TP size 1." diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_layout_kernels.py b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_layout_kernels.py index 2b624a1236f..867b2c9e171 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_layout_kernels.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_layout_kernels.py @@ -25,7 +25,7 @@ def _require_cute_cuda(): def _make_e2e_like_cu_seqlens(device: str = "cuda") -> torch.Tensor: - """Return the ragged THD prefix pattern used by DSv4 CP e2e tests. + """Return a representative ragged THD prefix pattern. The lengths sum to 4096, so CP4 owns 1024 local rows. They contain short, long, boundary-crossing, and padded-tail sequences without allocating diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py index 37bc61a87f1..2c92052128a 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py @@ -201,49 +201,51 @@ def fake_apply(hidden, boundary, cu, global_start, ratio, d_comp, c_cap): assert torch.equal(rank_rows, torch.tensor([0, 1, 2, 3, 10, 11, 12, 13], dtype=torch.int32)) -def test_compute_cp_indexer_topk_passes_valid_lengths(monkeypatch): +def test_compute_cp_indexer_topk_passes_offsets_without_repacking_k(monkeypatch): topk_calls = [] def fake_indexer_topk( - q, k, _weights, *, topk, cu_seqlens_q, cu_seqlens_kv, valid_k_lengths, **_ + q, k, _weights, *, topk, cu_seqlens_q, cu_seqlens_kv, q_causal_offsets, **_ ): topk_calls.append( - (k.clone(), cu_seqlens_q.clone(), cu_seqlens_kv.clone(), valid_k_lengths.clone()) + (k.clone(), cu_seqlens_q.clone(), cu_seqlens_kv.clone(), q_causal_offsets.clone()) ) return torch.full((q.shape[0], int(topk)), len(topk_calls), dtype=torch.int32), None monkeypatch.setattr(csa_cp_utils, "indexer_topk", fake_indexer_topk) - q = torch.randn(5, 2) - weights = torch.randn(5, 1) + q = torch.randn(8, 2) + weights = torch.randn(8, 1) k_seq = torch.arange(8, dtype=torch.float32).reshape(4, 2) - cu_q = torch.tensor([0, 10], dtype=torch.int32) - cu_comp = torch.tensor([0, 3], dtype=torch.int32) + cu_q = torch.tensor([0, 5, 13, 20], dtype=torch.int32) + cu_comp = torch.tensor([0, 1, 3, 4], dtype=torch.int32) - out = compute_cp_indexer_topk( + out, metadata = compute_cp_indexer_topk( q, weights, k_seq, cu_q, cu_comp, - 8, + 7, ratio=4, topk_width=2, indexer_softmax_scale=0.5, - max_seqlen_q=10, + max_seqlen_q=8, use_fused=True, ) - assert torch.equal(out, torch.ones(5, 2, dtype=torch.int32)) - assert torch.equal(topk_calls[0][0][:2], k_seq[:2]) - assert torch.count_nonzero(topk_calls[0][0][2:]) == 0 - assert torch.equal(topk_calls[0][1], torch.tensor([0, 2, 5], dtype=torch.int32)) - assert torch.equal(topk_calls[0][2], torch.tensor([0, 2, 2], dtype=torch.int32)) - assert torch.equal(topk_calls[0][3], torch.tensor([2, 2, 0, 0, 0], dtype=torch.int32)) - assert ( - compute_cp_indexer_topk(q, weights, k_seq[:0], cu_q, cu_comp, 2, 4, 2, 1.0, 10, True) - is None - ) + assert torch.equal(out, torch.ones(8, 2, dtype=torch.int32)) + assert torch.equal(topk_calls[0][0], k_seq) + assert torch.equal(topk_calls[0][1], torch.tensor([0, 0, 6, 8, 8], dtype=torch.int32)) + assert torch.equal(topk_calls[0][2], torch.tensor([0, 1, 3, 4, 4], dtype=torch.int32)) + assert torch.equal(topk_calls[0][3], torch.tensor([0, 2, 0, 0], dtype=torch.int32)) + metadata_q, metadata_k, metadata_offsets = metadata + assert torch.equal(metadata_q, topk_calls[0][1]) + assert torch.equal(metadata_k, topk_calls[0][2]) + assert torch.equal(metadata_offsets, topk_calls[0][3]) + assert compute_cp_indexer_topk( + q, weights, k_seq[:0], cu_q, cu_comp, 2, 4, 2, 1.0, 10, True + ) == (None, None) def test_compute_cp_indexer_topk_unfused_uses_exact_global_positions(monkeypatch): @@ -261,7 +263,7 @@ def fail_if_fused(*_args, **_kwargs): topk_width = 3 scale = 0.7 - actual = compute_cp_indexer_topk( + actual, _ = compute_cp_indexer_topk( q, weights, k, diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py index 096c73d1451..71d9f865216 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py @@ -2317,21 +2317,11 @@ def test_real_fused_dense_loss_matches_reference(self, reset_lazy_kernel_state): q_attn_bshd = query.permute(1, 0, 2, 3).contiguous().float() k_attn_bsd = kv_full[kv_offset:].permute(1, 0, 2).contiguous().float() - # PyTorch reference that mirrors the fused path's dense-loss math - # exactly. Two non-obvious requirements: - # * Use FlashMLA's emitted ``lse_indexer`` (logsumexp over the - # indexer-selected top-K positions, with the per-head sink term), - # not an analytical full-KV logsumexp. Otherwise the per-row LSE - # basis differs from the kernel by ~50x. - # * Do NOT apply the ratio-causal mask in the reference scores — - # the dense-score-recompute kernels emit values at every position - # (no internal masking). Masking the reference would shift the - # ``attn_score / attn_l1norm`` normalization and the indexer LSE - # basis, producing a different KL than the kernel's. + # The reference uses FlashMLA's actual normalization because it + # includes the selected top-k positions and the per-head sink term. from megatron.core.transformer.experimental_attention_variant.dsa_kernels import ( _dsa_fwd_flash_mla, _indexer_topk_core, - _kl_loss_from_dense_scores, ) # Run indexer + FlashMLA to capture the same ``lse_indexer`` the fused @@ -2363,23 +2353,17 @@ def test_real_fused_dense_loss_matches_reference(self, reset_lazy_kernel_state): ) lse_indexer_bsqh = lse_indexer.reshape(s['sq'], s['b'], s['np_']).permute(1, 0, 2) - # Attention path: exp(QK*scale - lse_indexer), head-summed. No mask. - qk_attn = torch.einsum('bqhd,bkd->bqhk', q_attn_bshd, k_attn_bsd) * s['softmax_scale'] - attn_score_ref = torch.exp(qk_attn - lse_indexer_bsqh.unsqueeze(-1)).sum(dim=2) - attn_l1norm_ref = attn_score_ref.sum(dim=-1) - - # Indexer path: ReLU(QK_indexer) * W head-summed, scaled by - # ``indexer_softmax_scale`` once. The fused path applies the scale - # via pre-scaled ``w_bsh_scaled`` only (not again inside the - # kernel), giving ``score = sum_h(relu(QK) * w_raw) * sm_scale``. - qk_idx = torch.einsum('bqhd,bkd->bqhk', q_idx_bshd, k_idx_bsd) - idx_score_ref = (torch.relu(qk_idx) * w_bsh.unsqueeze(-1)).sum(dim=2) * s[ - 'indexer_softmax_scale' - ] - idx_lse_ref = torch.logsumexp(idx_score_ref, dim=-1) - - loss_ref = _kl_loss_from_dense_scores( - attn_score_ref, attn_l1norm_ref, idx_score_ref, idx_lse_ref, loss_coeff + loss_ref = _ref_dense_indexer_loss( + q_idx_bshd, + k_idx_bsd, + w_bsh, + q_attn_bshd, + k_attn_bsd, + lse_indexer_bsqh, + indexer_softmax_scale=s['indexer_softmax_scale'], + attn_softmax_scale=s['softmax_scale'], + ratio=s['ratio'], + loss_coeff=loss_coeff, ) assert torch.allclose(indexer_loss, loss_ref, atol=5e-2, rtol=1e-1), ( f"actual = {indexer_loss.item():.6f}, ref = {loss_ref.item():.6f}, " @@ -2578,6 +2562,7 @@ def test_indexer_topk_thd_dispatch_calls_thd_kernel_path(self, reset_lazy_kernel q, k, w, cu_q, cu_kv = self._make_indexer_topk_thd_inputs() total_q, idx_nh, idx_hd = q.shape total_k = k.shape[0] + q_causal_offsets = torch.tensor([5, 0], dtype=torch.int32, device=q.device) fake_dsa = MagicMock(name='_DSA_thd_stub') @@ -2587,6 +2572,7 @@ def fake_indexer_forward(q_thd, k_thd, w_thd, ratio, **kwargs): assert 'cu_seqlens_k' in kwargs and kwargs['cu_seqlens_k'] is cu_kv assert kwargs['max_seqlen_q'] == 3 assert kwargs['max_seqlen_k'] == 2 + assert kwargs['q_causal_offsets'] is q_causal_offsets return {'scores': torch.zeros(total_q, 2, dtype=torch.float32, device=q_thd.device)} fake_dsa.indexer_forward_wrapper.side_effect = fake_indexer_forward @@ -2608,12 +2594,15 @@ def fake_indexer_forward(q_thd, k_thd, w_thd, ratio, **kwargs): cu_seqlens_kv=cu_kv, max_seqlen_q=3, max_seqlen_kv=2, + q_causal_offsets=q_causal_offsets, ) # THD return shape: (total_q, topk) + (total_q,). assert topk_idxs.shape == (total_q, 2) assert topk_len.shape == (total_q,) # Confirmed the kernel was called with THD kwargs. fake_dsa.indexer_forward_wrapper.assert_called_once() + seq_lens = fake_dsa.indexer_top_k_wrapper.call_args.args[1] + assert torch.equal(seq_lens, torch.tensor([1, 1, 2, 0, 0], device=q.device)) # ===================================================================== # dsa_sparse_attn(is_thd=True) @@ -3404,7 +3393,12 @@ def _sbhd_to_bshd(q_sbhd, k_sbd, w_sbh, sm_scale): idx_score_aligned = ( idx_score_analytical + (kernel_indexer_scores - idx_score_analytical).detach() ) - idx_lse_aligned = torch.logsumexp(idx_score_aligned, dim=-1) + position_valid = torch.isfinite(kernel_indexer_scores) + idx_score_aligned = idx_score_aligned.masked_fill(~position_valid, float('-inf')) + row_valid = position_valid.any(dim=-1) + idx_lse_aligned = torch.logsumexp( + torch.where(row_valid.unsqueeze(-1), idx_score_aligned, 0), dim=-1 + ).masked_fill(~row_valid, float('-inf')) loss_ref = _kl_loss_from_dense_scores( attn_score_const.detach(), attn_l1norm_const.detach(), diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py index 9bb0dfafc84..1126fc3d63a 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py @@ -57,10 +57,12 @@ _DSV4_CP_GRAPH_TIMING_MEASURE_REPLAYS = 5 _DSV4_CP_TEST_VARIANT = "flash" -# Padded total is 4096, divisible by CP2/CP4. Only the final segment has tail -# padding, so the intermediate sequence boundaries match the unpadded layout. -_DSV4_CP_RAGGED_SEG_LENS = (1, 127, 1000, 23, 129, 900, 55, 257, 800, 95, 509, 148) -_DSV4_CP_RAGGED_PADDED_SEG_LENS = (1, 127, 1000, 23, 129, 900, 55, 257, 800, 95, 509, 200) +# Keep every segment at least as long as the smallest compression ratio (4). +# The cuDNN dense-indexer backward kernel does not yet support a THD segment +# with Q rows but no compressed K rows. Padded total remains 4096, divisible +# by CP2/CP4, and only the final segment has tail padding. +_DSV4_CP_RAGGED_SEG_LENS = (5, 127, 1000, 23, 129, 900, 55, 257, 800, 95, 509, 148) +_DSV4_CP_RAGGED_PADDED_SEG_LENS = (5, 127, 1000, 23, 129, 900, 55, 257, 800, 95, 509, 196) # Same shape, padded total, and max padded sequence length as # _DSV4_CP_RAGGED_PADDED_SEG_LENS, but the padding is distributed across many # sequences instead of only the tail. CUDA graph replay must tolerate this value @@ -546,12 +548,19 @@ def _measure_cp1_cuda_graph_time(self, layer_number, packed, padded_tokens): return measured @pytest.mark.parametrize( - "layer_number", - [1, 2, 3], - ids=["ratio_0_window_only", "ratio_4_indexer", "ratio_128_compressor"], + ("layer_number", "sparse_loss"), + [(1, True), (2, True), (2, False), (3, True)], + ids=[ + "ratio_0_window_only", + "ratio_4_indexer_sparse", + "ratio_4_indexer_dense", + "ratio_128_compressor", + ], ) @pytest.mark.parametrize("apply_rope_fusion", [True, False], ids=["fused_rope", "unfused_rope"]) - def test_thd_cp_matches_full_reference_forward_backward(self, layer_number, apply_rope_fusion): + def test_thd_cp_matches_full_reference_forward_backward( + self, layer_number, sparse_loss, apply_rope_fusion + ): """CP path matches the full-sequence THD reference on ragged packed inputs using the DSv4 layer configuration. @@ -566,18 +575,16 @@ def test_thd_cp_matches_full_reference_forward_backward(self, layer_number, appl config_cp = _make_dsv4_cp_config( context_parallel_size=self.cp_size, dsa_indexer_loss_coeff=1.0, - dsa_indexer_use_sparse_loss=True, + dsa_indexer_use_sparse_loss=sparse_loss, apply_dsa_kernel_fusion=apply_dsa_kernel_fusion, apply_rope_fusion=apply_rope_fusion, ) config_ref = _make_dsv4_cp_config( context_parallel_size=1, dsa_indexer_loss_coeff=1.0, - dsa_indexer_use_sparse_loss=True, - # Numerical parity uses the unfused CP1 reference so failures point - # at the CP path instead of non-CP fused RoPE behavior. + dsa_indexer_use_sparse_loss=sparse_loss, apply_dsa_kernel_fusion=apply_dsa_kernel_fusion, - apply_rope_fusion=False, + apply_rope_fusion=apply_rope_fusion, ) cp_attn = _build_attention( config_cp, layer_number=layer_number, pg_collection=self.pg From 1ada640c9ec312cd2e4aa17d8bc3d0742bfe7c10 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Thu, 2 Jul 2026 03:29:18 -0700 Subject: [PATCH 18/23] Enable fused kernel test on hopper Signed-off-by: kunlunl --- docker/Dockerfile.ci.dev | 16 ++++- .../dsa_kernels.py | 2 +- .../core/transformer/transformer_config.py | 4 +- .../test_dsa_kernels.py | 28 ++++---- .../test_dsv4_hybrid_attention_cp.py | 19 ++++-- .../test_dsv4_hybrid_native_parity.py | 66 ++++++++++++++----- 6 files changed, 94 insertions(+), 41 deletions(-) diff --git a/docker/Dockerfile.ci.dev b/docker/Dockerfile.ci.dev index 7f3a5c0552a..33f0f2cd8a4 100644 --- a/docker/Dockerfile.ci.dev +++ b/docker/Dockerfile.ci.dev @@ -31,7 +31,7 @@ RUN bash -ex <<"EOF" curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | sh EOF -RUN ln -sf /usr/local/cuda/targets/x86_64-linux/include/cuda \ +RUN ln -sfn /usr/local/cuda/targets/x86_64-linux/include/cccl/cuda \ /usr/local/include/cuda RUN find /usr/local/cuda -name "utility" 2>/dev/null | head -5 && \ ls /usr/local/cuda/targets/x86_64-linux/include/ | head -20 @@ -40,6 +40,8 @@ COPY README.md pyproject.toml uv.lock /workspace/ COPY megatron/core/__init__.py /workspace/megatron/core/ COPY megatron/core/package_info.py /workspace/megatron/core/ ARG IMAGE_TYPE=dev +ARG CUDNN_FRONTEND_COMMIT=0a14b7181d129d30e7bad34b8c3ed0a0c995e23d +ARG FLASH_MLA_COMMIT=b7643bd54521f563b839b98289b5cd048c062ba2 RUN --mount=type=cache,target=/root/.cache/uv \ bash -ex <<"EOF" export NVTE_CUDA_ARCHS="80;90;100" @@ -61,7 +63,12 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --no-install-package nvidia-cusolver-cu12 \ --no-install-package nvidia-cusparse-cu12 \ --no-install-package nvidia-cusparselt-cu12 \ - --no-install-package nvidia-nccl-cu12 + --no-install-package nvidia-nccl-cu12 \ + --no-install-package nvidia-cudnn-frontend + uv pip install --no-build-isolation \ + "flash-mla @ git+https://github.com/deepseek-ai/FlashMLA.git@${FLASH_MLA_COMMIT}" \ + "nvidia-cudnn-frontend[cutedsl] @ git+https://github.com/NVIDIA/cudnn-frontend.git@${CUDNN_FRONTEND_COMMIT}" \ + "torch-c-dlpack-ext==0.1.5" EOF # Install DeepEP @@ -83,6 +90,11 @@ RUN bash -ex <<"EOF" EOF COPY assets/ /opt/data/ +# Reinstall CuTe DSL cleanly after upgrading the version provided by the base image. +RUN uv pip install --reinstall --no-deps \ + nvidia-cutlass-dsl==4.5.0 \ + nvidia-cutlass-dsl-libs-base==4.5.0 \ + nvidia-cutlass-dsl-libs-cu13==4.5.0 ENV UV_PYTHON=$UV_PROJECT_ENVIRONMENT/bin/python ##### For NVIDIANS only ##### diff --git a/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py b/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py index e3656c3971c..0fb747cce99 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py +++ b/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py @@ -334,7 +334,7 @@ def build_flat_topk_idxs( class SparseAttnFunc(torch.autograd.Function): - """SM100 sparse attention fwd + bwd on flat tensors. + """Sparse attention fwd + bwd on flat tensors. Forward uses :mod:`flash_mla`; backward uses cuDNN Frontend's :attr:`cudnn.DSA.sparse_attention_backward_wrapper`. diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 1e42a4a9438..2fe8475a2c1 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1555,8 +1555,8 @@ def __post_init__(self): torch.cuda.is_available() ), "apply_dsa_kernel_fusion requires a CUDA device, but none is available." sm = torch.cuda.get_device_capability() - assert sm[0] >= 10, ( - f"apply_dsa_kernel_fusion requires SM100+ (Blackwell or later), " + assert sm[0] >= 9, ( + f"apply_dsa_kernel_fusion requires SM90+ (Hopper or later), " f"but current device has compute capability {sm[0]}.{sm[1]}." ) diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py index 71d9f865216..16db0ed8777 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py @@ -1545,7 +1545,7 @@ def test_dense_path_fwd_kernel_calls_and_bwd_grads(self, reset_lazy_kernel_state # * CUDA is unavailable; # * cuDNN frontend is not installed (``import cudnn`` fails); # * ``cudnn.DSA`` namespace is missing; -# * SM is too low (sparse: SM90+; dense: SM100+); +# * SM is below SM90; # * for FlashMLA-needing tests, ``flash_mla`` is not installed. # --------------------------------------------------------------------------- @@ -1783,8 +1783,8 @@ def _ref_sparse_attn_forward( d=512, idx_nh=32, idx_hd=128, - # topk = lcm(64, 128) = 128 satisfies SparseScoreRecomputeSm100's - # `topk % n_block_size == 0` (64 for score_type=attention, 128 for indexer). + # topk = lcm(64, 128) = 128 satisfies both score block-size requirements + # (64 for score_type=attention, 128 for indexer). topk=128, ratio=1, softmax_scale=512**-0.5, @@ -1848,7 +1848,7 @@ class TestRealKernelScoreHelpers: ['sparse_indexer_predict', 'sparse_attn_target', 'dense_indexer_score', 'dense_attn_score'], ) def test_real_score_helper(self, case, reset_lazy_kernel_state): - _skip_if_real_kernels_unavailable(sm_min=10) + _skip_if_real_kernels_unavailable() s = _REAL_SHAPES_SPARSE # Each case needs a different combination of the input fixture. @@ -1991,7 +1991,7 @@ class TestRealKernelKLLossDense: @pytest.mark.parametrize("dummy", [None]) def test_real_dense_kl_loss_matches_reference(self, dummy, reset_lazy_kernel_state): - _skip_if_real_kernels_unavailable(sm_min=10) + _skip_if_real_kernels_unavailable() from megatron.core.transformer.experimental_attention_variant.dsa_kernels import ( _compute_dense_attn_score, _compute_dense_indexer_score, @@ -2053,7 +2053,7 @@ class TestRealKernelIndexerTopk: @pytest.mark.parametrize("dummy", [None]) def test_real_indexer_topk_set_matches_reference(self, dummy, reset_lazy_kernel_state): - _skip_if_real_kernels_unavailable(sm_min=10) # IndexerForward is SM100+ + _skip_if_real_kernels_unavailable() from megatron.core.transformer.experimental_attention_variant.dsa_kernels import ( indexer_topk, ) @@ -2173,7 +2173,7 @@ def test_real_dsa_sparse_attn_fwd_bwd_matches_reference(self, reset_lazy_kernel_ in one test halves cuDNN compile time vs running them separately, since they share the same kernel cache key. """ - _skip_if_real_kernels_unavailable(sm_min=10, need_flash_mla=True) + _skip_if_real_kernels_unavailable(need_flash_mla=True) s = self.SHAPES # ---- Real path: forward + backward via dsa_sparse_attn ---- @@ -2271,7 +2271,7 @@ def test_real_fused_dense_loss_matches_reference(self, reset_lazy_kernel_state): emits its own internal lse_indexer that differs slightly), so the tolerance is wider than for the kernel-only ``KLLossDense`` test. """ - _skip_if_real_kernels_unavailable(sm_min=10, need_flash_mla=True) + _skip_if_real_kernels_unavailable(need_flash_mla=True) s = self.SHAPES torch.manual_seed(0) dev = 'cuda' @@ -2747,7 +2747,7 @@ def test_thd_single_segment_matches_sbhd_b1(self, sparse_loss, reset_lazy_kernel on the same input tensors (just reshaped), for both dense-loss and sparse-loss Path B. """ - _skip_if_real_kernels_unavailable(sm_min=10, need_flash_mla=True) + _skip_if_real_kernels_unavailable(need_flash_mla=True) s = self.SHAPES torch.manual_seed(0) dev = 'cuda' @@ -3094,7 +3094,7 @@ def test_per_seg_padding_excluded_from_loss(self, sparse_loss, reset_lazy_kernel then expand each segment with intra-segment padding and supply cu_seqlens_q_unpadded. Losses should match. """ - _skip_if_real_kernels_unavailable(sm_min=10, need_flash_mla=True) + _skip_if_real_kernels_unavailable(need_flash_mla=True) dev = 'cuda' # Baseline: tightly packed (real lengths only, no padding). @@ -3120,7 +3120,7 @@ def test_per_seg_padding_unmasked_corrupts_loss(self, sparse_loss, reset_lazy_ke """Without cu_seqlens_q_unpadded, random per-segment padding rows DO corrupt the loss — confirming the masking is necessary. """ - _skip_if_real_kernels_unavailable(sm_min=10, need_flash_mla=True) + _skip_if_real_kernels_unavailable(need_flash_mla=True) dev = 'cuda' real = self._build_multi_seg_inputs(self.SEG_LENS_REAL, self.SHAPES, dev) @@ -3143,7 +3143,9 @@ def test_per_seg_padding_grads_are_zeroed(self, sparse_loss, reset_lazy_kernel_s """Indexer gradients at per-segment padding positions are zero, and gradients at real-token positions match the unpadded baseline. """ - _skip_if_real_kernels_unavailable(sm_min=10, need_flash_mla=True) + _skip_if_real_kernels_unavailable(need_flash_mla=True) + if torch.cuda.get_device_capability()[0] == 9 and not sparse_loss: + pytest.skip("cuDNN Frontend SM90 dense DSA has shape-incomplete compile caches") dev = 'cuda' # ---- Unpadded baseline (reference grads) ----------------------------- @@ -3259,7 +3261,7 @@ class TestRealKernelDenseIndexerBackward: SHAPES = TestRealKernelFusedIndexerSparseAttn.SHAPES def test_real_dense_backward_grad_matches_autograd(self, reset_lazy_kernel_state): - _skip_if_real_kernels_unavailable(sm_min=10, need_flash_mla=True) + _skip_if_real_kernels_unavailable(need_flash_mla=True) s = self.SHAPES torch.manual_seed(0) dev = 'cuda' diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py index 1126fc3d63a..f6f59ce129c 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py @@ -79,7 +79,7 @@ def _dsv4_cp_fused_kernels_available(): sm_major = torch.cuda.get_device_capability()[0] except RuntimeError: return False - if sm_major < 10: + if sm_major < 9: return False try: from cudnn import DSA # noqa: F401 @@ -90,7 +90,7 @@ def _dsv4_cp_fused_kernels_available(): _DSV4_CP_FUSED_KERNELS_UNAVAILABLE_REASON = ( - "DSv4 fused DSA cases require SM100+, flash_mla, and cudnn.DSA" + "DSv4 fused DSA cases require SM90+, flash_mla, and cudnn.DSA" ) @@ -567,6 +567,12 @@ def test_thd_cp_matches_full_reference_forward_backward( Verifies local CP output, hidden grad, and reduced parameter grads against the sliced full-reference tensors. """ + if ( + torch.cuda.get_device_capability()[0] == 9 + and self.fused_kernels_available + and not sparse_loss + ): + pytest.skip("cuDNN Frontend SM90 dense DSA has offset, cache, and stream bugs") packed, padded_tokens, local_idx = _make_ragged_cp_case(self.cp_size, self.cp_rank) torch.manual_seed(_SEED + layer_number) @@ -817,6 +823,11 @@ def test_thd_cp_cuda_graph_matches_eager_forward_backward( for name, param in graph_attn.named_parameters() if param.grad is not None } + # The unfused graph pool is large on H100. Its results are cloned, so + # release it before allocating the eager reference's backward buffers. + del graph, graph_output, graph_attn, full_hidden, test_hidden + del static_hidden, static_grad + _clear_cuda_test_state() eager_out, eager_hidden_grad, eager_param_grads = _run_dsv4_attention_forward_backward( eager_attn, eager_hidden, test_grad, packed @@ -838,9 +849,7 @@ def test_thd_cp_cuda_graph_matches_eager_forward_backward( f"layer={layer_number}:{mode}:param_grad:{name}", ) - del graph, graph_output - del graph_attn, eager_attn, full_hidden, test_hidden, test_grad, static_hidden - del eager_hidden, static_grad, graph_out, graph_hidden_grad + del eager_attn, test_grad, eager_hidden, graph_out, graph_hidden_grad _clear_cuda_test_state() @pytest.mark.parametrize( diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_native_parity.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_native_parity.py index 507b7879c21..c8816d6e18c 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_native_parity.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_native_parity.py @@ -304,23 +304,39 @@ def _native_sparse_attn( ) -> torch.Tensor: sq, batch_size, num_heads, head_dim = query.size() kv_t = kv_full.permute(1, 0, 2) - safe_indices = topk_indices.clamp(min=0).long() - gather_index = safe_indices.unsqueeze(-1).expand(-1, -1, -1, head_dim) - kv_gathered = torch.gather(kv_t.unsqueeze(1).expand(-1, sq, -1, -1), dim=2, index=gather_index) + chunk_size = 1024 if sq > 4096 else sq + outputs = [] + for query_chunk, topk_chunk in zip( + query.split(chunk_size), topk_indices.split(chunk_size, dim=1) + ): + safe_indices = topk_chunk.clamp(min=0).long() + chunk_len = query_chunk.size(0) + if sq > 4096: + # Expanded gather backward materializes [B, S, K, D], which is + # 80 GiB for pro/8192. Chunking also bounds einsum workspace. + batch = torch.arange(batch_size, device=kv_t.device)[:, None, None] + kv_gathered = kv_t[batch, safe_indices] + else: + gather_index = safe_indices.unsqueeze(-1).expand(-1, -1, -1, head_dim) + kv_gathered = torch.gather( + kv_t.unsqueeze(1).expand(-1, chunk_len, -1, -1), dim=2, index=gather_index + ) - q = query.permute(1, 2, 0, 3).float() - scores = torch.einsum("bnsh,bskh->bnsk", q, kv_gathered.float()) * softmax_scale - scores = scores.masked_fill((topk_indices < 0).unsqueeze(1), float("-inf")) + q = query_chunk.permute(1, 2, 0, 3).float() + kv_gathered_float = kv_gathered.float() + scores = torch.einsum("bnsh,bskh->bnsk", q, kv_gathered_float) * softmax_scale + scores = scores.masked_fill((topk_chunk < 0).unsqueeze(1), float("-inf")) - sink = attn_sink.view(1, num_heads, 1, 1).float() - scores_max = torch.max(scores.max(dim=-1, keepdim=True).values, sink) - exp_scores = torch.exp(scores - scores_max) - exp_sink = torch.exp(sink - scores_max) - attn_weights = exp_scores / (exp_scores.sum(dim=-1, keepdim=True) + exp_sink) + sink = attn_sink.view(1, num_heads, 1, 1).float() + scores_max = torch.max(scores.max(dim=-1, keepdim=True).values, sink) + exp_scores = torch.exp(scores - scores_max) + exp_sink = torch.exp(sink - scores_max) + attn_weights = exp_scores / (exp_scores.sum(dim=-1, keepdim=True) + exp_sink) - output = torch.einsum("bnsk,bskh->bnsh", attn_weights, kv_gathered.float()) - output = output.to(query.dtype).permute(2, 0, 1, 3).contiguous() - return output.reshape(sq, batch_size, num_heads * head_dim) + output = torch.einsum("bnsk,bskh->bnsh", attn_weights, kv_gathered_float) + output = output.to(query.dtype).permute(2, 0, 1, 3).contiguous() + outputs.append(output.reshape(chunk_len, batch_size, num_heads * head_dim)) + return outputs[0] if len(outputs) == 1 else torch.cat(outputs, dim=0) def _native_fused_sparse_indexer_loss( @@ -915,7 +931,7 @@ def test_attention_matches_native_reference( dsa_indexer_use_sparse_loss: bool, ): if apply_dsa_kernel_fusion: - _skip_if_real_kernels_unavailable(sm_min=10) + _skip_if_real_kernels_unavailable() major, _ = torch.cuda.get_device_capability() if major < 10 and not apply_dsa_kernel_fusion and seqlen > 4096: pytest.skip("seqlen > 4096 may OOM on Hopper with unfused DSA implementation") @@ -1021,8 +1037,15 @@ def test_thd_attention_matches_native_reference( gradients against the native reference. """ if apply_dsa_kernel_fusion: - _skip_if_real_kernels_unavailable(sm_min=10) + _skip_if_real_kernels_unavailable() major, _ = torch.cuda.get_device_capability() + if ( + major == 9 + and apply_dsa_kernel_fusion + and compress_ratio == 4 + and not dsa_indexer_use_sparse_loss + ): + pytest.skip("cuDNN Frontend SM90 THD dense DSA has cache and stream bugs") if major < 10 and not apply_dsa_kernel_fusion and seqlen > 4096: pytest.skip("seqlen > 4096 may OOM on Hopper with unfused DSA implementation") @@ -1143,7 +1166,7 @@ class of bug that single-segment and padding-invariance tests miss. max compress ratio) so compression is exact at every ratio. """ if apply_dsa_kernel_fusion: - _skip_if_real_kernels_unavailable(sm_min=10) + _skip_if_real_kernels_unavailable() major, _ = torch.cuda.get_device_capability() total_T = sum(seg_lens) if major < 10 and not apply_dsa_kernel_fusion and total_T > 4096: @@ -1273,7 +1296,14 @@ def test_thd_padded_attention_matches_unpadded( within tolerance. """ if apply_dsa_kernel_fusion: - _skip_if_real_kernels_unavailable(sm_min=10) + _skip_if_real_kernels_unavailable() + if ( + torch.cuda.get_device_capability()[0] == 9 + and apply_dsa_kernel_fusion + and compress_ratio == 4 + and not dsa_indexer_use_sparse_loss + ): + pytest.skip("cuDNN Frontend SM90 THD dense DSA has cache and stream bugs") actual_T = sum(seg_lens) assert actual_T <= pad_max_seqlen, "seg_lens must fit within pad_max_seqlen" From d4832dbf28a860a2ac9767127a409c8b95bd7167 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Thu, 2 Jul 2026 05:28:14 -0700 Subject: [PATCH 19/23] Fix dockerfile Signed-off-by: kunlunl --- docker/Dockerfile.ci.dev | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker/Dockerfile.ci.dev b/docker/Dockerfile.ci.dev index 33f0f2cd8a4..5411ca72acb 100644 --- a/docker/Dockerfile.ci.dev +++ b/docker/Dockerfile.ci.dev @@ -31,7 +31,7 @@ RUN bash -ex <<"EOF" curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | sh EOF -RUN ln -sfn /usr/local/cuda/targets/x86_64-linux/include/cccl/cuda \ +RUN ln -sf /usr/local/cuda/include/cccl/cuda \ /usr/local/include/cuda RUN find /usr/local/cuda -name "utility" 2>/dev/null | head -5 && \ ls /usr/local/cuda/targets/x86_64-linux/include/ | head -20 From e1198f6e9fd2ce3e4c8bb58c3fc3fa599d836201 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Thu, 2 Jul 2026 09:05:10 -0700 Subject: [PATCH 20/23] Fix CI Signed-off-by: kunlunl --- .../experimental_attention_variant/test_dsa_kernels.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py index 16db0ed8777..230300d0f3a 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py @@ -2626,7 +2626,7 @@ def test_dsa_sparse_attn_thd_wrong_ndim_raises(self, query_shape, kv_shape, matc dsa_sparse_attn(query, kv, attn_sink, topk, softmax_scale=0.125, is_thd=True) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - def test_dsa_sparse_attn_thd_pass_through_no_reshape(self): + def test_dsa_sparse_attn_thd_pass_through_no_reshape(self, reset_lazy_kernel_state): """THD inputs flow into ``SparseAttnFunc`` unchanged (no SBHD reshape) and the output is ``(total_q, np * d_v)``. """ From e8b0ff2eb7a93218477f60162d5b5e730ff75bc3 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Thu, 2 Jul 2026 10:49:49 -0700 Subject: [PATCH 21/23] Address comments Signed-off-by: kunlunl --- megatron/core/transformer/transformer_config.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 2fe8475a2c1..61e96d9d252 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1480,7 +1480,9 @@ def __post_init__(self): self.experimental_attention_variant != "dsv4_hybrid" and self.cp_partition_mode != "zigzag" ): - raise ValueError("cp_partition_mode='contiguous' is only supported with DSv4.") + raise ValueError( + "cp_partition_mode='contiguous' currently is only supported with dsv4_hybrid." + ) if self.experimental_attention_variant in ["gated_delta_net"]: assert ( From 3599a37c9b25860da6f49945ad4b65ef1eddf379 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Thu, 2 Jul 2026 21:32:57 -0700 Subject: [PATCH 22/23] Minor fix from claude code review comments Signed-off-by: kunlunl --- megatron/core/datasets/data_schedule.py | 2 +- .../experimental_attention_variant/csa.py | 4 +- .../csa_cp_layout_kernels.py | 7 +--- .../csa_cp_utils.py | 5 ++- .../core/transformer/transformer_config.py | 39 +++++++++++++++++++ .../test_csa_cp_utils.py | 4 ++ .../test_dsa_kernels.py | 7 +++- .../test_dsv4_hybrid_native_parity.py | 7 ++++ 8 files changed, 64 insertions(+), 11 deletions(-) diff --git a/megatron/core/datasets/data_schedule.py b/megatron/core/datasets/data_schedule.py index e359309970a..af87664b07b 100644 --- a/megatron/core/datasets/data_schedule.py +++ b/megatron/core/datasets/data_schedule.py @@ -630,7 +630,7 @@ def get_batch_on_this_rank_for_sequence_packing( group_size=local_cp_size_val ) - cp_partition_mode = config.cp_partition_mode if config is not None else "zigzag" + cp_partition_mode = getattr(config, "cp_partition_mode", "zigzag") contiguous_cp_local_target_len = None pad_alignment = ( getattr(config, 'pad_packed_seq_alignment', None) if config is not None else None diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index 0c3547d0b7d..0b8bc31e0be 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -701,7 +701,7 @@ def _unfused_indexer_sparse_attn_from_topk( ~valid.unsqueeze(1), float("-inf") ) target_logits = torch.where(row_valid.unsqueeze(1), target_logits, 0.0) - sink = attn_sink.view(1, np_, 1).float() + sink = attn_sink.detach().view(1, np_, 1).float() scores_max = torch.maximum(target_logits.max(dim=-1, keepdim=True).values, sink) exp_scores = torch.exp(target_logits - scores_max) exp_sink = torch.exp(sink - scores_max) @@ -758,7 +758,7 @@ def _unfused_indexer_sparse_attn_from_topk( attn_scores = torch.einsum("rhd,rkd->rhk", query.detach().float(), selected_kv.float()) attn_scores = attn_scores * softmax_scale attn_scores = attn_scores.masked_fill(~valid.unsqueeze(1), float("-inf")) - sink = attn_sink.view(1, np_, 1).float() + sink = attn_sink.detach().view(1, np_, 1).float() scores_max = torch.maximum(attn_scores.max(dim=-1, keepdim=True).values, sink) exp_scores = torch.exp(attn_scores - scores_max) exp_sink = torch.exp(sink - scores_max) diff --git a/megatron/core/transformer/experimental_attention_variant/csa_cp_layout_kernels.py b/megatron/core/transformer/experimental_attention_variant/csa_cp_layout_kernels.py index a77f72d56f0..ae3e4c7daf3 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa_cp_layout_kernels.py +++ b/megatron/core/transformer/experimental_attention_variant/csa_cp_layout_kernels.py @@ -8,7 +8,6 @@ """ import math -import os from typing import Optional, Tuple import torch @@ -591,11 +590,7 @@ def _build_attention_indices_launch( def _require_cute(message: str, *tensors: Optional[torch.Tensor]) -> None: """Raise ``RuntimeError`` when a wrapper cannot use CuTeDSL kernels.""" - if ( - not _CUTE_AVAILABLE - or os.environ.get("DSV4_CP_DISABLE_CUTE_KERNELS") - or not all(tensor is None or tensor.is_cuda for tensor in tensors) - ): + if not _CUTE_AVAILABLE or not all(tensor is None or tensor.is_cuda for tensor in tensors): raise RuntimeError(message) diff --git a/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py b/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py index dabe4f2d02b..3459a1954fb 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py +++ b/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py @@ -317,6 +317,9 @@ def compute_cp_indexer_topk( topk_width = int(topk_width) if topk_width == 0 or k_indexer_seq_major.shape[0] == 0: return None, None + max_seqlen_kv = int(max_seqlen_q) // int(ratio) + if max_seqlen_kv == 0: + return None, None global_start = int(global_start) l_local = q_indexer_local.shape[0] @@ -390,7 +393,7 @@ def compute_cp_indexer_topk( cu_seqlens_q=cu_q_topk, cu_seqlens_kv=cu_k_topk, max_seqlen_q=int(max_seqlen_q), - max_seqlen_kv=int(max_seqlen_q) // int(ratio), + max_seqlen_kv=int(max_seqlen_kv), q_causal_offsets=q_causal_offsets, ) return topk, (cu_q_topk, cu_k_topk, q_causal_offsets) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 61e96d9d252..f3c0775355d 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1,5 +1,6 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +import inspect import logging import math import warnings @@ -1561,6 +1562,19 @@ def __post_init__(self): f"apply_dsa_kernel_fusion requires SM90+ (Hopper or later), " f"but current device has compute capability {sm[0]}.{sm[1]}." ) + uses_ratio4_indexer = 4 in self.csa_compress_ratios and not self.csa_dense_mode + indexer_loss_enabled = (self.dsa_indexer_loss_coeff or 0.0) > 0 + if ( + sm[0] == 9 + and uses_ratio4_indexer + and indexer_loss_enabled + and not self.dsa_indexer_use_sparse_loss + ): + raise ValueError( + "DSv4 with fused DSA and dense indexer loss is not supported on SM90 " + "because the cuDNN Frontend SM90 dense DSA kernels are not reliable for " + "this path. Use sparse indexer loss or disable DSA kernel fusion." + ) _flash_mla_available = True try: @@ -1590,6 +1604,31 @@ def __post_init__(self): f"PyTorch fallback." ) + if ( + self.context_parallel_size > 1 or self.dynamic_context_parallel + ) and uses_ratio4_indexer: + required_wrappers = [DSA.indexer_forward_wrapper] + if indexer_loss_enabled and not self.dsa_indexer_use_sparse_loss: + required_wrappers.extend( + [ + DSA.dense_indexer_score_recompute_wrapper, + DSA.dense_attn_score_recompute_wrapper, + DSA.dense_indexer_backward_wrapper, + ] + ) + missing_offsets = [ + wrapper.__name__ + for wrapper in required_wrappers + if "q_causal_offsets" not in inspect.signature(wrapper).parameters + ] + if missing_offsets: + raise ValueError( + "DSv4 CP with ratio-4 fused DSA requires cuDNN Frontend wrappers " + "with q_causal_offsets support; missing from: " + f"{', '.join(missing_offsets)}. Install a compatible cuDNN Frontend " + "build or disable DSA kernel fusion." + ) + if ( self.gdn_pre_gated_delta_rule_fusion and self.experimental_attention_variant != "gated_delta_net" diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py index 2c92052128a..7a189e59ca6 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py @@ -246,6 +246,10 @@ def fake_indexer_topk( assert compute_cp_indexer_topk( q, weights, k_seq[:0], cu_q, cu_comp, 2, 4, 2, 1.0, 10, True ) == (None, None) + assert compute_cp_indexer_topk( + q, weights, k_seq, cu_q, torch.zeros_like(cu_comp), 2, 4, 2, 1.0, 3, True + ) == (None, None) + assert len(topk_calls) == 1 def test_compute_cp_indexer_topk_unfused_uses_exact_global_positions(monkeypatch): diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py index 230300d0f3a..beea1e26c1e 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py @@ -3395,7 +3395,12 @@ def _sbhd_to_bshd(q_sbhd, k_sbd, w_sbh, sm_scale): idx_score_aligned = ( idx_score_analytical + (kernel_indexer_scores - idx_score_analytical).detach() ) - position_valid = torch.isfinite(kernel_indexer_scores) + position_valid = ( + _ratio_causal_valid_mask(s['sq'], s['n_comp'], s['ratio'], dev) + .unsqueeze(0) + .expand(s['b'], -1, -1) + ) + assert torch.isfinite(kernel_indexer_scores[position_valid]).all() idx_score_aligned = idx_score_aligned.masked_fill(~position_valid, float('-inf')) row_valid = position_valid.any(dim=-1) idx_lse_aligned = torch.logsumexp( diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_native_parity.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_native_parity.py index c8816d6e18c..0831424c88d 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_native_parity.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_native_parity.py @@ -933,6 +933,13 @@ def test_attention_matches_native_reference( if apply_dsa_kernel_fusion: _skip_if_real_kernels_unavailable() major, _ = torch.cuda.get_device_capability() + if ( + major == 9 + and apply_dsa_kernel_fusion + and compress_ratio == 4 + and not dsa_indexer_use_sparse_loss + ): + pytest.skip("cuDNN Frontend SM90 dense DSA is not supported") if major < 10 and not apply_dsa_kernel_fusion and seqlen > 4096: pytest.skip("seqlen > 4096 may OOM on Hopper with unfused DSA implementation") From 15439504affc37b6dc47a247ef9c065c06fd812e Mon Sep 17 00:00:00 2001 From: kunlunl Date: Thu, 2 Jul 2026 23:53:11 -0700 Subject: [PATCH 23/23] Address CI comments Signed-off-by: kunlunl --- docker/Dockerfile.ci.dev | 14 +- pyproject.toml | 8 +- .../test_dsa_kernels.py | 2 + uv.lock | 479 ++++++++++++++---- 4 files changed, 389 insertions(+), 114 deletions(-) diff --git a/docker/Dockerfile.ci.dev b/docker/Dockerfile.ci.dev index 5411ca72acb..4ea7409b17f 100644 --- a/docker/Dockerfile.ci.dev +++ b/docker/Dockerfile.ci.dev @@ -40,8 +40,6 @@ COPY README.md pyproject.toml uv.lock /workspace/ COPY megatron/core/__init__.py /workspace/megatron/core/ COPY megatron/core/package_info.py /workspace/megatron/core/ ARG IMAGE_TYPE=dev -ARG CUDNN_FRONTEND_COMMIT=0a14b7181d129d30e7bad34b8c3ed0a0c995e23d -ARG FLASH_MLA_COMMIT=b7643bd54521f563b839b98289b5cd048c062ba2 RUN --mount=type=cache,target=/root/.cache/uv \ bash -ex <<"EOF" export NVTE_CUDA_ARCHS="80;90;100" @@ -63,12 +61,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --no-install-package nvidia-cusolver-cu12 \ --no-install-package nvidia-cusparse-cu12 \ --no-install-package nvidia-cusparselt-cu12 \ - --no-install-package nvidia-nccl-cu12 \ - --no-install-package nvidia-cudnn-frontend - uv pip install --no-build-isolation \ - "flash-mla @ git+https://github.com/deepseek-ai/FlashMLA.git@${FLASH_MLA_COMMIT}" \ - "nvidia-cudnn-frontend[cutedsl] @ git+https://github.com/NVIDIA/cudnn-frontend.git@${CUDNN_FRONTEND_COMMIT}" \ - "torch-c-dlpack-ext==0.1.5" + --no-install-package nvidia-nccl-cu12 EOF # Install DeepEP @@ -90,11 +83,6 @@ RUN bash -ex <<"EOF" EOF COPY assets/ /opt/data/ -# Reinstall CuTe DSL cleanly after upgrading the version provided by the base image. -RUN uv pip install --reinstall --no-deps \ - nvidia-cutlass-dsl==4.5.0 \ - nvidia-cutlass-dsl-libs-base==4.5.0 \ - nvidia-cutlass-dsl-libs-cu13==4.5.0 ENV UV_PYTHON=$UV_PROJECT_ENVIRONMENT/bin/python ##### For NVIDIANS only ##### diff --git a/pyproject.toml b/pyproject.toml index 95a965e959d..b7afa89b3cb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -98,7 +98,7 @@ dev = [ "megatron-energon[av_decode]~=6.0", "av", "flashinfer-python>=0.5.0,<0.7.0", - "nvidia-cudnn-frontend", + "nvidia-cudnn-frontend[cutedsl]", "wget", "onnxscript", "fastapi~=0.50", # Forcing a little bit more recent version of fastapi to be compatible with pydantic 2.0 @@ -169,12 +169,13 @@ linting = [ "pylint==3.2.6", ] ci = ["python-gitlab", "slack-sdk", "pandas"] -no_pypi_wheels = ["emerging_optimizers; python_version >= '3.12'", "fast-hadamard-transform"] +no_pypi_wheels = ["flash_mla", "emerging_optimizers; python_version >= '3.12'", "fast-hadamard-transform"] [tool.uv] default-groups = ["linting", "build", "test"] no-build-isolation-package = [ "causal-conv1d", + "flash_mla", "mamba-ssm", "transformer-engine", "transformer-engine-torch", @@ -194,8 +195,9 @@ override-dependencies = [ [tool.uv.sources] flash_mla = [ - { git = "https://github.com/deepseek-ai/FlashMLA", rev = "9edee0c022cd0938148a18e334203b0aab43aa19" }, + { git = "https://github.com/deepseek-ai/FlashMLA", rev = "nv_dev" }, ] +nvidia-cudnn-frontend = { git = "https://github.com/NVIDIA/cudnn-frontend.git", rev = "0a14b7181d129d30e7bad34b8c3ed0a0c995e23d" } transformer-engine = { git = "https://github.com/NVIDIA/TransformerEngine.git", rev = "f031cf87bd054c7558b887df7bed93975456667f" } nemo-run = { git = "https://github.com/NVIDIA-NeMo/Run.git", rev = "17ae86b64d7f75653351664f5d8c9e466faede00" } emerging_optimizers = { git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git", rev = "v0.3.0" } diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py index beea1e26c1e..5ae462f8af0 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py @@ -3095,6 +3095,8 @@ def test_per_seg_padding_excluded_from_loss(self, sparse_loss, reset_lazy_kernel cu_seqlens_q_unpadded. Losses should match. """ _skip_if_real_kernels_unavailable(need_flash_mla=True) + if torch.cuda.get_device_capability()[0] == 9 and not sparse_loss: + pytest.skip("cuDNN Frontend SM90 dense DSA has a default-stream race") dev = 'cuda' # Baseline: tightly packed (real lengths only, no padding). diff --git a/uv.lock b/uv.lock index 062d09d4217..1359f18ea06 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 2 +revision = 3 requires-python = ">=3.10" resolution-markers = [ "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform == 'win32' and extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts'", @@ -305,7 +305,7 @@ version = "1.4.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "frozenlist" }, - { name = "typing-extensions", marker = "python_full_version < '3.13'" }, + { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/61/62/06741b579156360248d1ec624842ad0edf697050bbaf7c3e46394e106ad1/aiosignal-1.4.0.tar.gz", hash = "sha256:f47eecd9468083c2029cc99945502cb7708b082c232f9aca65da147157b251c7", size = 25007, upload-time = "2025-07-03T22:54:43.528Z" } wheels = [ @@ -359,10 +359,10 @@ name = "anyio" version = "4.9.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "exceptiongroup", marker = "python_full_version < '3.11'" }, + { name = "exceptiongroup", marker = "python_full_version < '3.11' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "idna" }, { name = "sniffio" }, - { name = "typing-extensions", marker = "python_full_version < '3.13'" }, + { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/95/7d/4c1bd541d4dffa1b52bd83fb8527089e097a106fc90b467a7313b105f840/anyio-4.9.0.tar.gz", hash = "sha256:673c0c244e15788651a4ff38710fea9675823028a6f08a5eda409e0c9840a028", size = 190949, upload-time = "2025-03-17T00:02:54.77Z" } wheels = [ @@ -797,7 +797,7 @@ name = "cffi" version = "2.0.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pycparser", marker = "implementation_name != 'PyPy'" }, + { name = "pycparser", marker = "implementation_name != 'PyPy' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/eb/56/b1ba7935a17738ae8453301356628e8147c79dbb825bcbc73dc7401f9846/cffi-2.0.0.tar.gz", hash = "sha256:44d1b5909021139fe36001ae048dbdde8214afa20200eda0f64c068cac5d5529", size = 523588, upload-time = "2025-09-08T23:24:04.541Z" } wheels = [ @@ -984,7 +984,7 @@ name = "click" version = "8.3.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "colorama", marker = "sys_platform == 'win32' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/57/75/31212c6bf2503fdf920d87fee5d7a86a2e3bcf444984126f13d8e4016804/click-8.3.2.tar.gz", hash = "sha256:14162b8b3b3550a7d479eafa77dfd3c38d9dc8951f6f69c78913a8f9a7540fd5", size = 302856, upload-time = "2026-04-03T19:14:45.118Z" } wheels = [ @@ -1153,8 +1153,8 @@ name = "cryptography" version = "46.0.6" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cffi", marker = "platform_python_implementation != 'PyPy'" }, - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, + { name = "cffi", marker = "platform_python_implementation != 'PyPy' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "typing-extensions", marker = "python_full_version < '3.11' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/a4/ba/04b1bd4218cbc58dc90ce967106d51582371b898690f3ae0402876cc4f34/cryptography-46.0.6.tar.gz", hash = "sha256:27550628a518c5c6c903d84f637fbecf287f6cb9ced3804838a1295dc1fd0759", size = 750542, upload-time = "2026-03-25T23:34:53.396Z" } wheels = [ @@ -1213,7 +1213,7 @@ name = "cuda-bindings" version = "13.2.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cuda-pathfinder" }, + { name = "cuda-pathfinder", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32') or (sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev') or (sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev') or (sys_platform == 'win32' and extra == 'extra-13-megatron-core-lts')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/1a/fe/7351d7e586a8b4c9f89731bfe4cf0148223e8f9903ff09571f78b3fb0682/cuda_bindings-13.2.0-cp310-cp310-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:08b395f79cb89ce0cd8effff07c4a1e20101b873c256a1aeb286e8fd7bd0f556", size = 5744254, upload-time = "2026-03-11T00:12:29.798Z" }, @@ -1266,37 +1266,37 @@ wheels = [ [package.optional-dependencies] cublas = [ - { name = "nvidia-cublas", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or sys_platform == 'linux'" }, + { name = "nvidia-cublas", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or (python_full_version >= '3.11' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or sys_platform == 'linux' or (sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] cudart = [ - { name = "nvidia-cuda-runtime", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or sys_platform == 'linux'" }, + { name = "nvidia-cuda-runtime", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or (python_full_version >= '3.11' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or sys_platform == 'linux' or (sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] cufft = [ - { name = "nvidia-cufft", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or sys_platform == 'linux'" }, + { name = "nvidia-cufft", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or (python_full_version >= '3.11' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or sys_platform == 'linux' or (sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] cufile = [ - { name = "nvidia-cufile", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cufile", marker = "sys_platform == 'linux' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] cupti = [ - { name = "nvidia-cuda-cupti", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or sys_platform == 'linux'" }, + { name = "nvidia-cuda-cupti", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or (python_full_version >= '3.11' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or sys_platform == 'linux' or (sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] curand = [ - { name = "nvidia-curand", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or sys_platform == 'linux'" }, + { name = "nvidia-curand", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or (python_full_version >= '3.11' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or sys_platform == 'linux' or (sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] cusolver = [ - { name = "nvidia-cusolver", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or sys_platform == 'linux'" }, + { name = "nvidia-cusolver", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or (python_full_version >= '3.11' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or sys_platform == 'linux' or (sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] cusparse = [ - { name = "nvidia-cusparse", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or sys_platform == 'linux'" }, + { name = "nvidia-cusparse", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or (python_full_version >= '3.11' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or sys_platform == 'linux' or (sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] nvjitlink = [ - { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or (python_full_version >= '3.11' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or sys_platform == 'linux' or (sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] nvrtc = [ - { name = "nvidia-cuda-nvrtc", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or sys_platform == 'linux'" }, + { name = "nvidia-cuda-nvrtc", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or (python_full_version >= '3.11' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or sys_platform == 'linux' or (sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] nvtx = [ - { name = "nvidia-nvtx", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or sys_platform == 'linux'" }, + { name = "nvidia-nvtx", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or (python_full_version >= '3.11' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or sys_platform == 'linux' or (sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] [[package]] @@ -1564,8 +1564,8 @@ name = "emerging-optimizers" version = "0.3.0" source = { git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.3.0#b309e2f01cda75dc96a6dc1a2355a7b3b64b5e16" } dependencies = [ - { name = "absl-py", marker = "python_full_version >= '3.12'" }, - { name = "torch", marker = "python_full_version >= '3.12' and sys_platform == 'never'" }, + { name = "absl-py", marker = "python_full_version >= '3.12' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "torch", marker = "(python_full_version >= '3.12' and sys_platform == 'never') or (python_full_version < '3.12' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform != 'never' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] [[package]] @@ -1573,7 +1573,7 @@ name = "exceptiongroup" version = "1.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.13'" }, + { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" } wheels = [ @@ -1694,6 +1694,11 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/60/ee/a3cba17965482b35c4990af90bad108e82c32edcb59911c37f318b5f4198/flash_linear_attention-0.4.2-py3-none-any.whl", hash = "sha256:c08be006ce4dbe1be81f54938ee8e6fc7968cfba397c8d06c7669e97b8c44c0d", size = 284661, upload-time = "2026-03-12T14:45:44.905Z" }, ] +[[package]] +name = "flash-mla" +version = "1.0.0+b7643bd" +source = { git = "https://github.com/deepseek-ai/FlashMLA?rev=nv_dev#b7643bd54521f563b839b98289b5cd048c062ba2" } + [[package]] name = "flashinfer-python" version = "0.5.3" @@ -1705,8 +1710,10 @@ dependencies = [ { name = "ninja" }, { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and extra == 'extra-13-megatron-core-lts') or extra == 'extra-13-megatron-core-dev'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and extra == 'extra-13-megatron-core-lts') or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, - { name = "nvidia-cudnn-frontend" }, - { name = "nvidia-cutlass-dsl" }, + { name = "nvidia-cudnn-frontend", version = "1.25.0", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-13-megatron-core-lts'" }, + { name = "nvidia-cudnn-frontend", version = "1.25.0", source = { git = "https://github.com/NVIDIA/cudnn-frontend.git?rev=0a14b7181d129d30e7bad34b8c3ed0a0c995e23d#0a14b7181d129d30e7bad34b8c3ed0a0c995e23d" }, marker = "extra == 'extra-13-megatron-core-dev'" }, + { name = "nvidia-cutlass-dsl", version = "4.4.2", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-13-megatron-core-lts'" }, + { name = "nvidia-cutlass-dsl", version = "4.5.0", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-13-megatron-core-dev'" }, { name = "nvidia-ml-py" }, { name = "packaging" }, { name = "requests" }, @@ -2371,7 +2378,7 @@ dependencies = [ { name = "filelock" }, { name = "fsspec", version = "2026.2.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.14' or sys_platform != 'win32' or extra == 'extra-13-megatron-core-dev' or extra == 'extra-13-megatron-core-lts'" }, { name = "fsspec", version = "2026.3.0", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.14' and sys_platform == 'win32' and extra != 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts') or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, - { name = "hf-xet", marker = "platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64'" }, + { name = "hf-xet", marker = "platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "packaging" }, { name = "pyyaml" }, { name = "requests" }, @@ -2941,7 +2948,7 @@ resolution-markers = [ "python_full_version < '3.11' and extra != 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", ] dependencies = [ - { name = "mdurl", marker = "python_full_version < '3.11'" }, + { name = "mdurl", marker = "python_full_version < '3.11' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/38/71/3b932df36c1a044d397a1f92d1cf91ee0a503d91e470cbd670aa66b07ed0/markdown-it-py-3.0.0.tar.gz", hash = "sha256:e3f60a94fa066dc52ec76661e37c851cb232d92f9886b15cb560aaada2df8feb", size = 74596, upload-time = "2023-06-03T06:41:14.443Z" } wheels = [ @@ -3015,7 +3022,7 @@ resolution-markers = [ "python_full_version == '3.11.*' and sys_platform != 'emscripten' and sys_platform != 'win32' and extra != 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", ] dependencies = [ - { name = "mdurl", marker = "python_full_version >= '3.11'" }, + { name = "mdurl", marker = "python_full_version >= '3.11' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/5b/f5/4ec618ed16cc4f8fb3b701563655a69816155e79e24a17b651541804721d/markdown_it_py-4.0.0.tar.gz", hash = "sha256:cb0a2b4aa34f932c007117b194e945bd74e0ec24133ceb5bac59009cda1cb9f3", size = 73070, upload-time = "2025-08-11T12:57:52.854Z" } wheels = [ @@ -3187,7 +3194,7 @@ dev = [ { name = "mamba-ssm" }, { name = "megatron-energon", extra = ["av-decode"], marker = "extra == 'extra-13-megatron-core-dev'" }, { name = "multi-storage-client" }, - { name = "nvidia-cudnn-frontend" }, + { name = "nvidia-cudnn-frontend", version = "1.25.0", source = { git = "https://github.com/NVIDIA/cudnn-frontend.git?rev=0a14b7181d129d30e7bad34b8c3ed0a0c995e23d#0a14b7181d129d30e7bad34b8c3ed0a0c995e23d" }, extra = ["cutedsl"], marker = "extra == 'extra-13-megatron-core-dev'" }, { name = "nvidia-modelopt", marker = "(sys_platform != 'darwin' and extra == 'extra-13-megatron-core-dev') or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "nvidia-resiliency-ext" }, { name = "nvtx" }, @@ -3276,6 +3283,7 @@ linting = [ no-pypi-wheels = [ { name = "emerging-optimizers", marker = "python_full_version >= '3.12' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "fast-hadamard-transform" }, + { name = "flash-mla" }, ] test = [ { name = "coverage" }, @@ -3322,7 +3330,7 @@ requires-dist = [ { name = "multi-storage-client", marker = "extra == 'dev'", specifier = "~=0.27" }, { name = "multi-storage-client", marker = "extra == 'lts'", specifier = "~=0.27" }, { name = "numpy" }, - { name = "nvidia-cudnn-frontend", marker = "extra == 'dev'" }, + { name = "nvidia-cudnn-frontend", extras = ["cutedsl"], marker = "extra == 'dev'", git = "https://github.com/NVIDIA/cudnn-frontend.git?rev=0a14b7181d129d30e7bad34b8c3ed0a0c995e23d" }, { name = "nvidia-modelopt", extras = ["torch"], marker = "sys_platform != 'darwin' and extra == 'dev'" }, { name = "nvidia-resiliency-ext", marker = "extra == 'dev'", git = "https://github.com/NVIDIA/nvidia-resiliency-ext.git?rev=15a851565a4ce846c04431ecb0cf09903ab4837e" }, { name = "nvtx", marker = "extra == 'dev'", specifier = "~=0.2" }, @@ -3386,6 +3394,7 @@ linting = [ no-pypi-wheels = [ { name = "emerging-optimizers", marker = "python_full_version >= '3.12'", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.3.0" }, { name = "fast-hadamard-transform", git = "https://github.com/Dao-AILab/fast-hadamard-transform.git?rev=f134af63deb2df17e1171a9ec1ea4a7d8604d5ca" }, + { name = "flash-mla", git = "https://github.com/deepseek-ai/FlashMLA?rev=nv_dev" }, ] test = [ { name = "coverage" }, @@ -3655,7 +3664,7 @@ name = "multidict" version = "6.7.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, + { name = "typing-extensions", marker = "python_full_version < '3.11' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/1a/c2/c2d94cbe6ac1753f3fc980da97b3d930efe1da3af3c9f5125354436c073d/multidict-6.7.1.tar.gz", hash = "sha256:ec6652a1bee61c53a3e5776b6049172c53b6aaba34f18c9ad04f82712bac623d", size = 102010, upload-time = "2026-01-26T02:46:45.979Z" } wheels = [ @@ -4350,7 +4359,7 @@ name = "nvidia-cudnn-cu13" version = "9.19.0.56" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, + { name = "nvidia-cublas", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32') or (sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/f1/84/26025437c1e6b61a707442184fa0c03d083b661adf3a3eecfd6d21677740/nvidia_cudnn_cu13-9.19.0.56-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:6ed29ffaee1176c612daf442e4dd6cfeb6a0caa43ddcbeb59da94953030b1be4", size = 433781201, upload-time = "2026-02-03T20:40:53.805Z" }, @@ -4360,27 +4369,97 @@ wheels = [ [[package]] name = "nvidia-cudnn-frontend" -version = "1.24.0" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/71/05/bf31134b6e5d41c5a1e4abc65b5bb5461a73be74013da273b18fef7a2244/nvidia_cudnn_frontend-1.24.0-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e8833079f0283948cb5f99a2dc0c8fbff29d320e6a5635c4f77fa7eaf877043b", size = 3222069, upload-time = "2026-05-20T05:01:12.464Z" }, - { url = "https://files.pythonhosted.org/packages/50/4d/861c75757e688a5094da621871688d7583b52c7cb3ad75d0d5ab1dcdff68/nvidia_cudnn_frontend-1.24.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5736397ab8f29e06731960e055d27f65e8fda70c70f988e8bd91b671a3506999", size = 3370408, upload-time = "2026-05-20T05:01:36.098Z" }, - { url = "https://files.pythonhosted.org/packages/de/13/cee7f47acbb1d85a2019522f88238cbd4ac182a65fdea5d41e504b7c48be/nvidia_cudnn_frontend-1.24.0-cp310-cp310-win_amd64.whl", hash = "sha256:77bc9f3203c677f74b6cedf84125514b4881dc82f4177cc4ab33949693abe6aa", size = 2764295, upload-time = "2026-05-20T05:01:56.212Z" }, - { url = "https://files.pythonhosted.org/packages/81/75/767973a56b98d2a8fbc04c78fc28684cdb0df7c032ba4858090243df81cc/nvidia_cudnn_frontend-1.24.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9c0f39f211bb105798c7a8617b1d674e01dbd538b97714025628ea7127bfaf8a", size = 3223392, upload-time = "2026-05-20T05:02:56.092Z" }, - { url = "https://files.pythonhosted.org/packages/f5/8c/f5226ab5163dfffbe82ff41b9e1fbf649c908c077a5416bb16a1c3634cfa/nvidia_cudnn_frontend-1.24.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:dec07dcfea168792098b9a2652ec465d79e228d3d17a2f86a08404b487336530", size = 3372046, upload-time = "2026-05-20T05:03:20.091Z" }, - { url = "https://files.pythonhosted.org/packages/a7/bf/fa08a8bd953ad8db7ea3f64676b30f1cb2846f6518595a6faf7fb3558db4/nvidia_cudnn_frontend-1.24.0-cp311-cp311-win_amd64.whl", hash = "sha256:b461259b85b7a7e3a1c41b02c33ce4fde0dbcde7e0a227a968dedb74d311e2c3", size = 2764877, upload-time = "2026-05-20T05:04:21.965Z" }, - { url = "https://files.pythonhosted.org/packages/25/f4/a57efcbeb1dec0a047fa8f8aaa4defa4935db7d52e8afaf83224f1258ec6/nvidia_cudnn_frontend-1.24.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3b4398cecbaa555baa73a9b8716233632e3c16259c6ab999d83c51ca3b8fd09c", size = 3224420, upload-time = "2026-05-20T05:04:50.051Z" }, - { url = "https://files.pythonhosted.org/packages/19/b7/c1b8de8292b8fe28b2ffc95601a0b69392536a9515315bcfcdb0b07d2af8/nvidia_cudnn_frontend-1.24.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:144bfe09098d681d4c793c867fff53b1fb7c49f845324d5d0c52d824b14b74f5", size = 3376971, upload-time = "2026-05-20T05:05:19.998Z" }, - { url = "https://files.pythonhosted.org/packages/dc/00/284b8de99fbc1e8fd91c292a024943ee61d3361aa669c435ee44e14b6498/nvidia_cudnn_frontend-1.24.0-cp312-cp312-win_amd64.whl", hash = "sha256:5476d6a51ebaf5ef04e462e0052f1d9bca1af6274f738cb509715b4cf443a8df", size = 2765023, upload-time = "2026-05-20T05:05:40.269Z" }, - { url = "https://files.pythonhosted.org/packages/41/3e/430941e91a0c5234c79aa0bd6dcdb27c0e3a443f66953f881f8cb3428e93/nvidia_cudnn_frontend-1.24.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ccf1d2352f4b82fafbfb803512493c7a211224b424b6a78fbf45e42a94dcdb81", size = 3224661, upload-time = "2026-05-20T05:06:01.577Z" }, - { url = "https://files.pythonhosted.org/packages/d8/95/02691a6f0db4c2194899579a9c196df8990b140e0503fc6be1f09e7aa063/nvidia_cudnn_frontend-1.24.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ec45b08e0ab511f61532bc980c343c15eb5eda9ff14a3c80e75bc0aa3776860c", size = 3376128, upload-time = "2026-05-20T05:06:21.613Z" }, - { url = "https://files.pythonhosted.org/packages/d0/c5/e6c9ef439e167675e64fd1ca025df79211cafdf70ec9edbf6b99f876cafd/nvidia_cudnn_frontend-1.24.0-cp313-cp313-win_amd64.whl", hash = "sha256:2d02744a46726d262d80ed54299fae6491e4385e7d580eb6a027fb5b3b2c1db8", size = 2765241, upload-time = "2026-05-20T05:06:45.6Z" }, - { url = "https://files.pythonhosted.org/packages/f0/fc/c5a1d386f22dc8e17304874368ba213ff8c093093e9b0ccc3b2bd81a1d12/nvidia_cudnn_frontend-1.24.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:190857577a11d22b62da1863cb1a4b72692f98913d1c9c2e6a72224d08685575", size = 3226124, upload-time = "2026-05-20T05:07:14.426Z" }, - { url = "https://files.pythonhosted.org/packages/2e/39/e7f12c1a640174bdebbbf87be19819483c9da9f8f3af948f5dcb46cf501d/nvidia_cudnn_frontend-1.24.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b7e9aca4b6ce4d4bd484f01a1c4f530a9bea317cd69d21760d582b47e05514d8", size = 3375383, upload-time = "2026-05-20T05:07:38.976Z" }, - { url = "https://files.pythonhosted.org/packages/16/f8/cae1cffd3ab944301e3a736395f85925a887fc5255c100f561638d6b2fba/nvidia_cudnn_frontend-1.24.0-cp314-cp314-win_amd64.whl", hash = "sha256:b8968eb9dd9a71fe3d64b55d1e9731cffb7272578a9a39c9bf816f5e27c3f14e", size = 2765850, upload-time = "2026-05-20T05:07:59.264Z" }, - { url = "https://files.pythonhosted.org/packages/9e/27/c04b542fd2a882fabed1cb4e6778a8a6159f1123e00c2f0da8383690c713/nvidia_cudnn_frontend-1.24.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:65044724acc5fcb92ab829c2d2fcefc7fd4419881030e04358a66a4bcbce5d43", size = 3230149, upload-time = "2026-05-20T05:08:21.883Z" }, - { url = "https://files.pythonhosted.org/packages/fc/17/89a0eccbf5be9455c395f977022e299c910f1d7d5b0820661908f552ac70/nvidia_cudnn_frontend-1.24.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:37c6ef88c7cbc41eab6e36d5c6715ff0d6c639f76ca12c65ce7a24b453e184eb", size = 3382455, upload-time = "2026-05-20T05:08:48.412Z" }, - { url = "https://files.pythonhosted.org/packages/d7/d2/dac91d6a6fa2e6c07b59239f01fb07472f36d8368c29840647cfcd1c7dd0/nvidia_cudnn_frontend-1.24.0-cp314-cp314t-win_amd64.whl", hash = "sha256:79e902e124123d84d52fa06c0931415cccfa71ff8550c1c156dd3539bffeda5c", size = 2791017, upload-time = "2026-05-20T05:09:07.477Z" }, +version = "1.25.0" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform == 'win32'", + "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform == 'win32'", + "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform == 'emscripten'", + "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform == 'emscripten'", + "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform == 'win32'", + "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and platform_machine != 's390x' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform == 'win32'", + "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.12.*' and platform_machine != 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.11.*' and platform_machine != 's390x' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and platform_machine != 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.11.*' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version < '3.11' and platform_machine != 's390x'", + "python_full_version < '3.11' and platform_machine == 's390x'", +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/b3/58/5a853819023c1ea17b1e71363a1123bd9d9b1a31b41c80adff07a08d32d1/nvidia_cudnn_frontend-1.25.0-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:04c48329eb6918a92e83905981a02dc8f1817dc570720e5531adf053d04a4956", size = 3261918, upload-time = "2026-06-10T21:05:37.3Z" }, + { url = "https://files.pythonhosted.org/packages/63/46/fa9f68a9936d97498bc2c6d1f34e7dcbf2dd3a40bf7dd3f64461b9293f2e/nvidia_cudnn_frontend-1.25.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:92ccf81aad15764b67263901a733505d21772431675a56526e0a17c3f5a3674b", size = 3411836, upload-time = "2026-06-10T21:05:57.444Z" }, + { url = "https://files.pythonhosted.org/packages/02/19/4f36705f2c9a22733fc3dcd31d2e40cd422e614899a76030ada34c83255e/nvidia_cudnn_frontend-1.25.0-cp310-cp310-win_amd64.whl", hash = "sha256:2684f5f33a41f4cf3505f16c3adf59c0e43e8c4c79442c8c508da0e29ea3637c", size = 2796177, upload-time = "2026-06-10T21:06:18.372Z" }, + { url = "https://files.pythonhosted.org/packages/7c/57/5f2a32a40f7beeaec4020b7124ea854ba38ecb89663ba3449b42bb88ad54/nvidia_cudnn_frontend-1.25.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3ae5c281bcb23536c12b7fd2b28e2f599dd1e45e96d37b598175195eb75e8f1a", size = 3262531, upload-time = "2026-06-10T21:06:43.49Z" }, + { url = "https://files.pythonhosted.org/packages/a8/50/224ff36c5d9e02624f8d3c582982bfac74bec481cd331e704fb9a5ecd128/nvidia_cudnn_frontend-1.25.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:676d56062d3ade4ffb34315abe52ea766fa4488db1161b702d9ddd872fab4ddf", size = 3413687, upload-time = "2026-06-10T21:07:04.26Z" }, + { url = "https://files.pythonhosted.org/packages/fa/87/4716b610e0f5b695f76984cb7591944f2d72b10139ca952f3d0cd1cd9ea3/nvidia_cudnn_frontend-1.25.0-cp311-cp311-win_amd64.whl", hash = "sha256:05279eac512e923fc61154f5d463d9917f14d46aa7a507e2610458e1d2367f3b", size = 2797009, upload-time = "2026-06-10T21:07:27.112Z" }, + { url = "https://files.pythonhosted.org/packages/28/0f/df39a194f2529093db737d43cc4cbf594c6a79712a09aa104b999e4d95d4/nvidia_cudnn_frontend-1.25.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:09e6e1bc48ce1235743f89d8ea699c52b3008fd6dae7f2ecadb744bebf272a2b", size = 3263306, upload-time = "2026-06-10T21:07:48.093Z" }, + { url = "https://files.pythonhosted.org/packages/03/65/3b45941d8a22128b971e910f2e9af6bf5ef453e92cc329c56b6eb53c53de/nvidia_cudnn_frontend-1.25.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9a94a72d736bd79eb35f451aaf26d9493778e02ecabccc92c05425508c9e7a83", size = 3414884, upload-time = "2026-06-10T21:08:08.603Z" }, + { url = "https://files.pythonhosted.org/packages/2e/45/69517e8f028573a150e82b71205c920e78ebbe83ff0d073eaeee2ada18dc/nvidia_cudnn_frontend-1.25.0-cp312-cp312-win_amd64.whl", hash = "sha256:d1bfdc795a8bda570ca80ef2287e83f00974857a9a086c1653d2a28099496fee", size = 2798190, upload-time = "2026-06-10T21:08:30.506Z" }, + { url = "https://files.pythonhosted.org/packages/0e/37/ea07ff3578cb3cc847fbed4e7eb84bf02afa81bed437dc0a5e8b7f040c9f/nvidia_cudnn_frontend-1.25.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e8a223ef2e821bbe89fdc6a461cbce36b509e2daaa0a4425eb387060e0cc6ec7", size = 3263340, upload-time = "2026-06-10T21:08:53.528Z" }, + { url = "https://files.pythonhosted.org/packages/ff/ff/072bc1ed5e6d7efaf785498e6610a0b2152603c1e020faa6e6b8761d4505/nvidia_cudnn_frontend-1.25.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:95ab349dfc871a4e7e5b7ceaed649a22b968ce02bca9be1d0a0a577e80fa6832", size = 3414303, upload-time = "2026-06-10T21:09:15.951Z" }, + { url = "https://files.pythonhosted.org/packages/fd/7f/098f9d195f31e2d5ed5da91aa5dea693b582c70a8db0829ba9128a5f0cb2/nvidia_cudnn_frontend-1.25.0-cp313-cp313-win_amd64.whl", hash = "sha256:18cd90ff8429bc65888d7aeb0f9763a237667be652d286be7878ed0dc83216dd", size = 2798191, upload-time = "2026-06-10T21:09:39.559Z" }, + { url = "https://files.pythonhosted.org/packages/94/87/36d8027a6387a9db29f597bdd4070d0a277978331039d563fba104fc53a9/nvidia_cudnn_frontend-1.25.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:88edbf97e49d3b303af60c19e147a87bd0c34a7cdb14462009d488b7aa4bac41", size = 3265214, upload-time = "2026-06-10T21:10:01.283Z" }, + { url = "https://files.pythonhosted.org/packages/3a/97/5b41722547c894511bcb2064065534ab5e14404b673aff6a33098c12e223/nvidia_cudnn_frontend-1.25.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6ddfcff04af7193af05be8e8ef144bc62c99f9f78d7ddbe4e1977c52769c872f", size = 3414357, upload-time = "2026-06-10T21:10:21.76Z" }, + { url = "https://files.pythonhosted.org/packages/f3/24/0abdc3c6247e3691365200e3592dea1631baa1746944ad4d9aed5ed25411/nvidia_cudnn_frontend-1.25.0-cp314-cp314-win_amd64.whl", hash = "sha256:1cbbee1343b1c4037ea08058d16301e45e033916ce8908f583cef5f0905caf9d", size = 2799220, upload-time = "2026-06-10T21:10:42.824Z" }, + { url = "https://files.pythonhosted.org/packages/6e/ca/1a6fcdd672f27de3c0b14f5e19c29d2dabce09bfbe283df2fd04fed13a48/nvidia_cudnn_frontend-1.25.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5169603a011772b09d27f7652fb1c03e8ffe236e871c79cd6a3f63e74160aad7", size = 3269600, upload-time = "2026-06-10T21:11:04.691Z" }, + { url = "https://files.pythonhosted.org/packages/87/76/aa46edfa4a4ec8f90549340ff3738337691ff7c7064eb68c8b0821e62291/nvidia_cudnn_frontend-1.25.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c47b6df3bae7eef03040a62b38e7f8dc2323b16af5437729127488375d7fbba9", size = 3419590, upload-time = "2026-06-10T21:11:27.831Z" }, + { url = "https://files.pythonhosted.org/packages/a7/f0/d38cd0f8bd2412088ef3eb62e473a98d1d3baf7d4f0784f7645500bb4b5a/nvidia_cudnn_frontend-1.25.0-cp314-cp314t-win_amd64.whl", hash = "sha256:85852265fca494909467900a2cb739959da43686ce065d516c7cdb6e9add62e5", size = 2823738, upload-time = "2026-06-10T21:11:48.844Z" }, +] + +[[package]] +name = "nvidia-cudnn-frontend" +version = "1.25.0" +source = { git = "https://github.com/NVIDIA/cudnn-frontend.git?rev=0a14b7181d129d30e7bad34b8c3ed0a0c995e23d#0a14b7181d129d30e7bad34b8c3ed0a0c995e23d" } +resolution-markers = [ + "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform == 'win32'", + "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform == 'win32'", + "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform == 'emscripten'", + "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform == 'emscripten'", + "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and platform_machine != 's390x' and sys_platform == 'win32'", + "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform == 'win32'", + "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.12.*' and platform_machine != 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.11.*' and platform_machine != 's390x' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and platform_machine != 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.11.*' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version < '3.11' and platform_machine != 's390x'", + "python_full_version < '3.11' and platform_machine == 's390x'", +] + +[package.optional-dependencies] +cutedsl = [ + { name = "apache-tvm-ffi" }, + { name = "cuda-python" }, + { name = "nvidia-cutlass-dsl", version = "4.5.0", source = { registry = "https://pypi.org/simple" }, extra = ["cu13"], marker = "extra == 'extra-13-megatron-core-dev'" }, + { name = "torch", marker = "sys_platform == 'never'" }, + { name = "torch-c-dlpack-ext" }, ] [[package]] @@ -4388,7 +4467,7 @@ name = "nvidia-cufft" version = "12.0.0.61" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, + { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32') or (sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/8b/ae/f417a75c0259e85c1d2f83ca4e960289a5f814ed0cea74d18c353d3e989d/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5", size = 214053554, upload-time = "2025-09-04T08:31:38.196Z" }, @@ -4420,9 +4499,9 @@ name = "nvidia-cusolver" version = "12.0.4.66" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, - { name = "nvidia-cusparse", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, - { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, + { name = "nvidia-cublas", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32') or (sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "nvidia-cusparse", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32') or (sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32') or (sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/c8/c3/b30c9e935fc01e3da443ec0116ed1b2a009bb867f5324d3f2d7e533e776b/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2", size = 223467760, upload-time = "2025-09-04T08:33:04.222Z" }, @@ -4435,7 +4514,7 @@ name = "nvidia-cusparse" version = "12.6.3.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, + { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32') or (sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/f8/94/5c26f33738ae35276672f12615a64bd008ed5be6d1ebcb23579285d960a9/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c", size = 162155568, upload-time = "2025-09-04T08:33:42.864Z" }, @@ -4457,22 +4536,122 @@ wheels = [ name = "nvidia-cutlass-dsl" version = "4.4.2" source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform == 'win32'", + "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform == 'win32'", + "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform == 'emscripten'", + "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform == 'emscripten'", + "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform == 'win32'", + "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and platform_machine != 's390x' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform == 'win32'", + "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.12.*' and platform_machine != 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.11.*' and platform_machine != 's390x' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and platform_machine != 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.11.*' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version < '3.11' and platform_machine != 's390x'", + "python_full_version < '3.11' and platform_machine == 's390x'", +] dependencies = [ - { name = "nvidia-cutlass-dsl-libs-base" }, + { name = "nvidia-cutlass-dsl-libs-base", version = "4.4.2", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-13-megatron-core-lts'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/a9/03/678dab0383db1ddfc449da216220f40404189eb36eeed9d87a4fa4bdb0e6/nvidia_cutlass_dsl-4.4.2-py3-none-any.whl", hash = "sha256:7cfb9ef19062b055b9372c7a627004724e2755e4c8b16c3cc88807d64501a4ae", size = 10167, upload-time = "2026-03-16T02:18:59.043Z" }, ] +[[package]] +name = "nvidia-cutlass-dsl" +version = "4.5.0" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform == 'win32'", + "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform == 'win32'", + "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform == 'emscripten'", + "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform == 'emscripten'", + "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and platform_machine != 's390x' and sys_platform == 'win32'", + "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform == 'win32'", + "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.12.*' and platform_machine != 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.11.*' and platform_machine != 's390x' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and platform_machine != 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.11.*' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version < '3.11' and platform_machine != 's390x'", + "python_full_version < '3.11' and platform_machine == 's390x'", +] +dependencies = [ + { name = "nvidia-cutlass-dsl-libs-base", version = "4.5.0", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-13-megatron-core-dev'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/71/a3/46fdf77d373b06bc65a0eda6c921c746985fb3e496c90a09be476291ea80/nvidia_cutlass_dsl-4.5.0-py3-none-any.whl", hash = "sha256:3b051fe02ca69422ab840e64d9865667aba288a3984a7ca4ccd038a82aef1344", size = 10178, upload-time = "2026-05-06T01:17:33.592Z" }, +] + +[package.optional-dependencies] +cu13 = [ + { name = "nvidia-cutlass-dsl-libs-cu13" }, +] + [[package]] name = "nvidia-cutlass-dsl-libs-base" version = "4.4.2" source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform == 'win32'", + "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform == 'win32'", + "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform == 'emscripten'", + "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform == 'emscripten'", + "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform == 'win32'", + "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and platform_machine != 's390x' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform == 'win32'", + "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.12.*' and platform_machine != 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.11.*' and platform_machine != 's390x' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and platform_machine != 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.11.*' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version < '3.11' and platform_machine != 's390x'", + "python_full_version < '3.11' and platform_machine == 's390x'", +] dependencies = [ - { name = "cuda-python" }, - { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and extra == 'extra-13-megatron-core-lts') or extra == 'extra-13-megatron-core-dev'" }, + { name = "cuda-python", marker = "extra == 'extra-13-megatron-core-lts'" }, + { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and extra == 'extra-13-megatron-core-lts') or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and extra == 'extra-13-megatron-core-lts') or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, - { name = "typing-extensions" }, + { name = "typing-extensions", marker = "extra == 'extra-13-megatron-core-lts'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/5f/07/af1b456b5b6dd4a49e71a952a182a99fc863f70b9f78725324f89e0384e5/nvidia_cutlass_dsl_libs_base-4.4.2-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:06acb3acff3dcf4bf6630476efac7de94de30b988ded4fa00b647bbcec4224ff", size = 75471025, upload-time = "2026-03-16T02:23:49.61Z" }, @@ -4487,6 +4666,79 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/4b/ae/0998f328b28b956d7eb399d16f4ee681ca318b306007264444a623e86c64/nvidia_cutlass_dsl_libs_base-4.4.2-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:95db0c8d1d56992e2f5c2dcd5b3baab0297bedc0cbcefc1e70b57acd934e7b23", size = 74356280, upload-time = "2026-03-16T02:25:43.789Z" }, ] +[[package]] +name = "nvidia-cutlass-dsl-libs-base" +version = "4.5.0" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform == 'win32'", + "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform == 'win32'", + "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform == 'emscripten'", + "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform == 'emscripten'", + "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and platform_machine != 's390x' and sys_platform == 'win32'", + "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform == 'win32'", + "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.12.*' and platform_machine != 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.11.*' and platform_machine != 's390x' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and platform_machine != 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform == 'emscripten'", + "python_full_version == '3.11.*' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version < '3.11' and platform_machine != 's390x'", + "python_full_version < '3.11' and platform_machine == 's390x'", +] +dependencies = [ + { name = "cuda-python", marker = "extra == 'extra-13-megatron-core-dev'" }, + { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-13-megatron-core-dev'" }, + { name = "typing-extensions", marker = "extra == 'extra-13-megatron-core-dev'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/04/4d/1d0dc5f36f929885417acfff02af94f61d49e6d34acb480c080d4d887555/nvidia_cutlass_dsl_libs_base-4.5.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:c78b18f2b44ca10a91bc76380ebd65bb7b86aa97a9330bae9b73eb0a1bc51d55", size = 75636580, upload-time = "2026-05-06T01:23:19.17Z" }, + { url = "https://files.pythonhosted.org/packages/fe/81/1229637e8a14e1129475b8260a6ce66058148fa85faf10c94f9f95de5ef6/nvidia_cutlass_dsl_libs_base-4.5.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:5cfdf52bea8feede5e512a094484956693cb87adaafa310991d2876653b1a88e", size = 74505679, upload-time = "2026-05-06T01:26:23.41Z" }, + { url = "https://files.pythonhosted.org/packages/b4/39/155dcbcf942b2c170aa0d1115ef5f2d358d9916ddc7200ab6e70541b97a0/nvidia_cutlass_dsl_libs_base-4.5.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:f8635ad1e0a670323cc729f167067fa880cb56577ec2e79afb80a35ab371912e", size = 75635889, upload-time = "2026-05-06T01:25:20.572Z" }, + { url = "https://files.pythonhosted.org/packages/a4/36/2c2b3fc81a45a1bbbdcfd10c6d9793fd28848e6fefa6d4ed7c7c477f7d2a/nvidia_cutlass_dsl_libs_base-4.5.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:7bb6de91b00a2b392cd834fec174a1461bf0f10a9b6d28086c8f4885aed27218", size = 74505494, upload-time = "2026-05-06T01:27:29.616Z" }, + { url = "https://files.pythonhosted.org/packages/f8/d0/924048cfa43e1cb546735cb332b05a4fb92c63c1a1ac566f06445f9eca58/nvidia_cutlass_dsl_libs_base-4.5.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:3f7c133d31fa82ae7db697fd6943a5f9a2c97c8a40ee1056c67ef29fe00974d8", size = 75630723, upload-time = "2026-05-06T01:24:49.842Z" }, + { url = "https://files.pythonhosted.org/packages/c3/8b/2c187400d85f7d2acb328f20499b7b05745dca8485cf6ad247d5f2b434cf/nvidia_cutlass_dsl_libs_base-4.5.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:bd18322d9247f8c033a10ed4e519c4985ca6b4fb578ade382e5a264422ebd915", size = 74505487, upload-time = "2026-05-06T01:26:52.755Z" }, + { url = "https://files.pythonhosted.org/packages/ff/2c/21d5fc62e030a43c0f1a3dab6749fb632026a27d6a60f59975cd29a5d165/nvidia_cutlass_dsl_libs_base-4.5.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:90a4d802a03963fa36eb287fbc9b40a1374590fc7e8cc1b9673dee8872f75713", size = 75632646, upload-time = "2026-05-06T01:24:19.623Z" }, + { url = "https://files.pythonhosted.org/packages/1c/79/0dca3b465711ffb4c44b4252940cc5f51d2d4905e405707e5c6c2a83d3d6/nvidia_cutlass_dsl_libs_base-4.5.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:8e58b016da5bb09bd1d809d0c025433edb36b279adfbcd107e96361b214bd8bc", size = 74505936, upload-time = "2026-05-06T01:25:52.728Z" }, + { url = "https://files.pythonhosted.org/packages/59/85/2799e4de2fe7070cc4126ac501443d1cd7796b07ed880118e31956ae266a/nvidia_cutlass_dsl_libs_base-4.5.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:2e121b20f0a48122c9b48227d00a7d681189e1de2fd4d211f9661a4e1658f066", size = 75634241, upload-time = "2026-05-06T01:22:48.964Z" }, + { url = "https://files.pythonhosted.org/packages/04/c6/5aaa2dff6dfc615a83687df4462a91dad2ac1af85d6a9c91d9a6b9760a02/nvidia_cutlass_dsl_libs_base-4.5.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:0a60dfce3349984315306ef719ed1edf0e225527158f26019a5cf266e06cc45d", size = 74504851, upload-time = "2026-05-06T01:23:48.613Z" }, +] + +[[package]] +name = "nvidia-cutlass-dsl-libs-cu13" +version = "4.5.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuda-python" }, + { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" } }, + { name = "nvidia-cutlass-dsl-libs-base", version = "4.5.0", source = { registry = "https://pypi.org/simple" } }, + { name = "typing-extensions" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/50/32/6566f883444d4342f9b196ab4243c0ffa985d62bfdbc08f2d6f54b193a42/nvidia_cutlass_dsl_libs_cu13-4.5.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:b95cd78ca5f3440cd488ea03c30d394ca07e491e5bd6bb274f8f26d4f4713439", size = 79119607, upload-time = "2026-05-06T01:20:01.77Z" }, + { url = "https://files.pythonhosted.org/packages/04/6c/c7c32165d00dfa01d81a49626552192c3f3540d6cd5b8faf6507e161544a/nvidia_cutlass_dsl_libs_cu13-4.5.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:145cf5563b2d0d347c722d84a3b74541da5a1ec5b46821b54acc202693ce56ae", size = 78787095, upload-time = "2026-05-06T01:21:03.068Z" }, + { url = "https://files.pythonhosted.org/packages/a9/e0/d184ab4b764cdab3d2c09aa7726178c6cc294dc908bda80f0bec13c9bebc/nvidia_cutlass_dsl_libs_cu13-4.5.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:dadddab5bfe11973f8a0dfa66645a42ccb587287abb0766992e0e69da621dc82", size = 79119785, upload-time = "2026-05-06T01:18:00.059Z" }, + { url = "https://files.pythonhosted.org/packages/5c/0f/5ad5f2012fd00d683b2131d270300f45e386e55064a69047ee2989f6d480/nvidia_cutlass_dsl_libs_cu13-4.5.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:35726dddb6f6e63c767caf7533d583b4462e2c6b90df7fa67738d58b2eede742", size = 78786656, upload-time = "2026-05-06T01:20:33.039Z" }, + { url = "https://files.pythonhosted.org/packages/a1/bd/36767220a6ecd4284496708976017ef85053563971a13b95aa8b84d316d7/nvidia_cutlass_dsl_libs_cu13-4.5.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:a799eb9d65ba03095444907b8bf617b5a8bd7d03d1b95cec9637c558af6d0b60", size = 79118756, upload-time = "2026-05-06T01:18:58.118Z" }, + { url = "https://files.pythonhosted.org/packages/cd/f2/2c976759dff8836e41a8ff3716db4fc72b01969ea6ee062ae22877008030/nvidia_cutlass_dsl_libs_cu13-4.5.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:fc0b5a81ff591db72489134ca206ae886f0cce43f20863010a7f30fcfe484a7b", size = 78787917, upload-time = "2026-05-06T06:58:30.829Z" }, + { url = "https://files.pythonhosted.org/packages/e2/41/f8817a6ef5c93ae1009f7e6abe4e205f0014ae51faf8d0aa5817d7bf64f5/nvidia_cutlass_dsl_libs_cu13-4.5.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:f2b821671add2e69a1377e7fd87e6261995db281b2f8e516ddae2fd6b7a6c1d0", size = 79119918, upload-time = "2026-05-06T01:19:31.84Z" }, + { url = "https://files.pythonhosted.org/packages/08/48/13386e28bf2b724268d7ac95c41d0c718e91118bf89218b6b0471e5fa595/nvidia_cutlass_dsl_libs_cu13-4.5.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:554b775069d093f308949a65880bf9c9bfd48b1f4b2fd0e0d97aa2f608e6eea1", size = 78788617, upload-time = "2026-05-06T01:21:35.563Z" }, + { url = "https://files.pythonhosted.org/packages/4b/68/f303a5732763b72cee6f485c0fa1d4972b159770eb869f789a873f98f330/nvidia_cutlass_dsl_libs_cu13-4.5.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:652e34d8a78accceab321da0232157e493b113e417306269d700a297a7d62447", size = 79118174, upload-time = "2026-05-06T01:18:28.675Z" }, + { url = "https://files.pythonhosted.org/packages/d0/6e/6371065485ec91a75176b6b850a7bee31fd57bcb9045a3e07d7a8fd06c85/nvidia_cutlass_dsl_libs_cu13-4.5.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:3d460d03d6ea0a463e262ee07964ccafffde8eeefe058c9615d77da1ddd6d003", size = 78784876, upload-time = "2026-05-06T01:22:14.041Z" }, +] + [[package]] name = "nvidia-mathdx" version = "25.6.0" @@ -4730,11 +4982,11 @@ resolution-markers = [ "python_full_version < '3.11' and platform_machine == 's390x'", ] dependencies = [ - { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" } }, + { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (python_full_version >= '3.13' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and extra == 'extra-13-megatron-core-dev') or (python_full_version < '3.11' and extra == 'extra-13-megatron-core-lts') or (python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and extra == 'extra-13-megatron-core-lts') or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, - { name = "protobuf" }, - { name = "typing-extensions" }, + { name = "protobuf", marker = "(python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (python_full_version >= '3.13' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "typing-extensions", marker = "(python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (python_full_version >= '3.13' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/c5/93/942d2a0f6a70538eea042ce0445c8aefd46559ad153469986f29a743c01c/onnx-1.21.0.tar.gz", hash = "sha256:4d8b67d0aaec5864c87633188b91cc520877477ec0254eda122bef8be43cd764", size = 12074608, upload-time = "2026-03-27T21:33:36.118Z" } wheels = [ @@ -4829,12 +5081,12 @@ resolution-markers = [ "python_full_version < '3.11' and platform_machine == 's390x'", ] dependencies = [ - { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" } }, + { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (python_full_version >= '3.13' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and extra == 'extra-13-megatron-core-dev') or (python_full_version < '3.11' and extra == 'extra-13-megatron-core-lts') or (python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and extra == 'extra-13-megatron-core-lts') or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, - { name = "onnx", version = "1.21.0", source = { registry = "https://pypi.org/simple" } }, - { name = "sympy" }, - { name = "typing-extensions" }, + { name = "onnx", version = "1.21.0", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (python_full_version >= '3.13' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "sympy", marker = "(python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (python_full_version >= '3.13' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "typing-extensions", marker = "(python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (python_full_version >= '3.13' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/b2/a5/acc43c8fa6edbc584d127fb6bbd13ae9ebfc01b9675c74e0da2de15fa4a6/onnx_ir-0.2.0.tar.gz", hash = "sha256:8bad3906691987290789b26d05e0dbff467029a0b1e411e12e4cae02e43503e4", size = 141693, upload-time = "2026-02-24T02:31:10.998Z" } wheels = [ @@ -4905,13 +5157,13 @@ resolution-markers = [ "python_full_version < '3.11' and platform_machine == 's390x'", ] dependencies = [ - { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" } }, + { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (python_full_version >= '3.13' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and extra == 'extra-13-megatron-core-dev') or (python_full_version < '3.11' and extra == 'extra-13-megatron-core-lts') or (python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and extra == 'extra-13-megatron-core-lts') or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, - { name = "onnx", version = "1.21.0", source = { registry = "https://pypi.org/simple" } }, - { name = "onnx-ir", version = "0.2.0", source = { registry = "https://pypi.org/simple" } }, - { name = "packaging" }, - { name = "typing-extensions" }, + { name = "onnx", version = "1.21.0", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (python_full_version >= '3.13' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "onnx-ir", version = "0.2.0", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (python_full_version >= '3.13' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "packaging", marker = "(python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (python_full_version >= '3.13' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "typing-extensions", marker = "(python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (python_full_version >= '3.13' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/e7/2b/538fdeb0e25bed5d7e0f954af5710543e2629499fb74381afc3333f8a8ae/onnxscript-0.6.2.tar.gz", hash = "sha256:abb2e6f464db40c9b8c7fbb3e64cca04cf3f4495e67c4eda5eac17b784191ce3", size = 590865, upload-time = "2026-02-10T22:53:39.638Z" } wheels = [ @@ -5383,10 +5635,10 @@ resolution-markers = [ "python_full_version < '3.11' and extra != 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", ] dependencies = [ - { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "python-dateutil", marker = "python_full_version < '3.11'" }, - { name = "pytz", marker = "python_full_version < '3.11'" }, - { name = "tzdata", marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or extra == 'extra-13-megatron-core-dev'" }, + { name = "python-dateutil", marker = "python_full_version < '3.11' or extra == 'extra-13-megatron-core-dev'" }, + { name = "pytz", marker = "python_full_version < '3.11' or extra == 'extra-13-megatron-core-dev'" }, + { name = "tzdata", marker = "python_full_version < '3.11' or extra == 'extra-13-megatron-core-dev'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/33/01/d40b85317f86cf08d853a4f495195c73815fdf205eef3993821720274518/pandas-2.3.3.tar.gz", hash = "sha256:e05e1af93b977f7eafa636d043f9f94c7ee3ac81af99c13508215942e64c993b", size = 4495223, upload-time = "2025-09-29T23:34:51.853Z" } wheels = [ @@ -6246,7 +6498,7 @@ name = "pyjwt" version = "2.12.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, + { name = "typing-extensions", marker = "python_full_version < '3.11' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/c2/27/a3b6e5bf6ff856d2509292e95c8f57f0df7017cf5394921fc4e4ef40308a/pyjwt-2.12.1.tar.gz", hash = "sha256:c74a7a2adf861c04d002db713dd85f84beb242228e671280bf709d765b03672b", size = 102564, upload-time = "2026-03-13T19:27:37.25Z" } wheels = [ @@ -6648,7 +6900,7 @@ source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "attrs" }, { name = "rpds-py" }, - { name = "typing-extensions", marker = "python_full_version < '3.13'" }, + { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/22/f5/df4e9027acead3ecc63e50fe1e36aca1523e1719559c499951bb4b53188f/referencing-0.37.0.tar.gz", hash = "sha256:44aefc3142c5b842538163acb373e24cce6632bd54bdb01b21ad5863489f50d8", size = 78036, upload-time = "2025-10-13T15:30:48.871Z" } wheels = [ @@ -7864,7 +8116,7 @@ version = "0.52.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, - { name = "typing-extensions", marker = "python_full_version < '3.13'" }, + { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/c4/68/79977123bb7be889ad680d79a40f339082c1978b5cfcf62c2d8d196873ac/starlette-0.52.1.tar.gz", hash = "sha256:834edd1b0a23167694292e94f597773bc3f89f362be6effee198165a35d62933", size = 2653702, upload-time = "2026-01-18T13:34:11.062Z" } wheels = [ @@ -7876,7 +8128,7 @@ name = "sympy" version = "1.14.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "mpmath" }, + { name = "mpmath", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (python_full_version < '3.13' and sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev') or (python_full_version >= '3.13' and sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-lts') or (python_full_version < '3.13' and sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev') or (python_full_version >= '3.13' and sys_platform == 'win32' and extra == 'extra-13-megatron-core-lts') or (sys_platform != 'emscripten' and sys_platform != 'win32') or (sys_platform == 'emscripten' and extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'win32' and extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/83/d3/803453b36afefb7c2bb238361cd4ae6125a569b4db67cd9e79846ba2d68c/sympy-1.14.0.tar.gz", hash = "sha256:d3d3fe8df1e5a0b42f0e7bdf50541697dbe7d23746e894990c030e2b05e72517", size = 7793921, upload-time = "2025-04-27T18:05:01.611Z" } wheels = [ @@ -7968,7 +8220,7 @@ resolution-markers = [ dependencies = [ { name = "ml-dtypes", version = "0.4.1", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.13' and extra == 'extra-13-megatron-core-dev') or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and extra == 'extra-13-megatron-core-dev') or (python_full_version < '3.11' and extra == 'extra-13-megatron-core-lts') or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, - { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and extra == 'extra-13-megatron-core-dev') or (python_full_version < '3.11' and extra == 'extra-13-megatron-core-lts') or (python_full_version >= '3.13' and extra == 'extra-13-megatron-core-dev') or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/3c/b9/ea25aba62c688a87d7d7d9cc5926d602e2f9e84fa72586825486fb180b7e/tensorstore-0.1.74.tar.gz", hash = "sha256:a062875f27283d30ce4959c408c253ecb336fce8e3f9837c064e3d30cda79203", size = 6795605, upload-time = "2025-04-24T15:42:18.829Z" } wheels = [ @@ -8025,7 +8277,7 @@ resolution-markers = [ "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] dependencies = [ - { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (python_full_version >= '3.13' and extra == 'extra-13-megatron-core-lts') or (python_full_version < '3.11' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (python_full_version >= '3.11' and extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (python_full_version < '3.11' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (python_full_version >= '3.13' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and extra == 'extra-13-megatron-core-lts') or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] @@ -8263,21 +8515,21 @@ name = "torch" version = "2.11.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cuda-bindings", marker = "sys_platform == 'linux'" }, + { name = "cuda-bindings", marker = "sys_platform == 'linux' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "cuda-toolkit", extra = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "sys_platform == 'linux' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, - { name = "filelock", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, - { name = "fsspec", version = "2026.2.0", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, - { name = "jinja2", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, + { name = "filelock", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32') or (sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "fsspec", version = "2026.2.0", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32') or (sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "jinja2", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32') or (sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "networkx", version = "3.4.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "networkx", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and sys_platform != 'emscripten' and sys_platform != 'win32') or (python_full_version < '3.11' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, - { name = "nvidia-cudnn-cu13", marker = "sys_platform == 'linux'" }, - { name = "nvidia-cusparselt-cu13", marker = "sys_platform == 'linux'" }, - { name = "nvidia-nccl-cu13", marker = "sys_platform == 'linux'" }, - { name = "nvidia-nvshmem-cu13", marker = "sys_platform == 'linux'" }, - { name = "setuptools", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, - { name = "sympy", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, - { name = "triton", marker = "sys_platform == 'never'" }, - { name = "typing-extensions", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, + { name = "nvidia-cudnn-cu13", marker = "sys_platform == 'linux' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "nvidia-cusparselt-cu13", marker = "sys_platform == 'linux' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "nvidia-nccl-cu13", marker = "sys_platform == 'linux' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "nvidia-nvshmem-cu13", marker = "sys_platform == 'linux' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "setuptools", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32') or (sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "sympy", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32') or (sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "triton", marker = "sys_platform == 'never' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "typing-extensions", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32') or (sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/ac/f2/c1690994afe461aae2d0cac62251e6802a703dec0a6c549c02ecd0de92a9/torch-2.11.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:2c0d7fcfbc0c4e8bb5ebc3907cbc0c6a0da1b8f82b1fc6e14e914fa0b9baf74e", size = 80526521, upload-time = "2026-03-23T18:12:06.86Z" }, @@ -8310,6 +8562,37 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cf/bf/c8d12a2c86dbfd7f40fb2f56fbf5a505ccf2d9ce131eb559dfc7c51e1a04/torch-2.11.0-cp314-cp314t-win_amd64.whl", hash = "sha256:b2a43985ff5ef6ddd923bbcf99943e5f58059805787c5c9a2622bf05ca2965b0", size = 114792991, upload-time = "2026-03-23T18:08:19.216Z" }, ] +[[package]] +name = "torch-c-dlpack-ext" +version = "0.1.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "torch", marker = "sys_platform == 'never'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/37/de/921b6491efce5c389a5ef9bbed3d2d6660005840dae488124173180859ab/torch_c_dlpack_ext-0.1.5.tar.gz", hash = "sha256:d06f0357d575d22a168cc77acb9020fc4bae30968ceb6718a055dcbe92bacabe", size = 12913, upload-time = "2026-01-12T11:25:08.484Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b5/49/67a66932ab2fcdda3c5a4dcf606e713d86883a4a9a99a3bb832815b52b8e/torch_c_dlpack_ext-0.1.5-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:e0f6c197d5293884898b9ebf13d07501de39cb94799b374ed43f91731087d557", size = 7056755, upload-time = "2026-01-12T11:24:31.817Z" }, + { url = "https://files.pythonhosted.org/packages/ae/28/d2d6bf90e01a1f4da3277c9a56d9ecac648b6d6adaa8e20c17f802deb7fb/torch_c_dlpack_ext-0.1.5-cp310-cp310-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ba3d88f0f7d5e1d9c3d4a3179037fc8e261c3b77ac1fad23edc0d3a9214ef193", size = 432066, upload-time = "2026-01-12T11:24:33.619Z" }, + { url = "https://files.pythonhosted.org/packages/a1/e9/a1f9584a3af4ac6ae5ad5cf86927d8c3a9b6bb50d54e54d19313411216a0/torch_c_dlpack_ext-0.1.5-cp310-cp310-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c7468df84ec152d930fbc3acf460c44a60b3462b95af3d3a676d133629c7e176", size = 879488, upload-time = "2026-01-12T11:24:34.837Z" }, + { url = "https://files.pythonhosted.org/packages/6c/08/478cfcb5814e29f9b720111bdef315fc2fbc8b276e4b1183c8b9c9414a4f/torch_c_dlpack_ext-0.1.5-cp310-cp310-win_amd64.whl", hash = "sha256:78dd4904bd26170a2dd7c0eab56367756ee0a15672ce9b84146169e68f0c6ddc", size = 1461437, upload-time = "2026-01-12T11:24:36.385Z" }, + { url = "https://files.pythonhosted.org/packages/65/66/c12a9bb3a5ddc0962c00467891bf1ffdda39a4d4780bf0fbbf54523ff34e/torch_c_dlpack_ext-0.1.5-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:56bd25a2af19280bf8a06aa62cff5510106f43235b9327d8561b3e9a659c4d84", size = 5076782, upload-time = "2026-01-12T11:24:37.868Z" }, + { url = "https://files.pythonhosted.org/packages/20/e1/64e1e579d107064785549e70758e38a42376ab7e73d86897ed4beab10e74/torch_c_dlpack_ext-0.1.5-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:fba674110e1fab0b176bb5a28223e157db65c90767d4ba74abdbee9f537b0e9d", size = 440949, upload-time = "2026-01-12T11:24:39.716Z" }, + { url = "https://files.pythonhosted.org/packages/64/5c/3e1382a620824f92920ab3fae132d8fb4e85898284c99e0c6a7764e452ce/torch_c_dlpack_ext-0.1.5-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3448c4f0d64104d0b2e58080a7efa72304a04960c18f338024b80b13cd3eca26", size = 897768, upload-time = "2026-01-12T11:24:41.209Z" }, + { url = "https://files.pythonhosted.org/packages/54/4f/76ea1006b9038b496d01e916c91efd17cb782abde2491a261cf203f57e30/torch_c_dlpack_ext-0.1.5-cp311-cp311-win_amd64.whl", hash = "sha256:74676474e0afa9a4216c4755ea7cf05e8158be1d168f6bda669ba91097c263f2", size = 1479088, upload-time = "2026-01-12T11:24:42.436Z" }, + { url = "https://files.pythonhosted.org/packages/b1/67/10d236698525d7b7db4d74ec0a4b01f5b2db33968995fdd9ac6b4635e327/torch_c_dlpack_ext-0.1.5-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:c0f2bd51fcd99c0e5b50314e1985f2728c4941bfa821f065e6c30951d1f995ca", size = 5291237, upload-time = "2026-01-12T11:24:44.011Z" }, + { url = "https://files.pythonhosted.org/packages/87/06/8d760997307a5c3be4384424667bf31aae0a42060838c532c7d846516175/torch_c_dlpack_ext-0.1.5-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3562ee411258676f9c38b8ad39306d1c8d027b6a86f6a87c920d2d009a9d1510", size = 443069, upload-time = "2026-01-12T11:24:45.451Z" }, + { url = "https://files.pythonhosted.org/packages/e2/79/a914539b4785f3e44f891aa012a886edb8bc10fe081c440981c57543ce21/torch_c_dlpack_ext-0.1.5-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e6f9da4bb9af70e27facc777458be62e10dbbbddda7672d16138db0553c5a524", size = 897846, upload-time = "2026-01-12T11:24:48.168Z" }, + { url = "https://files.pythonhosted.org/packages/3a/e6/7d7a97a3953208d6d6ce749180c34d1dab48464ded9a76cecabe9d021ce6/torch_c_dlpack_ext-0.1.5-cp312-cp312-win_amd64.whl", hash = "sha256:670fbbab70123cc228bed41693a3720757af57a0ad22669063c9db25321e8f55", size = 1482855, upload-time = "2026-01-12T11:24:49.581Z" }, + { url = "https://files.pythonhosted.org/packages/ca/c6/65346a201d921b616731311fc9941f15137672b444cebdad702cb52ccee0/torch_c_dlpack_ext-0.1.5-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:74acea2ed395cadda63342845b9e9ee7cd4537846223dacfb4431b4610109265", size = 1993243, upload-time = "2026-01-12T11:24:51.079Z" }, + { url = "https://files.pythonhosted.org/packages/fd/ec/faf10be09a5812b1c5ec9922b53fb5def5fc4080b81a653b9347bb169ebb/torch_c_dlpack_ext-0.1.5-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:49f1e99d13c64e22dac0a34a1560e9e5a398a49a9fa81df83053e04fde6ec5bd", size = 443798, upload-time = "2026-01-12T11:24:52.754Z" }, + { url = "https://files.pythonhosted.org/packages/2d/68/f434b48700f3e04f33882f54d8d3910327b935f55e14ec49da7d607bf470/torch_c_dlpack_ext-0.1.5-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:debe62e5ef93e631065d6b9f6e60d3d39bae6b89fa1b25d9523f40b3efbf8aba", size = 755004, upload-time = "2026-01-12T11:24:54.004Z" }, + { url = "https://files.pythonhosted.org/packages/03/a8/cc64e563f05ea99bd79bdb43f71f0f46452d3acd734da4843ede5fc73a35/torch_c_dlpack_ext-0.1.5-cp313-cp313-win_amd64.whl", hash = "sha256:30e3eab616dbc81dfdb7492aca557be551a9163ba9b585f97394a42b336b113a", size = 999126, upload-time = "2026-01-12T11:24:55.44Z" }, + { url = "https://files.pythonhosted.org/packages/96/5e/449324ca8e81573e650b6851fc31c1038f750d1de85d0b185d788e1c7a3a/torch_c_dlpack_ext-0.1.5-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:cac94a4905d391889e679a8da31e46dc325af5d55d13b7c70c0ce3d71d1ced6d", size = 1982154, upload-time = "2026-01-12T11:24:58.038Z" }, + { url = "https://files.pythonhosted.org/packages/20/62/11c05b99f69aa5152bca0313e0dfa6d125a020cf890dc888ef009aa7891c/torch_c_dlpack_ext-0.1.5-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3a58fdf45fb0bda7bc459632cec891570f31c11636d5851c825cf308ec8b73c2", size = 163825, upload-time = "2026-01-12T11:24:59.474Z" }, + { url = "https://files.pythonhosted.org/packages/15/b5/be613cd8e71c9982bd07af530f86c5a7f30df7831d14cec5414857af7149/torch_c_dlpack_ext-0.1.5-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7b985a324c68241cf83a9474b28015524b66775b12a91930dd4c0760aa628d01", size = 171740, upload-time = "2026-01-12T11:25:00.776Z" }, + { url = "https://files.pythonhosted.org/packages/5c/11/52e291f1659e2ec70a09f5ca4ad27e015eb4f0a1371ae68d23a9fbd1c704/torch_c_dlpack_ext-0.1.5-cp314-cp314-win_amd64.whl", hash = "sha256:d794e19fa3f330ab7a29987c07e031fc08e4953aec516d35701d0827863e356b", size = 277086, upload-time = "2026-01-12T11:25:01.901Z" }, +] + [[package]] name = "torchx" version = "0.7.0" @@ -8336,7 +8619,7 @@ name = "tqdm" version = "4.67.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "colorama", marker = "sys_platform == 'win32' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/09/a9/6ba95a270c6f1fbcd8dac228323f2777d886cb206987444e4bce66338dd4/tqdm-4.67.3.tar.gz", hash = "sha256:7d825f03f89244ef73f1d4ce193cb1774a8179fd96f31d7e1dcde62092b960bb", size = 169598, upload-time = "2026-02-03T17:35:53.048Z" } wheels = [ @@ -8515,7 +8798,7 @@ source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "click" }, { name = "h11" }, - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, + { name = "typing-extensions", marker = "python_full_version < '3.11' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/62/f2/368268300fb8af33743508d738ef7bb4d56afdb46c6d9c0fa3dd515df171/uvicorn-0.43.0.tar.gz", hash = "sha256:ab1652d2fb23abf124f36ccc399828558880def222c3cb3d98d24021520dc6e8", size = 85686, upload-time = "2026-04-03T18:37:48.984Z" } wheels = [