diff --git a/docker/Dockerfile.ci.dev b/docker/Dockerfile.ci.dev index 7f3a5c0552a..4ea7409b17f 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 -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 diff --git a/megatron/core/datasets/data_schedule.py b/megatron/core/datasets/data_schedule.py index d5bf07053cc..af87664b07b 100644 --- a/megatron/core/datasets/data_schedule.py +++ b/megatron/core/datasets/data_schedule.py @@ -573,8 +573,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) @@ -630,6 +630,34 @@ def get_batch_on_this_rank_for_sequence_packing( group_size=local_cp_size_val ) + 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 + ) + 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 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: + # 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() + 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]. if is_tp_rank_0: @@ -644,7 +672,18 @@ 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) + partition_total_tokens = ( + 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, + cp_partition_mode=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. @@ -658,12 +697,17 @@ 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 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 + # 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) @@ -778,21 +822,13 @@ 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, ) # 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 190f898cfc3..87b46ac68bd 100644 --- a/megatron/core/datasets/data_schedule_utils.py +++ b/megatron/core/datasets/data_schedule_utils.py @@ -2,7 +2,7 @@ from functools import lru_cache from math import ceil, log2 -from typing import Dict, List, Optional, Sequence, Tuple +from typing import Dict, List, Literal, Optional, Sequence, Tuple import torch @@ -12,16 +12,26 @@ _DYNAMIC_CP_WORKLOAD_CAP_DELTA = 0.05 -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, + 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. + ``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. + 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. """ cp_size = cp_group.size() if cp_size <= 1: @@ -33,10 +43,44 @@ 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 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 " + 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 + + 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: batch[key] = batch[key].index_select(0, index) 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/fusions/fused_mla_yarn_rope_apply.py b/megatron/core/fusions/fused_mla_yarn_rope_apply.py index df81671397b..d4cdaa2bb7a 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,8 @@ def backward(ctx, grad): total_seqlen = grad.shape[0] else: seq_num = len(ctx.cu_seqlens_q) - 1 + if ctx.has_position_ids: + grad = grad.contiguous() total_seqlen, nheads, headdim = grad.shape assert grad.stride(-1) == 1 @@ -367,6 +385,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 +397,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 +412,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 +434,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 +452,7 @@ def fused_mla_rope_inplace( rotary_interleaved, inverse, remove_interleaving, + position_ids, ) diff --git a/megatron/core/packed_seq_params.py b/megatron/core/packed_seq_params.py index 3095b1b8464..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. @@ -198,6 +199,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 +207,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 +218,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 +236,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 +384,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 +409,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, @@ -474,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 2e57ec0b9f8..0b8bc31e0be 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -12,7 +12,10 @@ 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, DSAIndexerLossLoggingHelper, @@ -22,6 +25,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 +621,163 @@ 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, + 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 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 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.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) + 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: + 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() + + weights_scaled = weights.float() * float(indexer_softmax_scale) + 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) + 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.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) + attn_probs = exp_scores / (exp_scores.sum(dim=-1, keepdim=True) + exp_sink) + 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) + 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() + + indexer_loss = raw_local_loss * float(loss_coeff) / float(loss_divisor) + return output, indexer_loss + + # --------------------------------------------------------------------------- # Compressor # --------------------------------------------------------------------------- @@ -812,8 +973,9 @@ 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]: + ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: """THD per-segment compression — fully vectorized. Linear projections are token-wise on the flat input. The gated @@ -831,6 +993,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. @@ -841,25 +1005,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: + cu_seqlens_compressed = None + total_comp = compressed_group_ids.shape[0] + else: + # 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 total_comp == 0: return None, cu_seqlens_compressed @@ -868,29 +1038,38 @@ 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) + + kv_grouped = kv[gather_idx] # (total_comp, ratio, 1, coff * d) + 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) 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 +1082,49 @@ 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 pre_grouped: + 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 + 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 self.rotate: compressed_thd = rotate_activation(compressed_thd) @@ -1572,6 +1780,8 @@ def forward( 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. @@ -1589,7 +1799,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 +2241,286 @@ 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: Optional[torch.Tensor], + boundary_kv: Optional[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 = ( + 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 + ) + max_seqlen_q = int(packed_seq_params.max_seqlen_q) + + # ---- Step 2: local CP rows, local KV, and boundary tensors ------------ + 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( + "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] + # 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 + + # ``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 + 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 ---------------- + + # One compressed row per full ratio-sized group; tails are dropped. + compressed_lens = torch.div( + cu_seqlens[1:] - cu_seqlens[:-1], ratio, rounding_mode="floor" + ) + cu_seqlens_compressed = torch.cat( + ( + torch.zeros_like(cu_seqlens[:1]), + torch.cumsum(compressed_lens, dim=0, dtype=torch.int32), + ) + ) + # ``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, + boundary_hidden, + cu_seqlens, + cu_seqlens_compressed, + global_start, + cp_size, + ratio, + ) + ) + + if indexer is not None: + # ---- Step 4: optional indexer compressed path ----------------- + 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 + ) + 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, + indexer.index_head_dim - indexer.qk_pos_emb_head_dim, + indexer.qk_pos_emb_head_dim, + 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) + + indexer_compressed_local, _ = indexer.compressor._forward_thd( + hidden_compact.detach(), + cu_seqlens, + max_seqlen_q=max_seqlen_q, + compressed_group_ids=compressed_group_ids, + ) + # Gather indexer compressed K in rank-major order. + k_indexer_rank_major = gather_from_sequence_parallel_region( + indexer_compressed_local.squeeze(1), group=cp_group + ) + + # 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, seq_to_rank_row.clamp_min(0) + ) + # Each top-k entry is still a logical compressed id within that + # query's sequence here. + compressed_topk, indexer_layout = cp_utils.compute_cp_indexer_topk( + q_indexer_cp, + weights_indexer_cp, + k_indexer_seq_major, + cu_seqlens, + cu_seqlens_compressed, + global_start, + ratio, + 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 ------------------------- + compressed_kv_local, _ = self.compressor._forward_thd( + hidden_compact, + cu_seqlens, + max_seqlen_q=max_seqlen_q, + 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( + compressed_kv_local.squeeze(1), group=cp_group + ) + + # ---- 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) + 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 + 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, + global_start, + l_local, + d_window, + self.window_size, + ratio, + 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 ---------------------------------- + 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 + 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 = None + if cu_seqlens_q_unpadded is not None: + 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 - cu_seqlens[batch_ids] + q_padding_mask = positions >= real_seqlens[batch_ids] + 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, + q_indexer_cp, + k_indexer_for_loss, + weights_indexer_cp, + indexer_topk_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: + 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 ---------------------------------- + if self.apply_dsa_kernel_fusion: + output = dsa_sparse_attn( + query, + kv_full_thd, + self.attn_sink.float(), + 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, 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..ae3e4c7daf3 --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/csa_cp_layout_kernels.py @@ -0,0 +1,840 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +"""CuTeDSL kernels for DSv4 THD context parallel layout work. + +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 +from typing import Optional, Tuple + +import torch + +try: + 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 +except ImportError: + cuda = None + cutlass = None + cute = 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: + + 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 + + # 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). + # 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 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, + comp_ids: 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, + ): + 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) + 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_end = seq_end + if local_seq_end > range_end: + local_seq_end = range_end + + 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 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 + 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: + 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 + + # Launch contract for _compressor_input_compact_fwd_kernel. + # 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, + boundary_hidden: cute.Tensor, + hidden_compact: cute.Tensor, + cu_seqlens: cute.Tensor, + comp_ids: 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, + stream: cuda.CUstream, + ): + _launch_named( + _compressor_input_compact_fwd_kernel, + "dsv4_cp_compressor_input_compact_fwd", + ( + hidden_local, + boundary_hidden, + hidden_compact, + cu_seqlens, + comp_ids, + n_seq, + global_start, + l_local, + ratio, + d_comp, + d_window, + compact_len, + row_width, + ), + grid=(compact_len, 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, + ): + 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) + 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 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 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) + + 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, + ), + grid=(total_work, 1, 1), + block=(64, 1, 1), + stream=stream, + ) + + @cute.kernel + def _build_attention_indices_kernel( + cu_seqlens: cute.Tensor, + cu_seqlens_compressed: cute.Tensor, + indexer_topk: 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, + 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, + 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. + 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) + 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: + 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 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 + topk_length[row] = 0 + + if seq_start_found >= 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 < global_start: + topk_idxs[row, write_col] = pos - (global_start - d_window) + else: + topk_idxs[row, write_col] = d_window + pos - global_start + write_col = write_col + 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 < 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 + topk_length[row] = write_col + elif total_width > 0: + topk_idxs[row, 0] = 0 + topk_length[row] = 1 + + @cute.jit + def _build_attention_indices_launch( + cu_seqlens: cute.Tensor, + cu_seqlens_compressed: cute.Tensor, + indexer_topk: 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, + 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, + index_mode: cutlass.Constexpr, + launch_work: cutlass.Int32, + stream: cuda.CUstream, + ): + _launch_named( + _build_attention_indices_kernel, + "dsv4_cp_build_attention_indices", + ( + cu_seqlens, + cu_seqlens_compressed, + indexer_topk, + seq_to_rank_row, + topk_idxs, + topk_length, + indexer_rank_major, + n_seq, + global_start, + l_local, + d_window, + window_size, + ratio, + compressed_width, + seq_major_rows, + compressed_base, + total_width, + index_mode, + ), + grid=(cute.ceil_div(launch_work, 128), 1, 1), + block=(128, 1, 1), + stream=stream, + ) + + +# ============================================================================= +# Torch Wrapper Functions +# ============================================================================= +# 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. +# ============================================================================= + + +def _require_cute(message: str, *tensors: Optional[torch.Tensor]) -> None: + """Raise ``RuntimeError`` when a wrapper cannot use CuTeDSL kernels.""" + if not _CUTE_AVAILABLE or not all(tensor is None or tensor.is_cuda for tensor in tensors): + raise RuntimeError(message) + + +_COMPILED_LAUNCH_CACHE = {} + + +def _run_compiled_launch( + launch_fn, + tensor_args: Tuple[torch.Tensor, ...], + scalar_args: Tuple, + static_arg_indices: Tuple[int, ...] = (), +) -> None: + """Compile/cache a CuTe launch function and invoke it on the current stream.""" + 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 + ), + tuple((i, scalar_args[i]) for i in static_arg_indices), + ) + compiled = _COMPILED_LAUNCH_CACHE.get(key) + if compiled is None: + 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; " + "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) + 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 cutlass.Int32(arg) + for i, arg in enumerate(scalar_args) + ), + make_fake_stream(use_tvm_ffi_env_stream=False), + options=f"--enable-tvm-ffi --gpu-arch {arch}", + ) + _COMPILED_LAUNCH_CACHE[key] = compiled + compiled( + *tensor_args, + *(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 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: 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, + and ``comp_ids`` int32 shape ``(c_cap,)``. + The kernel copies the ratio source tokens for each visible compressed + group and emits each row's compressed group id. + """ + + @staticmethod + def forward( + ctx, + hidden_local: torch.Tensor, + boundary_hidden: torch.Tensor, + cu_seqlens: torch.Tensor, + global_start: int, + ratio: int, + d_comp: 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)) + 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:]) + _run_compiled_launch( + _compressor_input_compact_fwd_launch, + ( + 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, + comp_ids, + ), + ( + cu_seqlens.shape[0] - 1, + int(global_start), + int(l_local), + int(ratio), + int(d_comp), + int(d_window), + compact_len, + row_width, + ), + static_arg_indices=(3, 4, 5, 7), + ) + return hidden_compact, comp_ids + + @staticmethod + 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 + 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( + _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, + l_local + d_window, + ), + static_arg_indices=(3, 4, 5, 7), + ) + return (grad_hidden, grad_boundary, *([None] * 5)) + + +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, + compressed_topk: Optional[torch.Tensor] = None, + cu_seqlens_compressed: Optional[torch.Tensor] = None, + 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: 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. + compressed_topk: optional per-sequence compressed ids, shape + ``(l_local, compressed_width)``. + 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)``; + 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. + """ + 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 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) + 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: + 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 * 128 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, + compressed_topk, + seq_to_rank_row, + topk_idxs, + topk_length_kernel, + indexer_rank_major, + ), + ( + cu_seqlens.shape[0] - 1, + global_start, + l_local, + d_window, + window_size, + ratio, + compressed_width, + seq_to_rank_row.shape[0], + compressed_base, + total_width, + index_mode, + launch_work, + ), + static_arg_indices=(10,), + ) + 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 new file mode 100644 index 00000000000..3459a1954fb --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py @@ -0,0 +1,399 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +"""MCore-facing utilities for the DSv4 THD context-parallel path. + +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, 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.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 + +# ============================================================================= +# RoPE Wrappers +# ============================================================================= + + +def _thd_cp_position_ids( + cu_seqlens_padded: torch.Tensor, global_start: int, local_rows: int +) -> torch.Tensor: + """Map a consecutive CP row interval to positions within packed sequences.""" + global_rows = torch.arange( + int(global_start), + int(global_start) + int(local_rows), + dtype=cu_seqlens_padded.dtype, + device=cu_seqlens_padded.device, + ) + 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) + 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: + # 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, + 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 + + +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 +# ============================================================================= + + +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): + """Receive fixed left-boundary hidden rows needed by this CP 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 + if tensor.shape[0] < d_window: + raise RuntimeError( + "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:])) + + ops = [] + 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): + """Send boundary gradients back to ranks that own those hidden rows.""" + cp_group = ctx.cp_group + cp_size = cp_group.size() + cp_rank = cp_group.rank() + d_window = ctx.d_window + grad_input = grad_boundary.new_zeros(ctx.input_shape) + + ops = [] + 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, + compress_ratio: int, + csa_window_size: int, + cp_group: torch.distributed.ProcessGroup, +) -> torch.Tensor: + """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) + boundary_hidden = _LeftBoundaryExchange.apply(hidden_flat, d_window, cp_group) + return boundary_hidden.reshape((d_window,) + tuple(hidden_states.shape[1:])) + + +# ============================================================================= +# Compressed Metadata And Compressor Inputs +# ============================================================================= + + +def prepare_cp_compressor_input( + hidden_local: torch.Tensor, + boundary_hidden: torch.Tensor, + cu_seqlens: torch.Tensor, + cu_seqlens_compressed: torch.Tensor, + global_start: int, + cp_size: int, + ratio: int, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build fixed-capacity compressor input for this rank's token block. + + Returns: + ``hidden_compact``: rank-local compressor input, shape + ``(compact_group_capacity * ratio, ...)``. + ``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``. + ``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) + 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 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. + 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 + 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( + rank_starts, cu_seqlens[1:], out_int32=True, right=True + ).clamp_max(n_seq - 1) + first_comp_ids = torch.div( + (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[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, 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, + k_indexer_seq_major: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_compressed: torch.Tensor, + global_start: int, + ratio: int, + topk_width: int, + indexer_softmax_scale: float, + max_seqlen_q: int, + use_fused: bool, +) -> 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, 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] + 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]}." + ) + + 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, + 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, (cu_q_topk, cu_k_topk, q_causal_offsets) + + topk, _ = indexer_topk( + q_indexer_local, + k_indexer_seq_major, + 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_kv), + q_causal_offsets=q_causal_offsets, + ) + return topk, (cu_q_topk, cu_k_topk, q_causal_offsets) 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 f6f69ec065a..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, @@ -19,22 +23,13 @@ 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 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: @@ -257,27 +252,51 @@ 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 = 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.") + self.pg_collection.cp = cp_group + + boundary_hidden = None + if use_thd_cp: + boundary_hidden = cp_utils.exchange_cp_boundary_hidden( + hidden_states, + self._dsv4_compress_ratio, + self.config.csa_window_size, + 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( + qkv = self.get_query_key_value_tensors( hidden_states, key_value_states, position_ids, packed_seq_params, 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() @@ -300,9 +319,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': @@ -359,26 +383,50 @@ 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) - # 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( + if 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_fused( + core_attn_out, + rotary_pos_cos, + rotary_pos_sin, + nope_dim, + pos_dim, + cu_seqlens_kv, + global_start, + inverse=True, + ) + else: + if packed_seq: + core_attn_out = core_attn_out.squeeze(1) + # 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. + core_attn_out = fused_mla_rope_out_of_place( + 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) + 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_cos, - rotary_pos_sin, + rotary_pos_emb, nope_dim, pos_dim, cu_seqlens_kv, - self.pg_collection.cp.rank(), - self.pg_collection.cp.size(), + global_start, + self.config, 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 @@ -540,9 +588,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)``. 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] @@ -607,8 +660,11 @@ 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 if packed_seq_params is not None: # If sequence packing, TE expect [t, h, d] shaped qkv input. @@ -616,6 +672,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 @@ -629,7 +687,14 @@ 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, + cp_group, + boundary_kv_compressed=None, + ): """ Apply the up projection and RoPE to the query and key. When sequence packing enabled, the input tensors adopt a packed shape of [t, ...]; @@ -644,115 +709,202 @@ 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 = cp_group.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 and packed_seq: + 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( + q, + rotary_pos_cos, + rotary_pos_sin, + self.config.qk_head_dim, + self.config.qk_pos_emb_head_dim, + cu_seqlens_q, + global_start, + ) + 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, + global_start - boundary_rows, + ) + if boundary_kv_compressed is not None: + boundary_kv = kv[:boundary_rows] + kv = kv[boundary_rows:] + else: + cp_rank = cp_group.rank() + 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, + ) key = kv value = kv else: - q_len = q.size()[0] - if packed_seq_params is None or self.config.context_parallel_size == 1: + if packed_seq and cp_size > 1: + global_start = cp_group.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, + ) + 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) - - # 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 + # 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=cp_group, + 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=cp_group, + 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 = value = kv query = query.contiguous() key = key.contiguous() value = value.contiguous() + if boundary_kv is not None: + boundary_kv = boundary_kv.contiguous() - return query, key, value + 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 = self.qkv_up_checkpoint.checkpoint( - qkv_up_proj_and_rope_apply, q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb - ) + 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, + self.pg_collection.cp, + ) + 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, + self.pg_collection.cp, + boundary_kv_compressed, + ) else: - query, key, value = qkv_up_proj_and_rope_apply( - q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb - ) + 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, 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, + self.pg_collection.cp, + boundary_kv_compressed, + ) - return query, key, value, q_compressed, kv_compressed + 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.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..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`. @@ -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, + q_causal_offsets: Optional[Tensor] = None, ) -> Tuple[Tensor, Tensor, Tensor]: """Layout-agnostic core for :func:`indexer_topk`. @@ -513,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 @@ -525,18 +525,20 @@ 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, - 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). @@ -544,13 +546,12 @@ 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] + 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 = ( @@ -558,6 +559,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" @@ -584,7 +589,16 @@ 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,) + 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=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) + 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 +618,7 @@ def indexer_topk( 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]: """Score + top-K selection for inference (no KL loss, no backward). @@ -625,6 +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. + 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 @@ -638,6 +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 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), @@ -665,6 +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, + q_causal_offsets=q_causal_offsets, ) return topk_indices, topk_length @@ -828,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)``, @@ -843,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, @@ -854,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"] @@ -874,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). @@ -881,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, @@ -888,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"] @@ -1224,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( @@ -1302,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( @@ -1317,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, @@ -1459,6 +1497,248 @@ 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, + 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.""" + _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] + + 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, + ) + + bwd_loss_coeff = loss_coeff * total_q / loss_divisor + unit_grad_loss = torch.ones((), device=query.device, dtype=torch.float32) + + 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) + 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, + 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 8f01db519de..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 @@ -307,6 +308,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 #################### @@ -1457,6 +1461,30 @@ 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.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" + 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' currently is only supported with dsv4_hybrid." + ) + if self.experimental_attention_variant in ["gated_delta_net"]: assert ( self.linear_attention_freq is not None @@ -1530,10 +1558,23 @@ 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]}." ) + 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: @@ -1563,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/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index d3571152966..b283313b963 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 @@ -1321,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'), @@ -1897,18 +1903,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/megatron/training/training.py b/megatron/training/training.py index 77f90134d84..56fb0e71fc4 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -3147,6 +3147,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/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/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index acabd1ceccd..63df0a7e67e 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, diff --git a/tests/unit_tests/test_sequence_packing.py b/tests/unit_tests/test_sequence_packing.py index fe85b3ad8f9..942c15c8e29 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 ( DefaultDynamicCPScheduler, _build_thd_padding_mask, @@ -171,6 +172,143 @@ 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 + + batch = { + "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), + cp_partition_mode="contiguous", + 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"), + [(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( + cp_partition_mode="contiguous", + 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 + + 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 packed_seq_params.cp_partition_mode == "contiguous" + assert pad_input_lengths == [local_target] + + def test_next_hdp_group_packing_aware_can_use_larger_cp_group_for_short_sequences(): micro_batches, leftovers, exec_times, sample_ids = next_hdp_group_packing_aware( [(0, 6144), (1, 2048)], total_gpus=2, max_seq_len_per_rank=4096 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..867b2c9e171 --- /dev/null +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_layout_kernels.py @@ -0,0 +1,682 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from typing import List, Tuple + +import pytest +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. + +_E2E_RAGGED_PADDED_SEG_LENS = (1, 127, 1000, 23, 129, 900, 55, 257, 800, 95, 509, 200) +_E2E_CP_SIZE = 4 + + +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") -> torch.Tensor: + """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 + full e2e hidden sizes in these focused kernel tests. + """ + return torch.tensor( + [0] + list(torch.tensor(_E2E_RAGGED_PADDED_SEG_LENS).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_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 + return local, local + + +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, + ) + + comp_ids = torch.full((c_cap,), -1, dtype=torch.int32, device=hidden_local.device) + for slot, (_, comp_id) in enumerate(groups[:c_cap]): + comp_ids[slot] = comp_id + return hidden_compact, comp_ids + + +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, + seq_to_rank_row: torch.Tensor, + compressed_base: int, + indexer_topk: torch.Tensor = None, +): + """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) + lengths = torch.zeros((int(l_local),), 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()] + + for row in range(int(l_local)): + global_q = int(global_start) + row + 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_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) + 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: + 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 + 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: + 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 + lengths[row] = write_col + return topk, lengths + + +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, + seq_to_rank_row: torch.Tensor, + compressed_base: int, +): + """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) + 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()] + for row in range(int(l_local)): + global_q = int(global_start) + row + 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_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]) + if 0 <= comp_id < comp_len: + seq_major = cu_comp[seq_id] + comp_id + 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 + + 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: + topk[row, compressed_width + window_col] = pos - (int(global_start) - int(d_window)) + + return topk, rank_major + + +def test_compressor_input_compact_matches_native_forward_backward(): + _require_cute_cuda() + cu = _make_e2e_like_cu_seqlens() + global_start, l_local = _e2e_like_local_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, ratio, d_comp, c_cap + ) + fused[0].backward(grad) + assert torch.equal(fused[0], ref[0]) + assert torch.equal(fused[1], ref[1]) + assert torch.equal(hidden.grad, ref_hidden_grad) + assert torch.equal(boundary.grad, ref_boundary_grad) + + +def test_build_attention_indices_matches_native(): + _require_cute_cuda() + cu = _make_e2e_like_cu_seqlens() + 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) + 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, + 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(fused[0], expected[0]) + assert torch.equal(fused[1], expected[1]) + + 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, + global_start, + l_local, + d_window, + window, + ratio, + logical_ids.shape[1], + logical_ids, + cu_seqlens_compressed=cu_comp, + seq_to_rank_row=seq_to_rank_row, + ) + expected_indexer = _native_attention_indices( + cu, + cu_comp, + global_start, + l_local, + d_window, + window, + ratio, + logical_ids.shape[1], + seq_to_rank_row, + compressed_base, + logical_ids, + ) + 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 + ) + # 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")) + + +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_local_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 + 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, + 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(fused[0], expected[0]) + 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))], + 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 new file mode 100644 index 00000000000..7a189e59ca6 --- /dev/null +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py @@ -0,0 +1,300 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from types import SimpleNamespace + +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 ( + apply_thd_cp_local_rope_fused, + apply_thd_cp_local_rope_unfused, + compute_cp_indexer_topk, + exchange_cp_boundary_hidden, + prepare_cp_compressor_input, +) +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 _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 _sequence_positions(cu_seqlens, rows): + seq_ids = torch.bucketize(rows, cu_seqlens[1:], right=True).clamp_max(cu_seqlens.shape[0] - 2) + 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(): + """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, 0, d_window, 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() + + +@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 + ) + 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) + + +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) + 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) + 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): + calls = [] + + 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 + return hidden_compact, comp_ids + + monkeypatch.setattr( + 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(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 = prepare_cp_compressor_input( + hidden, boundary, cu, cu_comp, 0, cp_size=2, ratio=4 + ) + + assert calls == [(0, 16, 8, (2, 1))] + assert hidden_compact.shape == (32, 1) + 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_passes_offsets_without_repacking_k(monkeypatch): + topk_calls = [] + + def fake_indexer_topk( + 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(), 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(8, 2) + weights = torch.randn(8, 1) + k_seq = torch.arange(8, dtype=torch.float32).reshape(4, 2) + cu_q = torch.tensor([0, 5, 13, 20], dtype=torch.int32) + cu_comp = torch.tensor([0, 1, 3, 4], dtype=torch.int32) + + out, metadata = compute_cp_indexer_topk( + q, + weights, + k_seq, + cu_q, + cu_comp, + 7, + ratio=4, + topk_width=2, + indexer_softmax_scale=0.5, + max_seqlen_q=8, + use_fused=True, + ) + + 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) + 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): + 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_dsa_kernels.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py index 096c73d1451..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 @@ -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' @@ -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) @@ -2637,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)``. """ @@ -2758,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' @@ -3105,7 +3094,9 @@ 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) + 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). @@ -3131,7 +3122,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) @@ -3154,7 +3145,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) ----------------------------- @@ -3270,7 +3263,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' @@ -3404,7 +3397,17 @@ 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 = ( + _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( + 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 new file mode 100644 index 00000000000..f6f59ce129c --- /dev/null +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py @@ -0,0 +1,1078 @@ +# 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 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" + +# 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 +# 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) + + +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 < 9: + 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 SM90+, flash_mla, and cudnn.DSA" +) + + +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', + cp_partition_mode='contiguous', + ) + + +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, +): + """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, + 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, + 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, + ) + + +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_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) + 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): + """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, 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"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.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() + # 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", "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, sparse_loss, apply_rope_fusion + ): + """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. + """ + 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) + 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=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=sparse_loss, + apply_dsa_kernel_fusion=apply_dsa_kernel_fusion, + apply_rope_fusion=apply_rope_fusion, + ) + 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}:dsa={apply_dsa_kernel_fusion}:rope={apply_rope_fusion}: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}:dsa={apply_dsa_kernel_fusion}:rope={apply_rope_fusion}: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}: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() + + @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=2, pg_collection=self.ref_pg + ).cuda() + recompute_attn = _build_attention( + recompute_config, layer_number=2, 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) + + 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, + ) + 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), "layer=2:eval: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], + ids=["ratio_0_window_only", "ratio_4_indexer", "ratio_128_compressor"], + ) + @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 + 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. + """ + if (dsa_fused or rope_fused) and not self.fused_kernels_available: + pytest.skip(_DSV4_CP_FUSED_KERNELS_UNAVAILABLE_REASON) + + 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) + + 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=dsa_fused, + apply_rope_fusion=rope_fused, + ) + 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 + } + # 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 + ) + torch.cuda.synchronize() + assert graph_param_grads.keys() == eager_param_grads.keys() + 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 dsa_fused else _assert_cp_graph_bitwise_match + ) + bwd_match_fn( + 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}:param_grad:{name}", + ) + + del eager_attn, test_grad, eager_hidden, graph_out, graph_hidden_grad + _clear_cuda_test_state() + + @pytest.mark.parametrize( + "layer_number", [2, 3], ids=["ratio_4_indexer", "ratio_128_compressor"] + ) + 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 + 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. + """ + 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 + ) + 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_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) + 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, + ) + 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:output" + ) + _assert_cp_graph_fused_grad_match( + graph_hidden_grad, + eager_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: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"], + ) + 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 + 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(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, + ) + 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, + 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"], + ) + 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 + 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(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, + ) + 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, 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() 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..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 @@ -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,8 +931,15 @@ 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 == 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") @@ -1021,8 +1044,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 +1173,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 +1303,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" diff --git a/tests/unit_tests/transformer/test_thd_cuda_graph.py b/tests/unit_tests/transformer/test_thd_cuda_graph.py index 92a81b2fcdc..b8fe2de5896 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() @@ -474,6 +477,40 @@ 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): @@ -534,6 +571,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 @@ -541,6 +580,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) @@ -578,6 +618,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: 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 = [